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/package-lock.json b/package-lock.json index 54278e89a..1ec184176 100644 --- a/package-lock.json +++ b/package-lock.json @@ -14940,6 +14940,16 @@ "integrity": "sha512-PX1wu0AmAdPqOL1mWhqmlOd8kOIZQwGZw6rh7uby9fTc5lhaOWFLX3I6R1hrF9k3zUY40e6igsLGkDXK92LJNg==", "dev": true }, + "node_modules/ollama": { + "version": "0.5.14", + "resolved": "https://registry.npmjs.org/ollama/-/ollama-0.5.14.tgz", + "integrity": "sha512-pvOuEYa2WkkAumxzJP0RdEYHkbZ64AYyyUszXVX7ruLvk5L+EiO2G71da2GqEQ4IAk4j6eLoUbGk5arzFT1wJA==", + "dev": true, + "license": "MIT", + "dependencies": { + "whatwg-fetch": "^3.6.20" + } + }, "node_modules/once": { "version": "1.4.0", "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", @@ -19756,14 +19766,16 @@ "testcontainers": "^10.21.0" }, "devDependencies": { - "chromadb": "^1.8.1" + "chromadb": "^1.9.1", + "ollama": "^0.5.14" } }, "packages/modules/chromadb/node_modules/chromadb": { - "version": "1.8.1", - "resolved": "https://registry.npmjs.org/chromadb/-/chromadb-1.8.1.tgz", - "integrity": "sha512-NpbYydbg4Uqt/9BXKgkZXn0fqpsh2Z1yjhkhKH+rcHMoq0pwI18BFSU2QU7Fk/ZypwGefW2AvqyE/3ZJIgy4QA==", + "version": "1.10.5", + "resolved": "https://registry.npmjs.org/chromadb/-/chromadb-1.10.5.tgz", + "integrity": "sha512-+IeTjjf44pKUY3vp1BacwO2tFAPcWCd64zxPZZm98dVj/kbSBeaHKB2D6eX7iRLHS1PTVASuqoR6mAJ+nrsTBg==", "dev": true, + "license": "Apache-2.0", "dependencies": { "cliui": "^8.0.1", "isomorphic-fetch": "^3.0.0" @@ -19774,7 +19786,9 @@ "peerDependencies": { "@google/generative-ai": "^0.1.1", "cohere-ai": "^5.0.0 || ^6.0.0 || ^7.0.0", - "openai": "^3.0.0 || ^4.0.0" + "ollama": "^0.5.0", + "openai": "^3.0.0 || ^4.0.0", + "voyageai": "^0.0.3-1" }, "peerDependenciesMeta": { "@google/generative-ai": { @@ -19783,8 +19797,14 @@ "cohere-ai": { "optional": true }, + "ollama": { + "optional": true + }, "openai": { "optional": true + }, + "voyageai": { + "optional": true } } }, diff --git a/packages/modules/chromadb/package.json b/packages/modules/chromadb/package.json index 0e0e552e0..8e28ac025 100644 --- a/packages/modules/chromadb/package.json +++ b/packages/modules/chromadb/package.json @@ -29,7 +29,8 @@ "build": "tsc --project tsconfig.build.json" }, "devDependencies": { - "chromadb": "^1.8.1" + "chromadb": "^1.9.1", + "ollama": "^0.5.14" }, "dependencies": { "testcontainers": "^10.21.0" diff --git a/packages/modules/chromadb/src/chromadb-container.test.ts b/packages/modules/chromadb/src/chromadb-container.test.ts index a23f1200d..543b4f10b 100755 --- a/packages/modules/chromadb/src/chromadb-container.test.ts +++ b/packages/modules/chromadb/src/chromadb-container.test.ts @@ -1,21 +1,125 @@ -import { AdminClient, ChromaClient } from "chromadb"; -import { ChromaDBContainer } from "./chromadb-container"; +import { AdminClient, ChromaClient, OllamaEmbeddingFunction } from "chromadb"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { GenericContainer } from "testcontainers"; +import { ChromaDBContainer, StartedChromaDBContainer } from "./chromadb-container"; describe("ChromaDB", { timeout: 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(); + 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(); + 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(); + 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(), }); @@ -28,52 +132,14 @@ describe("ChromaDB", { timeout: 360_000 }, () => { 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"); }); // } }); diff --git a/packages/modules/chromadb/src/chromadb-container.ts b/packages/modules/chromadb/src/chromadb-container.ts index 2c4d58c0f..06d625152 100755 --- a/packages/modules/chromadb/src/chromadb-container.ts +++ b/packages/modules/chromadb/src/chromadb-container.ts @@ -3,10 +3,10 @@ import { AbstractStartedContainer, GenericContainer, StartedTestContainer, Wait const CHROMADB_PORT = 8000; export class ChromaDBContainer extends GenericContainer { - constructor(image = "chromadb/chroma:0.4.24") { + constructor(image = "chromadb/chroma:0.6.3") { super(image); this.withExposedPorts(CHROMADB_PORT) - .withWaitStrategy(Wait.forHttp("/api/v1/heartbeat", CHROMADB_PORT)) + .withWaitStrategy(Wait.forHttp("/api/v2/heartbeat", CHROMADB_PORT)) .withStartupTimeout(120_000); }