import { IEmbeddingFunction } from './embeddings/IEmbeddingFunction'; import { Configuration, ApiApi as DefaultApi } from "./generated"; import { handleSuccess, handleError } from "./utils"; import { Collection } from './Collection'; import { ChromaClientParams, CollectionMetadata, CollectionType, ConfigOptions, CreateCollectionParams, DeleteCollectionParams, GetCollectionParams, GetOrCreateCollectionParams, ListCollectionsParams } from './types'; import { AuthOptions, ClientAuthProtocolAdapter, IsomorphicFetchClientAuthProtocolAdapter } from "./auth"; import { DefaultEmbeddingFunction } from './embeddings/DefaultEmbeddingFunction'; import { AdminClient } from './AdminClient'; const DEFAULT_TENANT = "default_tenant" const DEFAULT_DATABASE = "default_database" export class ChromaClient { /** * @ignore */ private api: DefaultApi & ConfigOptions; private apiAdapter: ClientAuthProtocolAdapter|undefined; private tenant: string = DEFAULT_TENANT; private database: string = DEFAULT_DATABASE; private _adminClient?: AdminClient /** * Creates a new ChromaClient instance. * @param {Object} params - The parameters for creating a new client * @param {string} [params.path] - The base path for the Chroma API. * @returns {ChromaClient} A new ChromaClient instance. * * @example * ```typescript * const client = new ChromaClient({ * path: "http://localhost:8000" * }); * ``` */ constructor({ path, fetchOptions, auth, tenant = DEFAULT_TENANT, database = DEFAULT_DATABASE, }: ChromaClientParams = {}) { if (path === undefined) path = "http://localhost:8000"; this.tenant = tenant; this.database = database; const apiConfig: Configuration = new Configuration({ basePath: path, }); if (auth !== undefined) { this.apiAdapter = new IsomorphicFetchClientAuthProtocolAdapter(new DefaultApi(apiConfig), auth); this.api = this.apiAdapter.getApi(); } else { this.api = new DefaultApi(apiConfig); } this._adminClient = new AdminClient({ path: path, fetchOptions: fetchOptions, auth: auth, tenant: tenant, database: database }); // TODO: Validate tenant and database on client creation // this got tricky because: // - the constructor is sync but the generated api is async // - we need to inject auth information so a simple rewrite/fetch does not work this.api.options = fetchOptions ?? {}; } /** * Resets the state of the object by making an API call to the reset endpoint. * * @returns {Promise} A promise that resolves when the reset operation is complete. * @throws {Error} If there is an issue resetting the state. * * @example * ```typescript * await client.reset(); * ``` */ public async reset(): Promise { return await this.api.reset(this.api.options); } /** * Returns the version of the Chroma API. * @returns {Promise} A promise that resolves to the version of the Chroma API. * * @example * ```typescript * const version = await client.version(); * ``` */ public async version(): Promise { const response = await this.api.version(this.api.options); return await handleSuccess(response); } /** * Returns a heartbeat from the Chroma API. * @returns {Promise} A promise that resolves to the heartbeat from the Chroma API. * * @example * ```typescript * const heartbeat = await client.heartbeat(); * ``` */ public async heartbeat(): Promise { const response = await this.api.heartbeat(this.api.options); let ret = await handleSuccess(response); return ret["nanosecond heartbeat"] } /** * Creates a new collection with the specified properties. * * @param {Object} params - The parameters for creating a new collection. * @param {string} params.name - The name of the collection. * @param {CollectionMetadata} [params.metadata] - Optional metadata associated with the collection. * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. * * @returns {Promise} A promise that resolves to the created collection. * @throws {Error} If there is an issue creating the collection. * * @example * ```typescript * const collection = await client.createCollection({ * name: "my_collection", * metadata: { * "description": "My first collection" * } * }); * ``` */ public async createCollection({ name, metadata, embeddingFunction }: CreateCollectionParams): Promise { if (embeddingFunction === undefined) { embeddingFunction = new DefaultEmbeddingFunction(); } const newCollection = await this.api .createCollection(this.tenant, this.database, { name, metadata, }, this.api.options) .then(handleSuccess) .catch(handleError); if (newCollection.error) { throw new Error(newCollection.error); } return new Collection(name, newCollection.id, this.api, metadata, embeddingFunction); } /** * Gets or creates a collection with the specified properties. * * @param {Object} params - The parameters for creating a new collection. * @param {string} params.name - The name of the collection. * @param {CollectionMetadata} [params.metadata] - Optional metadata associated with the collection. * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. * * @returns {Promise} A promise that resolves to the got or created collection. * @throws {Error} If there is an issue getting or creating the collection. * * @example * ```typescript * const collection = await client.getOrCreateCollection({ * name: "my_collection", * metadata: { * "description": "My first collection" * } * }); * ``` */ public async getOrCreateCollection({ name, metadata, embeddingFunction }: GetOrCreateCollectionParams): Promise { if (embeddingFunction === undefined) { embeddingFunction = new DefaultEmbeddingFunction(); } const newCollection = await this.api .createCollection(this.tenant, this.database, { name, metadata, 'get_or_create': true }, this.api.options) .then(handleSuccess) .catch(handleError); if (newCollection.error) { throw new Error(newCollection.error); } return new Collection( name, newCollection.id, this.api, newCollection.metadata, embeddingFunction ); } /** * Lists all collections. * * @returns {Promise} A promise that resolves to a list of collection names. * @param {PositiveInteger} [params.limit] - Optional limit on the number of items to get. * @param {PositiveInteger} [params.offset] - Optional offset on the items to get. * @throws {Error} If there is an issue listing the collections. * * @example * ```typescript * const collections = await client.listCollections({ * limit: 10, * offset: 0, * }); * ``` */ public async listCollections({ limit, offset, }: ListCollectionsParams = {}): Promise { const response = await this.api.listCollections( this.tenant, this.database, limit, offset, this.api.options); return handleSuccess(response); } /** * Counts all collections. * * @returns {Promise} A promise that resolves to the number of collections. * @throws {Error} If there is an issue counting the collections. * * @example * ```typescript * const collections = await client.countCollections(); * ``` */ public async countCollections(): Promise { const response = await this.api.countCollections(this.tenant, this.database, this.api.options); return handleSuccess(response); } /** * Gets a collection with the specified name. * @param {Object} params - The parameters for getting a collection. * @param {string} params.name - The name of the collection. * @param {IEmbeddingFunction} [params.embeddingFunction] - Optional custom embedding function for the collection. * @returns {Promise} A promise that resolves to the collection. * @throws {Error} If there is an issue getting the collection. * * @example * ```typescript * const collection = await client.getCollection({ * name: "my_collection" * }); * ``` */ public async getCollection({ name, embeddingFunction }: GetCollectionParams): Promise { const response = await this.api .getCollection(name, this.tenant, this.database, this.api.options) .then(handleSuccess) .catch(handleError); if (response.error) { throw new Error(response.error); } return new Collection( response.name, response.id, this.api, response.metadata, embeddingFunction ); } /** * Deletes a collection with the specified name. * @param {Object} params - The parameters for deleting a collection. * @param {string} params.name - The name of the collection. * @returns {Promise} A promise that resolves when the collection is deleted. * @throws {Error} If there is an issue deleting the collection. * * @example * ```typescript * await client.deleteCollection({ * name: "my_collection" * }); * ``` */ public async deleteCollection({ name }: DeleteCollectionParams): Promise { return await this.api .deleteCollection(name, this.tenant, this.database, this.api.options) .then(handleSuccess) .catch(handleError); } }