chromadb 1.7.2 → 1.7.3-beta2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "chromadb",
-  "version": "1.7.2",
+  "version": "1.7.3-beta2",
   "description": "A JavaScript interface for chroma",
   "keywords": [],
   "author": "",
@@ -43,7 +43,7 @@
   "scripts": {
     "test": "run-s db:clean db:cleanauth db:run test:runfull db:clean test:runfull-authonly db:cleanauth",
     "testnoauth": "run-s db:clean db:run test:runfull db:clean",
-    "testauth": "run-s db:cleanauth test:runfull-authonly db:cleanauth",
+    "testauth": "run-s db:cleanauth db:run-auth test:runfull-authonly db:cleanauth",
     "test:set-port": "cross-env URL=localhost:8001",
     "test:run": "jest --runInBand --testPathIgnorePatterns=test/auth.*.test.ts",
     "test:run-auth-basic": "jest --runInBand --testPathPattern=test/auth.basic.test.ts",
@@ -76,14 +76,10 @@
     "cliui": "^8.0.1"
   },
   "peerDependencies": {
-    "@google/generative-ai": "^0.1.1",
     "cohere-ai": "^5.0.0 || ^6.0.0",
     "openai": "^3.0.0 || ^4.0.0"
   },
   "peerDependenciesMeta": {
-    "@google/generative-ai": {
-      "optional": true
-    },
     "cohere-ai": {
       "optional": true
     },

package/src/ChromaClient.ts

@@ -2,16 +2,14 @@ import { IEmbeddingFunction } from './embeddings/IEmbeddingFunction';
 import { Configuration, ApiApi as DefaultApi } from "./generated";
 import { handleSuccess, handleError } from "./utils";
 import { Collection } from './Collection';
-import { ChromaClientParams, CollectionMetadata, CollectionType, ConfigOptions, CreateCollectionParams, DeleteCollectionParams, GetCollectionParams, GetOrCreateCollectionParams, ListCollectionsParams } from './types';
+import { CollectionMetadata, CollectionType, ConfigOptions } from './types';
 import {
     AuthOptions,
     ClientAuthProtocolAdapter,
     IsomorphicFetchClientAuthProtocolAdapter
 } from "./auth";
-import { AdminClient } from './AdminClient';
+import { DefaultEmbeddingFunction } from './embeddings/DefaultEmbeddingFunction';
 
-const DEFAULT_TENANT = "default_tenant"
-const DEFAULT_DATABASE = "default_database"
 
 export class ChromaClient {
     /**
@@ -19,9 +17,6 @@ export class ChromaClient {
      */
     private api: DefaultApi & ConfigOptions;
     private apiAdapter: ClientAuthProtocolAdapter<any>|undefined;
-    private tenant: string = DEFAULT_TENANT;
-    private database: string = DEFAULT_DATABASE;
-    private _adminClient?: AdminClient
 
     /**
      * Creates a new ChromaClient instance.
@@ -40,17 +35,15 @@ export class ChromaClient {
         path,
         fetchOptions,
         auth,
-        tenant = DEFAULT_TENANT,
-        database = DEFAULT_DATABASE,
-    }: ChromaClientParams = {}) {
+    }: {
+        path?: string,
+        fetchOptions?: RequestInit,
+        auth?: AuthOptions,
+    } = {}) {
         if (path === undefined) path = "http://localhost:8000";
-        this.tenant = tenant;
-        this.database = database;
-
         const apiConfig: Configuration = new Configuration({
             basePath: path,
         });
-
         if (auth !== undefined) {
             this.apiAdapter = new IsomorphicFetchClientAuthProtocolAdapter(new DefaultApi(apiConfig), auth);
             this.api = this.apiAdapter.getApi();
@@ -58,19 +51,6 @@ export class ChromaClient {
             this.api = new DefaultApi(apiConfig);
         }
 
-        this._adminClient = new AdminClient({
-            path: path,
-            fetchOptions: fetchOptions,
-            auth: auth,
-            tenant: tenant,
-            database: database
-        });
-
-        // TODO: Validate tenant and database on client creation
-        // this got tricky because:
-        // - the constructor is sync but the generated api is async
-        // - we need to inject auth information so a simple rewrite/fetch does not work
-
         this.api.options = fetchOptions ?? {};
     }
 
@@ -143,9 +123,18 @@ export class ChromaClient {
         name,
         metadata,
         embeddingFunction
-    }: CreateCollectionParams): Promise<Collection> {
+    }: {
+        name: string,
+        metadata?: CollectionMetadata,
+        embeddingFunction?: IEmbeddingFunction
+}): Promise<Collection> {
+
+        if (embeddingFunction === undefined) {
+            embeddingFunction = new DefaultEmbeddingFunction();
+        }
+
         const newCollection = await this.api
-            .createCollection(this.tenant, this.database, {
+            .createCollection({
                 name,
                 metadata,
             }, this.api.options)
@@ -184,9 +173,18 @@ export class ChromaClient {
         name,
         metadata,
         embeddingFunction
-    }: GetOrCreateCollectionParams): Promise<Collection> {
+    }: {
+        name: string,
+        metadata?: CollectionMetadata,
+        embeddingFunction?: IEmbeddingFunction
+    }): Promise<Collection> {
+
+        if (embeddingFunction === undefined) {
+            embeddingFunction = new DefaultEmbeddingFunction();
+        }
+
         const newCollection = await this.api
-            .createCollection(this.tenant, this.database, {
+            .createCollection({
                 name,
                 metadata,
                 'get_or_create': true
@@ -211,44 +209,15 @@ export class ChromaClient {
      * Lists all collections.
      *
      * @returns {Promise<CollectionType[]>} A promise that resolves to a list of collection names.
-     * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get.
-     * @param {PositiveInteger} [params.offset] - Optional offset on the items to get.
      * @throws {Error} If there is an issue listing the collections.
      *
      * @example
      * ```typescript
-     * const collections = await client.listCollections({
-     *     limit: 10,
-     *     offset: 0,
-     * });
-     * ```
-     */
-    public async listCollections({
-        limit,
-        offset,
-    }: ListCollectionsParams = {}): Promise<CollectionType[]> {
-        const response = await this.api.listCollections(
-            this.tenant,
-            this.database,
-            limit,
-            offset,
-            this.api.options);
-        return handleSuccess(response);
-    }
-
-    /**
-     * Counts all collections.
-     *
-     * @returns {Promise<number>} A promise that resolves to the number of collections.
-     * @throws {Error} If there is an issue counting the collections.
-     *
-     * @example
-     * ```typescript
-     * const collections = await client.countCollections();
+     * const collections = await client.listCollections();
      * ```
      */
-    public async countCollections(): Promise<number> {
-        const response = await this.api.countCollections(this.tenant, this.database, this.api.options);
+    public async listCollections(): Promise<CollectionType[]> {
+        const response = await this.api.listCollections(this.api.options);
         return handleSuccess(response);
     }
 
@@ -270,9 +239,12 @@ export class ChromaClient {
     public async getCollection({
         name,
         embeddingFunction
-    }: GetCollectionParams): Promise<Collection> {
+    }: {
+        name: string;
+        embeddingFunction?: IEmbeddingFunction
+    }): Promise<Collection> {
         const response = await this.api
-            .getCollection(name, this.tenant, this.database, this.api.options)
+            .getCollection(name, this.api.options)
             .then(handleSuccess)
             .catch(handleError);
 
@@ -306,9 +278,11 @@ export class ChromaClient {
      */
     public async deleteCollection({
         name
-    }: DeleteCollectionParams): Promise<void> {
+    }: {
+        name: string
+    }): Promise<void> {
         return await this.api
-            .deleteCollection(name, this.tenant, this.database, this.api.options)
+            .deleteCollection(name, this.api.options)
             .then(handleSuccess)
             .catch(handleError);
     }

package/src/Collection.ts

@@ -1,17 +1,21 @@
 import {
+    IncludeEnum,
+    Metadata,
+    Metadatas,
+    Embedding,
+    Embeddings,
+    Document,
+    Documents,
+    Where,
+    WhereDocument,
+    ID,
+    IDs,
+    PositiveInteger,
     GetResponse,
     QueryResponse,
     AddResponse,
     CollectionMetadata,
-    ConfigOptions,
-    GetParams,
-    AddParams,
-    UpsertParams,
-    ModifyCollectionParams,
-    UpdateParams,
-    QueryParams,
-    PeekParams,
-    DeleteParams
+    ConfigOptions
 } from "./types";
 import { IEmbeddingFunction } from './embeddings/IEmbeddingFunction';
 import { ApiApi as DefaultApi } from "./generated";
@@ -169,7 +173,12 @@ export class Collection {
         embeddings,
         metadatas,
         documents,
-    }: AddParams): Promise<AddResponse> {
+    }: {
+        ids: ID | IDs,
+        embeddings?: Embedding | Embeddings,
+        metadatas?: Metadata | Metadatas,
+        documents?: Document | Documents,
+    }): Promise<AddResponse> {
 
         const [idsArray, embeddingsArray, metadatasArray, documentsArray] = await this.validate(
             true,
@@ -219,7 +228,12 @@ export class Collection {
         embeddings,
         metadatas,
         documents,
-    }: UpsertParams): Promise<boolean> {
+    }: {
+        ids: ID | IDs,
+        embeddings?: Embedding | Embeddings,
+        metadatas?: Metadata | Metadatas,
+        documents?: Document | Documents,
+    }): Promise<boolean> {
         const [idsArray, embeddingsArray, metadatasArray, documentsArray] = await this.validate(
             true,
             ids,
@@ -279,7 +293,10 @@ export class Collection {
     public async modify({
         name,
         metadata
-    }: ModifyCollectionParams = {}): Promise<void> {
+    }: {
+        name?: string,
+        metadata?: CollectionMetadata
+    } = {}): Promise<void> {
         const response = await this.api
             .updateCollection(
                 this.id,
@@ -296,6 +313,7 @@ export class Collection {
         this.setMetadata(metadata || this.metadata);
 
         return response;
+
     }
 
     /**
@@ -328,7 +346,14 @@ export class Collection {
         offset,
         include,
         whereDocument,
-    }: GetParams = {}): Promise<GetResponse> {
+    }: {
+        ids?: ID | IDs,
+        where?: Where,
+        limit?: PositiveInteger,
+        offset?: PositiveInteger,
+        include?: IncludeEnum[],
+        whereDocument?: WhereDocument
+    } = {}): Promise<GetResponse> {
         let idsArray = undefined;
         if (ids !== undefined) idsArray = toArray(ids);
 
@@ -370,7 +395,12 @@ export class Collection {
         embeddings,
         metadatas,
         documents,
-    }: UpdateParams): Promise<boolean> {
+    }: {
+        ids: ID | IDs,
+        embeddings?: Embedding | Embeddings,
+        metadatas?: Metadata | Metadatas,
+        documents?: Document | Documents,
+    }): Promise<boolean> {
         if (
             embeddings === undefined &&
             documents === undefined &&
@@ -451,7 +481,14 @@ export class Collection {
         queryTexts,
         whereDocument,
         include,
-    }: QueryParams): Promise<QueryResponse> {
+    }: {
+        queryEmbeddings?: Embedding | Embeddings,
+        nResults?: PositiveInteger,
+        where?: Where,
+        queryTexts?: string | string[],
+        whereDocument?: WhereDocument, // {"$contains":"search_string"}
+        include?: IncludeEnum[] // ["metadata", "document"]
+    }): Promise<QueryResponse> {
         if (nResults === undefined) nResults = 10
         if (queryEmbeddings === undefined && queryTexts === undefined) {
             throw new Error(
@@ -499,7 +536,7 @@ export class Collection {
      * });
      * ```
      */
-    public async peek({ limit }: PeekParams = {}): Promise<GetResponse> {
+    public async peek({ limit }: { limit?: PositiveInteger } = {}): Promise<GetResponse> {
         if (limit === undefined) limit = 10;
         const response = await this.api.aGet(this.id, {
             limit: limit,
@@ -529,7 +566,11 @@ export class Collection {
         ids,
         where,
         whereDocument
-    }: DeleteParams = {}): Promise<string[]> {
+    }: {
+        ids?: ID | IDs,
+        where?: Where,
+        whereDocument?: WhereDocument
+    } = {}): Promise<string[]> {
         let idsArray = undefined;
         if (ids !== undefined) idsArray = toArray(ids);
         return await this.api

package/src/embeddings/DefaultEmbeddingFunction.ts

@@ -0,0 +1,99 @@
+import { IEmbeddingFunction } from "./IEmbeddingFunction";
+
+// Dynamically import module
+let TransformersApi: Promise<any>;
+
+export class DefaultEmbeddingFunction implements IEmbeddingFunction {
+  private pipelinePromise?: Promise<any> | null;
+  private transformersApi: any;
+  private model: string;
+  private revision: string;
+  private quantized: boolean;
+  private progress_callback: Function | null;
+
+  /**
+   * DefaultEmbeddingFunction constructor.
+   * @param options The configuration options.
+   * @param options.model The model to use to calculate embeddings. Defaults to 'Xenova/all-MiniLM-L6-v2', which is an ONNX port of `sentence-transformers/all-MiniLM-L6-v2`.
+   * @param options.revision The specific model version to use (can be a branch, tag name, or commit id). Defaults to 'main'.
+   * @param options.quantized Whether to load the 8-bit quantized version of the model. Defaults to `false`.
+   * @param options.progress_callback If specified, this function will be called during model construction, to provide the user with progress updates.
+   */
+  constructor({
+    model = "Xenova/all-MiniLM-L6-v2",
+    revision = "main",
+    quantized = false,
+    progress_callback = null,
+  }: {
+    model?: string;
+    revision?: string;
+    quantized?: boolean;
+    progress_callback?: Function | null;
+  } = {}) {
+    this.model = model;
+    this.revision = revision;
+    this.quantized = quantized;
+    this.progress_callback = progress_callback;
+  }
+
+  public async generate(texts: string[]): Promise<number[][]> {
+    await this.loadClient();
+
+     // Store a promise that resolves to the pipeline
+    this.pipelinePromise = new Promise(async (resolve, reject) => {
+      try {
+        const pipeline = this.transformersApi
+
+        const quantized = this.quantized
+        const revision = this.revision
+        const progress_callback = this.progress_callback
+
+        resolve(
+          await pipeline("feature-extraction", this.model, {
+            quantized,
+            revision,
+            progress_callback,
+          })
+        );
+      } catch (e) {
+        reject(e);
+      }
+    });
+
+    let pipe = await this.pipelinePromise;
+    let output = await pipe(texts, { pooling: "mean", normalize: true });
+    return output.tolist();
+  }
+
+  private async loadClient() {
+      if(this.transformersApi) return;
+      try {
+          // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+          let { pipeline } = await DefaultEmbeddingFunction.import();
+          TransformersApi = pipeline;
+      } catch (_a) {
+          // @ts-ignore
+          if (_a.code === 'MODULE_NOT_FOUND') {
+              throw new Error("Please install the chromadb-default-embed package to use the DefaultEmbeddingFunction, `npm install -S chromadb-default-embed`");
+          }
+          throw _a; // Re-throw other errors
+      }
+      this.transformersApi = TransformersApi;
+  }
+
+  /** @ignore */
+  static async import(): Promise<{
+      // @ts-ignore
+      pipeline: typeof import("chromadb-default-embed");
+  }> {
+      try {
+          // @ts-ignore
+          const { pipeline } = await import("chromadb-default-embed");
+          return { pipeline };
+      } catch (e) {
+          throw new Error(
+              "Please install chromadb-default-embed as a dependency with, e.g. `yarn add chromadb-default-embed`"
+          );
+      }
+  }
+}

package/src/embeddings/TransformersEmbeddingFunction.ts

@@ -0,0 +1,99 @@
+import { IEmbeddingFunction } from "./IEmbeddingFunction";
+
+// Dynamically import module
+let TransformersApi: Promise<any>;
+
+export class TransformersEmbeddingFunction implements IEmbeddingFunction {
+  private pipelinePromise?: Promise<any> | null;
+  private transformersApi: any;
+  private model: string;
+  private revision: string;
+  private quantized: boolean;
+  private progress_callback: Function | null;
+
+  /**
+   * TransformersEmbeddingFunction constructor.
+   * @param options The configuration options.
+   * @param options.model The model to use to calculate embeddings. Defaults to 'Xenova/all-MiniLM-L6-v2', which is an ONNX port of `sentence-transformers/all-MiniLM-L6-v2`.
+   * @param options.revision The specific model version to use (can be a branch, tag name, or commit id). Defaults to 'main'.
+   * @param options.quantized Whether to load the 8-bit quantized version of the model. Defaults to `false`.
+   * @param options.progress_callback If specified, this function will be called during model construction, to provide the user with progress updates.
+   */
+  constructor({
+    model = "Xenova/all-MiniLM-L6-v2",
+    revision = "main",
+    quantized = false,
+    progress_callback = null,
+  }: {
+    model?: string;
+    revision?: string;
+    quantized?: boolean;
+    progress_callback?: Function | null;
+  } = {}) {
+    this.model = model;
+    this.revision = revision;
+    this.quantized = quantized;
+    this.progress_callback = progress_callback;
+  }
+
+  public async generate(texts: string[]): Promise<number[][]> {
+    await this.loadClient();
+
+     // Store a promise that resolves to the pipeline
+    this.pipelinePromise = new Promise(async (resolve, reject) => {
+      try {
+        const pipeline = this.transformersApi
+
+        const quantized = this.quantized
+        const revision = this.revision
+        const progress_callback = this.progress_callback
+
+        resolve(
+          await pipeline("feature-extraction", this.model, {
+            quantized,
+            revision,
+            progress_callback,
+          })
+        );
+      } catch (e) {
+        reject(e);
+      }
+    });
+
+    let pipe = await this.pipelinePromise;
+    let output = await pipe(texts, { pooling: "mean", normalize: true });
+    return output.tolist();
+  }
+
+  private async loadClient() {
+      if(this.transformersApi) return;
+      try {
+          // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+          let { pipeline } = await TransformersEmbeddingFunction.import();
+          TransformersApi = pipeline;
+      } catch (_a) {
+          // @ts-ignore
+          if (_a.code === 'MODULE_NOT_FOUND') {
+              throw new Error("Please install the @xenova/transformers package to use the TransformersEmbeddingFunction, `npm install -S @xenova/transformers`");
+          }
+          throw _a; // Re-throw other errors
+      }
+      this.transformersApi = TransformersApi;
+  }
+
+  /** @ignore */
+  static async import(): Promise<{
+      // @ts-ignore
+      pipeline: typeof import("@xenova/transformers");
+  }> {
+      try {
+          // @ts-ignore
+          const { pipeline } = await import("@xenova/transformers");
+          return { pipeline };
+      } catch (e) {
+          throw new Error(
+              "Please install @xenova/transformers as a dependency with, e.g. `yarn add @xenova/transformers`"
+          );
+      }
+  }
+}