chromadb 1.2.1 → 1.3.0

package/dist/main/index.d.ts

@@ -1,23 +1,41 @@
 import { DefaultApi } from "./generated/api";
+export declare class OpenAIEmbeddingFunction {
+    private api_key;
+    private org_id;
+    private model;
+    constructor(openai_api_key: string, openai_model: string, openai_organization_id: string);
+    generate(texts: string[]): Promise<any[]>;
+}
+export declare class CohereEmbeddingFunction {
+    private api_key;
+    constructor(cohere_api_key: string);
+    generate(texts: string[]): Promise<any>;
+}
+type CallableFunction = {
+    new (): CallableFunction;
+    generate(texts: string[]): number[][];
+};
 export declare class Collection {
-    private name;
+    name: string;
     private api;
-    constructor(name: string, api: DefaultApi);
-    add(ids: string | Array<any>, embeddings: Array<any>, metadatas?: Array<any> | object, documents?: string | Array<any>, increment_index?: boolean): Promise<import("axios").AxiosResponse<any, any>>;
+    embeddingFunction: CallableFunction | undefined;
+    constructor(name: string, api: DefaultApi, embeddingFunction?: CallableFunction);
+    add(ids: string | string[], embeddings: number[] | number[][], metadatas?: object | object[], documents?: string | string[], increment_index?: boolean): Promise<any>;
     count(): Promise<any>;
-    get(ids?: string[], where?: object, limit?: number, offset?: number): Promise<import("axios").AxiosResponse<any, any>>;
-    query(query_embeddings: number[], n_results?: number, where?: object): Promise<any>;
-    peek(limit?: number): Promise<import("axios").AxiosResponse<any, any>>;
+    get(ids?: string[], where?: object, limit?: number, offset?: number): Promise<any>;
+    query(query_embeddings: number[] | number[][], n_results?: number, where?: object, query_text?: string | string[]): Promise<any>;
+    peek(limit?: number): Promise<any>;
     createIndex(): Promise<import("axios").AxiosResponse<any, any>>;
-    delete(ids?: string[], where?: object): Promise<import("axios").AxiosResponse<any, any>>;
+    delete(ids?: string[], where?: object): Promise<any>;
 }
 export declare class ChromaClient {
     private api;
-    constructor(basePath: string);
+    constructor(basePath?: string);
     reset(): Promise<import("axios").AxiosResponse<any, any>>;
-    createCollection(name: string, metadata?: object): Promise<Collection>;
+    createCollection(name: string, metadata?: object, embeddingFunction?: CallableFunction): Promise<Collection>;
     listCollections(): Promise<any>;
-    getCollection(name: string): Promise<Collection>;
-    deleteCollection(name: string): Promise<import("axios").AxiosResponse<any, any>>;
+    getCollection(name: string, embeddingFunction?: CallableFunction): Promise<Collection>;
+    deleteCollection(name: string): Promise<any>;
 }
+export {};
 //# sourceMappingURL=index.d.ts.map

package/dist/main/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAqB7C,qBAAa,UAAU;IACrB,OAAO,CAAC,IAAI,CAAS;IACrB,OAAO,CAAC,GAAG,CAAa;gBAEZ,IAAI,EAAE,MAAM,EAAE,GAAG,EAAE,UAAU;IAK5B,GAAG,CACd,GAAG,EAAE,MAAM,GAAG,KAAK,CAAC,GAAG,CAAC,EACxB,UAAU,EAAE,KAAK,CAAC,GAAG,CAAC,EACtB,SAAS,CAAC,EAAE,KAAK,CAAC,GAAG,CAAC,GAAG,MAAM,EAC/B,SAAS,CAAC,EAAE,MAAM,GAAG,KAAK,CAAC,GAAG,CAAC,EAC/B,eAAe,GAAE,OAAc;IA0CpB,KAAK;IAKL,GAAG,CACd,GAAG,CAAC,EAAE,MAAM,EAAE,EACd,KAAK,CAAC,EAAE,MAAM,EACd,KAAK,CAAC,EAAE,MAAM,EACd,MAAM,CAAC,EAAE,MAAM;IAeJ,KAAK,CAChB,gBAAgB,EAAE,MAAM,EAAE,EAC1B,SAAS,GAAE,MAAW,EACtB,KAAK,CAAC,EAAE,MAAM;IAeH,IAAI,CAAC,KAAK,GAAE,MAAW;IAOvB,WAAW;IAIX,MAAM,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,EAAE,KAAK,CAAC,EAAE,MAAM;CAOnD;AAED,qBAAa,YAAY;IACvB,OAAO,CAAC,GAAG,CAAa;gBAEZ,QAAQ,EAAE,MAAM;IAOf,KAAK;IAIL,gBAAgB,CAAC,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,EAAE,MAAM;IAOhD,eAAe;IAKf,aAAa,CAAC,IAAI,EAAE,MAAM;IAI1B,gBAAgB,CAAC,IAAI,EAAE,MAAM;CAI3C"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAyB7C,qBAAa,uBAAuB;IAClC,OAAO,CAAC,OAAO,CAAS;IACxB,OAAO,CAAC,MAAM,CAAS;IACvB,OAAO,CAAC,KAAK,CAAS;gBAEV,cAAc,EAAE,MAAM,EAAE,YAAY,EAAE,MAAM,EAAE,sBAAsB,EAAE,MAAM;IAc3E,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE;CAiBtC;AAID,qBAAa,uBAAuB;IAClC,OAAO,CAAC,OAAO,CAAS;gBAEZ,cAAc,EAAE,MAAM;IAYrB,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE;CAQtC;AAED,KAAK,gBAAgB,GAAG;IACtB,QAAO,gBAAgB,CAAC;IACxB,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE,GAAG,MAAM,EAAE,EAAE,CAAC;CACvC,CAAC;AAEF,qBAAa,UAAU;IACd,IAAI,EAAE,MAAM,CAAC;IACpB,OAAO,CAAC,GAAG,CAAa;IACjB,iBAAiB,EAAE,gBAAgB,GAAG,SAAS,CAAC;gBAE3C,IAAI,EAAE,MAAM,EAAE,GAAG,EAAE,UAAU,EAAE,iBAAiB,CAAC,EAAE,gBAAgB;IAOlE,GAAG,CACd,GAAG,EAAE,MAAM,GAAG,MAAM,EAAE,EACtB,UAAU,EAAE,MAAM,EAAE,GAAG,MAAM,EAAE,EAAE,EACjC,SAAS,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE,EAC7B,SAAS,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE,EAC7B,eAAe,GAAE,OAAc;IA8DpB,KAAK;IAKL,GAAG,CACd,GAAG,CAAC,EAAE,MAAM,EAAE,EACd,KAAK,CAAC,EAAE,MAAM,EACd,KAAK,CAAC,EAAE,MAAM,EACd,MAAM,CAAC,EAAE,MAAM;IAuBJ,KAAK,CAChB,gBAAgB,EAAE,MAAM,EAAE,GAAG,MAAM,EAAE,EAAE,EACvC,SAAS,GAAE,MAAW,EACtB,KAAK,CAAC,EAAE,MAAM,EACd,UAAU,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE;IAmCnB,IAAI,CAAC,KAAK,GAAE,MAAW;IAQvB,WAAW;IAIX,MAAM,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,EAAE,KAAK,CAAC,EAAE,MAAM;CAanD;AAED,qBAAa,YAAY;IACvB,OAAO,CAAC,GAAG,CAAa;gBAEZ,QAAQ,CAAC,EAAE,MAAM;IAQhB,KAAK;IAIL,gBAAgB,CAAC,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,EAAE,MAAM,EAAE,iBAAiB,CAAC,EAAE,gBAAgB;IAgBtF,eAAe;IAKf,aAAa,CAAC,IAAI,EAAE,MAAM,EAAE,iBAAiB,CAAC,EAAE,gBAAgB;IAIhE,gBAAgB,CAAC,IAAI,EAAE,MAAM;CAU3C"}

package/dist/main/index.js

@@ -1,6 +1,6 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.ChromaClient = exports.Collection = void 0;
+exports.ChromaClient = exports.Collection = exports.CohereEmbeddingFunction = exports.OpenAIEmbeddingFunction = void 0;
 const api_1 = require("./generated/api");
 const configuration_1 = require("./generated/configuration");
 // a function to convert a non-Array object to an Array
@@ -21,12 +21,83 @@ function toArrayOfArrays(obj) {
         return [obj];
     }
 }
+class EmbeddingFunction {
+}
+let OpenAIApi;
+class OpenAIEmbeddingFunction {
+    constructor(openai_api_key, openai_model, openai_organization_id) {
+        try {
+            // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+            OpenAIApi = require("openai");
+        }
+        catch (_a) {
+            throw new Error("Please install the openai package to use the OpenAIEmbeddingFunction, `npm install -S openai`");
+        }
+        this.api_key = openai_api_key;
+        this.org_id = openai_organization_id || "";
+        this.model = openai_model || "text-embedding-ada-002";
+    }
+    async generate(texts) {
+        const configuration = new OpenAIApi.Configuration({
+            organization: this.org_id,
+            apiKey: this.api_key,
+        });
+        const openai = new OpenAIApi.OpenAIApi(configuration);
+        const embeddings = [];
+        const response = await openai.createEmbedding({
+            model: "text-embedding-ada-002",
+            input: texts,
+        });
+        const data = response.data['data'];
+        for (let i = 0; i < data.length; i += 1) {
+            embeddings.push(data[i]['embedding']);
+        }
+        return embeddings;
+    }
+}
+exports.OpenAIEmbeddingFunction = OpenAIEmbeddingFunction;
+let CohereAiApi;
+class CohereEmbeddingFunction {
+    constructor(cohere_api_key) {
+        try {
+            // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+            CohereAiApi = require("cohere-ai");
+        }
+        catch (_a) {
+            throw new Error("Please install the cohere-ai package to use the CohereEmbeddingFunction, `npm install -S cohere-ai`");
+        }
+        this.api_key = cohere_api_key;
+    }
+    async generate(texts) {
+        const cohere = CohereAiApi.init(this.api_key);
+        const embeddings = [];
+        const response = await CohereAiApi.embed({
+            texts: texts,
+        });
+        return response.body.embeddings;
+    }
+}
+exports.CohereEmbeddingFunction = CohereEmbeddingFunction;
 class Collection {
-    constructor(name, api) {
+    constructor(name, api, embeddingFunction) {
         this.name = name;
         this.api = api;
+        if (embeddingFunction !== undefined)
+            this.embeddingFunction = embeddingFunction;
     }
     async add(ids, embeddings, metadatas, documents, increment_index = true) {
+        if ((embeddings === undefined) && (documents === undefined)) {
+            throw new Error("embeddings and documents cannot both be undefined");
+        }
+        else if ((embeddings === undefined) && (documents !== undefined)) {
+            const documentsArray = toArray(documents);
+            if (this.embeddingFunction !== undefined) {
+                embeddings = await this.embeddingFunction.generate(documentsArray);
+            }
+            else {
+                throw new Error("embeddingFunction is undefined. Please configure an embedding function");
+            }
+        }
         const idsArray = toArray(ids);
         const embeddingsArray = toArrayOfArrays(embeddings);
         let metadatasArray;
@@ -37,18 +108,18 @@ class Collection {
             metadatasArray = toArray(metadatas);
         }
         let documentsArray;
-        if (metadatas === undefined) {
+        if (documents === undefined) {
             documentsArray = undefined;
         }
         else {
-            documentsArray = toArray(metadatas);
+            documentsArray = toArray(documents);
         }
-        if (idsArray.length !== embeddingsArray.length ||
+        if (((embeddingsArray !== undefined) && idsArray.length !== embeddingsArray.length) ||
             ((metadatasArray !== undefined) && idsArray.length !== metadatasArray.length) ||
             ((documentsArray !== undefined) && idsArray.length !== documentsArray.length)) {
             throw new Error("ids, embeddings, metadatas, and documents must all be the same length");
         }
-        return await this.api.add({
+        const response = await this.api.add({
             collectionName: this.name,
             addEmbedding: {
                 ids: idsArray,
@@ -57,15 +128,22 @@ class Collection {
                 metadatas: metadatasArray,
                 increment_index: increment_index,
             },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
+        return response;
     }
     async count() {
         const response = await this.api.count({ collectionName: this.name });
         return response.data;
     }
     async get(ids, where, limit, offset) {
-        const idsArray = toArray(ids);
-        return await this.api.get({
+        let idsArray = undefined;
+        if (ids !== undefined)
+            idsArray = toArray(ids);
+        var resp = await this.api.get({
             collectionName: this.name,
             getEmbedding: {
                 ids: idsArray,
@@ -73,9 +151,26 @@ class Collection {
                 limit,
                 offset,
             },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
+        return resp;
     }
-    async query(query_embeddings, n_results = 10, where) {
+    async query(query_embeddings, n_results = 10, where, query_text) {
+        if ((query_embeddings === undefined) && (query_text === undefined)) {
+            throw new Error("query_embeddings and query_text cannot both be undefined");
+        }
+        else if ((query_embeddings === undefined) && (query_text !== undefined)) {
+            const query_texts = toArray(query_text);
+            if (this.embeddingFunction !== undefined) {
+                query_embeddings = await this.embeddingFunction.generate(query_texts);
+            }
+            else {
+                throw new Error("embeddingFunction is undefined. Please configure an embedding function");
+            }
+        }
         const query_embeddingsArray = toArrayOfArrays(query_embeddings);
         const response = await this.api.getNearestNeighbors({
             collectionName: this.name,
@@ -84,28 +179,40 @@ class Collection {
                 where,
                 n_results,
             },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
-        return response.data;
+        return response;
     }
     async peek(limit = 10) {
-        return await this.api.get({
+        const response = await this.api.get({
             collectionName: this.name,
             getEmbedding: { limit: limit },
         });
+        return response.data;
     }
     async createIndex() {
         return await this.api.createIndex({ collectionName: this.name });
     }
     async delete(ids, where) {
-        return await this.api._delete({
+        var response = await this.api._delete({
             collectionName: this.name,
             deleteEmbedding: { ids: ids, where: where },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
+        return response;
     }
 }
 exports.Collection = Collection;
 class ChromaClient {
     constructor(basePath) {
+        if (basePath === undefined)
+            basePath = "http://localhost:8000";
         const apiConfig = new configuration_1.Configuration({
             basePath,
         });
@@ -114,21 +221,33 @@ class ChromaClient {
     async reset() {
         return await this.api.reset();
     }
-    async createCollection(name, metadata) {
+    async createCollection(name, metadata, embeddingFunction) {
         const newCollection = await this.api.createCollection({
             createCollection: { name, metadata },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
-        return new Collection(name, this.api);
+        if (newCollection.error) {
+            throw new Error(newCollection.error);
+        }
+        return new Collection(name, this.api, embeddingFunction);
     }
     async listCollections() {
         const response = await this.api.listCollections();
         return response.data;
     }
-    async getCollection(name) {
-        return new Collection(name, this.api);
+    async getCollection(name, embeddingFunction) {
+        return new Collection(name, this.api, embeddingFunction);
     }
     async deleteCollection(name) {
-        return await this.api.deleteCollection({ collectionName: name });
+        const response = await this.api.deleteCollection({ collectionName: name }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
+        });
+        return response;
     }
 }
 exports.ChromaClient = ChromaClient;

package/dist/main/index.js.map

@@ -1 +1 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":";;;AAAA,yCAA6C;AAC7C,6DAA0D;AAE1D,uDAAuD;AACvD,SAAS,OAAO,CAAI,GAAiB;IACnC,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE;QACtB,OAAO,GAAG,CAAC;KACZ;SAAM;QACL,OAAO,CAAC,GAAG,CAAC,CAAC;KACd;AACH,CAAC;AAED,oDAAoD;AACpD,SAAS,eAAe,CAAI,GAA+B;IACzD,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE;QACzB,OAAO,GAAsB,CAAC;KAC/B;SAAM;QACL,OAAO,CAAC,GAAG,CAAoB,CAAC;KACjC;AACH,CAAC;AAED,MAAa,UAAU;IAIrB,YAAY,IAAY,EAAE,GAAe;QACvC,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;QACjB,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;IACjB,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAwB,EACxB,UAAsB,EACtB,SAA+B,EAC/B,SAA+B,EAC/B,kBAA2B,IAAI;QAG/B,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAC9B,MAAM,eAAe,GAAG,eAAe,CAAC,UAAU,CAAC,CAAC;QAEpD,IAAI,cAAc,CAAC;QACnB,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IAAI,cAAc,CAAC;QACnB,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IACE,QAAQ,CAAC,MAAM,KAAK,eAAe,CAAC,MAAM;YAC1C,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC;YAC7E,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC,EAC7E;YACA,MAAM,IAAI,KAAK,CACb,uEAAuE,CACxE,CAAC;SACH;QAED,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YACxB,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,UAAU,EAAE,eAAe;gBAC3B,SAAS,EAAE,cAAc;gBACzB,SAAS,EAAE,cAAc;gBACzB,eAAe,EAAE,eAAe;aACjC;SACF,CAAC,CAAC;IACL,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;QACrE,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAc,EACd,KAAc,EACd,KAAc,EACd,MAAe;QAEf,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAE9B,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YACxB,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,KAAK;gBACL,KAAK;gBACL,MAAM;aACP;SACF,CAAC,CAAC;IACL,CAAC;IAEM,KAAK,CAAC,KAAK,CAChB,gBAA0B,EAC1B,YAAoB,EAAE,EACtB,KAAc;QAEd,MAAM,qBAAqB,GAAG,eAAe,CAAC,gBAAgB,CAAC,CAAC;QAEhE,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,mBAAmB,CAAC;YAClD,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,cAAc,EAAE;gBACd,gBAAgB,EAAE,qBAAqB;gBACvC,KAAK;gBACL,SAAS;aACV;SACF,CAAC,CAAC;QACH,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,IAAI,CAAC,QAAgB,EAAE;QAClC,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YACxB,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE;SAC/B,CAAC,CAAC;IACL,CAAC;IAEM,KAAK,CAAC,WAAW;QACtB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;IACnE,CAAC;IAEM,KAAK,CAAC,MAAM,CAAC,GAAc,EAAE,KAAc;QAChD,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC;YAC5B,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,eAAe,EAAE,EAAE,GAAG,EAAE,GAAG,EAAE,KAAK,EAAE,KAAK,EAAE;SAC5C,CAAC,CAAC;IACL,CAAC;CAEF;AApHD,gCAoHC;AAED,MAAa,YAAY;IAGvB,YAAY,QAAgB;QAC1B,MAAM,SAAS,GAAkB,IAAI,6BAAa,CAAC;YACjD,QAAQ;SACT,CAAC,CAAC;QACH,IAAI,CAAC,GAAG,GAAG,IAAI,gBAAU,CAAC,SAAS,CAAC,CAAC;IACvC,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,CAAC;IAChC,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY,EAAE,QAAiB;QAC3D,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC;YACpD,gBAAgB,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE;SACrC,CAAC,CAAC;QACH,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC;IACxC,CAAC;IAEM,KAAK,CAAC,eAAe;QAC1B,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,eAAe,EAAE,CAAC;QAClD,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,aAAa,CAAC,IAAY;QACrC,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC;IACxC,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY;QACxC,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC,EAAE,cAAc,EAAE,IAAI,EAAE,CAAC,CAAC;IACnE,CAAC;CAEF;AAlCD,oCAkCC"}
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":";;;AAAA,yCAA6C;AAC7C,6DAA0D;AAE1D,uDAAuD;AACvD,SAAS,OAAO,CAAI,GAAiB;IACnC,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE;QACtB,OAAO,GAAG,CAAC;KACZ;SAAM;QACL,OAAO,CAAC,GAAG,CAAC,CAAC;KACd;AACH,CAAC;AAED,oDAAoD;AACpD,SAAS,eAAe,CAAI,GAA+B;IACzD,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE;QACzB,OAAO,GAAsB,CAAC;KAC/B;SAAM;QACL,OAAO,CAAC,GAAG,CAAoB,CAAC;KACjC;AACH,CAAC;AAED,MAAM,iBAAiB;CAAI;AAE3B,IAAI,SAAc,CAAC;AAEnB,MAAa,uBAAuB;IAKlC,YAAY,cAAsB,EAAE,YAAoB,EAAE,sBAA8B;QACtF,IAAI;YACF,4EAA4E;YAC5E,SAAS,GAAG,OAAO,CAAC,QAAQ,CAAC,CAAC;SAC/B;QAAC,WAAM;YACN,MAAM,IAAI,KAAK,CACb,+FAA+F,CAChG,CAAC;SACH;QACD,IAAI,CAAC,OAAO,GAAG,cAAc,CAAC;QAC9B,IAAI,CAAC,MAAM,GAAG,sBAAsB,IAAI,EAAE,CAAC;QAC3C,IAAI,CAAC,KAAK,GAAG,YAAY,IAAI,wBAAwB,CAAC;IACxD,CAAC;IAEM,KAAK,CAAC,QAAQ,CAAC,KAAe;QACnC,MAAM,aAAa,GAAG,IAAI,SAAS,CAAC,aAAa,CAAC;YAChD,YAAY,EAAE,IAAI,CAAC,MAAM;YACzB,MAAM,EAAE,IAAI,CAAC,OAAO;SACrB,CAAC,CAAC;QACH,MAAM,MAAM,GAAG,IAAI,SAAS,CAAC,SAAS,CAAC,aAAa,CAAC,CAAC;QACtD,MAAM,UAAU,GAAG,EAAE,CAAC;QACtB,MAAM,QAAQ,GAAG,MAAM,MAAM,CAAC,eAAe,CAAC;YAC5C,KAAK,EAAE,wBAAwB;YAC/B,KAAK,EAAE,KAAK;SACb,CAAC,CAAC;QACH,MAAM,IAAI,GAAG,QAAQ,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QACnC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC,IAAI,CAAC,EAAE;YACvC,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,WAAW,CAAC,CAAC,CAAC;SACvC;QACD,OAAO,UAAU,CAAC;IACpB,CAAC;CACF;AApCD,0DAoCC;AAED,IAAI,WAAgB,CAAC;AAErB,MAAa,uBAAuB;IAGlC,YAAY,cAAsB;QAChC,IAAI;YACF,4EAA4E;YAC5E,WAAW,GAAG,OAAO,CAAC,WAAW,CAAC,CAAC;SACpC;QAAC,WAAM;YACN,MAAM,IAAI,KAAK,CACb,qGAAqG,CACtG,CAAC;SACH;QACD,IAAI,CAAC,OAAO,GAAG,cAAc,CAAC;IAChC,CAAC;IAEM,KAAK,CAAC,QAAQ,CAAC,KAAe;QACnC,MAAM,MAAM,GAAG,WAAW,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QAC9C,MAAM,UAAU,GAAG,EAAE,CAAC;QACtB,MAAM,QAAQ,GAAG,MAAM,WAAW,CAAC,KAAK,CAAC;YACvC,KAAK,EAAE,KAAK;SACb,CAAC,CAAC;QACH,OAAO,QAAQ,CAAC,IAAI,CAAC,UAAU,CAAC;IAClC,CAAC;CACF;AAvBD,0DAuBC;AAOD,MAAa,UAAU;IAKrB,YAAY,IAAY,EAAE,GAAe,EAAE,iBAAoC;QAC7E,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;QACjB,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;QACf,IAAI,iBAAiB,KAAK,SAAS;YACjC,IAAI,CAAC,iBAAiB,GAAG,iBAAiB,CAAC;IAC/C,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAsB,EACtB,UAAiC,EACjC,SAA6B,EAC7B,SAA6B,EAC7B,kBAA2B,IAAI;QAE/B,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,IAAI,CAAC,SAAS,KAAK,SAAS,CAAC,EAAE;YAC3D,MAAM,IAAI,KAAK,CACb,mDAAmD,CACpD,CAAC;SACH;aAAM,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,IAAI,CAAC,SAAS,KAAK,SAAS,CAAC,EAAE;YAClE,MAAM,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;YAC1C,IAAI,IAAI,CAAC,iBAAiB,KAAK,SAAS,EAAE;gBACxC,UAAU,GAAG,MAAM,IAAI,CAAC,iBAAiB,CAAC,QAAQ,CAAC,cAAc,CAAC,CAAA;aACnE;iBAAM;gBACL,MAAM,IAAI,KAAK,CACb,wEAAwE,CACzE,CAAC;aACH;SACF;QAED,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAC9B,MAAM,eAAe,GAAG,eAAe,CAAC,UAAU,CAAC,CAAC;QAEpD,IAAI,cAAoC,CAAC;QACzC,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IAAI,cAAkD,CAAC;QACvD,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IACE,CAAC,CAAC,eAAe,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,eAAe,CAAC,MAAM,CAAC;YAC/E,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC;YAC7E,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC,EAC7E;YACA,MAAM,IAAI,KAAK,CACb,uEAAuE,CACxE,CAAC;SACH;QAED,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YAClC,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,UAAU,EAAE,eAAe;gBAC3B,SAAS,EAAE,cAAc;gBACzB,SAAS,EAAE,cAAc;gBACzB,eAAe,EAAE,eAAe;aACjC;SACF,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAA;IACjB,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;QACrE,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAc,EACd,KAAc,EACd,KAAc,EACd,MAAe;QAEf,IAAI,QAAQ,GAAG,SAAS,CAAA;QACxB,IAAI,GAAG,KAAK,SAAS;YAAE,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAE/C,IAAI,IAAI,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YAC5B,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,KAAK;gBACL,KAAK;gBACL,MAAM;aACP;SACF,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,IAAI,CAAA;IAEb,CAAC;IAEM,KAAK,CAAC,KAAK,CAChB,gBAAuC,EACvC,YAAoB,EAAE,EACtB,KAAc,EACd,UAA8B;QAE9B,IAAI,CAAC,gBAAgB,KAAK,SAAS,CAAC,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,EAAE;YAClE,MAAM,IAAI,KAAK,CACb,0DAA0D,CAC3D,CAAC;SACH;aAAM,IAAI,CAAC,gBAAgB,KAAK,SAAS,CAAC,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,EAAE;YACzE,MAAM,WAAW,GAAG,OAAO,CAAC,UAAU,CAAC,CAAC;YACxC,IAAI,IAAI,CAAC,iBAAiB,KAAK,SAAS,EAAE;gBACxC,gBAAgB,GAAG,MAAM,IAAI,CAAC,iBAAiB,CAAC,QAAQ,CAAC,WAAW,CAAC,CAAA;aACtE;iBAAM;gBACL,MAAM,IAAI,KAAK,CACb,wEAAwE,CACzE,CAAC;aACH;SACF;QAED,MAAM,qBAAqB,GAAG,eAAe,CAAC,gBAAgB,CAAC,CAAC;QAEhE,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,mBAAmB,CAAC;YAClD,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,cAAc,EAAE;gBACd,gBAAgB,EAAE,qBAAqB;gBACvC,KAAK;gBACL,SAAS;aACV;SACF,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAC;IAClB,CAAC;IAEM,KAAK,CAAC,IAAI,CAAC,QAAgB,EAAE;QAClC,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YAClC,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE;SAC/B,CAAC,CAAC;QACH,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,WAAW;QACtB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;IACnE,CAAC;IAEM,KAAK,CAAC,MAAM,CAAC,GAAc,EAAE,KAAc;QAChD,IAAI,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC;YACpC,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,eAAe,EAAE,EAAE,GAAG,EAAE,GAAG,EAAE,KAAK,EAAE,KAAK,EAAE;SAC5C,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAA;IACjB,CAAC;CAEF;AA/KD,gCA+KC;AAED,MAAa,YAAY;IAGvB,YAAY,QAAiB;QAC3B,IAAI,QAAQ,KAAK,SAAS;YAAE,QAAQ,GAAG,uBAAuB,CAAC;QAC/D,MAAM,SAAS,GAAkB,IAAI,6BAAa,CAAC;YACjD,QAAQ;SACT,CAAC,CAAC;QACH,IAAI,CAAC,GAAG,GAAG,IAAI,gBAAU,CAAC,SAAS,CAAC,CAAC;IACvC,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,CAAC;IAChC,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY,EAAE,QAAiB,EAAE,iBAAoC;QACjG,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC;YACpD,gBAAgB,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE;SACrC,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,IAAI,aAAa,CAAC,KAAK,EAAE;YACvB,MAAM,IAAI,KAAK,CAAC,aAAa,CAAC,KAAK,CAAC,CAAC;SACtC;QAED,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,EAAE,iBAAiB,CAAC,CAAC;IAC3D,CAAC;IAEM,KAAK,CAAC,eAAe;QAC1B,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,eAAe,EAAE,CAAC;QAClD,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,aAAa,CAAC,IAAY,EAAE,iBAAoC;QAC3E,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,EAAE,iBAAiB,CAAC,CAAC;IAC3D,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY;QACxC,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC,EAAE,cAAc,EAAE,IAAI,EAAE,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YAChG,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAA;IACjB,CAAC;CAEF;AAlDD,oCAkDC"}

package/dist/module/index.d.ts

@@ -1,23 +1,41 @@
 import { DefaultApi } from "./generated/api";
+export declare class OpenAIEmbeddingFunction {
+    private api_key;
+    private org_id;
+    private model;
+    constructor(openai_api_key: string, openai_model: string, openai_organization_id: string);
+    generate(texts: string[]): Promise<any[]>;
+}
+export declare class CohereEmbeddingFunction {
+    private api_key;
+    constructor(cohere_api_key: string);
+    generate(texts: string[]): Promise<any>;
+}
+type CallableFunction = {
+    new (): CallableFunction;
+    generate(texts: string[]): number[][];
+};
 export declare class Collection {
-    private name;
+    name: string;
     private api;
-    constructor(name: string, api: DefaultApi);
-    add(ids: string | Array<any>, embeddings: Array<any>, metadatas?: Array<any> | object, documents?: string | Array<any>, increment_index?: boolean): Promise<import("axios").AxiosResponse<any, any>>;
+    embeddingFunction: CallableFunction | undefined;
+    constructor(name: string, api: DefaultApi, embeddingFunction?: CallableFunction);
+    add(ids: string | string[], embeddings: number[] | number[][], metadatas?: object | object[], documents?: string | string[], increment_index?: boolean): Promise<any>;
     count(): Promise<any>;
-    get(ids?: string[], where?: object, limit?: number, offset?: number): Promise<import("axios").AxiosResponse<any, any>>;
-    query(query_embeddings: number[], n_results?: number, where?: object): Promise<any>;
-    peek(limit?: number): Promise<import("axios").AxiosResponse<any, any>>;
+    get(ids?: string[], where?: object, limit?: number, offset?: number): Promise<any>;
+    query(query_embeddings: number[] | number[][], n_results?: number, where?: object, query_text?: string | string[]): Promise<any>;
+    peek(limit?: number): Promise<any>;
     createIndex(): Promise<import("axios").AxiosResponse<any, any>>;
-    delete(ids?: string[], where?: object): Promise<import("axios").AxiosResponse<any, any>>;
+    delete(ids?: string[], where?: object): Promise<any>;
 }
 export declare class ChromaClient {
     private api;
-    constructor(basePath: string);
+    constructor(basePath?: string);
     reset(): Promise<import("axios").AxiosResponse<any, any>>;
-    createCollection(name: string, metadata?: object): Promise<Collection>;
+    createCollection(name: string, metadata?: object, embeddingFunction?: CallableFunction): Promise<Collection>;
     listCollections(): Promise<any>;
-    getCollection(name: string): Promise<Collection>;
-    deleteCollection(name: string): Promise<import("axios").AxiosResponse<any, any>>;
+    getCollection(name: string, embeddingFunction?: CallableFunction): Promise<Collection>;
+    deleteCollection(name: string): Promise<any>;
 }
+export {};
 //# sourceMappingURL=index.d.ts.map

package/dist/module/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAqB7C,qBAAa,UAAU;IACrB,OAAO,CAAC,IAAI,CAAS;IACrB,OAAO,CAAC,GAAG,CAAa;gBAEZ,IAAI,EAAE,MAAM,EAAE,GAAG,EAAE,UAAU;IAK5B,GAAG,CACd,GAAG,EAAE,MAAM,GAAG,KAAK,CAAC,GAAG,CAAC,EACxB,UAAU,EAAE,KAAK,CAAC,GAAG,CAAC,EACtB,SAAS,CAAC,EAAE,KAAK,CAAC,GAAG,CAAC,GAAG,MAAM,EAC/B,SAAS,CAAC,EAAE,MAAM,GAAG,KAAK,CAAC,GAAG,CAAC,EAC/B,eAAe,GAAE,OAAc;IA0CpB,KAAK;IAKL,GAAG,CACd,GAAG,CAAC,EAAE,MAAM,EAAE,EACd,KAAK,CAAC,EAAE,MAAM,EACd,KAAK,CAAC,EAAE,MAAM,EACd,MAAM,CAAC,EAAE,MAAM;IAeJ,KAAK,CAChB,gBAAgB,EAAE,MAAM,EAAE,EAC1B,SAAS,GAAE,MAAW,EACtB,KAAK,CAAC,EAAE,MAAM;IAeH,IAAI,CAAC,KAAK,GAAE,MAAW;IAOvB,WAAW;IAIX,MAAM,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,EAAE,KAAK,CAAC,EAAE,MAAM;CAOnD;AAED,qBAAa,YAAY;IACvB,OAAO,CAAC,GAAG,CAAa;gBAEZ,QAAQ,EAAE,MAAM;IAOf,KAAK;IAIL,gBAAgB,CAAC,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,EAAE,MAAM;IAOhD,eAAe;IAKf,aAAa,CAAC,IAAI,EAAE,MAAM;IAI1B,gBAAgB,CAAC,IAAI,EAAE,MAAM;CAI3C"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAyB7C,qBAAa,uBAAuB;IAClC,OAAO,CAAC,OAAO,CAAS;IACxB,OAAO,CAAC,MAAM,CAAS;IACvB,OAAO,CAAC,KAAK,CAAS;gBAEV,cAAc,EAAE,MAAM,EAAE,YAAY,EAAE,MAAM,EAAE,sBAAsB,EAAE,MAAM;IAc3E,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE;CAiBtC;AAID,qBAAa,uBAAuB;IAClC,OAAO,CAAC,OAAO,CAAS;gBAEZ,cAAc,EAAE,MAAM;IAYrB,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE;CAQtC;AAED,KAAK,gBAAgB,GAAG;IACtB,QAAO,gBAAgB,CAAC;IACxB,QAAQ,CAAC,KAAK,EAAE,MAAM,EAAE,GAAG,MAAM,EAAE,EAAE,CAAC;CACvC,CAAC;AAEF,qBAAa,UAAU;IACd,IAAI,EAAE,MAAM,CAAC;IACpB,OAAO,CAAC,GAAG,CAAa;IACjB,iBAAiB,EAAE,gBAAgB,GAAG,SAAS,CAAC;gBAE3C,IAAI,EAAE,MAAM,EAAE,GAAG,EAAE,UAAU,EAAE,iBAAiB,CAAC,EAAE,gBAAgB;IAOlE,GAAG,CACd,GAAG,EAAE,MAAM,GAAG,MAAM,EAAE,EACtB,UAAU,EAAE,MAAM,EAAE,GAAG,MAAM,EAAE,EAAE,EACjC,SAAS,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE,EAC7B,SAAS,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE,EAC7B,eAAe,GAAE,OAAc;IA8DpB,KAAK;IAKL,GAAG,CACd,GAAG,CAAC,EAAE,MAAM,EAAE,EACd,KAAK,CAAC,EAAE,MAAM,EACd,KAAK,CAAC,EAAE,MAAM,EACd,MAAM,CAAC,EAAE,MAAM;IAuBJ,KAAK,CAChB,gBAAgB,EAAE,MAAM,EAAE,GAAG,MAAM,EAAE,EAAE,EACvC,SAAS,GAAE,MAAW,EACtB,KAAK,CAAC,EAAE,MAAM,EACd,UAAU,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE;IAmCnB,IAAI,CAAC,KAAK,GAAE,MAAW;IAQvB,WAAW;IAIX,MAAM,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,EAAE,KAAK,CAAC,EAAE,MAAM;CAanD;AAED,qBAAa,YAAY;IACvB,OAAO,CAAC,GAAG,CAAa;gBAEZ,QAAQ,CAAC,EAAE,MAAM;IAQhB,KAAK;IAIL,gBAAgB,CAAC,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,EAAE,MAAM,EAAE,iBAAiB,CAAC,EAAE,gBAAgB;IAgBtF,eAAe;IAKf,aAAa,CAAC,IAAI,EAAE,MAAM,EAAE,iBAAiB,CAAC,EAAE,gBAAgB;IAIhE,gBAAgB,CAAC,IAAI,EAAE,MAAM;CAU3C"}

package/dist/module/index.js

@@ -18,12 +18,81 @@ function toArrayOfArrays(obj) {
         return [obj];
     }
 }
+class EmbeddingFunction {
+}
+let OpenAIApi;
+export class OpenAIEmbeddingFunction {
+    constructor(openai_api_key, openai_model, openai_organization_id) {
+        try {
+            // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+            OpenAIApi = require("openai");
+        }
+        catch (_a) {
+            throw new Error("Please install the openai package to use the OpenAIEmbeddingFunction, `npm install -S openai`");
+        }
+        this.api_key = openai_api_key;
+        this.org_id = openai_organization_id || "";
+        this.model = openai_model || "text-embedding-ada-002";
+    }
+    async generate(texts) {
+        const configuration = new OpenAIApi.Configuration({
+            organization: this.org_id,
+            apiKey: this.api_key,
+        });
+        const openai = new OpenAIApi.OpenAIApi(configuration);
+        const embeddings = [];
+        const response = await openai.createEmbedding({
+            model: "text-embedding-ada-002",
+            input: texts,
+        });
+        const data = response.data['data'];
+        for (let i = 0; i < data.length; i += 1) {
+            embeddings.push(data[i]['embedding']);
+        }
+        return embeddings;
+    }
+}
+let CohereAiApi;
+export class CohereEmbeddingFunction {
+    constructor(cohere_api_key) {
+        try {
+            // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+            CohereAiApi = require("cohere-ai");
+        }
+        catch (_a) {
+            throw new Error("Please install the cohere-ai package to use the CohereEmbeddingFunction, `npm install -S cohere-ai`");
+        }
+        this.api_key = cohere_api_key;
+    }
+    async generate(texts) {
+        const cohere = CohereAiApi.init(this.api_key);
+        const embeddings = [];
+        const response = await CohereAiApi.embed({
+            texts: texts,
+        });
+        return response.body.embeddings;
+    }
+}
 export class Collection {
-    constructor(name, api) {
+    constructor(name, api, embeddingFunction) {
         this.name = name;
         this.api = api;
+        if (embeddingFunction !== undefined)
+            this.embeddingFunction = embeddingFunction;
     }
     async add(ids, embeddings, metadatas, documents, increment_index = true) {
+        if ((embeddings === undefined) && (documents === undefined)) {
+            throw new Error("embeddings and documents cannot both be undefined");
+        }
+        else if ((embeddings === undefined) && (documents !== undefined)) {
+            const documentsArray = toArray(documents);
+            if (this.embeddingFunction !== undefined) {
+                embeddings = await this.embeddingFunction.generate(documentsArray);
+            }
+            else {
+                throw new Error("embeddingFunction is undefined. Please configure an embedding function");
+            }
+        }
         const idsArray = toArray(ids);
         const embeddingsArray = toArrayOfArrays(embeddings);
         let metadatasArray;
@@ -34,18 +103,18 @@ export class Collection {
             metadatasArray = toArray(metadatas);
         }
         let documentsArray;
-        if (metadatas === undefined) {
+        if (documents === undefined) {
             documentsArray = undefined;
         }
         else {
-            documentsArray = toArray(metadatas);
+            documentsArray = toArray(documents);
         }
-        if (idsArray.length !== embeddingsArray.length ||
+        if (((embeddingsArray !== undefined) && idsArray.length !== embeddingsArray.length) ||
             ((metadatasArray !== undefined) && idsArray.length !== metadatasArray.length) ||
             ((documentsArray !== undefined) && idsArray.length !== documentsArray.length)) {
             throw new Error("ids, embeddings, metadatas, and documents must all be the same length");
         }
-        return await this.api.add({
+        const response = await this.api.add({
             collectionName: this.name,
             addEmbedding: {
                 ids: idsArray,
@@ -54,15 +123,22 @@ export class Collection {
                 metadatas: metadatasArray,
                 increment_index: increment_index,
             },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
+        return response;
     }
     async count() {
         const response = await this.api.count({ collectionName: this.name });
         return response.data;
     }
     async get(ids, where, limit, offset) {
-        const idsArray = toArray(ids);
-        return await this.api.get({
+        let idsArray = undefined;
+        if (ids !== undefined)
+            idsArray = toArray(ids);
+        var resp = await this.api.get({
             collectionName: this.name,
             getEmbedding: {
                 ids: idsArray,
@@ -70,9 +146,26 @@ export class Collection {
                 limit,
                 offset,
             },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
+        return resp;
     }
-    async query(query_embeddings, n_results = 10, where) {
+    async query(query_embeddings, n_results = 10, where, query_text) {
+        if ((query_embeddings === undefined) && (query_text === undefined)) {
+            throw new Error("query_embeddings and query_text cannot both be undefined");
+        }
+        else if ((query_embeddings === undefined) && (query_text !== undefined)) {
+            const query_texts = toArray(query_text);
+            if (this.embeddingFunction !== undefined) {
+                query_embeddings = await this.embeddingFunction.generate(query_texts);
+            }
+            else {
+                throw new Error("embeddingFunction is undefined. Please configure an embedding function");
+            }
+        }
         const query_embeddingsArray = toArrayOfArrays(query_embeddings);
         const response = await this.api.getNearestNeighbors({
             collectionName: this.name,
@@ -81,27 +174,39 @@ export class Collection {
                 where,
                 n_results,
             },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
-        return response.data;
+        return response;
     }
     async peek(limit = 10) {
-        return await this.api.get({
+        const response = await this.api.get({
             collectionName: this.name,
             getEmbedding: { limit: limit },
         });
+        return response.data;
     }
     async createIndex() {
         return await this.api.createIndex({ collectionName: this.name });
     }
     async delete(ids, where) {
-        return await this.api._delete({
+        var response = await this.api._delete({
             collectionName: this.name,
             deleteEmbedding: { ids: ids, where: where },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
+        return response;
     }
 }
 export class ChromaClient {
     constructor(basePath) {
+        if (basePath === undefined)
+            basePath = "http://localhost:8000";
         const apiConfig = new Configuration({
             basePath,
         });
@@ -110,21 +215,33 @@ export class ChromaClient {
     async reset() {
         return await this.api.reset();
     }
-    async createCollection(name, metadata) {
+    async createCollection(name, metadata, embeddingFunction) {
         const newCollection = await this.api.createCollection({
             createCollection: { name, metadata },
+        }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
         });
-        return new Collection(name, this.api);
+        if (newCollection.error) {
+            throw new Error(newCollection.error);
+        }
+        return new Collection(name, this.api, embeddingFunction);
     }
     async listCollections() {
         const response = await this.api.listCollections();
         return response.data;
     }
-    async getCollection(name) {
-        return new Collection(name, this.api);
+    async getCollection(name, embeddingFunction) {
+        return new Collection(name, this.api, embeddingFunction);
     }
     async deleteCollection(name) {
-        return await this.api.deleteCollection({ collectionName: name });
+        const response = await this.api.deleteCollection({ collectionName: name }).then(function (response) {
+            return response.data;
+        }).catch(function ({ response }) {
+            return response.data;
+        });
+        return response;
     }
 }
 //# sourceMappingURL=index.js.map

package/dist/module/index.js.map

@@ -1 +1 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAC7C,OAAO,EAAE,aAAa,EAAE,MAAM,2BAA2B,CAAC;AAE1D,uDAAuD;AACvD,SAAS,OAAO,CAAI,GAAiB;IACnC,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE;QACtB,OAAO,GAAG,CAAC;KACZ;SAAM;QACL,OAAO,CAAC,GAAG,CAAC,CAAC;KACd;AACH,CAAC;AAED,oDAAoD;AACpD,SAAS,eAAe,CAAI,GAA+B;IACzD,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE;QACzB,OAAO,GAAsB,CAAC;KAC/B;SAAM;QACL,OAAO,CAAC,GAAG,CAAoB,CAAC;KACjC;AACH,CAAC;AAED,MAAM,OAAO,UAAU;IAIrB,YAAY,IAAY,EAAE,GAAe;QACvC,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;QACjB,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;IACjB,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAwB,EACxB,UAAsB,EACtB,SAA+B,EAC/B,SAA+B,EAC/B,kBAA2B,IAAI;QAG/B,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAC9B,MAAM,eAAe,GAAG,eAAe,CAAC,UAAU,CAAC,CAAC;QAEpD,IAAI,cAAc,CAAC;QACnB,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IAAI,cAAc,CAAC;QACnB,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IACE,QAAQ,CAAC,MAAM,KAAK,eAAe,CAAC,MAAM;YAC1C,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC;YAC7E,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC,EAC7E;YACA,MAAM,IAAI,KAAK,CACb,uEAAuE,CACxE,CAAC;SACH;QAED,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YACxB,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,UAAU,EAAE,eAAe;gBAC3B,SAAS,EAAE,cAAc;gBACzB,SAAS,EAAE,cAAc;gBACzB,eAAe,EAAE,eAAe;aACjC;SACF,CAAC,CAAC;IACL,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;QACrE,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAc,EACd,KAAc,EACd,KAAc,EACd,MAAe;QAEf,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAE9B,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YACxB,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,KAAK;gBACL,KAAK;gBACL,MAAM;aACP;SACF,CAAC,CAAC;IACL,CAAC;IAEM,KAAK,CAAC,KAAK,CAChB,gBAA0B,EAC1B,YAAoB,EAAE,EACtB,KAAc;QAEd,MAAM,qBAAqB,GAAG,eAAe,CAAC,gBAAgB,CAAC,CAAC;QAEhE,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,mBAAmB,CAAC;YAClD,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,cAAc,EAAE;gBACd,gBAAgB,EAAE,qBAAqB;gBACvC,KAAK;gBACL,SAAS;aACV;SACF,CAAC,CAAC;QACH,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,IAAI,CAAC,QAAgB,EAAE;QAClC,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YACxB,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE;SAC/B,CAAC,CAAC;IACL,CAAC;IAEM,KAAK,CAAC,WAAW;QACtB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;IACnE,CAAC;IAEM,KAAK,CAAC,MAAM,CAAC,GAAc,EAAE,KAAc;QAChD,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC;YAC5B,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,eAAe,EAAE,EAAE,GAAG,EAAE,GAAG,EAAE,KAAK,EAAE,KAAK,EAAE;SAC5C,CAAC,CAAC;IACL,CAAC;CAEF;AAED,MAAM,OAAO,YAAY;IAGvB,YAAY,QAAgB;QAC1B,MAAM,SAAS,GAAkB,IAAI,aAAa,CAAC;YACjD,QAAQ;SACT,CAAC,CAAC;QACH,IAAI,CAAC,GAAG,GAAG,IAAI,UAAU,CAAC,SAAS,CAAC,CAAC;IACvC,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,CAAC;IAChC,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY,EAAE,QAAiB;QAC3D,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC;YACpD,gBAAgB,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE;SACrC,CAAC,CAAC;QACH,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC;IACxC,CAAC;IAEM,KAAK,CAAC,eAAe;QAC1B,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,eAAe,EAAE,CAAC;QAClD,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,aAAa,CAAC,IAAY;QACrC,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC;IACxC,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY;QACxC,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC,EAAE,cAAc,EAAE,IAAI,EAAE,CAAC,CAAC;IACnE,CAAC;CAEF"}
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC;AAC7C,OAAO,EAAE,aAAa,EAAE,MAAM,2BAA2B,CAAC;AAE1D,uDAAuD;AACvD,SAAS,OAAO,CAAI,GAAiB;IACnC,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE;QACtB,OAAO,GAAG,CAAC;KACZ;SAAM;QACL,OAAO,CAAC,GAAG,CAAC,CAAC;KACd;AACH,CAAC;AAED,oDAAoD;AACpD,SAAS,eAAe,CAAI,GAA+B;IACzD,IAAI,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,EAAE;QACzB,OAAO,GAAsB,CAAC;KAC/B;SAAM;QACL,OAAO,CAAC,GAAG,CAAoB,CAAC;KACjC;AACH,CAAC;AAED,MAAM,iBAAiB;CAAI;AAE3B,IAAI,SAAc,CAAC;AAEnB,MAAM,OAAO,uBAAuB;IAKlC,YAAY,cAAsB,EAAE,YAAoB,EAAE,sBAA8B;QACtF,IAAI;YACF,4EAA4E;YAC5E,SAAS,GAAG,OAAO,CAAC,QAAQ,CAAC,CAAC;SAC/B;QAAC,WAAM;YACN,MAAM,IAAI,KAAK,CACb,+FAA+F,CAChG,CAAC;SACH;QACD,IAAI,CAAC,OAAO,GAAG,cAAc,CAAC;QAC9B,IAAI,CAAC,MAAM,GAAG,sBAAsB,IAAI,EAAE,CAAC;QAC3C,IAAI,CAAC,KAAK,GAAG,YAAY,IAAI,wBAAwB,CAAC;IACxD,CAAC;IAEM,KAAK,CAAC,QAAQ,CAAC,KAAe;QACnC,MAAM,aAAa,GAAG,IAAI,SAAS,CAAC,aAAa,CAAC;YAChD,YAAY,EAAE,IAAI,CAAC,MAAM;YACzB,MAAM,EAAE,IAAI,CAAC,OAAO;SACrB,CAAC,CAAC;QACH,MAAM,MAAM,GAAG,IAAI,SAAS,CAAC,SAAS,CAAC,aAAa,CAAC,CAAC;QACtD,MAAM,UAAU,GAAG,EAAE,CAAC;QACtB,MAAM,QAAQ,GAAG,MAAM,MAAM,CAAC,eAAe,CAAC;YAC5C,KAAK,EAAE,wBAAwB;YAC/B,KAAK,EAAE,KAAK;SACb,CAAC,CAAC;QACH,MAAM,IAAI,GAAG,QAAQ,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QACnC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,MAAM,EAAE,CAAC,IAAI,CAAC,EAAE;YACvC,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,WAAW,CAAC,CAAC,CAAC;SACvC;QACD,OAAO,UAAU,CAAC;IACpB,CAAC;CACF;AAED,IAAI,WAAgB,CAAC;AAErB,MAAM,OAAO,uBAAuB;IAGlC,YAAY,cAAsB;QAChC,IAAI;YACF,4EAA4E;YAC5E,WAAW,GAAG,OAAO,CAAC,WAAW,CAAC,CAAC;SACpC;QAAC,WAAM;YACN,MAAM,IAAI,KAAK,CACb,qGAAqG,CACtG,CAAC;SACH;QACD,IAAI,CAAC,OAAO,GAAG,cAAc,CAAC;IAChC,CAAC;IAEM,KAAK,CAAC,QAAQ,CAAC,KAAe;QACnC,MAAM,MAAM,GAAG,WAAW,CAAC,IAAI,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QAC9C,MAAM,UAAU,GAAG,EAAE,CAAC;QACtB,MAAM,QAAQ,GAAG,MAAM,WAAW,CAAC,KAAK,CAAC;YACvC,KAAK,EAAE,KAAK;SACb,CAAC,CAAC;QACH,OAAO,QAAQ,CAAC,IAAI,CAAC,UAAU,CAAC;IAClC,CAAC;CACF;AAOD,MAAM,OAAO,UAAU;IAKrB,YAAY,IAAY,EAAE,GAAe,EAAE,iBAAoC;QAC7E,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;QACjB,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;QACf,IAAI,iBAAiB,KAAK,SAAS;YACjC,IAAI,CAAC,iBAAiB,GAAG,iBAAiB,CAAC;IAC/C,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAsB,EACtB,UAAiC,EACjC,SAA6B,EAC7B,SAA6B,EAC7B,kBAA2B,IAAI;QAE/B,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,IAAI,CAAC,SAAS,KAAK,SAAS,CAAC,EAAE;YAC3D,MAAM,IAAI,KAAK,CACb,mDAAmD,CACpD,CAAC;SACH;aAAM,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,IAAI,CAAC,SAAS,KAAK,SAAS,CAAC,EAAE;YAClE,MAAM,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;YAC1C,IAAI,IAAI,CAAC,iBAAiB,KAAK,SAAS,EAAE;gBACxC,UAAU,GAAG,MAAM,IAAI,CAAC,iBAAiB,CAAC,QAAQ,CAAC,cAAc,CAAC,CAAA;aACnE;iBAAM;gBACL,MAAM,IAAI,KAAK,CACb,wEAAwE,CACzE,CAAC;aACH;SACF;QAED,MAAM,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAC9B,MAAM,eAAe,GAAG,eAAe,CAAC,UAAU,CAAC,CAAC;QAEpD,IAAI,cAAoC,CAAC;QACzC,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IAAI,cAAkD,CAAC;QACvD,IAAI,SAAS,KAAK,SAAS,EAAE;YAC3B,cAAc,GAAG,SAAS,CAAA;SAC3B;aAAM;YACL,cAAc,GAAG,OAAO,CAAC,SAAS,CAAC,CAAC;SACrC;QAED,IACE,CAAC,CAAC,eAAe,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,eAAe,CAAC,MAAM,CAAC;YAC/E,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC;YAC7E,CAAC,CAAC,cAAc,KAAK,SAAS,CAAC,IAAI,QAAQ,CAAC,MAAM,KAAK,cAAc,CAAC,MAAM,CAAC,EAC7E;YACA,MAAM,IAAI,KAAK,CACb,uEAAuE,CACxE,CAAC;SACH;QAED,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YAClC,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,UAAU,EAAE,eAAe;gBAC3B,SAAS,EAAE,cAAc;gBACzB,SAAS,EAAE,cAAc;gBACzB,eAAe,EAAE,eAAe;aACjC;SACF,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAA;IACjB,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;QACrE,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,GAAG,CACd,GAAc,EACd,KAAc,EACd,KAAc,EACd,MAAe;QAEf,IAAI,QAAQ,GAAG,SAAS,CAAA;QACxB,IAAI,GAAG,KAAK,SAAS;YAAE,QAAQ,GAAG,OAAO,CAAC,GAAG,CAAC,CAAC;QAE/C,IAAI,IAAI,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YAC5B,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE;gBACZ,GAAG,EAAE,QAAQ;gBACb,KAAK;gBACL,KAAK;gBACL,MAAM;aACP;SACF,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,IAAI,CAAA;IAEb,CAAC;IAEM,KAAK,CAAC,KAAK,CAChB,gBAAuC,EACvC,YAAoB,EAAE,EACtB,KAAc,EACd,UAA8B;QAE9B,IAAI,CAAC,gBAAgB,KAAK,SAAS,CAAC,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,EAAE;YAClE,MAAM,IAAI,KAAK,CACb,0DAA0D,CAC3D,CAAC;SACH;aAAM,IAAI,CAAC,gBAAgB,KAAK,SAAS,CAAC,IAAI,CAAC,UAAU,KAAK,SAAS,CAAC,EAAE;YACzE,MAAM,WAAW,GAAG,OAAO,CAAC,UAAU,CAAC,CAAC;YACxC,IAAI,IAAI,CAAC,iBAAiB,KAAK,SAAS,EAAE;gBACxC,gBAAgB,GAAG,MAAM,IAAI,CAAC,iBAAiB,CAAC,QAAQ,CAAC,WAAW,CAAC,CAAA;aACtE;iBAAM;gBACL,MAAM,IAAI,KAAK,CACb,wEAAwE,CACzE,CAAC;aACH;SACF;QAED,MAAM,qBAAqB,GAAG,eAAe,CAAC,gBAAgB,CAAC,CAAC;QAEhE,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,mBAAmB,CAAC;YAClD,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,cAAc,EAAE;gBACd,gBAAgB,EAAE,qBAAqB;gBACvC,KAAK;gBACL,SAAS;aACV;SACF,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAC;IAClB,CAAC;IAEM,KAAK,CAAC,IAAI,CAAC,QAAgB,EAAE;QAClC,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC;YAClC,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,YAAY,EAAE,EAAE,KAAK,EAAE,KAAK,EAAE;SAC/B,CAAC,CAAC;QACH,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,WAAW;QACtB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,WAAW,CAAC,EAAE,cAAc,EAAE,IAAI,CAAC,IAAI,EAAE,CAAC,CAAC;IACnE,CAAC;IAEM,KAAK,CAAC,MAAM,CAAC,GAAc,EAAE,KAAc;QAChD,IAAI,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,OAAO,CAAC;YACpC,cAAc,EAAE,IAAI,CAAC,IAAI;YACzB,eAAe,EAAE,EAAE,GAAG,EAAE,GAAG,EAAE,KAAK,EAAE,KAAK,EAAE;SAC5C,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAA;IACjB,CAAC;CAEF;AAED,MAAM,OAAO,YAAY;IAGvB,YAAY,QAAiB;QAC3B,IAAI,QAAQ,KAAK,SAAS;YAAE,QAAQ,GAAG,uBAAuB,CAAC;QAC/D,MAAM,SAAS,GAAkB,IAAI,aAAa,CAAC;YACjD,QAAQ;SACT,CAAC,CAAC;QACH,IAAI,CAAC,GAAG,GAAG,IAAI,UAAU,CAAC,SAAS,CAAC,CAAC;IACvC,CAAC;IAEM,KAAK,CAAC,KAAK;QAChB,OAAO,MAAM,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,CAAC;IAChC,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY,EAAE,QAAiB,EAAE,iBAAoC;QACjG,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC;YACpD,gBAAgB,EAAE,EAAE,IAAI,EAAE,QAAQ,EAAE;SACrC,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YACxB,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,IAAI,aAAa,CAAC,KAAK,EAAE;YACvB,MAAM,IAAI,KAAK,CAAC,aAAa,CAAC,KAAK,CAAC,CAAC;SACtC;QAED,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,EAAE,iBAAiB,CAAC,CAAC;IAC3D,CAAC;IAEM,KAAK,CAAC,eAAe;QAC1B,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,eAAe,EAAE,CAAC;QAClD,OAAO,QAAQ,CAAC,IAAI,CAAC;IACvB,CAAC;IAEM,KAAK,CAAC,aAAa,CAAC,IAAY,EAAE,iBAAoC;QAC3E,OAAO,IAAI,UAAU,CAAC,IAAI,EAAE,IAAI,CAAC,GAAG,EAAE,iBAAiB,CAAC,CAAC;IAC3D,CAAC;IAEM,KAAK,CAAC,gBAAgB,CAAC,IAAY;QACxC,MAAM,QAAQ,GAAG,MAAM,IAAI,CAAC,GAAG,CAAC,gBAAgB,CAAC,EAAE,cAAc,EAAE,IAAI,EAAE,CAAC,CAAC,IAAI,CAAC,UAAU,QAAQ;YAChG,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC,KAAK,CAAC,UAAU,EAAE,QAAQ,EAAE;YAC7B,OAAO,QAAQ,CAAC,IAAI,CAAC;QACvB,CAAC,CAAC,CAAC;QAEH,OAAO,QAAQ,CAAA;IACjB,CAAC;CAEF"}

package/package.json

@@ -1,14 +1,20 @@
 {
   "name": "chromadb",
-  "version": "1.2.1",
+  "version": "1.3.0",
   "description": "A JavaScript interface for chroma",
   "keywords": [],
   "author": "",
   "license": "Apache-2.0",
   "devDependencies": {
     "@openapitools/openapi-generator-cli": "^2.5.2",
+    "@types/jest": "^29.4.0",
+    "jest": "^29.4.3",
     "npm-run-all": "^4.1.5",
-    "rimraf": "^3.0.2"
+    "rimraf": "^3.0.2",
+    "ts-jest": "^29.0.5",
+    "ts-node": "^10.9.1",
+    "tsd": "^0.24.1",
+    "typescript": "^4.5.5"
   },
   "dependencies": {
     "axios": "^0.26.0"
@@ -20,6 +26,13 @@
     "dist"
   ],
   "scripts": {
+    "test": "run-s db:clean db:run test:runfull db:clean",
+    "test:set-port": "cross-env URL=localhost:8001",
+    "test:run": "jest --runInBand",
+    "test:runfull": "PORT=8001 jest --runInBand",
+    "test:update": "run-s db:clean db:run && jest --runInBand --updateSnapshot && run-s db:clean",
+    "db:clean": "cd ../.. && docker-compose -f docker-compose-js-tests.yml down --volumes",
+    "db:run": "cd ../.. && docker-compose -f docker-compose-js-tests.yml up --detach && sleep 5",
     "clean": "rimraf dist",
     "build": "run-s clean build:*",
     "build:main": "tsc -p tsconfig.json",

package/src/index.ts

@@ -19,42 +19,133 @@ function toArrayOfArrays<T>(obj: Array<Array<T>> | Array<T>): Array<Array<T>> {
   }
 }
 
+class EmbeddingFunction { }
+
+let OpenAIApi: any;
+
+export class OpenAIEmbeddingFunction {
+  private api_key: string;
+  private org_id: string;
+  private model: string;
+
+  constructor(openai_api_key: string, openai_model: string, openai_organization_id: string) {
+    try {
+      // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+      OpenAIApi = require("openai");
+    } catch {
+      throw new Error(
+        "Please install the openai package to use the OpenAIEmbeddingFunction, `npm install -S openai`",
+      );
+    }
+    this.api_key = openai_api_key;
+    this.org_id = openai_organization_id || "";
+    this.model = openai_model || "text-embedding-ada-002";
+  }
+
+  public async generate(texts: string[]) {
+    const configuration = new OpenAIApi.Configuration({
+      organization: this.org_id,
+      apiKey: this.api_key,
+    });
+    const openai = new OpenAIApi.OpenAIApi(configuration);
+    const embeddings = [];
+    const response = await openai.createEmbedding({
+      model: "text-embedding-ada-002",
+      input: texts,
+    });
+    const data = response.data['data'];
+    for (let i = 0; i < data.length; i += 1) {
+      embeddings.push(data[i]['embedding']);
+    }
+    return embeddings;
+  }
+}
+
+let CohereAiApi: any;
+
+export class CohereEmbeddingFunction {
+  private api_key: string;
+
+  constructor(cohere_api_key: string) {
+    try {
+      // eslint-disable-next-line global-require,import/no-extraneous-dependencies
+      CohereAiApi = require("cohere-ai");
+    } catch {
+      throw new Error(
+        "Please install the cohere-ai package to use the CohereEmbeddingFunction, `npm install -S cohere-ai`",
+      );
+    }
+    this.api_key = cohere_api_key;
+  }
+
+  public async generate(texts: string[]) {
+    const cohere = CohereAiApi.init(this.api_key);
+    const embeddings = [];
+    const response = await CohereAiApi.embed({
+      texts: texts,
+    });
+    return response.body.embeddings;
+  }
+}
+
+type CallableFunction = {
+  new(): CallableFunction;
+  generate(texts: string[]): number[][];
+};
+
 export class Collection {
-  private name: string;
+  public name: string;
   private api: DefaultApi;
+  public embeddingFunction: CallableFunction | undefined;
 
-  constructor(name: string, api: DefaultApi) {
+  constructor(name: string, api: DefaultApi, embeddingFunction?: CallableFunction) {
     this.name = name;
     this.api = api;
+    if (embeddingFunction !== undefined)
+      this.embeddingFunction = embeddingFunction;
   }
 
   public async add(
-    ids: string | Array<any>,
-    embeddings: Array<any>,
-    metadatas?: Array<any> | object,
-    documents?: string | Array<any>,
+    ids: string | string[],
+    embeddings: number[] | number[][],
+    metadatas?: object | object[],
+    documents?: string | string[],
     increment_index: boolean = true,
   ) {
+    if ((embeddings === undefined) && (documents === undefined)) {
+      throw new Error(
+        "embeddings and documents cannot both be undefined",
+      );
+    } else if ((embeddings === undefined) && (documents !== undefined)) {
+      const documentsArray = toArray(documents);
+      if (this.embeddingFunction !== undefined) {
+        embeddings = await this.embeddingFunction.generate(documentsArray)
+      } else {
+        throw new Error(
+          "embeddingFunction is undefined. Please configure an embedding function",
+        );
+      }
+    }
 
     const idsArray = toArray(ids);
     const embeddingsArray = toArrayOfArrays(embeddings);
 
-    let metadatasArray;
+    let metadatasArray: object[] | undefined;
     if (metadatas === undefined) {
       metadatasArray = undefined
     } else {
       metadatasArray = toArray(metadatas);
     }
 
-    let documentsArray;
-    if (metadatas === undefined) {
+    let documentsArray: (string | undefined)[] | undefined;
+    if (documents === undefined) {
       documentsArray = undefined
     } else {
-      documentsArray = toArray(metadatas);
+      documentsArray = toArray(documents);
     }
 
     if (
-      idsArray.length !== embeddingsArray.length ||
+      ((embeddingsArray !== undefined) && idsArray.length !== embeddingsArray.length) ||
       ((metadatasArray !== undefined) && idsArray.length !== metadatasArray.length) ||
       ((documentsArray !== undefined) && idsArray.length !== documentsArray.length)
     ) {
@@ -63,7 +154,7 @@ export class Collection {
       );
     }
 
-    return await this.api.add({
+    const response = await this.api.add({
       collectionName: this.name,
       addEmbedding: {
         ids: idsArray,
@@ -72,7 +163,13 @@ export class Collection {
         metadatas: metadatasArray,
         increment_index: increment_index,
       },
+    }).then(function (response) {
+      return response.data;
+    }).catch(function ({ response }) {
+      return response.data;
     });
+
+    return response
   }
 
   public async count() {
@@ -86,9 +183,10 @@ export class Collection {
     limit?: number,
     offset?: number,
   ) {
-    const idsArray = toArray(ids);
+    let idsArray = undefined
+    if (ids !== undefined) idsArray = toArray(ids);
 
-    return await this.api.get({
+    var resp = await this.api.get({
       collectionName: this.name,
       getEmbedding: {
         ids: idsArray,
@@ -96,14 +194,37 @@ export class Collection {
         limit,
         offset,
       },
+    }).then(function (response) {
+      return response.data;
+    }).catch(function ({ response }) {
+      return response.data;
     });
+
+    return resp
+
   }
 
   public async query(
-    query_embeddings: number[],
+    query_embeddings: number[] | number[][],
     n_results: number = 10,
     where?: object,
+    query_text?: string | string[],
   ) {
+    if ((query_embeddings === undefined) && (query_text === undefined)) {
+      throw new Error(
+        "query_embeddings and query_text cannot both be undefined",
+      );
+    } else if ((query_embeddings === undefined) && (query_text !== undefined)) {
+      const query_texts = toArray(query_text);
+      if (this.embeddingFunction !== undefined) {
+        query_embeddings = await this.embeddingFunction.generate(query_texts)
+      } else {
+        throw new Error(
+          "embeddingFunction is undefined. Please configure an embedding function",
+        );
+      }
+    }
+
     const query_embeddingsArray = toArrayOfArrays(query_embeddings);
 
     const response = await this.api.getNearestNeighbors({
@@ -113,15 +234,21 @@ export class Collection {
         where,
         n_results,
       },
+    }).then(function (response) {
+      return response.data;
+    }).catch(function ({ response }) {
+      return response.data;
     });
-    return response.data;
+
+    return response;
   }
 
   public async peek(limit: number = 10) {
-    return await this.api.get({
+    const response = await this.api.get({
       collectionName: this.name,
       getEmbedding: { limit: limit },
     });
+    return response.data;
   }
 
   public async createIndex() {
@@ -129,10 +256,16 @@ export class Collection {
   }
 
   public async delete(ids?: string[], where?: object) {
-    return await this.api._delete({
+    var response = await this.api._delete({
       collectionName: this.name,
       deleteEmbedding: { ids: ids, where: where },
+    }).then(function (response) {
+      return response.data;
+    }).catch(function ({ response }) {
+      return response.data;
     });
+
+    return response
   }
 
 }
@@ -140,7 +273,8 @@ export class Collection {
 export class ChromaClient {
   private api: DefaultApi;
 
-  constructor(basePath: string) {
+  constructor(basePath?: string) {
+    if (basePath === undefined) basePath = "http://localhost:8000";
     const apiConfig: Configuration = new Configuration({
       basePath,
     });
@@ -151,11 +285,20 @@ export class ChromaClient {
     return await this.api.reset();
   }
 
-  public async createCollection(name: string, metadata?: object) {
+  public async createCollection(name: string, metadata?: object, embeddingFunction?: CallableFunction) {
     const newCollection = await this.api.createCollection({
       createCollection: { name, metadata },
+    }).then(function (response) {
+      return response.data;
+    }).catch(function ({ response }) {
+      return response.data;
     });
-    return new Collection(name, this.api);
+
+    if (newCollection.error) {
+      throw new Error(newCollection.error);
+    }
+
+    return new Collection(name, this.api, embeddingFunction);
   }
 
   public async listCollections() {
@@ -163,12 +306,18 @@ export class ChromaClient {
     return response.data;
   }
 
-  public async getCollection(name: string) {
-    return new Collection(name, this.api);
+  public async getCollection(name: string, embeddingFunction?: CallableFunction) {
+    return new Collection(name, this.api, embeddingFunction);
   }
 
   public async deleteCollection(name: string) {
-    return await this.api.deleteCollection({ collectionName: name });
+    const response = await this.api.deleteCollection({ collectionName: name }).then(function (response) {
+      return response.data;
+    }).catch(function ({ response }) {
+      return response.data;
+    });
+
+    return response
   }
 
 }