Update ChromaDB + improve docs and tests (#940)

Author: cristianrgreco

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
+
+<!--codeinclude-->
+[Connect to Chroma:](../../packages/modules/chromadb/src/chromadb-container.test.ts)
+inside_block:simpleConnect
+<!--/codeinclude-->
+
+<!--codeinclude-->
+[Create Collection:](../../packages/modules/chromadb/src/chromadb-container.test.ts)
+inside_block:createCollection
+<!--/codeinclude-->
+
+<!--codeinclude-->
+[Query Collection with Embedding Function:](../../packages/modules/chromadb/src/chromadb-container.test.ts)
+inside_block:queryCollectionWithEmbeddingFunction
+<!--/codeinclude-->
+
+<!--codeinclude-->
+[Work with persistent directory:](../../packages/modules/chromadb/src/chromadb-container.test.ts)
+inside_block:persistentData
+<!--/codeinclude-->
 
 <!--codeinclude-->
-[](../../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
 <!--/codeinclude-->
 

package-lock.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"

packages/modules/chromadb/package.json

@@ -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");
   });
   // }
 });

packages/modules/chromadb/src/chromadb-container.test.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);
   }