[ENH] Add collection config to js client

Author: jairad26

Dockerfile

@@ -11,57 +11,55 @@ RUN apt-get update --fix-missing && apt-get install -y --fix-missing \
     unzip \
     curl \
     make && \
+    curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain stable && \
     rm -rf /var/lib/apt/lists/* && \
     mkdir /install
-
+ENV PATH="/root/.cargo/bin:$PATH"
 # Install specific Protobuf compiler (v28.2)
 RUN curl -LO https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-x86_64.zip && \
     unzip protoc-${PROTOBUF_VERSION}-linux-x86_64.zip -d /usr/local/ && \
     rm protoc-${PROTOBUF_VERSION}-linux-x86_64.zip && \
     chmod +x /usr/local/bin/protoc && \
     protoc --version  # Verify installed version
-
 WORKDIR /install
 
 COPY ./requirements.txt requirements.txt
 
+RUN --mount=type=cache,target=/root/.cache/pip pip install maturin cffi patchelf
 RUN --mount=type=cache,target=/root/.cache/pip pip install --upgrade --prefix="/install" -r requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip if [ "$REBUILD_HNSWLIB" = "true" ]; then pip install --no-binary :all: --force-reinstall --prefix="/install" chroma-hnswlib; fi
 
 # Install gRPC tools for Python with fixed version
 RUN pip install grpcio==1.58.0 grpcio-tools==1.58.0
-
 # Copy source files to build Protobufs
 COPY ./ /chroma
 
 # Generate Protobufs
 WORKDIR /chroma
 RUN make -C idl proto_python
+RUN python3 -m maturin build
+RUN pip uninstall chromadb -y
+RUN pip install --prefix="/install" --find-links target/wheels/ --upgrade  chromadb
 
 FROM python:3.11-slim-bookworm AS final
 
 # Create working directory
 RUN mkdir /chroma
 WORKDIR /chroma
-
 # Copy entrypoint
 COPY ./bin/docker_entrypoint.sh /docker_entrypoint.sh
-
 RUN apt-get update --fix-missing && apt-get install -y curl && \
     chmod +x /docker_entrypoint.sh && \
     rm -rf /var/lib/apt/lists/*
-
 # Copy built dependencies and generated Protobufs
 COPY --from=builder /install /usr/local
 COPY --from=builder /chroma /chroma
-
 ENV CHROMA_HOST_ADDR="0.0.0.0"
 ENV CHROMA_HOST_PORT=8000
 ENV CHROMA_WORKERS=1
 ENV CHROMA_LOG_CONFIG="chromadb/log_config.yml"
 ENV CHROMA_TIMEOUT_KEEP_ALIVE=30
-
 EXPOSE 8000
 
 ENTRYPOINT ["/docker_entrypoint.sh"]
-CMD [ "--workers ${CHROMA_WORKERS} --host ${CHROMA_HOST_ADDR} --port ${CHROMA_HOST_PORT} --proxy-headers --log-config ${CHROMA_LOG_CONFIG} --timeout-keep-alive ${CHROMA_TIMEOUT_KEEP_ALIVE}"]
+CMD [ "--workers ${CHROMA_WORKERS} --host ${CHROMA_HOST_ADDR} --port ${CHROMA_HOST_PORT} --proxy-headers --reload --log-config ${CHROMA_LOG_CONFIG} --timeout-keep-alive ${CHROMA_TIMEOUT_KEEP_ALIVE}"]

clients/js/package.json

@@ -65,14 +65,15 @@
   },
   "dependencies": {
     "cliui": "^8.0.1",
-    "isomorphic-fetch": "^3.0.0"
+    "isomorphic-fetch": "^3.0.0",
+    "os": "^0.1.2"
   },
   "peerDependencies": {
     "@google/generative-ai": "^0.1.1",
     "cohere-ai": "^5.0.0 || ^6.0.0 || ^7.0.0",
+    "ollama": "^0.5.0",
     "openai": "^3.0.0 || ^4.0.0",
-    "voyageai": "^0.0.3-1",
-    "ollama": "^0.5.0"
+    "voyageai": "^0.0.3-1"
   },
   "peerDependenciesMeta": {
     "@google/generative-ai": {

clients/js/packages/chromadb-core/src/ChromaClient.ts

@@ -4,7 +4,7 @@ import { authOptionsToAuthProvider, ClientAuthProvider } from "./auth";
 import { chromaFetch } from "./ChromaFetch";
 import { Collection } from "./Collection";
 import { DefaultEmbeddingFunction } from "./embeddings/DefaultEmbeddingFunction";
-import { Configuration, ApiApi as DefaultApi } from "./generated";
+import { Configuration, ApiApi as DefaultApi, Api } from "./generated";
 import type {
   ChromaClientParams,
   CollectionMetadata,
@@ -18,7 +18,7 @@ import type {
   UserIdentity,
 } from "./types";
 import { validateTenantDatabase, wrapCollection } from "./utils";
-
+import { loadApiCollectionConfigurationFromCreateCollectionConfiguration } from "./CollectionConfiguration";
 const DEFAULT_TENANT = "default_tenant";
 const DEFAULT_DATABASE = "default_database";
 
@@ -224,25 +224,39 @@ export class ChromaClient {
     name,
     metadata,
     embeddingFunction = new DefaultEmbeddingFunction(),
+    configuration,
   }: CreateCollectionParams): Promise<Collection> {
     await this.init();
+    let collectionConfiguration: Api.CollectionConfiguration | undefined =
+      undefined;
+    if (configuration) {
+      collectionConfiguration =
+        loadApiCollectionConfigurationFromCreateCollectionConfiguration(
+          configuration,
+        );
+    }
     const newCollection = await this.api.createCollection(
       this.tenant,
       this.database,
       {
         name,
-        // @ts-ignore: we need to generate the client libraries again
-        configuration: null, //TODO: Configuration type in JavaScript
+        configuration: collectionConfiguration,
         metadata: metadata,
       },
       this.api.options,
     );
 
+    console.log(
+      "NEW CREATE COLLECTION CONFIGURATION: ",
+      newCollection.configuration_json,
+    );
+
     return wrapCollection(this, {
       name: newCollection.name,
       id: newCollection.id,
       metadata: newCollection.metadata as CollectionMetadata | undefined,
       embeddingFunction,
+      configuration: newCollection.configuration_json,
     });
   }
 
@@ -286,11 +300,18 @@ export class ChromaClient {
       this.api.options,
     );
 
+    console.log(
+      "NEW GET_OR_CREATE COLLECTION CONFIGURATION: ",
+      newCollection.configuration_json,
+    );
+
     return wrapCollection(this, {
       name: newCollection.name,
       id: newCollection.id,
       metadata: newCollection.metadata as CollectionMetadata | undefined,
       embeddingFunction,
+      // TODO: once server returns configuration, add it here
+      configuration: newCollection.configuration_json,
     });
   }
 
@@ -412,6 +433,7 @@ export class ChromaClient {
       this.api.options,
     );
 
+    console.log("GET COLLECTION CONFIGURATION: ", response.configuration_json);
     return wrapCollection(this, {
       id: response.id,
       name: response.name,
@@ -420,6 +442,7 @@ export class ChromaClient {
         embeddingFunction !== undefined
           ? embeddingFunction
           : new DefaultEmbeddingFunction(),
+      configuration: response.configuration_json,
     });
   }
 

clients/js/packages/chromadb-core/src/Collection.ts

@@ -16,7 +16,11 @@ import {
   CollectionParams,
 } from "./types";
 import { prepareRecordRequest, toArray, toArrayOfArrays } from "./utils";
-import { Api as GeneratedApi } from "./generated";
+import { Api as GeneratedApi, Api } from "./generated";
+import {
+  UpdateCollectionConfiguration,
+  loadApiUpdateCollectionConfigurationFromUpdateCollectionConfiguration,
+} from "./CollectionConfiguration";
 
 export class Collection {
   public name: string;
@@ -31,6 +35,8 @@ export class Collection {
    */
   public embeddingFunction: IEmbeddingFunction;
 
+  public configuration: Api.CollectionConfiguration | undefined;
+
   /**
    * @ignore
    */
@@ -40,12 +46,14 @@ export class Collection {
     client: ChromaClient,
     embeddingFunction: IEmbeddingFunction,
     metadata?: CollectionMetadata,
+    configuration?: Api.CollectionConfiguration,
   ) {
     this.name = name;
     this.id = id;
     this.metadata = metadata;
     this.client = client;
     this.embeddingFunction = embeddingFunction;
+    this.configuration = configuration;
   }
 
   /**
@@ -318,19 +326,30 @@ export class Collection {
   async modify({
     name,
     metadata,
+    configuration,
   }: {
     name?: string;
     metadata?: CollectionMetadata;
+    configuration?: UpdateCollectionConfiguration;
   }): Promise<CollectionParams> {
     await this.client.init();
-
+    let updateCollectionConfiguration:
+      | Api.UpdateCollectionConfiguration
+      | undefined = undefined;
+    if (configuration) {
+      updateCollectionConfiguration =
+        loadApiUpdateCollectionConfigurationFromUpdateCollectionConfiguration(
+          configuration,
+        );
+    }
     const resp = (await this.client.api.updateCollection(
       this.client.tenant,
       this.client.database,
       this.id,
       {
         new_name: name,
         new_metadata: metadata,
+        new_configuration: updateCollectionConfiguration,
       },
       this.client.api.options,
     )) as CollectionParams;