Mintplex-Labs/chromadb-extended
1
0
Fork 0
mirror of https://github.com/Mintplex-Labs/chromadb-extended.git synced 2026-07-01 20:44:00 -04:00
Code Issues Packages Projects Releases Wiki Activity

bump patch version

Browse Source 
add support for collection methods
This commit is contained in:
timothycarambat
2023-06-22 11:19:11 -07:00
parent ee484e74a6
commit 572de5c220
7 changed files with 1280 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dist/index.d.ts
Vendored
+248 -4
View File
@@ -1,13 +1,11 @@
import { IncludeEnum, Metadata, Metadatas, Embedding, Embeddings, Document, Documents, Where, WhereDocument, ID, IDs, PositiveInteger, GetResponse, QueryResponse, AddResponse, CollectionMetadata, CollectionType } from "chromadb/dist/main/types";
import { IEmbeddingFunction } from "chromadb/dist/main//embeddings/IEmbeddingFunction";
import { Api } from "chromadb/dist/main/generated";
import { Collection } from "chromadb/dist/main/Collection";
import { CollectionMetadata, CollectionType } from "chromadb/dist/main/types";
import { ApiApi as DefaultApi, Api } from "chromadb/dist/main/generated";
export declare class ChromaClientExtended {
/**
* @ignore
*/
private api;
private options;
/**
* Creates a new ChromaClient instance.
* @param {Object} params - The parameters for creating a new client
@@ -162,4 +160,250 @@ export declare class ChromaClientExtended {
name: string;
}): Promise<void>;
}
declare class Collection {
name: string;
id: string;
metadata: CollectionMetadata | undefined;
/**
* @ignore
*/
private api;
/**
* @ignore
*/
embeddingFunction: IEmbeddingFunction | undefined;
/**
* @ignore
*/
constructor(name: string, id: string, api: DefaultApi, metadata?: CollectionMetadata, embeddingFunction?: IEmbeddingFunction);
/**
* @ignore
*/
private setName;
/**
* @ignore
*/
private setMetadata;
/**
* @ignore
*/
private validate;
/**
* Add items to the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - IDs of the items to add.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
* @param {Document | Documents} [params.documents] - Optional documents of the items to add.
* @returns {Promise<AddResponse>} - The response from the API. True if successful.
*
* @example
* ```typescript
* const response = await collection.add({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["document1", "document2"]
* });
* ```
*/
add({ ids, embeddings, metadatas, documents, }: {
ids: ID | IDs;
embeddings?: Embedding | Embeddings;
metadatas?: Metadata | Metadatas;
documents?: Document | Documents;
}): Promise<AddResponse>;
/**
* Upsert items to the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - IDs of the items to add.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
* @param {Document | Documents} [params.documents] - Optional documents of the items to add.
* @returns {Promise<boolean>} - The response from the API. True if successful.
*
* @example
* ```typescript
* const response = await collection.upsert({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["document1", "document2"],
* });
* ```
*/
upsert({ ids, embeddings, metadatas, documents, }: {
ids: ID | IDs;
embeddings?: Embedding | Embeddings;
metadatas?: Metadata | Metadatas;
documents?: Document | Documents;
}): Promise<boolean>;
/**
* Count the number of items in the collection
* @returns {Promise<number>} - The response from the API.
*
* @example
* ```typescript
* const response = await collection.count();
* ```
*/
count(): Promise<number>;
/**
* Modify the collection name or metadata
* @param {Object} params - The parameters for the query.
* @param {string} [params.name] - Optional new name for the collection.
* @param {CollectionMetadata} [params.metadata] - Optional new metadata for the collection.
* @returns {Promise<void>} - The response from the API.
*
* @example
* ```typescript
* const response = await collection.modify({
* name: "new name",
* metadata: { "key": "value" },
* });
* ```
*/
modify({ name, metadata, }?: {
name?: string;
metadata?: CollectionMetadata;
}): Promise<void>;
/**
* Get items from the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - Optional IDs of the items to get.
* @param {Where} [params.where] - Optional where clause to filter items by.
* @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get.
* @param {PositiveInteger} [params.offset] - Optional offset on the items to get.
* @param {IncludeEnum[]} [params.include] - Optional list of items to include in the response.
* @param {WhereDocument} [params.whereDocument] - Optional where clause to filter items by.
* @returns {Promise<GetResponse>} - The response from the server.
*
* @example
* ```typescript
* const response = await collection.get({
* ids: ["id1", "id2"],
* where: { "key": "value" },
* limit: 10,
* offset: 0,
* include: ["embeddings", "metadatas", "documents"],
* whereDocument: { $contains: "value" },
* });
* ```
*/
get({ ids, where, limit, offset, include, whereDocument, }?: {
ids?: ID | IDs;
where?: Where;
limit?: PositiveInteger;
offset?: PositiveInteger;
include?: IncludeEnum[];
whereDocument?: WhereDocument;
}): Promise<GetResponse>;
/**
* Update the embeddings, documents, and/or metadatas of existing items
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - The IDs of the items to update.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings to update.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadatas to update.
* @param {Document | Documents} [params.documents] - Optional documents to update.
* @returns {Promise<boolean>} - The API Response. True if successful. Else, error.
*
* @example
* ```typescript
* const response = await collection.update({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["new document 1", "new document 2"],
* });
* ```
*/
update({ ids, embeddings, metadatas, documents, }: {
ids: ID | IDs;
embeddings?: Embedding | Embeddings;
metadatas?: Metadata | Metadatas;
documents?: Document | Documents;
}): Promise<boolean>;
/**
* Performs a query on the collection using the specified parameters.
*
* @param {Object} params - The parameters for the query.
* @param {Embedding | Embeddings} [params.queryEmbeddings] - Optional query embeddings to use for the search.
* @param {PositiveInteger} [params.nResults] - Optional number of results to return (default is 10).
* @param {Where} [params.where] - Optional query condition to filter results based on metadata values.
* @param {string | string[]} [params.queryTexts] - Optional query text(s) to search for in the collection.
* @param {WhereDocument} [params.whereDocument] - Optional query condition to filter results based on document content.
* @param {IncludeEnum[]} [params.include] - Optional array of fields to include in the result, such as "metadata" and "document".
*
* @returns {Promise<QueryResponse>} A promise that resolves to the query results.
* @throws {Error} If there is an issue executing the query.
* @example
* // Query the collection using embeddings
* const results = await collection.query({
* queryEmbeddings: [[0.1, 0.2, ...], ...],
* nResults: 10,
* where: {"name": {"$eq": "John Doe"}},
* include: ["metadata", "document"]
* });
* @example
* ```js
* // Query the collection using query text
* const results = await collection.query({
* queryTexts: "some text",
* nResults: 10,
* where: {"name": {"$eq": "John Doe"}},
* include: ["metadata", "document"]
* });
* ```
*
*/
query({ queryEmbeddings, nResults, where, queryTexts, whereDocument, include, }: {
queryEmbeddings?: Embedding | Embeddings;
nResults?: PositiveInteger;
where?: Where;
queryTexts?: string | string[];
whereDocument?: WhereDocument;
include?: IncludeEnum[];
}): Promise<QueryResponse>;
/**
* Peek inside the collection
* @param {Object} params - The parameters for the query.
* @param {PositiveInteger} [params.limit] - Optional number of results to return (default is 10).
* @returns {Promise<GetResponse>} A promise that resolves to the query results.
* @throws {Error} If there is an issue executing the query.
*
* @example
* ```typescript
* const results = await collection.peek({
* limit: 10
* });
* ```
*/
peek({ limit, }?: {
limit?: PositiveInteger;
}): Promise<GetResponse>;
/**
* Deletes items from the collection.
* @param {Object} params - The parameters for deleting items from the collection.
* @param {ID | IDs} [params.ids] - Optional ID or array of IDs of items to delete.
* @param {Where} [params.where] - Optional query condition to filter items to delete based on metadata values.
* @param {WhereDocument} [params.whereDocument] - Optional query condition to filter items to delete based on document content.
* @returns {Promise<string[]>} A promise that resolves to the IDs of the deleted items.
* @throws {Error} If there is an issue deleting items from the collection.
*
* @example
* ```typescript
* const results = await collection.delete({
* ids: "some_id",
* where: {"name": {"$eq": "John Doe"}},
* whereDocument: {"$contains":"search_string"}
* });
* ```
*/
delete({ ids, where, whereDocument, }?: {
ids?: ID | IDs;
where?: Where;
whereDocument?: WhereDocument;
}): Promise<string[]>;
}
export {};
//# sourceMappingURL=index.d.ts.map
dist/index.d.ts.map
Vendored
+1 -1
View File
@@ -1 +1 @@
{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,kBAAkB,EAAE,MAAM,mDAAmD,CAAC;AACvF,OAAO,EAGL,GAAG,EACJ,MAAM,8BAA8B,CAAC;AAEtC,OAAO,EAAE,UAAU,EAAE,MAAM,+BAA+B,CAAC;AAC3D,OAAO,EAAE,kBAAkB,EAAE,cAAc,EAAE,MAAM,0BAA0B,CAAC;AAE9E,qBAAa,oBAAoB;IAC/B;;OAEG;IACH,OAAO,CAAC,GAAG,CAAa;IACxB,OAAO,CAAC,OAAO,CAAc;IAE7B;;;;;;;;;;;;OAYG;gBACS,EACV,IAAI,EACJ,YAAY,GACb,GAAE;QAAE,IAAI,CAAC,EAAE,MAAM,CAAC;QAAC,YAAY,CAAC,EAAE,WAAW,CAAA;KAAO;IASrD;;;;;;;;;;OAUG;IACU,KAAK,IAAI,OAAO,CAAC,GAAG,CAAC,gBAAgB,CAAC;IAInD;;;;;;;;OAQG;IACU,OAAO,IAAI,OAAO,CAAC,MAAM,CAAC;IAKvC;;;;;;;;OAQG;IACU,SAAS,IAAI,OAAO,CAAC,MAAM,CAAC;IAMzC;;OAEG;IACU,OAAO,IAAI,OAAO,CAAC,KAAK,CAAC;IAItC;;;;;;;;;;;;;;;;;;;;OAoBG;IACU,gBAAgB,CAAC,EAC5B,IAAI,EACJ,QAAQ,EACR,iBAAiB,GAClB,EAAE;QACD,IAAI,EAAE,MAAM,CAAC;QACb,QAAQ,CAAC,EAAE,kBAAkB,CAAC;QAC9B,iBAAiB,CAAC,EAAE,kBAAkB,CAAC;KACxC,GAAG,OAAO,CAAC,UAAU,CAAC;IAyBvB;;;;;;;;;;;;;;;;;;;;OAoBG;IACU,qBAAqB,CAAC,EACjC,IAAI,EACJ,QAAQ,EACR,iBAAiB,GAClB,EAAE;QACD,IAAI,EAAE,MAAM,CAAC;QACb,QAAQ,CAAC,EAAE,kBAAkB,CAAC;QAC9B,iBAAiB,CAAC,EAAE,kBAAkB,CAAC;KACxC,GAAG,OAAO,CAAC,UAAU,CAAC;IA0BvB;;;;;;;;;;OAUG;IACU,eAAe,IAAI,OAAO,CAAC,cAAc,EAAE,CAAC;IAKzD;;;;;;;;;;;;;;OAcG;IACU,aAAa,CAAC,EACzB,IAAI,EACJ,iBAAiB,GAClB,EAAE;QACD,IAAI,EAAE,MAAM,CAAC;QACb,iBAAiB,CAAC,EAAE,kBAAkB,CAAC;KACxC,GAAG,OAAO,CAAC,UAAU,CAAC;IAmBvB;;;;;;;;;;;;;OAaG;IACU,gBAAgB,CAAC,EAAE,IAAI,EAAE,EAAE;QAAE,IAAI,EAAE,MAAM,CAAA;KAAE,GAAG,OAAO,CAAC,IAAI,CAAC;CAMzE"}
{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EACL,WAAW,EACX,QAAQ,EACR,SAAS,EACT,SAAS,EACT,UAAU,EACV,QAAQ,EACR,SAAS,EACT,KAAK,EACL,aAAa,EACb,EAAE,EACF,GAAG,EACH,eAAe,EACf,WAAW,EACX,aAAa,EACb,WAAW,EACX,kBAAkB,EAClB,cAAc,EACf,MAAM,0BAA0B,CAAC;AAClC,OAAO,EAAE,kBAAkB,EAAE,MAAM,mDAAmD,CAAC;AACvF,OAAO,EAEL,MAAM,IAAI,UAAU,EACpB,GAAG,EACJ,MAAM,8BAA8B,CAAC;AAYtC,qBAAa,oBAAoB;IAC/B;;OAEG;IACH,OAAO,CAAC,GAAG,CAAqB;IAEhC;;;;;;;;;;;;OAYG;gBACS,EACV,IAAI,EACJ,YAAY,GACb,GAAE;QAAE,IAAI,CAAC,EAAE,MAAM,CAAC;QAAC,YAAY,CAAC,EAAE,WAAW,CAAA;KAAO;IASrD;;;;;;;;;;OAUG;IACU,KAAK,IAAI,OAAO,CAAC,GAAG,CAAC,gBAAgB,CAAC;IAInD;;;;;;;;OAQG;IACU,OAAO,IAAI,OAAO,CAAC,MAAM,CAAC;IAKvC;;;;;;;;OAQG;IACU,SAAS,IAAI,OAAO,CAAC,MAAM,CAAC;IAMzC;;OAEG;IACU,OAAO,IAAI,OAAO,CAAC,KAAK,CAAC;IAItC;;;;;;;;;;;;;;;;;;;;OAoBG;IACU,gBAAgB,CAAC,EAC5B,IAAI,EACJ,QAAQ,EACR,iBAAiB,GAClB,EAAE;QACD,IAAI,EAAE,MAAM,CAAC;QACb,QAAQ,CAAC,EAAE,kBAAkB,CAAC;QAC9B,iBAAiB,CAAC,EAAE,kBAAkB,CAAC;KACxC,GAAG,OAAO,CAAC,UAAU,CAAC;IAyBvB;;;;;;;;;;;;;;;;;;;;OAoBG;IACU,qBAAqB,CAAC,EACjC,IAAI,EACJ,QAAQ,EACR,iBAAiB,GAClB,EAAE;QACD,IAAI,EAAE,MAAM,CAAC;QACb,QAAQ,CAAC,EAAE,kBAAkB,CAAC;QAC9B,iBAAiB,CAAC,EAAE,kBAAkB,CAAC;KACxC,GAAG,OAAO,CAAC,UAAU,CAAC;IA0BvB;;;;;;;;;;OAUG;IACU,eAAe,IAAI,OAAO,CAAC,cAAc,EAAE,CAAC;IAKzD;;;;;;;;;;;;;;OAcG;IACU,aAAa,CAAC,EACzB,IAAI,EACJ,iBAAiB,GAClB,EAAE;QACD,IAAI,EAAE,MAAM,CAAC;QACb,iBAAiB,CAAC,EAAE,kBAAkB,CAAC;KACxC,GAAG,OAAO,CAAC,UAAU,CAAC;IAmBvB;;;;;;;;;;;;;OAaG;IACU,gBAAgB,CAAC,EAAE,IAAI,EAAE,EAAE;QAAE,IAAI,EAAE,MAAM,CAAA;KAAE,GAAG,OAAO,CAAC,IAAI,CAAC;CAMzE;AAED,cAAM,UAAU;IACP,IAAI,EAAE,MAAM,CAAC;IACb,EAAE,EAAE,MAAM,CAAC;IACX,QAAQ,EAAE,kBAAkB,GAAG,SAAS,CAAC;IAChD;;OAEG;IACH,OAAO,CAAC,GAAG,CAAqB;IAChC;;OAEG;IACI,iBAAiB,EAAE,kBAAkB,GAAG,SAAS,CAAC;IAEzD;;OAEG;gBAED,IAAI,EAAE,MAAM,EACZ,EAAE,EAAE,MAAM,EACV,GAAG,EAAE,UAAU,EACf,QAAQ,CAAC,EAAE,kBAAkB,EAC7B,iBAAiB,CAAC,EAAE,kBAAkB;IAUxC;;OAEG;IACH,OAAO,CAAC,OAAO;IAGf;;OAEG;IACH,OAAO,CAAC,WAAW;IAInB;;OAEG;YACW,QAAQ;IAgFtB;;;;;;;;;;;;;;;;;;OAkBG;IACU,GAAG,CAAC,EACf,GAAG,EACH,UAAU,EACV,SAAS,EACT,SAAS,GACV,EAAE;QACD,GAAG,EAAE,EAAE,GAAG,GAAG,CAAC;QACd,UAAU,CAAC,EAAE,SAAS,GAAG,UAAU,CAAC;QACpC,SAAS,CAAC,EAAE,QAAQ,GAAG,SAAS,CAAC;QACjC,SAAS,CAAC,EAAE,QAAQ,GAAG,SAAS,CAAC;KAClC,GAAG,OAAO,CAAC,WAAW,CAAC;IAuBxB;;;;;;;;;;;;;;;;;;OAkBG;IACU,MAAM,CAAC,EAClB,GAAG,EACH,UAAU,EACV,SAAS,EACT,SAAS,GACV,EAAE;QACD,GAAG,EAAE,EAAE,GAAG,GAAG,CAAC;QACd,UAAU,CAAC,EAAE,SAAS,GAAG,UAAU,CAAC;QACpC,SAAS,CAAC,EAAE,QAAQ,GAAG,SAAS,CAAC;QACjC,SAAS,CAAC,EAAE,QAAQ,GAAG,SAAS,CAAC;KAClC,GAAG,OAAO,CAAC,OAAO,CAAC;IAuBpB;;;;;;;;OAQG;IACU,KAAK,IAAI,OAAO,CAAC,MAAM,CAAC;IAKrC;;;;;;;;;;;;;;OAcG;IACU,MAAM,CAAC,EAClB,IAAI,EACJ,QAAQ,GACT,GAAE;QACD,IAAI,CAAC,EAAE,MAAM,CAAC;QACd,QAAQ,CAAC,EAAE,kBAAkB,CAAC;KAC1B,GAAG,OAAO,CAAC,IAAI,CAAC;IAmBtB;;;;;;;;;;;;;;;;;;;;;;OAsBG;IACU,GAAG,CAAC,EACf,GAAG,EACH,KAAK,EACL,KAAK,EACL,MAAM,EACN,OAAO,EACP,aAAa,GACd,GAAE;QACD,GAAG,CAAC,EAAE,EAAE,GAAG,GAAG,CAAC;QACf,KAAK,CAAC,EAAE,KAAK,CAAC;QACd,KAAK,CAAC,EAAE,eAAe,CAAC;QACxB,MAAM,CAAC,EAAE,eAAe,CAAC;QACzB,OAAO,CAAC,EAAE,WAAW,EAAE,CAAC;QACxB,aAAa,CAAC,EAAE,aAAa,CAAC;KAC1B,GAAG,OAAO,CAAC,WAAW,CAAC;IAqB7B;;;;;;;;;;;;;;;;;;OAkBG;IACU,MAAM,CAAC,EAClB,GAAG,EACH,UAAU,EACV,SAAS,EACT,SAAS,GACV,EAAE;QACD,GAAG,EAAE,EAAE,GAAG,GAAG,CAAC;QACd,UAAU,CAAC,EAAE,SAAS,GAAG,UAAU,CAAC;QACpC,SAAS,CAAC,EAAE,QAAQ,GAAG,SAAS,CAAC;QACjC,SAAS,CAAC,EAAE,QAAQ,GAAG,SAAS,CAAC;KAClC,GAAG,OAAO,CAAC,OAAO,CAAC;IAyCpB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;OAgCG;IACU,KAAK,CAAC,EACjB,eAAe,EACf,QAAQ,EACR,KAAK,EACL,UAAU,EACV,aAAa,EACb,OAAO,GACR,EAAE;QACD,eAAe,CAAC,EAAE,SAAS,GAAG,UAAU,CAAC;QACzC,QAAQ,CAAC,EAAE,eAAe,CAAC;QAC3B,KAAK,CAAC,EAAE,KAAK,CAAC;QACd,UAAU,CAAC,EAAE,MAAM,GAAG,MAAM,EAAE,CAAC;QAC/B,aAAa,CAAC,EAAE,aAAa,CAAC;QAC9B,OAAO,CAAC,EAAE,WAAW,EAAE,CAAC;KACzB,GAAG,OAAO,CAAC,aAAa,CAAC;IAuC1B;;;;;;;;;;;;;OAaG;IACU,IAAI,CAAC,EAChB,KAAK,GACN,GAAE;QAAE,KAAK,CAAC,EAAE,eAAe,CAAA;KAAO,GAAG,OAAO,CAAC,WAAW,CAAC;IAY1D;;;;;;;;;;;;;;;;;OAiBG;IACU,MAAM,CAAC,EAClB,GAAG,EACH,KAAK,EACL,aAAa,GACd,GAAE;QACD,GAAG,CAAC,EAAE,EAAE,GAAG,GAAG,CAAC;QACf,KAAK,CAAC,EAAE,KAAK,CAAC;QACd,aAAa,CAAC,EAAE,aAAa,CAAC;KAC1B,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC;CAY3B"}
dist/index.js
Vendored
+403 -13
View File
@@ -3,7 +3,6 @@ Object.defineProperty(exports, "__esModule", { value: true });
exports.ChromaClientExtended = void 0;
const generated_1 = require("chromadb/dist/main/generated");
const utils_1 = require("chromadb/dist/main/utils");
const Collection_1 = require("chromadb/dist/main/Collection");
class ChromaClientExtended {
/**
* Creates a new ChromaClient instance.
@@ -25,7 +24,7 @@ class ChromaClientExtended {
basePath: path,
});
this.api = new generated_1.ApiApi(apiConfig);
this.options = fetchOptions !== null && fetchOptions !== void 0 ? fetchOptions : {};
this.api.options = fetchOptions !== null && fetchOptions !== void 0 ? fetchOptions : {};
}
/**
* Resets the state of the object by making an API call to the reset endpoint.
@@ -39,7 +38,7 @@ class ChromaClientExtended {
* ```
*/
async reset() {
return await this.api.reset(this.options);
return await this.api.reset(this.api.options);
}
/**
* Returns the version of the Chroma API.
@@ -51,7 +50,7 @@ class ChromaClientExtended {
* ```
*/
async version() {
const response = await this.api.version(this.options);
const response = await this.api.version(this.api.options);
return await (0, utils_1.handleSuccess)(response);
}
/**
@@ -64,7 +63,7 @@ class ChromaClientExtended {
* ```
*/
async heartbeat() {
const response = await this.api.heartbeat(this.options);
const response = await this.api.heartbeat(this.api.options);
let ret = await (0, utils_1.handleSuccess)(response);
return ret["nanosecond heartbeat"];
}
@@ -100,13 +99,13 @@ class ChromaClientExtended {
.createCollection({
name,
metadata,
}, this.options)
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
if (newCollection.error) {
throw new Error(newCollection.error);
}
return new Collection_1.Collection(name, newCollection.id, this.api, metadata, embeddingFunction);
return new Collection(name, newCollection.id, this.api, metadata, embeddingFunction);
}
/**
* Gets or creates a collection with the specified properties.
@@ -135,13 +134,13 @@ class ChromaClientExtended {
name,
metadata,
get_or_create: true,
}, this.options)
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
if (newCollection.error) {
throw new Error(newCollection.error);
}
return new Collection_1.Collection(name, newCollection.id, this.api, newCollection.metadata, embeddingFunction);
return new Collection(name, newCollection.id, this.api, newCollection.metadata, embeddingFunction);
}
/**
* Lists all collections.
@@ -155,7 +154,7 @@ class ChromaClientExtended {
* ```
*/
async listCollections() {
const response = await this.api.listCollections(this.options);
const response = await this.api.listCollections(this.api.options);
return (0, utils_1.handleSuccess)(response);
}
/**
@@ -175,13 +174,13 @@ class ChromaClientExtended {
*/
async getCollection({ name, embeddingFunction, }) {
const response = await this.api
.getCollection(name, this.options)
.getCollection(name, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
if (response.error) {
throw new Error(response.error);
}
return new Collection_1.Collection(response.name, response.id, this.api, response.metadata, embeddingFunction);
return new Collection(response.name, response.id, this.api, response.metadata, embeddingFunction);
}
/**
* Deletes a collection with the specified name.
@@ -199,10 +198,401 @@ class ChromaClientExtended {
*/
async deleteCollection({ name }) {
return await this.api
.deleteCollection(name, this.options)
.deleteCollection(name, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
}
}
exports.ChromaClientExtended = ChromaClientExtended;
class Collection {
/**
* @ignore
*/
constructor(name, id, api, metadata, embeddingFunction) {
this.name = name;
this.id = id;
this.metadata = metadata;
this.api = api;
if (embeddingFunction !== undefined)
this.embeddingFunction = embeddingFunction;
}
/**
* @ignore
*/
setName(name) {
this.name = name;
}
/**
* @ignore
*/
setMetadata(metadata) {
this.metadata = metadata;
}
/**
* @ignore
*/
async validate(require_embeddings_or_documents, // set to false in the case of Update
ids, embeddings, metadatas, documents) {
if (require_embeddings_or_documents) {
if (embeddings === undefined && documents === undefined) {
throw new Error("embeddings and documents cannot both be undefined");
}
}
if (embeddings === undefined && documents !== undefined) {
const documentsArray = (0, utils_1.toArray)(documents);
if (this.embeddingFunction !== undefined) {
embeddings = await this.embeddingFunction.generate(documentsArray);
}
else {
throw new Error("embeddingFunction is undefined. Please configure an embedding function");
}
}
if (embeddings === undefined)
throw new Error("embeddings is undefined but shouldnt be");
const idsArray = (0, utils_1.toArray)(ids);
const embeddingsArray = (0, utils_1.toArrayOfArrays)(embeddings);
let metadatasArray;
if (metadatas === undefined) {
metadatasArray = undefined;
}
else {
metadatasArray = (0, utils_1.toArray)(metadatas);
}
let documentsArray;
if (documents === undefined) {
documentsArray = undefined;
}
else {
documentsArray = (0, utils_1.toArray)(documents);
}
// validate all ids are strings
for (let i = 0; i < idsArray.length; i += 1) {
if (typeof idsArray[i] !== "string") {
throw new Error(`Expected ids to be strings, found ${typeof idsArray[i]} at index ${i}`);
}
}
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");
}
const uniqueIds = new Set(idsArray);
if (uniqueIds.size !== idsArray.length) {
const duplicateIds = idsArray.filter((item, index) => idsArray.indexOf(item) !== index);
throw new Error(`Expected IDs to be unique, found duplicates for: ${duplicateIds}`);
}
return [idsArray, embeddingsArray, metadatasArray, documentsArray];
}
/**
* Add items to the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - IDs of the items to add.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
* @param {Document | Documents} [params.documents] - Optional documents of the items to add.
* @returns {Promise<AddResponse>} - The response from the API. True if successful.
*
* @example
* ```typescript
* const response = await collection.add({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["document1", "document2"]
* });
* ```
*/
async add({ ids, embeddings, metadatas, documents, }) {
const [idsArray, embeddingsArray, metadatasArray, documentsArray] = await this.validate(true, ids, embeddings, metadatas, documents);
const response = await this.api
.add(this.id, {
// @ts-ignore
ids: idsArray,
embeddings: embeddingsArray,
// @ts-ignore
documents: documentsArray,
metadatas: metadatasArray,
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
return response;
}
/**
* Upsert items to the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - IDs of the items to add.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
* @param {Document | Documents} [params.documents] - Optional documents of the items to add.
* @returns {Promise<boolean>} - The response from the API. True if successful.
*
* @example
* ```typescript
* const response = await collection.upsert({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["document1", "document2"],
* });
* ```
*/
async upsert({ ids, embeddings, metadatas, documents, }) {
const [idsArray, embeddingsArray, metadatasArray, documentsArray] = await this.validate(true, ids, embeddings, metadatas, documents);
const response = await this.api
.upsert(this.id, {
//@ts-ignore
ids: idsArray,
embeddings: embeddingsArray,
//@ts-ignore
documents: documentsArray,
metadatas: metadatasArray,
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
return response;
}
/**
* Count the number of items in the collection
* @returns {Promise<number>} - The response from the API.
*
* @example
* ```typescript
* const response = await collection.count();
* ```
*/
async count() {
const response = await this.api.count(this.id, this.api.options);
return (0, utils_1.handleSuccess)(response);
}
/**
* Modify the collection name or metadata
* @param {Object} params - The parameters for the query.
* @param {string} [params.name] - Optional new name for the collection.
* @param {CollectionMetadata} [params.metadata] - Optional new metadata for the collection.
* @returns {Promise<void>} - The response from the API.
*
* @example
* ```typescript
* const response = await collection.modify({
* name: "new name",
* metadata: { "key": "value" },
* });
* ```
*/
async modify({ name, metadata, } = {}) {
const response = await this.api
.updateCollection(this.id, {
new_name: name,
new_metadata: metadata,
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
this.setName(name || this.name);
this.setMetadata(metadata || this.metadata);
return response;
}
/**
* Get items from the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - Optional IDs of the items to get.
* @param {Where} [params.where] - Optional where clause to filter items by.
* @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get.
* @param {PositiveInteger} [params.offset] - Optional offset on the items to get.
* @param {IncludeEnum[]} [params.include] - Optional list of items to include in the response.
* @param {WhereDocument} [params.whereDocument] - Optional where clause to filter items by.
* @returns {Promise<GetResponse>} - The response from the server.
*
* @example
* ```typescript
* const response = await collection.get({
* ids: ["id1", "id2"],
* where: { "key": "value" },
* limit: 10,
* offset: 0,
* include: ["embeddings", "metadatas", "documents"],
* whereDocument: { $contains: "value" },
* });
* ```
*/
async get({ ids, where, limit, offset, include, whereDocument, } = {}) {
let idsArray = undefined;
if (ids !== undefined)
idsArray = (0, utils_1.toArray)(ids);
return await this.api
.aGet(this.id, {
ids: idsArray,
where,
limit,
offset,
include,
where_document: whereDocument,
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
}
/**
* Update the embeddings, documents, and/or metadatas of existing items
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - The IDs of the items to update.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings to update.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadatas to update.
* @param {Document | Documents} [params.documents] - Optional documents to update.
* @returns {Promise<boolean>} - The API Response. True if successful. Else, error.
*
* @example
* ```typescript
* const response = await collection.update({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["new document 1", "new document 2"],
* });
* ```
*/
async update({ ids, embeddings, metadatas, documents, }) {
if (embeddings === undefined &&
documents === undefined &&
metadatas === undefined) {
throw new Error("embeddings, documents, and metadatas cannot all be undefined");
}
else if (embeddings === undefined && documents !== undefined) {
const documentsArray = (0, utils_1.toArray)(documents);
if (this.embeddingFunction !== undefined) {
embeddings = await this.embeddingFunction.generate(documentsArray);
}
else {
throw new Error("embeddingFunction is undefined. Please configure an embedding function");
}
}
// backend expects None if metadatas is undefined
if (metadatas !== undefined)
metadatas = (0, utils_1.toArray)(metadatas);
if (documents !== undefined)
documents = (0, utils_1.toArray)(documents);
var resp = await this.api
.update(this.id, {
ids: (0, utils_1.toArray)(ids),
embeddings: embeddings ? (0, utils_1.toArrayOfArrays)(embeddings) : undefined,
documents: documents,
metadatas: metadatas,
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
return resp;
}
/**
* Performs a query on the collection using the specified parameters.
*
* @param {Object} params - The parameters for the query.
* @param {Embedding | Embeddings} [params.queryEmbeddings] - Optional query embeddings to use for the search.
* @param {PositiveInteger} [params.nResults] - Optional number of results to return (default is 10).
* @param {Where} [params.where] - Optional query condition to filter results based on metadata values.
* @param {string | string[]} [params.queryTexts] - Optional query text(s) to search for in the collection.
* @param {WhereDocument} [params.whereDocument] - Optional query condition to filter results based on document content.
* @param {IncludeEnum[]} [params.include] - Optional array of fields to include in the result, such as "metadata" and "document".
*
* @returns {Promise<QueryResponse>} A promise that resolves to the query results.
* @throws {Error} If there is an issue executing the query.
* @example
* // Query the collection using embeddings
* const results = await collection.query({
* queryEmbeddings: [[0.1, 0.2, ...], ...],
* nResults: 10,
* where: {"name": {"$eq": "John Doe"}},
* include: ["metadata", "document"]
* });
* @example
* ```js
* // Query the collection using query text
* const results = await collection.query({
* queryTexts: "some text",
* nResults: 10,
* where: {"name": {"$eq": "John Doe"}},
* include: ["metadata", "document"]
* });
* ```
*
*/
async query({ queryEmbeddings, nResults, where, queryTexts, whereDocument, include, }) {
if (nResults === undefined)
nResults = 10;
if (queryEmbeddings === undefined && queryTexts === undefined) {
throw new Error("queryEmbeddings and queryTexts cannot both be undefined");
}
else if (queryEmbeddings === undefined && queryTexts !== undefined) {
const queryTextsArray = (0, utils_1.toArray)(queryTexts);
if (this.embeddingFunction !== undefined) {
queryEmbeddings = await this.embeddingFunction.generate(queryTextsArray);
}
else {
throw new Error("embeddingFunction is undefined. Please configure an embedding function");
}
}
if (queryEmbeddings === undefined)
throw new Error("embeddings is undefined but shouldnt be");
const query_embeddingsArray = (0, utils_1.toArrayOfArrays)(queryEmbeddings);
return await this.api
.getNearestNeighbors(this.id, {
query_embeddings: query_embeddingsArray,
where,
n_results: nResults,
where_document: whereDocument,
include: include,
}, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
}
/**
* Peek inside the collection
* @param {Object} params - The parameters for the query.
* @param {PositiveInteger} [params.limit] - Optional number of results to return (default is 10).
* @returns {Promise<GetResponse>} A promise that resolves to the query results.
* @throws {Error} If there is an issue executing the query.
*
* @example
* ```typescript
* const results = await collection.peek({
* limit: 10
* });
* ```
*/
async peek({ limit, } = {}) {
if (limit === undefined)
limit = 10;
const response = await this.api.aGet(this.id, {
limit: limit,
}, this.api.options);
return (0, utils_1.handleSuccess)(response);
}
/**
* Deletes items from the collection.
* @param {Object} params - The parameters for deleting items from the collection.
* @param {ID | IDs} [params.ids] - Optional ID or array of IDs of items to delete.
* @param {Where} [params.where] - Optional query condition to filter items to delete based on metadata values.
* @param {WhereDocument} [params.whereDocument] - Optional query condition to filter items to delete based on document content.
* @returns {Promise<string[]>} A promise that resolves to the IDs of the deleted items.
* @throws {Error} If there is an issue deleting items from the collection.
*
* @example
* ```typescript
* const results = await collection.delete({
* ids: "some_id",
* where: {"name": {"$eq": "John Doe"}},
* whereDocument: {"$contains":"search_string"}
* });
* ```
*/
async delete({ ids, where, whereDocument, } = {}) {
let idsArray = undefined;
if (ids !== undefined)
idsArray = (0, utils_1.toArray)(ids);
return await this.api
.aDelete(this.id, { ids: idsArray, where: where, where_document: whereDocument }, this.api.options)
.then(utils_1.handleSuccess)
.catch(utils_1.handleError);
}
}
//# sourceMappingURL=index.js.map
dist/index.js.map
Vendored
+1 -1
View File
File diff suppressed because one or more lines are too long
package.json
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@mintplex-labs/chromadb-extended",
"version": "1.0.0",
"version": "1.0.1",
"description": "A chromadb NodeJS extended client that enables you to pass custom authorization headers and options to your chromaDB instance",
"main": "dist/index.js",
"files": [
src/index.ts
+606 -14
View File
@@ -1,19 +1,44 @@
import {
IncludeEnum,
Metadata,
Metadatas,
Embedding,
Embeddings,
Document,
Documents,
Where,
WhereDocument,
ID,
IDs,
PositiveInteger,
GetResponse,
QueryResponse,
AddResponse,
CollectionMetadata,
CollectionType,
} from "chromadb/dist/main/types";
import { IEmbeddingFunction } from "chromadb/dist/main//embeddings/IEmbeddingFunction";
import {
Configuration,
ApiApi as DefaultApi,
Api,
} from "chromadb/dist/main/generated";
import { handleSuccess, handleError } from "chromadb/dist/main/utils";
import { Collection } from "chromadb/dist/main/Collection";
import { CollectionMetadata, CollectionType } from "chromadb/dist/main/types";
import {
handleSuccess,
handleError,
toArray,
toArrayOfArrays,
} from "chromadb/dist/main/utils";
interface ExtendedDefaultAPI extends DefaultApi {
options?: RequestInit;
}
export class ChromaClientExtended {
/**
* @ignore
*/
private api: DefaultApi;
private options: RequestInit;
private api: ExtendedDefaultAPI;
/**
* Creates a new ChromaClient instance.
@@ -37,7 +62,7 @@ export class ChromaClientExtended {
basePath: path,
});
this.api = new DefaultApi(apiConfig);
this.options = fetchOptions ?? {};
this.api.options = fetchOptions ?? {};
}
/**
@@ -52,7 +77,7 @@ export class ChromaClientExtended {
* ```
*/
public async reset(): Promise<Api.Reset200Response> {
return await this.api.reset(this.options);
return await this.api.reset(this.api.options);
}
/**
@@ -65,7 +90,7 @@ export class ChromaClientExtended {
* ```
*/
public async version(): Promise<string> {
const response = await this.api.version(this.options);
const response = await this.api.version(this.api.options);
return await handleSuccess(response);
}
@@ -79,7 +104,7 @@ export class ChromaClientExtended {
* ```
*/
public async heartbeat(): Promise<number> {
const response = await this.api.heartbeat(this.options);
const response = await this.api.heartbeat(this.api.options);
let ret = await handleSuccess(response);
return ret["nanosecond heartbeat"];
}
@@ -127,7 +152,7 @@ export class ChromaClientExtended {
name,
metadata,
},
this.options
this.api.options
)
.then(handleSuccess)
.catch(handleError);
@@ -182,7 +207,7 @@ export class ChromaClientExtended {
metadata,
get_or_create: true,
},
this.options
this.api.options
)
.then(handleSuccess)
.catch(handleError);
@@ -212,7 +237,7 @@ export class ChromaClientExtended {
* ```
*/
public async listCollections(): Promise<CollectionType[]> {
const response = await this.api.listCollections(this.options);
const response = await this.api.listCollections(this.api.options);
return handleSuccess(response);
}
@@ -239,7 +264,7 @@ export class ChromaClientExtended {
embeddingFunction?: IEmbeddingFunction;
}): Promise<Collection> {
const response = await this.api
.getCollection(name, this.options)
.getCollection(name, this.api.options)
.then(handleSuccess)
.catch(handleError);
@@ -272,7 +297,574 @@ export class ChromaClientExtended {
*/
public async deleteCollection({ name }: { name: string }): Promise<void> {
return await this.api
.deleteCollection(name, this.options)
.deleteCollection(name, this.api.options)
.then(handleSuccess)
.catch(handleError);
}
}
class Collection {
public name: string;
public id: string;
public metadata: CollectionMetadata | undefined;
/**
* @ignore
*/
private api: ExtendedDefaultAPI;
/**
* @ignore
*/
public embeddingFunction: IEmbeddingFunction | undefined;
/**
* @ignore
*/
constructor(
name: string,
id: string,
api: DefaultApi,
metadata?: CollectionMetadata,
embeddingFunction?: IEmbeddingFunction
) {
this.name = name;
this.id = id;
this.metadata = metadata;
this.api = api;
if (embeddingFunction !== undefined)
this.embeddingFunction = embeddingFunction;
}
/**
* @ignore
*/
private setName(name: string): void {
this.name = name;
}
/**
* @ignore
*/
private setMetadata(metadata: CollectionMetadata | undefined): void {
this.metadata = metadata;
}
/**
* @ignore
*/
private async validate(
require_embeddings_or_documents: boolean, // set to false in the case of Update
ids: string | string[],
embeddings: number[] | number[][] | undefined,
metadatas?: object | object[],
documents?: string | string[]
) {
if (require_embeddings_or_documents) {
if (embeddings === undefined && documents === undefined) {
throw new Error("embeddings and documents cannot both be undefined");
}
}
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"
);
}
}
if (embeddings === undefined)
throw new Error("embeddings is undefined but shouldnt be");
const idsArray = toArray(ids);
const embeddingsArray: number[][] = toArrayOfArrays(embeddings);
let metadatasArray: object[] | undefined;
if (metadatas === undefined) {
metadatasArray = undefined;
} else {
metadatasArray = toArray(metadatas);
}
let documentsArray: (string | undefined)[] | undefined;
if (documents === undefined) {
documentsArray = undefined;
} else {
documentsArray = toArray(documents);
}
// validate all ids are strings
for (let i = 0; i < idsArray.length; i += 1) {
if (typeof idsArray[i] !== "string") {
throw new Error(
`Expected ids to be strings, found ${typeof idsArray[
i
]} at index ${i}`
);
}
}
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"
);
}
const uniqueIds = new Set(idsArray);
if (uniqueIds.size !== idsArray.length) {
const duplicateIds = idsArray.filter(
(item, index) => idsArray.indexOf(item) !== index
);
throw new Error(
`Expected IDs to be unique, found duplicates for: ${duplicateIds}`
);
}
return [idsArray, embeddingsArray, metadatasArray, documentsArray];
}
/**
* Add items to the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - IDs of the items to add.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
* @param {Document | Documents} [params.documents] - Optional documents of the items to add.
* @returns {Promise<AddResponse>} - The response from the API. True if successful.
*
* @example
* ```typescript
* const response = await collection.add({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["document1", "document2"]
* });
* ```
*/
public async add({
ids,
embeddings,
metadatas,
documents,
}: {
ids: ID | IDs;
embeddings?: Embedding | Embeddings;
metadatas?: Metadata | Metadatas;
documents?: Document | Documents;
}): Promise<AddResponse> {
const [idsArray, embeddingsArray, metadatasArray, documentsArray] =
await this.validate(true, ids, embeddings, metadatas, documents);
const response = await this.api
.add(
this.id,
{
// @ts-ignore
ids: idsArray,
embeddings: embeddingsArray as number[][], // We know this is defined because of the validate function
// @ts-ignore
documents: documentsArray,
metadatas: metadatasArray,
},
this.api.options
)
.then(handleSuccess)
.catch(handleError);
return response;
}
/**
* Upsert items to the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - IDs of the items to add.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings of the items to add.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadata of the items to add.
* @param {Document | Documents} [params.documents] - Optional documents of the items to add.
* @returns {Promise<boolean>} - The response from the API. True if successful.
*
* @example
* ```typescript
* const response = await collection.upsert({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["document1", "document2"],
* });
* ```
*/
public async upsert({
ids,
embeddings,
metadatas,
documents,
}: {
ids: ID | IDs;
embeddings?: Embedding | Embeddings;
metadatas?: Metadata | Metadatas;
documents?: Document | Documents;
}): Promise<boolean> {
const [idsArray, embeddingsArray, metadatasArray, documentsArray] =
await this.validate(true, ids, embeddings, metadatas, documents);
const response = await this.api
.upsert(
this.id,
{
//@ts-ignore
ids: idsArray,
embeddings: embeddingsArray as number[][], // We know this is defined because of the validate function
//@ts-ignore
documents: documentsArray,
metadatas: metadatasArray,
},
this.api.options
)
.then(handleSuccess)
.catch(handleError);
return response;
}
/**
* Count the number of items in the collection
* @returns {Promise<number>} - The response from the API.
*
* @example
* ```typescript
* const response = await collection.count();
* ```
*/
public async count(): Promise<number> {
const response = await this.api.count(this.id, this.api.options);
return handleSuccess(response);
}
/**
* Modify the collection name or metadata
* @param {Object} params - The parameters for the query.
* @param {string} [params.name] - Optional new name for the collection.
* @param {CollectionMetadata} [params.metadata] - Optional new metadata for the collection.
* @returns {Promise<void>} - The response from the API.
*
* @example
* ```typescript
* const response = await collection.modify({
* name: "new name",
* metadata: { "key": "value" },
* });
* ```
*/
public async modify({
name,
metadata,
}: {
name?: string;
metadata?: CollectionMetadata;
} = {}): Promise<void> {
const response = await this.api
.updateCollection(
this.id,
{
new_name: name,
new_metadata: metadata,
},
this.api.options
)
.then(handleSuccess)
.catch(handleError);
this.setName(name || this.name);
this.setMetadata(metadata || this.metadata);
return response;
}
/**
* Get items from the collection
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - Optional IDs of the items to get.
* @param {Where} [params.where] - Optional where clause to filter items by.
* @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get.
* @param {PositiveInteger} [params.offset] - Optional offset on the items to get.
* @param {IncludeEnum[]} [params.include] - Optional list of items to include in the response.
* @param {WhereDocument} [params.whereDocument] - Optional where clause to filter items by.
* @returns {Promise<GetResponse>} - The response from the server.
*
* @example
* ```typescript
* const response = await collection.get({
* ids: ["id1", "id2"],
* where: { "key": "value" },
* limit: 10,
* offset: 0,
* include: ["embeddings", "metadatas", "documents"],
* whereDocument: { $contains: "value" },
* });
* ```
*/
public async get({
ids,
where,
limit,
offset,
include,
whereDocument,
}: {
ids?: ID | IDs;
where?: Where;
limit?: PositiveInteger;
offset?: PositiveInteger;
include?: IncludeEnum[];
whereDocument?: WhereDocument;
} = {}): Promise<GetResponse> {
let idsArray = undefined;
if (ids !== undefined) idsArray = toArray(ids);
return await this.api
.aGet(
this.id,
{
ids: idsArray,
where,
limit,
offset,
include,
where_document: whereDocument,
},
this.api.options
)
.then(handleSuccess)
.catch(handleError);
}
/**
* Update the embeddings, documents, and/or metadatas of existing items
* @param {Object} params - The parameters for the query.
* @param {ID | IDs} [params.ids] - The IDs of the items to update.
* @param {Embedding | Embeddings} [params.embeddings] - Optional embeddings to update.
* @param {Metadata | Metadatas} [params.metadatas] - Optional metadatas to update.
* @param {Document | Documents} [params.documents] - Optional documents to update.
* @returns {Promise<boolean>} - The API Response. True if successful. Else, error.
*
* @example
* ```typescript
* const response = await collection.update({
* ids: ["id1", "id2"],
* embeddings: [[1, 2, 3], [4, 5, 6]],
* metadatas: [{ "key": "value" }, { "key": "value" }],
* documents: ["new document 1", "new document 2"],
* });
* ```
*/
public async update({
ids,
embeddings,
metadatas,
documents,
}: {
ids: ID | IDs;
embeddings?: Embedding | Embeddings;
metadatas?: Metadata | Metadatas;
documents?: Document | Documents;
}): Promise<boolean> {
if (
embeddings === undefined &&
documents === undefined &&
metadatas === undefined
) {
throw new Error(
"embeddings, documents, and metadatas cannot all 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"
);
}
}
// backend expects None if metadatas is undefined
if (metadatas !== undefined) metadatas = toArray(metadatas);
if (documents !== undefined) documents = toArray(documents);
var resp = await this.api
.update(
this.id,
{
ids: toArray(ids),
embeddings: embeddings ? toArrayOfArrays(embeddings) : undefined,
documents: documents,
metadatas: metadatas,
},
this.api.options
)
.then(handleSuccess)
.catch(handleError);
return resp;
}
/**
* Performs a query on the collection using the specified parameters.
*
* @param {Object} params - The parameters for the query.
* @param {Embedding | Embeddings} [params.queryEmbeddings] - Optional query embeddings to use for the search.
* @param {PositiveInteger} [params.nResults] - Optional number of results to return (default is 10).
* @param {Where} [params.where] - Optional query condition to filter results based on metadata values.
* @param {string | string[]} [params.queryTexts] - Optional query text(s) to search for in the collection.
* @param {WhereDocument} [params.whereDocument] - Optional query condition to filter results based on document content.
* @param {IncludeEnum[]} [params.include] - Optional array of fields to include in the result, such as "metadata" and "document".
*
* @returns {Promise<QueryResponse>} A promise that resolves to the query results.
* @throws {Error} If there is an issue executing the query.
* @example
* // Query the collection using embeddings
* const results = await collection.query({
* queryEmbeddings: [[0.1, 0.2, ...], ...],
* nResults: 10,
* where: {"name": {"$eq": "John Doe"}},
* include: ["metadata", "document"]
* });
* @example
* ```js
* // Query the collection using query text
* const results = await collection.query({
* queryTexts: "some text",
* nResults: 10,
* where: {"name": {"$eq": "John Doe"}},
* include: ["metadata", "document"]
* });
* ```
*
*/
public async query({
queryEmbeddings,
nResults,
where,
queryTexts,
whereDocument,
include,
}: {
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(
"queryEmbeddings and queryTexts cannot both be undefined"
);
} else if (queryEmbeddings === undefined && queryTexts !== undefined) {
const queryTextsArray = toArray(queryTexts);
if (this.embeddingFunction !== undefined) {
queryEmbeddings = await this.embeddingFunction.generate(
queryTextsArray
);
} else {
throw new Error(
"embeddingFunction is undefined. Please configure an embedding function"
);
}
}
if (queryEmbeddings === undefined)
throw new Error("embeddings is undefined but shouldnt be");
const query_embeddingsArray = toArrayOfArrays(queryEmbeddings);
return await this.api
.getNearestNeighbors(
this.id,
{
query_embeddings: query_embeddingsArray,
where,
n_results: nResults,
where_document: whereDocument,
include: include,
},
this.api.options
)
.then(handleSuccess)
.catch(handleError);
}
/**
* Peek inside the collection
* @param {Object} params - The parameters for the query.
* @param {PositiveInteger} [params.limit] - Optional number of results to return (default is 10).
* @returns {Promise<GetResponse>} A promise that resolves to the query results.
* @throws {Error} If there is an issue executing the query.
*
* @example
* ```typescript
* const results = await collection.peek({
* limit: 10
* });
* ```
*/
public async peek({
limit,
}: { limit?: PositiveInteger } = {}): Promise<GetResponse> {
if (limit === undefined) limit = 10;
const response = await this.api.aGet(
this.id,
{
limit: limit,
},
this.api.options
);
return handleSuccess(response);
}
/**
* Deletes items from the collection.
* @param {Object} params - The parameters for deleting items from the collection.
* @param {ID | IDs} [params.ids] - Optional ID or array of IDs of items to delete.
* @param {Where} [params.where] - Optional query condition to filter items to delete based on metadata values.
* @param {WhereDocument} [params.whereDocument] - Optional query condition to filter items to delete based on document content.
* @returns {Promise<string[]>} A promise that resolves to the IDs of the deleted items.
* @throws {Error} If there is an issue deleting items from the collection.
*
* @example
* ```typescript
* const results = await collection.delete({
* ids: "some_id",
* where: {"name": {"$eq": "John Doe"}},
* whereDocument: {"$contains":"search_string"}
* });
* ```
*/
public async delete({
ids,
where,
whereDocument,
}: {
ids?: ID | IDs;
where?: Where;
whereDocument?: WhereDocument;
} = {}): Promise<string[]> {
let idsArray = undefined;
if (ids !== undefined) idsArray = toArray(ids);
return await this.api
.aDelete(
this.id,
{ ids: idsArray, where: where, where_document: whereDocument },
this.api.options
)
.then(handleSuccess)
.catch(handleError);
}
test/test.js
+20 -8
View File
@@ -1,8 +1,20 @@
const { ChromaClientExtended } = require("../dist/index.js");
function makeid(length) {
let result = '';
const characters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
const charactersLength = characters.length;
let counter = 0;
while (counter < length) {
result += characters.charAt(Math.floor(Math.random() * charactersLength));
counter += 1;
}
return result;
}
async function run() {
const PATH = 'http://localhost:8000'
const reachable = await fetch(`${PATH}/api/v1`).then((res) => res.ok).catch(() => false);
const PATH = 'https://sgrh4jb9s5.execute-api.us-west-1.amazonaws.com/dev'
const reachable = await fetch(`${PATH}/api/v1`).then((res) => ![404, 500].includes(res.status)).catch(() => false);
if (!reachable) {
console.log(`\nService is not online at ${PATH} - cannot test. Exiting\n`)
process.exit(1);
@@ -12,7 +24,7 @@ async function run() {
path: PATH,
fetchOptions: {
headers: {
'X-Api-Token': "sk-live-Hunt3r2"
'x-api-key': "Wl2WNki9DT2r434RXMHbCpE8ponXsXnaxNn9tYxa"
}
}
});
@@ -20,17 +32,17 @@ async function run() {
console.log(await client.reset())
console.log(await client.heartbeat())
await client.createCollection({ name: "test" });
await client.getOrCreateCollection({ name: "test" });
console.log(await client.listCollections())
const collection = await client.getCollection({ name: 'test' })
console.log({ count: await collection.count() })
console.log({ currentCount: await collection.count() })
const ids = 'test1'
const embeddings = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
const ids = makeid(20)
const embeddings = Array.from({ length: 10 }, () => Math.floor(Math.random() * 99));
const metadatas = { test: 'test' }
await collection.add({ ids, embeddings, metadatas })
console.log(await collection.count())
console.log({ updatedCount: await collection.count() })
console.log('Tests completed.')
}