[ENH] Add Schema to js client

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;