[ENH] Add Schema to js client (#5621)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
- This PR adds schema support to the js client, along with tests, and
logic to embed sparse vectors using efs, using the schema for dense
vecs, and tests to ensure serialization and deserialization work
- New functionality
  - ...

## Test plan

_How are these changes tested?_

added schema unit tests matching python ones

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

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility
changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: jairad26

clients/new-js/packages/chromadb/src/chroma-client.ts

@@ -22,6 +22,23 @@ import {
   CreateCollectionConfiguration,
   processCreateCollectionConfig,
 } from "./collection-configuration";
+import { EMBEDDING_KEY, Schema } from "./schema";
+
+const resolveSchemaEmbeddingFunction = (
+  schema: Schema | undefined,
+): EmbeddingFunction | undefined => {
+  if (!schema) {
+    return undefined;
+  }
+
+  const embeddingOverride =
+    schema.keys[EMBEDDING_KEY]?.floatList?.vectorIndex?.config.embeddingFunction ?? undefined;
+  if (embeddingOverride) {
+    return embeddingOverride;
+  }
+
+  return schema.defaults.floatList?.vectorIndex?.config.embeddingFunction ?? undefined;
+};
 
 /**
  * Configuration options for the ChromaClient.
@@ -217,22 +234,27 @@ export class ChromaClient {
     });
 
     return Promise.all(
-      data.map(
-        async (collection) =>
-          new CollectionImpl({
-            chromaClient: this,
-            apiClient: this.apiClient,
-            name: collection.name,
-            id: collection.id,
-            embeddingFunction: await getEmbeddingFunction(
-              collection.name,
-              collection.configuration_json.embedding_function ?? undefined,
-            ),
-            configuration: collection.configuration_json,
-            metadata:
-              deserializeMetadata(collection.metadata ?? undefined) ?? undefined,
-          }),
-      ),
+      data.map(async (collection) => {
+        const schema = Schema.deserializeFromJSON(collection.schema ?? undefined);
+        const schemaEmbeddingFunction = resolveSchemaEmbeddingFunction(schema);
+        const resolvedEmbeddingFunction =
+          getEmbeddingFunction(
+            collection.name,
+            collection.configuration_json.embedding_function ?? undefined,
+          ) ?? schemaEmbeddingFunction;
+
+        return new CollectionImpl({
+          chromaClient: this,
+          apiClient: this.apiClient,
+          name: collection.name,
+          id: collection.id,
+          embeddingFunction: resolvedEmbeddingFunction,
+          configuration: collection.configuration_json,
+          metadata:
+            deserializeMetadata(collection.metadata ?? undefined) ?? undefined,
+          schema,
+        });
+      }),
     );
   }
 
@@ -264,11 +286,13 @@ export class ChromaClient {
     configuration,
     metadata,
     embeddingFunction,
+    schema,
   }: {
     name: string;
     configuration?: CreateCollectionConfiguration;
     metadata?: CollectionMetadata;
     embeddingFunction?: EmbeddingFunction | null;
+    schema?: Schema;
   }): Promise<Collection> {
     const collectionConfig = await processCreateCollectionConfig({
       configuration,
@@ -284,22 +308,29 @@ export class ChromaClient {
         configuration: collectionConfig,
         metadata: serializeMetadata(metadata),
         get_or_create: false,
+        schema: schema ? schema.serializeToJSON() : undefined,
       },
     });
 
+    const serverSchema = Schema.deserializeFromJSON(data.schema ?? undefined);
+    const schemaEmbeddingFunction = resolveSchemaEmbeddingFunction(serverSchema);
+    const resolvedEmbeddingFunction =
+      embeddingFunction ??
+      getEmbeddingFunction(
+        data.name,
+        data.configuration_json.embedding_function ?? undefined,
+      ) ??
+      schemaEmbeddingFunction;
+
     return new CollectionImpl({
       chromaClient: this,
       apiClient: this.apiClient,
       name,
       configuration: data.configuration_json,
       metadata: deserializeMetadata(data.metadata ?? undefined) ?? undefined,
-      embeddingFunction:
-        embeddingFunction ??
-        (await getEmbeddingFunction(
-          data.name,
-          data.configuration_json.embedding_function ?? undefined,
-        )),
+      embeddingFunction: resolvedEmbeddingFunction,
       id: data.id,
+      schema: serverSchema,
     });
   }
 
@@ -323,19 +354,25 @@ export class ChromaClient {
       path: { ...(await this._path()), collection_id: name },
     });
 
+    const schema = Schema.deserializeFromJSON(data.schema ?? undefined);
+    const schemaEmbeddingFunction = resolveSchemaEmbeddingFunction(schema);
+    const resolvedEmbeddingFunction =
+      embeddingFunction ??
+      getEmbeddingFunction(
+        data.name,
+        data.configuration_json.embedding_function ?? undefined,
+      ) ??
+      schemaEmbeddingFunction;
+
     return new CollectionImpl({
       chromaClient: this,
       apiClient: this.apiClient,
       name,
       configuration: data.configuration_json,
       metadata: deserializeMetadata(data.metadata ?? undefined) ?? undefined,
-      embeddingFunction: embeddingFunction
-        ? embeddingFunction
-        : await getEmbeddingFunction(
-          data.name,
-          data.configuration_json.embedding_function ?? undefined,
-        ),
+      embeddingFunction: resolvedEmbeddingFunction,
       id: data.id,
+      schema,
     });
   }
 
@@ -382,11 +419,13 @@ export class ChromaClient {
     configuration,
     metadata,
     embeddingFunction,
+    schema,
   }: {
     name: string;
     configuration?: CreateCollectionConfiguration;
     metadata?: CollectionMetadata;
     embeddingFunction?: EmbeddingFunction | null;
+    schema?: Schema;
   }): Promise<Collection> {
     const collectionConfig = await processCreateCollectionConfig({
       configuration,
@@ -402,22 +441,29 @@ export class ChromaClient {
         configuration: collectionConfig,
         metadata: serializeMetadata(metadata),
         get_or_create: true,
+        schema: schema ? schema.serializeToJSON() : undefined,
       },
     });
 
+    const serverSchema = Schema.deserializeFromJSON(data.schema ?? undefined);
+    const schemaEmbeddingFunction = resolveSchemaEmbeddingFunction(serverSchema);
+    const resolvedEmbeddingFunction =
+      embeddingFunction ??
+      getEmbeddingFunction(
+        name,
+        data.configuration_json.embedding_function ?? undefined,
+      ) ??
+      schemaEmbeddingFunction;
+
     return new CollectionImpl({
       chromaClient: this,
       apiClient: this.apiClient,
       name,
       configuration: data.configuration_json,
       metadata: deserializeMetadata(data.metadata ?? undefined) ?? undefined,
-      embeddingFunction:
-        embeddingFunction ??
-        (await getEmbeddingFunction(
-          name,
-          data.configuration_json.embedding_function ?? undefined,
-        )),
+      embeddingFunction: resolvedEmbeddingFunction,
       id: data.id,
+      schema: serverSchema,
     });
   }
 

clients/new-js/packages/chromadb/src/collection-configuration.ts

@@ -164,10 +164,10 @@ export const processUpdateCollectionConfig = async ({
 
   const embeddingFunction =
     currentEmbeddingFunction ||
-    (await getEmbeddingFunction(
+    getEmbeddingFunction(
       collectionName,
       currentConfiguration.embeddingFunction ?? undefined,
-    ));
+    );
 
   const newEmbeddingFunction = newConfiguration.embeddingFunction;