From a41e4e6cea7eadeb2de63a86214bb7ed3833f8fc Mon Sep 17 00:00:00 2001 From: Drew Kim Date: Tue, 15 Oct 2024 11:44:35 -0700 Subject: [PATCH] Update JS client --- clients/js/src/ChromaClient.ts | 50 +++++++++++++++++++++++++---- clients/js/src/Collection.ts | 18 +++++++++++ clients/js/src/generated/api.ts | 2 +- clients/js/src/generated/runtime.ts | 1 + clients/js/src/types.ts | 5 +++ 5 files changed, 68 insertions(+), 8 deletions(-) diff --git a/clients/js/src/ChromaClient.ts b/clients/js/src/ChromaClient.ts index a8b87552869e..dfb068fd3750 100644 --- a/clients/js/src/ChromaClient.ts +++ b/clients/js/src/ChromaClient.ts @@ -13,6 +13,7 @@ import type { GetCollectionParams, GetOrCreateCollectionParams, ListCollectionsParams, + TenantAndDatabases, } from "./types"; import { validateTenantDatabase, wrapCollection } from "./utils"; @@ -27,11 +28,11 @@ export class ChromaClient { /** * @ignore */ - private tenant: string; + public tenant: string; /** * @ignore */ - private database: string; + public database: string; /** * @ignore */ @@ -39,7 +40,7 @@ export class ChromaClient { /** * @ignore */ - private authProvider: ClientAuthProvider | undefined; + public authProvider: ClientAuthProvider | undefined; /** * @ignore */ @@ -94,7 +95,9 @@ export class ChromaClient { } /** @ignore */ - init(): Promise { + async init(): Promise { + await this.resolveTenantAndDatabases(); + if (!this._initPromise) { this._initPromise = validateTenantDatabase( this._adminClient, @@ -106,6 +109,39 @@ export class ChromaClient { return this._initPromise; } + /** + * Tries to set the tenant and database for the client. + * + * @returns {Promise} A promise that resolves when the tenant/database is resolved. + * @throws {Error} If there is an issue resolving the tenant and database. + * + */ + async resolveTenantAndDatabases(): Promise { + const response = (await this.api.resolveTenantAndDatabases( + this.api.options, + )) as TenantAndDatabases; + const user_identity = response as TenantAndDatabases; + const user_tenant = user_identity.tenant; + const user_databases = user_identity.databases; + + if ( + user_tenant !== null && + user_tenant !== undefined && + user_tenant !== "*" + ) { + this.tenant = user_tenant; + } + + if ( + user_databases !== null && + user_databases !== undefined && + user_databases.length == 1 && + user_databases[0] !== "*" + ) { + this.database = user_databases[0]; + } + } + /** * Resets the state of the object by making an API call to the reset endpoint. * @@ -269,10 +305,10 @@ export class ChromaClient { > { await this.init(); return (await this.api.listCollections( - limit, - offset, this.tenant, this.database, + limit, + offset, this.api.options, )) as CollectionParams[]; } @@ -320,9 +356,9 @@ export class ChromaClient { await this.init(); const response = (await this.api.getCollection( - name, this.tenant, this.database, + name, this.api.options, )) as CollectionParams; diff --git a/clients/js/src/Collection.ts b/clients/js/src/Collection.ts index 5282be9714f5..c5bc89c8dcd0 100644 --- a/clients/js/src/Collection.ts +++ b/clients/js/src/Collection.ts @@ -71,6 +71,8 @@ export class Collection { await this.client.init(); await this.client.api.add( + this.client.tenant, + this.client.database, this.id, // TODO: For some reason the auto generated code requires metadata to be defined here. (await prepareRecordRequest( @@ -104,6 +106,8 @@ export class Collection { await this.client.init(); await this.client.api.upsert( + this.client.tenant, + this.client.database, this.id, // TODO: For some reason the auto generated code requires metadata to be defined here. (await prepareRecordRequest( @@ -126,6 +130,8 @@ export class Collection { async count(): Promise { await this.client.init(); return (await this.client.api.count( + this.client.tenant, + this.client.database, this.id, this.client.api.options, )) as number; @@ -168,6 +174,8 @@ export class Collection { const resp = (await this.client.api.aGet( this.id, + this.client.tenant, + this.client.database, { ids: idsArray, where, @@ -205,6 +213,8 @@ export class Collection { await this.client.init(); await this.client.api.update( + this.client.tenant, + this.client.database, this.id, await prepareRecordRequest(params, this.embeddingFunction, true), this.client.api.options, @@ -266,6 +276,8 @@ export class Collection { : toArrayOfArrays(queryEmbeddings); return (await this.client.api.getNearestNeighbors( + this.client.tenant, + this.client.database, this.id, { query_embeddings: arrayQueryEmbeddings, @@ -303,6 +315,8 @@ export class Collection { await this.client.init(); return this.client.api .updateCollection( + this.client.tenant, + this.client.database, this.id, { new_name: name, @@ -342,6 +356,8 @@ export class Collection { await this.client.init(); return (await this.client.api.aGet( this.id, + this.client.tenant, + this.client.database, { limit, }, @@ -377,6 +393,8 @@ export class Collection { if (ids !== undefined) idsArray = toArray(ids); await this.client.api.aDelete( this.id, + this.client.tenant, + this.client.database, { ids: idsArray, where: where, where_document: whereDocument }, this.client.api.options, ); diff --git a/clients/js/src/generated/api.ts b/clients/js/src/generated/api.ts index e4a87a8824ae..7d0806b8c8c5 100644 --- a/clients/js/src/generated/api.ts +++ b/clients/js/src/generated/api.ts @@ -989,7 +989,7 @@ export const ApiApiFetchParamCreator = function (configuration?: Configuration) * ApiApi - functional programming interface * @export */ -export const ApiApiFp = function(configuration?: Configuration) { +export const ApiApiFp = function (configuration?: Configuration) { return { /** * @summary Add diff --git a/clients/js/src/generated/runtime.ts b/clients/js/src/generated/runtime.ts index d73b079b1b4c..8fa17d91fc55 100644 --- a/clients/js/src/generated/runtime.ts +++ b/clients/js/src/generated/runtime.ts @@ -1,4 +1,5 @@ import 'isomorphic-fetch'; +import 'isomorphic-fetch'; /* eslint-disable */ // tslint:disable /** diff --git a/clients/js/src/types.ts b/clients/js/src/types.ts index 5e9e6cc34fd8..3479835e4d83 100644 --- a/clients/js/src/types.ts +++ b/clients/js/src/types.ts @@ -258,3 +258,8 @@ export type DeleteParams = { where?: Where; whereDocument?: WhereDocument; }; + +export type TenantAndDatabases = { + tenant: string; + databases: string[]; +};