diff --git a/docs/modules/chromadb.md b/docs/modules/chromadb.md index 0177da671..3d80ffc2e 100644 --- a/docs/modules/chromadb.md +++ b/docs/modules/chromadb.md @@ -8,9 +8,37 @@ npm install @testcontainers/chromadb --save-dev ``` -## Example +## Resources + +* [GitHub](https://github.com/chroma-core/chroma) +* [Node.js Client](https://www.npmjs.com/package/chromadb) +* [Docs](https://docs.trychroma.com) +* [Discord](https://discord.gg/MMeYNTmh3x) +* [Cookbook](https://cookbook.chromadb.dev) + +## Examples + + +[Connect to Chroma:](../../packages/modules/chromadb/src/chromadb-container.test.ts) +inside_block:simpleConnect + + + +[Create Collection:](../../packages/modules/chromadb/src/chromadb-container.test.ts) +inside_block:createCollection + + + +[Query Collection with Embedding Function:](../../packages/modules/chromadb/src/chromadb-container.test.ts) +inside_block:queryCollectionWithEmbeddingFunction + + + +[Work with persistent directory:](../../packages/modules/chromadb/src/chromadb-container.test.ts) +inside_block:persistentData + -[](../../packages/modules/chromadb/src/chromadb-container.test.ts) inside_block:docs +[Work with authentication:](../../packages/modules/chromadb/src/chromadb-container.test.ts) inside_block:auth diff --git a/packages/modules/chromadb/src/chromadb-container.test.ts b/packages/modules/chromadb/src/chromadb-container.test.ts index 7368ded4b..74bdbe6c0 100755 --- a/packages/modules/chromadb/src/chromadb-container.test.ts +++ b/packages/modules/chromadb/src/chromadb-container.test.ts @@ -1,23 +1,133 @@ -import { ChromaClient, AdminClient } from "chromadb"; -import { ChromaDBContainer } from "./chromadb-container"; - +import { ChromaClient, AdminClient, OllamaEmbeddingFunction } from "chromadb"; +import { ChromaDBContainer, StartedChromaDBContainer } from "./chromadb-container"; +import * as path from "node:path"; +import { GenericContainer } from "testcontainers"; +import * as os from "node:os"; +import * as fs from "node:fs"; +// run tests with NODE_OPTIONS=--experimental-vm-modules jest packages/modules/chromadb/src/chromadb-container.test.ts describe("ChromaDB", () => { jest.setTimeout(360_000); - // docs { - it("should connect and return a query result", async () => { + // startContainer { + it("should connect", async () => { + const container = await new ChromaDBContainer().start(); + const client = await connectTo(container); + expect(await client.heartbeat()).toBeDefined(); + // Do something with the client + await container.stop(); + }); + // } + + // simpleConnect { + async function connectTo(container: StartedChromaDBContainer) { + const client = new ChromaClient({ + path: container.getHttpUrl(), + }); + const hb = await client.heartbeat(); + expect(hb).toBeDefined(); + return client; + } + // } + + // createCollection { + it("should create collection and get data", async () => { const container = await new ChromaDBContainer().start(); + const client = await connectTo(container); + const collection = await client.createCollection({ name: "test", metadata: { "hnsw:space": "cosine" } }); + expect(collection.name).toBe("test"); + expect(collection.metadata).toBeDefined(); + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + expect(collection.metadata["hnsw:space"]).toBe("cosine"); + await collection.add({ ids: ["1"], embeddings: [[1, 2, 3]], documents: ["my doc"], metadatas: [{ key: "value" }] }); + const getResults = await collection.get({ ids: ["1"] }); + expect(getResults.ids[0]).toBe("1"); + expect(getResults.documents[0]).toStrictEqual("my doc"); + expect(getResults.metadatas).toBeDefined(); + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + expect(getResults.metadatas[0].key).toStrictEqual("value"); + await container.stop(); + }); + // } + + // queryCollectionWithEmbeddingFunction { + it("should create collection and query", async () => { + const container = await new ChromaDBContainer().start(); + const ollama = await new GenericContainer("ollama/ollama").withExposedPorts(11434).start(); + await ollama.exec(["ollama", "pull", "nomic-embed-text"]); + const client = await connectTo(container); + const embedder = new OllamaEmbeddingFunction({ + url: `http://${ollama.getHost()}:${ollama.getMappedPort(11434)}/api/embeddings`, + model: "nomic-embed-text", + }); + const collection = await client.createCollection({ + name: "test", + metadata: { "hnsw:space": "cosine" }, + embeddingFunction: embedder, + }); + expect(collection.name).toBe("test"); + await collection.add({ + ids: ["1", "2"], + documents: [ + "This is a document about dogs. Dogs are awesome.", + "This is a document about cats. Cats are awesome.", + ], + }); + const results = await collection.query({ queryTexts: ["Tell me about dogs"], nResults: 1 }); + expect(results).toBeDefined(); + expect(results.ids[0]).toEqual(["1"]); + expect(results.ids[0][0]).toBe("1"); + await container.stop(); + }); + + // persistentData { + it("should reconnect with volume and persistence data", async () => { + const sourcePath = fs.mkdtempSync(path.join(os.tmpdir(), "chroma-temp")); + const container = await new ChromaDBContainer() + .withBindMounts([{ source: sourcePath, target: "/chroma/chroma" }]) + .start(); + const client = await connectTo(container); + const collection = await client.createCollection({ name: "test", metadata: { "hnsw:space": "cosine" } }); + expect(collection.name).toBe("test"); + expect(collection.metadata).toBeDefined(); + // eslint-disable-next-line @typescript-eslint/ban-ts-comment + // @ts-ignore + expect(collection.metadata["hnsw:space"]).toBe("cosine"); + await collection.add({ ids: ["1"], embeddings: [[1, 2, 3]], documents: ["my doc"] }); + const getResults = await collection.get({ ids: ["1"] }); + expect(getResults.ids[0]).toBe("1"); + expect(getResults.documents[0]).toStrictEqual("my doc"); + await container.stop(); + expect(fs.existsSync(`${sourcePath}/chroma.sqlite3`)).toBe(true); + try { + fs.rmSync(sourcePath, { force: true, recursive: true }); + } catch (e) { + //Ignore clean up, when have no access on fs. + console.log(e); + } + }); + // } + + // auth { + it("should use auth", async () => { const tenant = "test-tenant"; const key = "test-key"; const database = "test-db"; + const container = await new ChromaDBContainer() + .withEnvironment({ + CHROMA_SERVER_AUTHN_CREDENTIALS: key, + CHROMA_SERVER_AUTHN_PROVIDER: "chromadb.auth.token_authn.TokenAuthenticationServerProvider", + CHROMA_AUTH_TOKEN_TRANSPORT_HEADER: "X_CHROMA_TOKEN", + }) + .start(); + const adminClient = new AdminClient({ tenant: tenant, auth: { provider: "token", credentials: key, - providerOptions: { - headerType: "X_CHROMA_TOKEN", - }, + tokenHeaderType: "X_CHROMA_TOKEN", }, path: container.getHttpUrl(), }); @@ -30,52 +140,14 @@ describe("ChromaDB", () => { auth: { provider: "token", credentials: key, - providerOptions: { - headerType: "X_CHROMA_TOKEN", - }, + tokenHeaderType: "X_CHROMA_TOKEN", }, path: container.getHttpUrl(), database, }); const collection = await dbClient.createCollection({ name: "test-collection" }); - - await collection.add({ - ids: ["1", "2", "3"], - documents: ["apple", "oranges", "pineapple"], - embeddings: [ - [1, 2, 3], - [4, 5, 6], - [7, 8, 9], - ], - }); - - const result = await collection.get({ ids: ["1", "2", "3"] }); - - expect(result).toMatchInlineSnapshot(` - { - "data": null, - "documents": [ - "apple", - "oranges", - "pineapple", - ], - "embeddings": null, - "ids": [ - "1", - "2", - "3", - ], - "metadatas": [ - null, - null, - null, - ], - "uris": null, - } - `); - - await container.stop(); + expect(collection.name).toBe("test-collection"); }); // } });