diff --git a/.github/workflows/chroma-test.yml b/.github/workflows/chroma-test.yml index e78c028285b..ec1c7e37c3c 100644 --- a/.github/workflows/chroma-test.yml +++ b/.github/workflows/chroma-test.yml @@ -27,4 +27,4 @@ jobs: - name: Test run: python -m pytest - name: Integration Test - run: bin/integration-test + run: bin/integration-test \ No newline at end of file diff --git a/README.md b/README.md index 23f5821f2ab..984c883458c 100644 --- a/README.md +++ b/README.md @@ -90,3 +90,4 @@ Chroma is a rapidly developing project. We welcome PR contributors and ideas for ## License [Apache 2.0](./LICENSE) + diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py index cba6e1ad7fc..d24515ef8bb 100644 --- a/chromadb/server/fastapi/__init__.py +++ b/chromadb/server/fastapi/__init__.py @@ -70,6 +70,7 @@ def __init__(self, settings): self.router.add_api_route("/api/v1", self.root, methods=["GET"]) self.router.add_api_route("/api/v1/reset", self.reset, methods=["POST"]) self.router.add_api_route("/api/v1/version", self.version, methods=["GET"]) + self.router.add_api_route("/api/v1/heartbeat", self.heartbeat, methods=["GET"]) self.router.add_api_route("/api/v1/persist", self.persist, methods=["POST"]) self.router.add_api_route("/api/v1/raw_sql", self.raw_sql, methods=["POST"]) @@ -124,6 +125,9 @@ def app(self): def root(self): return {"nanosecond heartbeat": self._api.heartbeat()} + def heartbeat(self): + return self.root() + def persist(self): self._api.persist() diff --git a/clients/js/package.json b/clients/js/package.json index f7b301f98b3..666e6d4c8e1 100644 --- a/clients/js/package.json +++ b/clients/js/package.json @@ -1,6 +1,6 @@ { "name": "chromadb", - "version": "1.3.1", + "version": "1.4.0", "description": "A JavaScript interface for chroma", "keywords": [], "author": "", @@ -31,8 +31,8 @@ "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", + "db:clean": "cd ../.. && docker-compose -f docker-compose.test.yml down --volumes", + "db:run": "cd ../.. && docker-compose -f docker-compose.test.yml up --detach && sleep 5", "clean": "rimraf dist", "build": "run-s clean build:*", "build:main": "tsc -p tsconfig.json", diff --git a/clients/js/src/generated/api/default-api.ts b/clients/js/src/generated/api/default-api.ts index 05476bdcecd..fefcd568328 100644 --- a/clients/js/src/generated/api/default-api.ts +++ b/clients/js/src/generated/api/default-api.ts @@ -376,6 +376,36 @@ export const DefaultApiAxiosParamCreator = function (configuration?: Configurati options: localVarRequestOptions, }; }, + /** + * + * @summary Heartbeat + * @param {*} [options] Override http request option. + * @throws {RequiredError} + */ + heartbeat: async (options: AxiosRequestConfig = {}): Promise => { + const localVarPath = `/api/v1/heartbeat`; + // use dummy base URL string because the URL constructor only accepts absolute URLs. + const localVarUrlObj = new URL(localVarPath, DUMMY_BASE_URL); + let baseOptions; + if (configuration) { + baseOptions = configuration.baseOptions; + } + + const localVarRequestOptions = { method: 'GET', ...baseOptions, ...options}; + const localVarHeaderParameter = {} as any; + const localVarQueryParameter = {} as any; + + + + setSearchParams(localVarUrlObj, localVarQueryParameter); + let headersFromBaseOptions = baseOptions && baseOptions.headers ? baseOptions.headers : {}; + localVarRequestOptions.headers = {...localVarHeaderParameter, ...headersFromBaseOptions, ...options.headers}; + + return { + url: toPathString(localVarUrlObj), + options: localVarRequestOptions, + }; + }, /** * * @summary List Collections @@ -607,6 +637,36 @@ export const DefaultApiAxiosParamCreator = function (configuration?: Configurati localVarRequestOptions.headers = {...localVarHeaderParameter, ...headersFromBaseOptions, ...options.headers}; localVarRequestOptions.data = serializeDataIfNeeded(updateCollection, localVarRequestOptions, configuration) + return { + url: toPathString(localVarUrlObj), + options: localVarRequestOptions, + }; + }, + /** + * + * @summary Version + * @param {*} [options] Override http request option. + * @throws {RequiredError} + */ + version: async (options: AxiosRequestConfig = {}): Promise => { + const localVarPath = `/api/v1/version`; + // use dummy base URL string because the URL constructor only accepts absolute URLs. + const localVarUrlObj = new URL(localVarPath, DUMMY_BASE_URL); + let baseOptions; + if (configuration) { + baseOptions = configuration.baseOptions; + } + + const localVarRequestOptions = { method: 'GET', ...baseOptions, ...options}; + const localVarHeaderParameter = {} as any; + const localVarQueryParameter = {} as any; + + + + setSearchParams(localVarUrlObj, localVarQueryParameter); + let headersFromBaseOptions = baseOptions && baseOptions.headers ? baseOptions.headers : {}; + localVarRequestOptions.headers = {...localVarHeaderParameter, ...headersFromBaseOptions, ...options.headers}; + return { url: toPathString(localVarUrlObj), options: localVarRequestOptions, @@ -725,6 +785,16 @@ export const DefaultApiFp = function(configuration?: Configuration) { const localVarAxiosArgs = await localVarAxiosParamCreator.getNearestNeighbors(collectionName, queryEmbedding, options); return createRequestFunction(localVarAxiosArgs, globalAxios, BASE_PATH, configuration); }, + /** + * + * @summary Heartbeat + * @param {*} [options] Override http request option. + * @throws {RequiredError} + */ + async heartbeat(options?: AxiosRequestConfig): Promise<(axios?: AxiosInstance, basePath?: string) => AxiosPromise> { + const localVarAxiosArgs = await localVarAxiosParamCreator.heartbeat(options); + return createRequestFunction(localVarAxiosArgs, globalAxios, BASE_PATH, configuration); + }, /** * * @summary List Collections @@ -800,6 +870,16 @@ export const DefaultApiFp = function(configuration?: Configuration) { const localVarAxiosArgs = await localVarAxiosParamCreator.updateCollection(collectionName, updateCollection, options); return createRequestFunction(localVarAxiosArgs, globalAxios, BASE_PATH, configuration); }, + /** + * + * @summary Version + * @param {*} [options] Override http request option. + * @throws {RequiredError} + */ + async version(options?: AxiosRequestConfig): Promise<(axios?: AxiosInstance, basePath?: string) => AxiosPromise> { + const localVarAxiosArgs = await localVarAxiosParamCreator.version(options); + return createRequestFunction(localVarAxiosArgs, globalAxios, BASE_PATH, configuration); + }, } }; @@ -904,6 +984,15 @@ export const DefaultApiFactory = function (configuration?: Configuration, basePa getNearestNeighbors(collectionName: any, queryEmbedding: QueryEmbedding, options?: any): AxiosPromise { return localVarFp.getNearestNeighbors(collectionName, queryEmbedding, options).then((request) => request(axios, basePath)); }, + /** + * + * @summary Heartbeat + * @param {*} [options] Override http request option. + * @throws {RequiredError} + */ + heartbeat(options?: any): AxiosPromise { + return localVarFp.heartbeat(options).then((request) => request(axios, basePath)); + }, /** * * @summary List Collections @@ -972,6 +1061,15 @@ export const DefaultApiFactory = function (configuration?: Configuration, basePa updateCollection(collectionName: any, updateCollection: UpdateCollection, options?: any): AxiosPromise { return localVarFp.updateCollection(collectionName, updateCollection, options).then((request) => request(axios, basePath)); }, + /** + * + * @summary Version + * @param {*} [options] Override http request option. + * @throws {RequiredError} + */ + version(options?: any): AxiosPromise { + return localVarFp.version(options).then((request) => request(axios, basePath)); + }, }; }; @@ -1300,6 +1398,17 @@ export class DefaultApi extends BaseAPI { return DefaultApiFp(this.configuration).getNearestNeighbors(requestParameters.collectionName, requestParameters.queryEmbedding, options).then((request) => request(this.axios, this.basePath)); } + /** + * + * @summary Heartbeat + * @param {*} [options] Override http request option. + * @throws {RequiredError} + * @memberof DefaultApi + */ + public heartbeat(options?: AxiosRequestConfig) { + return DefaultApiFp(this.configuration).heartbeat(options).then((request) => request(this.axios, this.basePath)); + } + /** * * @summary List Collections @@ -1379,4 +1488,15 @@ export class DefaultApi extends BaseAPI { public updateCollection(requestParameters: DefaultApiUpdateCollectionRequest, options?: AxiosRequestConfig) { return DefaultApiFp(this.configuration).updateCollection(requestParameters.collectionName, requestParameters.updateCollection, options).then((request) => request(this.axios, this.basePath)); } + + /** + * + * @summary Version + * @param {*} [options] Override http request option. + * @throws {RequiredError} + * @memberof DefaultApi + */ + public version(options?: AxiosRequestConfig) { + return DefaultApiFp(this.configuration).version(options).then((request) => request(this.axios, this.basePath)); + } } diff --git a/clients/js/src/index.ts b/clients/js/src/index.ts index 4ec2a0dd973..70687387009 100644 --- a/clients/js/src/index.ts +++ b/clients/js/src/index.ts @@ -1,3 +1,4 @@ +import { GetEmbeddingIncludeEnum, QueryEmbeddingIncludeEnum } from "./generated"; import { DefaultApi } from "./generated/api"; import { Configuration } from "./generated/configuration"; @@ -94,16 +95,25 @@ type CallableFunction = { export class Collection { public name: string; + public metadata: object | undefined; private api: DefaultApi; public embeddingFunction: CallableFunction | undefined; - constructor(name: string, api: DefaultApi, embeddingFunction?: CallableFunction) { + constructor(name: string, api: DefaultApi, metadata?: object, embeddingFunction?: CallableFunction) { this.name = name; + this.metadata = metadata; this.api = api; if (embeddingFunction !== undefined) this.embeddingFunction = embeddingFunction; } + private setName(name: string) { + this.name = name; + } + private setMetadata(metadata: object | undefined) { + this.metadata = metadata; + } + public async add( ids: string | string[], embeddings: number[] | number[][] | undefined, @@ -177,11 +187,35 @@ export class Collection { return response.data; } + public async modify( + name?: string, + metadata?: object, + ) { + const response = await this.api.updateCollection({ + collectionName: this.name, + updateCollection: { + new_name: name, + new_metadata: metadata, + }, + }).then(function (response) { + return response.data; + }).catch(function ({ response }) { + return response.data; + }); + + this.setName(name || this.name) + this.setMetadata(metadata || this.metadata) + + return response + } + public async get( ids?: string[], where?: object, limit?: number, offset?: number, + include?: GetEmbeddingIncludeEnum[], + where_document?: object, ) { let idsArray = undefined if (ids !== undefined) idsArray = toArray(ids); @@ -193,6 +227,8 @@ export class Collection { where, limit, offset, + include, + where_document, }, }).then(function (response) { return response.data; @@ -204,11 +240,51 @@ export class Collection { } + public async update( + ids: string | string[], + embeddings?: number[] | number[][], + metadatas?: object | object[], + documents?: string | string[], + ) { + 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", + ); + } + } + + var resp = await this.api.update({ + collectionName: this.name, + updateEmbedding: { + ids: toArray(ids), + embeddings: (embeddings ? toArrayOfArrays(embeddings) : undefined), + documents: toArray(documents), + metadatas: toArray(metadatas), + }, + }).then(function (response) { + return response.data; + }).catch(function ({ response }) { + return response.data; + }); + + return resp + } + public async query( query_embeddings: number[] | number[][] | undefined, n_results: number = 10, where?: object, - query_text?: string | string[], + query_text?: string | string[], // TODO: should be named query_texts to match python API + where_document?: object, // {"$contains":"search_string"} + include?: QueryEmbeddingIncludeEnum[], // ["metadata", "document"] ) { if ((query_embeddings === undefined) && (query_text === undefined)) { throw new Error( @@ -234,6 +310,8 @@ export class Collection { query_embeddings: query_embeddingsArray, where, n_results, + where_document: where_document, + include: include }, }).then(function (response) { return response.data; @@ -256,10 +334,10 @@ export class Collection { return await this.api.createIndex({ collectionName: this.name }); } - public async delete(ids?: string[], where?: object) { + public async delete(ids?: string[], where?: object, where_document?: object) { var response = await this.api._delete({ collectionName: this.name, - deleteEmbedding: { ids: ids, where: where }, + deleteEmbedding: { ids: ids, where: where, where_document: where_document }, }).then(function (response) { return response.data; }).catch(function ({ response }) { @@ -286,6 +364,20 @@ export class ChromaClient { return await this.api.reset(); } + public async version() { + const response = await this.api.version(); + return response.data; + } + + public async heartbeat() { + const response = await this.api.heartbeat(); + return response.data["nanosecond heartbeat"]; + } + + public async persist() { + throw new Error("Not implemented in JS client") + } + public async createCollection(name: string, metadata?: object, embeddingFunction?: CallableFunction) { const newCollection = await this.api.createCollection({ createCollection: { name, metadata }, @@ -299,7 +391,24 @@ export class ChromaClient { throw new Error(newCollection.error); } - return new Collection(name, this.api, embeddingFunction); + return new Collection(name, this.api, metadata, embeddingFunction); + } + + public async getOrCreateCollection(name: string, metadata?: object, embeddingFunction?: CallableFunction) { + const newCollection = await this.api.createCollection({ + createCollection: { name, metadata, get_or_create: true }, + + }).then(function (response) { + return response.data; + }).catch(function ({ response }) { + return response.data; + }); + + if (newCollection.error) { + throw new Error(newCollection.error); + } + + return new Collection(name, this.api, newCollection.metadata, embeddingFunction); } public async listCollections() { @@ -308,7 +417,13 @@ export class ChromaClient { } public async getCollection(name: string, embeddingFunction?: CallableFunction) { - return new Collection(name, this.api, embeddingFunction); + const response = await this.api.getCollection({ collectionName: name }).then(function (response) { + return response.data; + }).catch(function ({ response }) { + return response.data; + }); + + return new Collection(response.name, this.api, response.metadata, embeddingFunction); } public async deleteCollection(name: string) { diff --git a/clients/js/test/add.collections.test.ts b/clients/js/test/add.collections.test.ts new file mode 100644 index 00000000000..6fa1729417e --- /dev/null +++ b/clients/js/test/add.collections.test.ts @@ -0,0 +1,50 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' +import { DOCUMENTS, EMBEDDINGS, IDS } from './data'; +import { GetEmbeddingIncludeEnum } from '../src/generated'; + +test('it should add single embeddings to a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + const id = 'test1' + const embedding = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + const metadata = { test: 'test' } + await collection.add(id, embedding, metadata) + const count = await collection.count() + expect(count).toBe(1) + var res = await collection.get([id], undefined, undefined, undefined, [GetEmbeddingIncludeEnum.Embeddings]) + expect(res.embeddings[0]).toEqual(embedding) +}) + +test('it should add batch embeddings to a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS) + const count = await collection.count() + expect(count).toBe(3) + var res = await collection.get(IDS, undefined, undefined, undefined, [GetEmbeddingIncludeEnum.Embeddings]) + expect(res.embeddings).toEqual(EMBEDDINGS) // reverse because of the order of the ids +}) + +test('add documents', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, undefined, DOCUMENTS) + const results = await collection.get(["test1"]) + expect(results.documents[0]).toBe("This is a test") +}) + +test('test skipping indexing and manually doing it', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, undefined, DOCUMENTS, false) + + // expect collection.query to throw an error + const result = await collection.query([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 3) + expect(result.error).toContain("NoIndexException") + + await collection.createIndex() + const result2 = await collection.query([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 3) + expect(result2.error).toBeUndefined() + expect(result2.ids[0].length).toBe(3) +}) \ No newline at end of file diff --git a/clients/js/test/client.test.ts b/clients/js/test/client.test.ts index b631eef3c35..568e0ca7b96 100644 --- a/clients/js/test/client.test.ts +++ b/clients/js/test/client.test.ts @@ -1,199 +1,40 @@ import { expect, test } from '@jest/globals'; import { ChromaClient } from '../src/index' - -const PORT = process.env.PORT || '8000' -const URL = 'http://localhost:' + PORT -const chroma = new ChromaClient(URL) -console.log('using URL: ' + URL) - -// sleep for 10 seconds - to allow sentence transformers to download -// test('await1', async () => { -// await chroma.reset() -// let collections = await chroma.listCollections() -// await new Promise(r => setTimeout(r, 4500)); -// }) -// test('await2', async () => { -// await new Promise(r => setTimeout(r, 4500)); -// }) -// test('await3', async () => { -// await new Promise(r => setTimeout(r, 4500)); -// }) +import chroma from './initClient' test('it should create the client connection', async () => { expect(chroma).toBeDefined() expect(chroma).toBeInstanceOf(ChromaClient) }) -test('it should reset the database', async () => { - await chroma.reset() - let collections = await chroma.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) - const collection = await chroma.createCollection('test') - await chroma.reset() - collections = await chroma.listCollections() - expect(collections).toBeDefined() - expect(collections).toBeInstanceOf(Array) - expect(collections.length).toBe(0) +test('it should get the version', async () => { + const version = await chroma.version() + expect(version).toBeDefined() + expect(version).toMatch(/^[0-9]+\.[0-9]+\.[0-9]+$/) }) -test('it should create a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - expect(collection).toBeDefined() - expect(collection).toHaveProperty('name') - let collections = await chroma.listCollections() - expect([{ name: 'test', metadata: null }]).toEqual(expect.arrayContaining(collections)); - expect([{ name: 'test2', metadata: null }]).not.toEqual(expect.arrayContaining(collections)); +test('it should get the heartbeat', async () => { + const heartbeat = await chroma.heartbeat() + expect(heartbeat).toBeDefined() + expect(heartbeat).toBeGreaterThan(0) }) -test('it should list collections', async () => { +test('it should reset the database', async () => { await chroma.reset() - let collections = await chroma.listCollections() + const collections = await chroma.listCollections() expect(collections).toBeDefined() expect(collections).toBeInstanceOf(Array) expect(collections.length).toBe(0) + const collection = await chroma.createCollection('test') - collections = await chroma.listCollections() - expect(collections.length).toBe(1) -}) - -test('it should get a collection', async () => { + const collections2 = await chroma.listCollections() + expect(collections2).toBeDefined() + expect(collections2).toBeInstanceOf(Array) + expect(collections2.length).toBe(1) + await chroma.reset() - const collection = await chroma.createCollection('test') - const collection2 = await chroma.getCollection('test') - expect(collection).toBeDefined() - expect(collection2).toBeDefined() - expect(collection).toHaveProperty('name') - expect(collection2).toHaveProperty('name') - expect(collection.name).toBe(collection2.name) + const collections3 = await chroma.listCollections() + expect(collections3).toBeDefined() + expect(collections3).toBeInstanceOf(Array) + expect(collections3.length).toBe(0) }) - -test('it should delete a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - let collections = await chroma.listCollections() - expect(collections.length).toBe(1) - var resp = await chroma.deleteCollection('test') - collections = await chroma.listCollections() - expect(collections.length).toBe(0) -}) - -test('it should add single embeddings to a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const id = 'test1' - const embedding = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] - const metadata = { test: 'test' } - await collection.add(id, embedding, metadata) - const count = await collection.count() - expect(count).toBe(1) -}) - -test('it should add batch embeddings to a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add(ids, embeddings) - const count = await collection.count() - expect(count).toBe(3) -}) - -test('it should query a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add(ids, embeddings) - const results = await collection.query([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 2) - expect(results).toBeDefined() - expect(results).toBeInstanceOf(Object) - // expect(results.embeddings[0].length).toBe(2) - expect(['test1', 'test2']).toEqual(expect.arrayContaining(results.ids[0])); - expect(['test3']).not.toEqual(expect.arrayContaining(results.ids[0])); -}) - -test('it should peek a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - await collection.add(ids, embeddings) - const results = await collection.peek(2) - expect(results).toBeDefined() - expect(results).toBeInstanceOf(Object) - expect(results.ids.length).toBe(2) - expect(['test1', 'test2']).toEqual(expect.arrayContaining(results.ids)); -}) - -test('it should get a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - const metadatas = [{ test: 'test1' }, { test: 'test2' }, { test: 'test3' }] - await collection.add(ids, embeddings, metadatas) - const results = await collection.get(['test1']) - expect(results).toBeDefined() - expect(results).toBeInstanceOf(Object) - expect(results.ids.length).toBe(1) - expect(['test1']).toEqual(expect.arrayContaining(results.ids)); - expect(['test2']).not.toEqual(expect.arrayContaining(results.ids)); - - const results2 = await collection.get(undefined, { 'test': 'test1' }) - expect(results2).toBeDefined() - expect(results2).toBeInstanceOf(Object) - expect(results2.ids.length).toBe(1) -}) - -test('it should delete a collection', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - const metadatas = [{ test: 'test1' }, { test: 'test2' }, { test: 'test3' }] - await collection.add(ids, embeddings, metadatas) - let count = await collection.count() - expect(count).toBe(3) - var resp = await collection.delete(undefined, { 'test': 'test1' }) - count = await collection.count() - expect(count).toBe(2) -}) - -test('wrong code returns an error', async () => { - await chroma.reset() - const collection = await chroma.createCollection('test') - const ids = ['test1', 'test2', 'test3'] - const embeddings = [ - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] - ] - const metadatas = [{ test: 'test1' }, { test: 'test2' }, { test: 'test3' }] - await collection.add(ids, embeddings, metadatas) - const results = await collection.get(undefined, { "test": { "$contains": "hello" } }); - expect(results.error).toBeDefined() - expect(results.error).toBe("ValueError('Expected one of $gt, $lt, $gte, $lte, $ne, $eq, got $contains')") -}) \ No newline at end of file diff --git a/clients/js/test/collection.client.test.ts b/clients/js/test/collection.client.test.ts new file mode 100644 index 00000000000..2542f7a5117 --- /dev/null +++ b/clients/js/test/collection.client.test.ts @@ -0,0 +1,80 @@ +import { expect, test } from '@jest/globals'; +import { ChromaClient } from '../src/index' +import chroma from './initClient' + +test('it should list collections', async () => { + await chroma.reset() + let collections = await chroma.listCollections() + expect(collections).toBeDefined() + expect(collections).toBeInstanceOf(Array) + expect(collections.length).toBe(0) + const collection = await chroma.createCollection('test') + collections = await chroma.listCollections() + expect(collections.length).toBe(1) +}) + +test('it should create a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + expect(collection).toBeDefined() + expect(collection).toHaveProperty('name') + expect(collection.name).toBe('test') + let collections = await chroma.listCollections() + expect([{ name: 'test', metadata: null }]).toEqual(expect.arrayContaining(collections)); + expect([{ name: 'test2', metadata: null }]).not.toEqual(expect.arrayContaining(collections)); + + await chroma.reset() + const collection2 = await chroma.createCollection('test2', { test: 'test' }) + expect(collection2).toBeDefined() + expect(collection2).toHaveProperty('name') + expect(collection2.name).toBe('test2') + expect(collection2).toHaveProperty('metadata') + expect(collection2.metadata).toHaveProperty('test') + expect(collection2.metadata).toEqual({ test: 'test' }) + let collections2 = await chroma.listCollections() + expect([{ name: 'test2', metadata: { test: 'test' } }]).toEqual(expect.arrayContaining(collections2)); + +}) + +test('it should get a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + const collection2 = await chroma.getCollection('test') + expect(collection).toBeDefined() + expect(collection2).toBeDefined() + expect(collection).toHaveProperty('name') + expect(collection2).toHaveProperty('name') + expect(collection.name).toBe(collection2.name) +}) + +test('it should get or create a collection', async () => { + await chroma.reset() + await chroma.createCollection('test') + + const collection2 = await chroma.getOrCreateCollection('test') + expect(collection2).toBeDefined() + expect(collection2).toHaveProperty('name') + expect(collection2.name).toBe('test') + + const collection3 = await chroma.getOrCreateCollection('test3') + expect(collection3).toBeDefined() + expect(collection3).toHaveProperty('name') + expect(collection3.name).toBe('test3') +}) + +test('it should delete a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + let collections = await chroma.listCollections() + expect(collections.length).toBe(1) + await chroma.deleteCollection('test') + collections = await chroma.listCollections() + expect(collections.length).toBe(0) +}) + +// TODO: I want to test this, but I am not sure how to +// test('custom index params', async () => { +// throw new Error('not implemented') +// await chroma.reset() +// const collection = await chroma.createCollection('test', {"hnsw:space": "cosine"}) +// }) \ No newline at end of file diff --git a/clients/js/test/collection.test.ts b/clients/js/test/collection.test.ts new file mode 100644 index 00000000000..5d47c492a13 --- /dev/null +++ b/clients/js/test/collection.test.ts @@ -0,0 +1,66 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' + +test('it should modify collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + expect(collection.name).toBe('test') + expect(collection.metadata).toBeUndefined() + + await collection.modify('test2') + expect(collection.name).toBe('test2') + expect(collection.metadata).toBeUndefined() + + const collection2 = await chroma.getCollection('test2') + expect(collection2.name).toBe('test2') + expect(collection2.metadata).toBeNull() + + // test changing name and metadata independently + // and verify there are no side effects + const original_name = 'test3' + const new_name = 'test4' + const original_metadata = { test: 'test' } + const new_metadata = { test: 'test2' } + + const collection3 = await chroma.createCollection(original_name, original_metadata) + expect(collection3.name).toBe(original_name) + expect(collection3.metadata).toEqual(original_metadata) + + await collection3.modify(new_name) + expect(collection3.name).toBe(new_name) + expect(collection3.metadata).toEqual(original_metadata) + + const collection4 = await chroma.getCollection(new_name) + expect(collection4.name).toBe(new_name) + expect(collection4.metadata).toEqual(original_metadata) + + await collection3.modify(undefined, new_metadata) + expect(collection3.name).toBe(new_name) + expect(collection3.metadata).toEqual(new_metadata) + + const collection5 = await chroma.getCollection(new_name) + expect(collection5.name).toBe(new_name) + expect(collection5.metadata).toEqual(new_metadata) +}) + +test('it should store metadata', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test', { test: 'test' }) + expect(collection.metadata).toEqual({ test: 'test' }) + + // get the collection + const collection2 = await chroma.getCollection('test') + expect(collection2.metadata).toEqual({ test: 'test' }) + + // get or create the collection + const collection3 = await chroma.getOrCreateCollection('test') + expect(collection3.metadata).toEqual({ test: 'test' }) + + // modify + await collection3.modify(undefined, { test: 'test2' }) + expect(collection3.metadata).toEqual({ test: 'test2' }) + + // get it again + const collection4 = await chroma.getCollection('test') + expect(collection4.metadata).toEqual({ test: 'test2' }) +}) \ No newline at end of file diff --git a/clients/js/test/data.ts b/clients/js/test/data.ts new file mode 100644 index 00000000000..a26f0b7ba1b --- /dev/null +++ b/clients/js/test/data.ts @@ -0,0 +1,10 @@ +const IDS = ['test1', 'test2', 'test3'] +const EMBEDDINGS = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] +] +const METADATAS = [{ test: 'test1', 'float_value': -2 }, { test: 'test2', 'float_value': 0 }, { test: 'test3', 'float_value': 2 }] +const DOCUMENTS = ["This is a test", "This is another test", "This is a third test"] + +export { IDS, EMBEDDINGS, METADATAS, DOCUMENTS } \ No newline at end of file diff --git a/clients/js/test/delete.collection.test.ts b/clients/js/test/delete.collection.test.ts new file mode 100644 index 00000000000..fe4273b0f01 --- /dev/null +++ b/clients/js/test/delete.collection.test.ts @@ -0,0 +1,17 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' +import { EMBEDDINGS, IDS, METADATAS } from './data'; + +test('it should delete a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, METADATAS) + let count = await collection.count() + expect(count).toBe(3) + var resp = await collection.delete(undefined, { 'test': 'test1' }) + count = await collection.count() + expect(count).toBe(2) + + var remainingEmbeddings = await collection.get() + expect(['test2', 'test3']).toEqual(expect.arrayContaining(remainingEmbeddings.ids)); +}) \ No newline at end of file diff --git a/clients/js/test/get.collection.test.ts b/clients/js/test/get.collection.test.ts new file mode 100644 index 00000000000..8e85b6e7175 --- /dev/null +++ b/clients/js/test/get.collection.test.ts @@ -0,0 +1,39 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' +import { EMBEDDINGS, IDS, METADATAS } from './data'; + +test('it should get a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, METADATAS) + const results = await collection.get(['test1']) + expect(results).toBeDefined() + expect(results).toBeInstanceOf(Object) + expect(results.ids.length).toBe(1) + expect(['test1']).toEqual(expect.arrayContaining(results.ids)); + expect(['test2']).not.toEqual(expect.arrayContaining(results.ids)); + + const results2 = await collection.get(undefined, { 'test': 'test1' }) + expect(results2).toBeDefined() + expect(results2).toBeInstanceOf(Object) + expect(results2.ids.length).toBe(1) + expect(['test1']).toEqual(expect.arrayContaining(results2.ids)); +}) + +test('wrong code returns an error', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, METADATAS) + const results = await collection.get(undefined, { "test": { "$contains": "hello" } }); + expect(results.error).toBeDefined() + expect(results.error).toContain("ValueError") +}) + +test('test gt, lt, in a simple small way', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, METADATAS) + const items = await collection.get(undefined, {"float_value": {"$gt": -1.4}}) + expect(items.ids.length).toBe(2) + expect(['test2', 'test3']).toEqual(expect.arrayContaining(items.ids)); +}) \ No newline at end of file diff --git a/clients/js/test/initClient.ts b/clients/js/test/initClient.ts new file mode 100644 index 00000000000..a12a60c4c1f --- /dev/null +++ b/clients/js/test/initClient.ts @@ -0,0 +1,7 @@ +import { ChromaClient } from '../src/index' + +const PORT = process.env.PORT || '8000' +const URL = 'http://localhost:' + PORT +const chroma = new ChromaClient(URL) + +export default chroma \ No newline at end of file diff --git a/clients/js/test/peek.collection.test.ts b/clients/js/test/peek.collection.test.ts new file mode 100644 index 00000000000..5c5b6346d28 --- /dev/null +++ b/clients/js/test/peek.collection.test.ts @@ -0,0 +1,14 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' +import { IDS, EMBEDDINGS } from './data'; + +test('it should peek a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS) + const results = await collection.peek(2) + expect(results).toBeDefined() + expect(results).toBeInstanceOf(Object) + expect(results.ids.length).toBe(2) + expect(['test1', 'test2']).toEqual(expect.arrayContaining(results.ids)); +}) diff --git a/clients/js/test/query.collection.test.ts b/clients/js/test/query.collection.test.ts new file mode 100644 index 00000000000..472e5801b12 --- /dev/null +++ b/clients/js/test/query.collection.test.ts @@ -0,0 +1,39 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' +import { QueryEmbeddingIncludeEnum } from '../src/generated'; +import { EMBEDDINGS, IDS, METADATAS, DOCUMENTS } from './data'; + +test('it should query a collection', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS) + const results = await collection.query([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 2) + expect(results).toBeDefined() + expect(results).toBeInstanceOf(Object) + expect(['test1', 'test2']).toEqual(expect.arrayContaining(results.ids[0])); + expect(['test3']).not.toEqual(expect.arrayContaining(results.ids[0])); +}) + +// test where_document +test('it should get embedding with matching documents', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, METADATAS, DOCUMENTS) + + const results = await collection.query([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 3, undefined, undefined, { "$contains": "This is a test" }) + + // it should only return doc1 + expect(results).toBeDefined() + expect(results).toBeInstanceOf(Object) + expect(results.ids.length).toBe(1) + expect(['test1']).toEqual(expect.arrayContaining(results.ids[0])); + expect(['test2']).not.toEqual(expect.arrayContaining(results.ids[0])); + expect(['This is a test']).toEqual(expect.arrayContaining(results.documents[0])); + + const results2 = await collection.query([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 3, undefined, undefined, { "$contains": "This is a test" }, [QueryEmbeddingIncludeEnum.Embeddings]) + + expect(results2.embeddings[0][0]).toBeInstanceOf(Array) + expect(results2.embeddings[0].length).toBe(1) + expect(results2.embeddings[0][0]).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +}) + diff --git a/clients/js/test/update.collection.test.ts b/clients/js/test/update.collection.test.ts new file mode 100644 index 00000000000..c609253035c --- /dev/null +++ b/clients/js/test/update.collection.test.ts @@ -0,0 +1,29 @@ +import { expect, test } from '@jest/globals'; +import chroma from './initClient' +import { GetEmbeddingIncludeEnum } from '../src/generated'; +import { IDS, DOCUMENTS, EMBEDDINGS, METADATAS } from './data'; + +test('it should get embedding with matching documents', async () => { + await chroma.reset() + const collection = await chroma.createCollection('test') + await collection.add(IDS, EMBEDDINGS, METADATAS, DOCUMENTS) + + const results = await collection.get(['test1'], undefined, undefined, undefined, [GetEmbeddingIncludeEnum.Embeddings, GetEmbeddingIncludeEnum.Metadatas, GetEmbeddingIncludeEnum.Documents]) + expect(results).toBeDefined() + expect(results).toBeInstanceOf(Object) + expect(results.embeddings[0]).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + + await collection.update( + ['test1'], + [[1, 2, 3, 4, 5, 6, 7, 8, 9, 11]], + [{ test: 'test1new' }], + ["doc1new"] + ) + + const results2 = await collection.get(['test1'], undefined, undefined, undefined, [GetEmbeddingIncludeEnum.Embeddings, GetEmbeddingIncludeEnum.Metadatas, GetEmbeddingIncludeEnum.Documents]) + expect(results2).toBeDefined() + expect(results2).toBeInstanceOf(Object) + expect(results2.embeddings[0]).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 11]) + expect(results2.metadatas[0]).toEqual({ test: 'test1new' }) + expect(results2.documents[0]).toEqual('doc1new') +}) \ No newline at end of file