[ENH] Add collection config to js client

Author: jairad26

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;