[ENH] Add Schema to js client

Author: jairad26

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

@@ -22,6 +22,7 @@ import {
   CreateCollectionConfiguration,
   processCreateCollectionConfig,
 } from "./collection-configuration";
+import { Schema } from "./schema";
 
 /**
  * Configuration options for the ChromaClient.
@@ -217,21 +218,21 @@ 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) =>
+        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,
+          schema: Schema.deserializeFromJSON(collection.schema ?? undefined),
+        }),
       ),
     );
   }
@@ -264,11 +265,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,9 +287,12 @@ export class ChromaClient {
         configuration: collectionConfig,
         metadata: serializeMetadata(metadata),
         get_or_create: false,
+        schema: schema ? schema.serializeToJSON() : undefined,
       },
     });
 
+    const serverSchema = Schema.deserializeFromJSON(data.schema ?? undefined);
+
     return new CollectionImpl({
       chromaClient: this,
       apiClient: this.apiClient,
@@ -300,6 +306,7 @@ export class ChromaClient {
           data.configuration_json.embedding_function ?? undefined,
         )),
       id: data.id,
+      schema: serverSchema,
     });
   }
 
@@ -323,6 +330,8 @@ export class ChromaClient {
       path: { ...(await this._path()), collection_id: name },
     });
 
+    const schema = Schema.deserializeFromJSON(data.schema ?? undefined);
+
     return new CollectionImpl({
       chromaClient: this,
       apiClient: this.apiClient,
@@ -336,6 +345,7 @@ export class ChromaClient {
           data.configuration_json.embedding_function ?? undefined,
         ),
       id: data.id,
+      schema,
     });
   }
 
@@ -382,11 +392,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,9 +414,12 @@ export class ChromaClient {
         configuration: collectionConfig,
         metadata: serializeMetadata(metadata),
         get_or_create: true,
+        schema: schema ? schema.serializeToJSON() : undefined,
       },
     });
 
+    const serverSchema = Schema.deserializeFromJSON(data.schema ?? undefined);
+
     return new CollectionImpl({
       chromaClient: this,
       apiClient: this.apiClient,
@@ -418,6 +433,7 @@ export class ChromaClient {
           data.configuration_json.embedding_function ?? undefined,
         )),
       id: data.id,
+      schema: serverSchema,
     });
   }
 

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

@@ -40,6 +40,7 @@ import {
   UpdateCollectionConfiguration,
 } from "./collection-configuration";
 import { SearchLike, SearchResult, toSearch } from "./execution/expression";
+import { Schema, EMBEDDING_KEY } from "./schema";
 
 /**
  * Interface for collection operations using collection ID.
@@ -56,6 +57,8 @@ export interface Collection {
   configuration: CollectionConfiguration;
   /** Optional embedding function. Must match the one used to create the collection. */
   embeddingFunction?: EmbeddingFunction;
+  /** Collection schema describing index configuration */
+  schema?: Schema;
   /** Gets the total number of records in the collection */
   count(): Promise<number>;
   /**
@@ -212,6 +215,8 @@ export interface CollectionArgs {
   configuration: CollectionConfiguration;
   /** Optional collection metadata */
   metadata?: CollectionMetadata;
+  /** Optional schema returned by the server */
+  schema?: Schema;
 }
 
 /**
@@ -226,6 +231,7 @@ export class CollectionImpl implements Collection {
   private _metadata: CollectionMetadata | undefined;
   private _configuration: CollectionConfiguration;
   protected _embeddingFunction: EmbeddingFunction | undefined;
+  protected _schema: Schema | undefined;
 
   /**
    * Creates a new CollectionAPIImpl instance.
@@ -239,6 +245,7 @@ export class CollectionImpl implements Collection {
     metadata,
     configuration,
     embeddingFunction,
+    schema,
   }: CollectionArgs) {
     this.chromaClient = chromaClient;
     this.apiClient = apiClient;
@@ -247,6 +254,7 @@ export class CollectionImpl implements Collection {
     this._metadata = metadata;
     this._configuration = configuration;
     this._embeddingFunction = embeddingFunction;
+    this._schema = schema;
   }
 
   public get name(): string {
@@ -283,6 +291,14 @@ export class CollectionImpl implements Collection {
     this._embeddingFunction = embeddingFunction;
   }
 
+  public get schema(): Schema | undefined {
+    return this._schema;
+  }
+
+  protected set schema(schema: Schema | undefined) {
+    this._schema = schema;
+  }
+
   protected async path(): Promise<{
     tenant: string;
     database: string;
@@ -296,17 +312,36 @@ export class CollectionImpl implements Collection {
   }
 
   private async embed(inputs: string[], isQuery: boolean): Promise<number[][]> {
-    if (!this._embeddingFunction) {
+    const embeddingFunction =
+      this._embeddingFunction ?? this.getSchemaEmbeddingFunction();
+
+    if (!embeddingFunction) {
       throw new ChromaValueError(
         "Embedding function must be defined for operations requiring embeddings.",
       );
     }
 
-	if (this._embeddingFunction.generateForQueries && isQuery) {
-		return await this._embeddingFunction.generateForQueries(inputs);
-	} else {
-    	return await this._embeddingFunction.generate(inputs);
-	}
+    if (isQuery && embeddingFunction.generateForQueries) {
+      return await embeddingFunction.generateForQueries(inputs);
+    }
+
+    return await embeddingFunction.generate(inputs);
+  }
+
+  private getSchemaEmbeddingFunction(): EmbeddingFunction | undefined {
+    const schema = this._schema;
+    if (!schema) return undefined;
+
+    const schemaOverride = schema.keys[EMBEDDING_KEY];
+    const overrideFunction = schemaOverride?.floatList?.vectorIndex?.config
+      .embeddingFunction;
+    if (overrideFunction) {
+      return overrideFunction;
+    }
+
+    const defaultFunction = schema.defaults.floatList?.vectorIndex?.config
+      .embeddingFunction;
+    return defaultFunction ?? undefined;
   }
 
   private async prepareRecords<T extends boolean = false>({

clients/new-js/packages/chromadb/src/embedding-function.ts

@@ -130,6 +130,11 @@ export const knownSparseEmbeddingFunctions = new Map<
   SparseEmbeddingFunctionClass
 >();
 
+/**
+ * Union type covering both dense and sparse embedding functions.
+ */
+export type AnyEmbeddingFunction = EmbeddingFunction | SparseEmbeddingFunction;
+
 /**
  * Registers an embedding function in the global registry.
  * @param name - Unique name for the embedding function
@@ -170,9 +175,9 @@ export const registerSparseEmbeddingFunction = (
  * Retrieves and instantiates an embedding function from configuration.
  * @param collectionName - Name of the collection (for error messages)
  * @param efConfig - Configuration for the embedding function
- * @returns Promise resolving to an EmbeddingFunction instance
+ * @returns EmbeddingFunction instance or undefined if it cannot be constructed
  */
-export const getEmbeddingFunction = async (
+export const getEmbeddingFunction = (
   collectionName: string,
   efConfig?: EmbeddingFunctionConfiguration,
 ) => {
@@ -235,9 +240,9 @@ export const getEmbeddingFunction = async (
  * Retrieves and instantiates a sparse embedding function from configuration.
  * @param collectionName - Name of the collection (for error messages)
  * @param efConfig - Configuration for the sparse embedding function
- * @returns Promise resolving to a SparseEmbeddingFunction instance
+ * @returns SparseEmbeddingFunction instance or undefined if it cannot be constructed
  */
-export const getSparseEmbeddingFunction = async (
+export const getSparseEmbeddingFunction = (
   collectionName: string,
   efConfig?: EmbeddingFunctionConfiguration,
 ) => {

clients/new-js/packages/chromadb/src/index.ts

@@ -16,3 +16,4 @@ export * from "./cloud-client";
 export * from "./errors";
 export * from "./collection-configuration";
 export * from "./execution";
+export * from "./schema";