[ENH] Add collection config to js client (#4116)

## Description of changes

This PR adds support for collection configuration in the javascript client. it updates tests, functionality, and introduces a breaking change involving how the rust server responds with include

## Test plan
*How are these changes tested?*

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?*

Author: jairad26

chromadb/types.py

@@ -148,7 +148,7 @@ def get_configuration(self) -> CollectionConfiguration:
             return load_collection_configuration_from_json(self.configuration_json)
         except Exception as e:
             warnings.warn(
-                f"server does not respond with configuration_json. Please update server: {e}",
+                f"Server does not respond with configuration_json. Please update server: {e}",
                 DeprecationWarning,
                 stacklevel=2,
             )
@@ -181,7 +181,7 @@ def from_json(cls, json_map: Dict[str, Any]) -> Self:
             configuration = load_collection_configuration_from_json(configuration_json)
         except Exception as e:
             warnings.warn(
-                f"server does not respond with configuration_json. Please update server: {e}",
+                f"Server does not respond with configuration_json. Please update server: {e}",
                 DeprecationWarning,
                 stacklevel=2,
             )

clients/js/README.md

@@ -4,6 +4,8 @@ Chroma is the open-source embedding database. Chroma makes it easy to build LLM
 
 This package gives you a JS/TS interface to talk to a backend Chroma DB over REST.
 
+**Note:** JS client version 3._ is only compatible with chromadb v1.0.6 and newer or Chroma Cloud. For prior version compatiblity, please use JS client version 2._.
+
 [Learn more about Chroma](https://github.com/chroma-core/chroma)
 
 - [💬 Community Discord](https://discord.gg/MMeYNTmh3x)

clients/js/package.json

@@ -1,7 +1,7 @@
 {
   "name": "chromadb-root",
   "private": true,
-  "version": "2.2.1",
+  "version": "3.0.0",
   "description": "A JavaScript interface for chroma",
   "keywords": [],
   "author": "",
@@ -70,9 +70,9 @@
   "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-client/README.md

@@ -2,6 +2,8 @@
 
 Chroma is the open-source embedding database. Chroma makes it easy to build LLM apps by making knowledge, facts, and skills pluggable for LLMs.
 
+**Note:** JS client version 3._ is only compatible with chromadb v1.0.6 and newer or Chroma Cloud. For prior version compatiblity, please use JS client version 2._.
+
 **This package provides embedding libraries as peer dependencies**, allowing you to manage your own versions of embedding libraries and keep your dependency tree lean by not bundling dependencies you don't use. For a thick client with bundled embedding functions, install `chromadb`.
 
 ## Features

clients/js/packages/chromadb-client/package.json

@@ -1,6 +1,6 @@
 {
   "name": "chromadb-client",
-  "version": "2.2.1",
+  "version": "3.0.0",
   "description": "A JavaScript interface for chroma with embedding functions as peer dependencies",
   "license": "Apache-2.0",
   "type": "module",

clients/js/packages/chromadb-core/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@internal/chromadb-core",
-  "version": "2.2.1",
+  "version": "3.0.0",
   "private": true,
   "description": "Core functionality for ChromaDB JavaScript client",
   "license": "Apache-2.0",

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,8 @@ import type {
   UserIdentity,
 } from "./types";
 import { validateTenantDatabase, wrapCollection } from "./utils";
-
+import { loadApiCollectionConfigurationFromCreateCollectionConfiguration } from "./CollectionConfiguration";
+import { warn } from "console";
 const DEFAULT_TENANT = "default_tenant";
 const DEFAULT_DATABASE = "default_database";
 
@@ -224,25 +225,43 @@ 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,
     );
 
+    let config: Api.CollectionConfiguration = {};
+    try {
+      config = newCollection.configuration_json;
+    } catch {
+      warn(
+        "Server does not respond with configuration_json. Please update server",
+      );
+    }
+
     return wrapCollection(this, {
       name: newCollection.name,
       id: newCollection.id,
       metadata: newCollection.metadata as CollectionMetadata | undefined,
       embeddingFunction,
+      configuration: config,
     });
   }
 
@@ -271,26 +290,47 @@ export class ChromaClient {
     name,
     metadata,
     embeddingFunction = new DefaultEmbeddingFunction(),
+    configuration,
   }: GetOrCreateCollectionParams): 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,
         get_or_create: true,
       },
       this.api.options,
     );
 
+    let config: Api.CollectionConfiguration = {};
+    try {
+      config = newCollection.configuration_json;
+    } catch {
+      warn(
+        "Server does not respond with configuration_json. Please update server",
+      );
+    }
+
     return wrapCollection(this, {
       name: newCollection.name,
       id: newCollection.id,
       metadata: newCollection.metadata as CollectionMetadata | undefined,
       embeddingFunction,
+      configuration: config,
     });
   }
 
@@ -412,6 +452,15 @@ export class ChromaClient {
       this.api.options,
     );
 
+    let config: Api.CollectionConfiguration = {};
+    try {
+      config = response.configuration_json;
+    } catch {
+      warn(
+        "Server does not respond with configuration_json. Please update server",
+      );
+    }
+
     return wrapCollection(this, {
       id: response.id,
       name: response.name,
@@ -420,6 +469,7 @@ export class ChromaClient {
         embeddingFunction !== undefined
           ? embeddingFunction
           : new DefaultEmbeddingFunction(),
+      configuration: config,
     });
   }
 

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

@@ -14,9 +14,15 @@ import {
   DeleteParams,
   Embeddings,
   CollectionParams,
+  IncludeEnum,
+  Metadata,
 } 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 +37,8 @@ export class Collection {
    */
   public embeddingFunction: IEmbeddingFunction;
 
+  public configuration: Api.CollectionConfiguration | undefined;
+
   /**
    * @ignore
    */
@@ -40,12 +48,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;
   }
 
   /**
@@ -172,7 +182,7 @@ export class Collection {
 
     const idsArray = ids ? toArray(ids) : undefined;
 
-    const resp = (await this.client.api.collectionGet(
+    const resp = await this.client.api.collectionGet(
       this.client.tenant,
       this.client.database,
       this.id,
@@ -185,9 +195,17 @@ export class Collection {
         where_document: whereDocument,
       },
       this.client.api.options,
-    )) as unknown as GetResponse;
+    );
 
-    return resp;
+    const finalResp: GetResponse = {
+      ...resp,
+      metadatas: resp.metadatas as (Metadata | null)[],
+      documents: resp.documents as (string | null)[],
+      embeddings: resp.embeddings as Embeddings | null,
+      included: resp.include as unknown as IncludeEnum[],
+    };
+
+    return finalResp;
   }
 
   /**
@@ -281,7 +299,7 @@ export class Collection {
       );
     }
 
-    const resp = (await this.client.api.collectionQuery(
+    const resp = await this.client.api.collectionQuery(
       this.client.tenant,
       this.client.database,
       this.id,
@@ -295,9 +313,18 @@ export class Collection {
         include: include as GeneratedApi.Include[] | undefined,
       },
       this.client.api.options,
-    )) as unknown as MultiQueryResponse;
+    );
 
-    return resp;
+    const finalResp: MultiQueryResponse = {
+      ...resp,
+      metadatas: resp.metadatas as (Metadata | null)[][],
+      documents: resp.documents as (string | null)[][],
+      embeddings: resp.embeddings as Embeddings[] | null,
+      distances: resp.distances as number[][] | null,
+      included: resp.include as unknown as IncludeEnum[],
+    };
+
+    return finalResp;
   }
 
   /**
@@ -318,19 +345,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;