From 22f70d50e156445eab9ac13449d0f22f5911dd37 Mon Sep 17 00:00:00 2001 From: sahithreddy05 <45757789+sahithreddy05@users.noreply.github.com> Date: Wed, 8 Jul 2026 06:44:05 -0700 Subject: [PATCH] feat(ts-sdk): add Chroma vector store provider (#6145) Co-authored-by: kartik-mem0 --- docs/components/vectordbs/dbs/chroma.mdx | 58 ++- mem0-ts/package.json | 1 + mem0-ts/pnpm-lock.yaml | 63 +++ mem0-ts/src/oss/src/tests/chroma.test.ts | 457 ++++++++++++++++++++ mem0-ts/src/oss/src/utils/factory.ts | 3 + mem0-ts/src/oss/src/vector_stores/chroma.ts | 370 ++++++++++++++++ mem0-ts/tsup.config.ts | 1 + 7 files changed, 949 insertions(+), 4 deletions(-) create mode 100644 mem0-ts/src/oss/src/tests/chroma.test.ts create mode 100644 mem0-ts/src/oss/src/vector_stores/chroma.ts diff --git a/docs/components/vectordbs/dbs/chroma.mdx b/docs/components/vectordbs/dbs/chroma.mdx index ef6c4c2ca..cedb9f3fa 100644 --- a/docs/components/vectordbs/dbs/chroma.mdx +++ b/docs/components/vectordbs/dbs/chroma.mdx @@ -6,9 +6,8 @@ description: "Use Chroma as a vector database in Mem0 for local or cloud-hosted ### Usage -#### Local Installation - -```python + +```python Python import os from mem0 import Memory @@ -37,10 +36,46 @@ messages = [ m.add(messages, user_id="alice", metadata={"category": "movies"}) ``` +```typescript TypeScript +import { Memory } from 'mem0ai/oss'; + +// The Node.js client connects to a running Chroma server. +// Start one locally with: chroma run --host localhost --port 8000 +const config = { + vectorStore: { + provider: 'chroma', + config: { + collectionName: 'memories', + host: 'localhost', + port: 8000, + // Optional: ChromaDB Cloud configuration + // apiKey: 'your-chroma-cloud-api-key', + // tenant: 'your-chroma-cloud-tenant-id', + }, + }, +}; + +const memory = new Memory(config); +const messages = [ + {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"}, + {"role": "assistant", "content": "How about thriller movies? They can be quite engaging."}, + {"role": "user", "content": "I’m not a big fan of thriller movies but I love sci-fi movies."}, + {"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."} +] +await memory.add(messages, { userId: "alice", metadata: { category: "movies" } }); +``` + + + +The Node.js SDK uses the `chromadb` v3 client, which talks to a Chroma server over HTTP (local server or ChromaDB Cloud). Install it with `npm install chromadb`. Mem0 supplies the embeddings, so the collection is created without an embedding function. + + ### Config Here are the parameters available for configuring Chroma: + + | Parameter | Description | Default Value | | --- | --- | --- | | `collection_name` | The name of the collection | `mem0` | @@ -49,4 +84,19 @@ Here are the parameters available for configuring Chroma: | `host` | The host where the Chroma server is running | `None` | | `port` | The port where the Chroma server is running | `None` | | `api_key` | ChromaDB Cloud API key (for cloud usage) | `None` | -| `tenant` | ChromaDB Cloud tenant ID (for cloud usage) | `None` | \ No newline at end of file +| `tenant` | ChromaDB Cloud tenant ID (for cloud usage) | `None` | + + +| Parameter | Description | Default Value | +| --- | --- | --- | +| `collectionName` | The name of the collection | `mem0` | +| `client` | Pre-configured `ChromaClient` or `CloudClient` instance | `None` | +| `host` | The host where the Chroma server is running | `None` | +| `port` | The port where the Chroma server is running | `None` | +| `ssl` | Whether to use SSL when connecting to the Chroma server | `false` | +| `path` | Full URL of a Chroma server, e.g. `http://localhost:8000` (alternative to `host` and `port`) | `None` | +| `apiKey` | ChromaDB Cloud API key (for cloud usage) | `None` | +| `tenant` | ChromaDB Cloud tenant ID (for cloud usage) | `None` | +| `database` | ChromaDB Cloud database name (for cloud usage) | `mem0` | + + diff --git a/mem0-ts/package.json b/mem0-ts/package.json index d75177f4d..461bb7f26 100644 --- a/mem0-ts/package.json +++ b/mem0-ts/package.json @@ -126,6 +126,7 @@ "@upstash/vector": "^1.2.3", "better-sqlite3": "^12.6.2", "cassandra-driver": "4.8.0", + "chromadb": "^3.5.0", "cloudflare": "^4.2.0", "fastembed": "^2.1.0", "groq-sdk": "0.3.0", diff --git a/mem0-ts/pnpm-lock.yaml b/mem0-ts/pnpm-lock.yaml index 4338db33b..8df3854b1 100644 --- a/mem0-ts/pnpm-lock.yaml +++ b/mem0-ts/pnpm-lock.yaml @@ -94,6 +94,9 @@ importers: cassandra-driver: specifier: 4.8.0 version: 4.8.0 + chromadb: + specifier: ^3.5.0 + version: 3.5.0 cloudflare: specifier: ^4.2.0 version: 4.5.0 @@ -1682,6 +1685,41 @@ packages: resolution: {integrity: sha512-+IxzY9BZOQd/XuYPRmrvEVjF/nqj5kgT4kEq7VofrDoM1MxoRjEWkrCC3EtLi59TVawxTAn+orJwFQcrqEN1+g==} engines: {node: '>=18'} + chromadb-js-bindings-darwin-arm64@1.3.4: + resolution: {integrity: sha512-pDdWCcR0hCz3pSDiIijBito7YG6FumewfzYE2mIjSwjrO1CKSfQKLwDdTVK4ZNpVGQa16ePoMX+pg7ShSUARJg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [darwin] + + chromadb-js-bindings-darwin-x64@1.3.4: + resolution: {integrity: sha512-5vKVFXSFo+S9qC9by/7oofjcXqQK4IGHiUnPrvTfc7B52gNlPSKeyDO1wha556COc3KtXSZjICEH0nYBJ8UXng==} + engines: {node: '>= 10'} + cpu: [x64] + os: [darwin] + + chromadb-js-bindings-linux-arm64-gnu@1.3.4: + resolution: {integrity: sha512-ELiqDZzU5mZd1BvZugGrhoMkVeNyQaa+PGizd06WIpIJVoHY+S+9ZBjdeM6ZkRnL3vTxD79XBrosxNWhzqy/kg==} + engines: {node: '>= 10'} + cpu: [arm64] + os: [linux] + + chromadb-js-bindings-linux-x64-gnu@1.3.4: + resolution: {integrity: sha512-HmTUe7PEIaBngCQ70Yyle81z2wYyg9OMgJYyRUNBYCBp4ED6zTmdaOro41Vs/c/h2UyQeBXFIKNmHDeD2Nr2fw==} + engines: {node: '>= 10'} + cpu: [x64] + os: [linux] + + chromadb-js-bindings-win32-x64-msvc@1.3.4: + resolution: {integrity: sha512-punrofQRzKspNMxHmv24qoPeRycV2T3KfedekGV4lOuPPG14i6qggS4x/OWLSAK7ozPfBbrzCejvCa5wfvujhQ==} + engines: {node: '>= 10'} + cpu: [x64] + os: [win32] + + chromadb@3.5.0: + resolution: {integrity: sha512-A68f2RQjY3R95Pzy0j3ykOu6fi+WWCN1tjbvzrZXZoYQ1d7ZjgWr1FyzitZCeaz/0qmc8y2LEK7QmlxloOAs9Q==} + engines: {node: '>=20'} + hasBin: true + ci-info@3.9.0: resolution: {integrity: sha512-NIxF55hv4nSqQswkAeiOi1r83xy8JldOFDTWiug55KBu9Jnblncd2U6ViHmYgHf01TPZS77NJBhBMKdWj9HQMQ==} engines: {node: '>=8'} @@ -5602,6 +5640,31 @@ snapshots: chownr@3.0.0: {} + chromadb-js-bindings-darwin-arm64@1.3.4: + optional: true + + chromadb-js-bindings-darwin-x64@1.3.4: + optional: true + + chromadb-js-bindings-linux-arm64-gnu@1.3.4: + optional: true + + chromadb-js-bindings-linux-x64-gnu@1.3.4: + optional: true + + chromadb-js-bindings-win32-x64-msvc@1.3.4: + optional: true + + chromadb@3.5.0: + dependencies: + semver: 7.8.4 + optionalDependencies: + chromadb-js-bindings-darwin-arm64: 1.3.4 + chromadb-js-bindings-darwin-x64: 1.3.4 + chromadb-js-bindings-linux-arm64-gnu: 1.3.4 + chromadb-js-bindings-linux-x64-gnu: 1.3.4 + chromadb-js-bindings-win32-x64-msvc: 1.3.4 + ci-info@3.9.0: {} cjs-module-lexer@1.4.3: {} diff --git a/mem0-ts/src/oss/src/tests/chroma.test.ts b/mem0-ts/src/oss/src/tests/chroma.test.ts new file mode 100644 index 000000000..65874d402 --- /dev/null +++ b/mem0-ts/src/oss/src/tests/chroma.test.ts @@ -0,0 +1,457 @@ +// jest.mock is hoisted above variable declarations, so shared mock functions +// are attached to the module-level `__mocks__` object populated inside the +// factory. The hoisted mock reaches them through this stable reference. + +const __mocks__: { + add: jest.Mock; + query: jest.Mock; + get: jest.Mock; + update: jest.Mock; + upsert: jest.Mock; + deleteFn: jest.Mock; + getOrCreateCollection: jest.Mock; + deleteCollection: jest.Mock; + ChromaClient: jest.Mock; + CloudClient: jest.Mock; +} = {} as any; + +jest.mock("chromadb", () => { + const add = jest.fn().mockResolvedValue(undefined); + const query = jest + .fn() + .mockResolvedValue({ ids: [[]], distances: [[]], metadatas: [[]] }); + const get = jest.fn().mockResolvedValue({ ids: [], metadatas: [] }); + const update = jest.fn().mockResolvedValue(undefined); + const upsert = jest.fn().mockResolvedValue(undefined); + const deleteFn = jest.fn().mockResolvedValue(undefined); + + const collectionHandle = { + add, + query, + get, + update, + upsert, + delete: deleteFn, + }; + const getOrCreateCollection = jest.fn().mockResolvedValue(collectionHandle); + const deleteCollection = jest.fn().mockResolvedValue(undefined); + + const clientImpl = () => ({ getOrCreateCollection, deleteCollection }); + const ChromaClient = jest.fn().mockImplementation(clientImpl); + const CloudClient = jest.fn().mockImplementation(clientImpl); + + Object.assign(__mocks__, { + add, + query, + get, + update, + upsert, + deleteFn, + getOrCreateCollection, + deleteCollection, + ChromaClient, + CloudClient, + }); + + return { ChromaClient, CloudClient }; +}); + +import { ChromaDB } from "../vector_stores/chroma"; +import { VectorStoreFactory } from "../utils/factory"; + +// --- Helpers --- + +function makeDb(overrides: Record = {}): ChromaDB { + return new ChromaDB({ + collectionName: "test-collection", + host: "localhost", + port: 8000, + ...overrides, + } as any); +} + +async function initDb(overrides: Record = {}): Promise { + const db = makeDb(overrides); + await db.initialize(); + return db; +} + +// --- Reset mocks between tests --- + +beforeEach(() => { + jest.clearAllMocks(); + + __mocks__.add.mockResolvedValue(undefined); + __mocks__.query.mockResolvedValue({ + ids: [[]], + distances: [[]], + metadatas: [[]], + }); + __mocks__.get.mockResolvedValue({ ids: [], metadatas: [] }); + __mocks__.update.mockResolvedValue(undefined); + __mocks__.upsert.mockResolvedValue(undefined); + __mocks__.deleteFn.mockResolvedValue(undefined); + + const collectionHandle = { + add: __mocks__.add, + query: __mocks__.query, + get: __mocks__.get, + update: __mocks__.update, + upsert: __mocks__.upsert, + delete: __mocks__.deleteFn, + }; + __mocks__.getOrCreateCollection.mockResolvedValue(collectionHandle); + __mocks__.deleteCollection.mockResolvedValue(undefined); + __mocks__.ChromaClient.mockImplementation(() => ({ + getOrCreateCollection: __mocks__.getOrCreateCollection, + deleteCollection: __mocks__.deleteCollection, + })); + __mocks__.CloudClient.mockImplementation(() => ({ + getOrCreateCollection: __mocks__.getOrCreateCollection, + deleteCollection: __mocks__.deleteCollection, + })); +}); + +// --- Test suites --- + +describe("VectorStoreFactory", () => { + it("returns a ChromaDB instance for provider 'chroma'", async () => { + const db = VectorStoreFactory.create("chroma", { + collectionName: "x", + host: "localhost", + port: 8000, + } as any); + expect(db).toBeInstanceOf(ChromaDB); + await (db as any).initialize(); + }); +}); + +describe("Constructor", () => { + it("builds a local ChromaClient with host and port", () => { + makeDb({ host: "localhost", port: 8000 }); + expect(__mocks__.ChromaClient).toHaveBeenCalledWith({ + host: "localhost", + port: 8000, + }); + expect(__mocks__.CloudClient).not.toHaveBeenCalled(); + }); + + it("passes ssl and path through to ChromaClient when provided", () => { + makeDb({ host: "example.com", port: 443, ssl: true, path: "/db" }); + expect(__mocks__.ChromaClient).toHaveBeenCalledWith({ + host: "example.com", + port: 443, + ssl: true, + path: "/db", + }); + }); + + it("builds a CloudClient when apiKey and tenant are set", () => { + new ChromaDB({ + collectionName: "test-collection", + apiKey: "key-123", + tenant: "tenant-abc", + } as any); + expect(__mocks__.CloudClient).toHaveBeenCalledWith({ + apiKey: "key-123", + tenant: "tenant-abc", + database: "mem0", + }); + expect(__mocks__.ChromaClient).not.toHaveBeenCalled(); + }); + + it("honors an explicit cloud database name", () => { + new ChromaDB({ + collectionName: "test-collection", + apiKey: "key-123", + tenant: "tenant-abc", + database: "custom-db", + } as any); + expect(__mocks__.CloudClient).toHaveBeenCalledWith( + expect.objectContaining({ database: "custom-db" }), + ); + }); + + it("accepts a pre-built client via config.client", async () => { + const fakeClient = { + getOrCreateCollection: __mocks__.getOrCreateCollection, + deleteCollection: __mocks__.deleteCollection, + }; + const db = new ChromaDB({ + collectionName: "test-collection", + client: fakeClient, + } as any); + await db.initialize(); + expect(__mocks__.ChromaClient).not.toHaveBeenCalled(); + expect(__mocks__.CloudClient).not.toHaveBeenCalled(); + expect(__mocks__.getOrCreateCollection).toHaveBeenCalled(); + }); +}); + +describe("initialize", () => { + it("creates the collection with embeddingFunction null (mem0 supplies embeddings)", async () => { + await initDb(); + expect(__mocks__.getOrCreateCollection).toHaveBeenCalledWith({ + name: "test-collection", + embeddingFunction: null, + }); + }); + + it("also creates the migrations collection", async () => { + await initDb(); + expect(__mocks__.getOrCreateCollection).toHaveBeenCalledWith({ + name: "memory_migrations", + embeddingFunction: null, + }); + }); +}); + +describe("insert", () => { + it("adds records with ids, embeddings, and metadatas", async () => { + const db = await initDb(); + await db.insert([[1, 2, 3]], ["id-1"], [{ text: "hello" }]); + expect(__mocks__.add).toHaveBeenCalledWith({ + ids: ["id-1"], + embeddings: [[1, 2, 3]], + metadatas: [{ text: "hello" }], + }); + }); +}); + +describe("search", () => { + it("queries with embeddings, nResults, and where", async () => { + const db = await initDb(); + await db.search([1, 2, 3], 10, { user_id: "alice" }); + expect(__mocks__.query).toHaveBeenCalledWith({ + queryEmbeddings: [[1, 2, 3]], + nResults: 10, + where: { user_id: { $eq: "alice" } }, + }); + }); + + it("maps a nested query response to VectorStoreResult with 1/(1+distance) scores", async () => { + __mocks__.query.mockResolvedValue({ + ids: [["a", "b"]], + distances: [[0, 1]], + metadatas: [[{ k: "v1" }, { k: "v2" }]], + }); + const db = await initDb(); + const results = await db.search([1, 2, 3]); + expect(results).toEqual([ + { id: "a", payload: { k: "v1" }, score: 1 }, + { id: "b", payload: { k: "v2" }, score: 0.5 }, + ]); + }); + + it("returns [] when the query yields no matches", async () => { + __mocks__.query.mockResolvedValue({ + ids: [[]], + distances: [[]], + metadatas: [[]], + }); + const db = await initDb(); + const results = await db.search([1, 2, 3]); + expect(results).toEqual([]); + }); + + it("passes where undefined when there are no filters", async () => { + const db = await initDb(); + await db.search([1, 2, 3], 5); + expect(__mocks__.query).toHaveBeenCalledWith( + expect.objectContaining({ where: undefined }), + ); + }); +}); + +describe("get", () => { + it("fetches by id and returns the first parsed result", async () => { + __mocks__.get.mockResolvedValue({ + ids: ["vec-1"], + metadatas: [{ text: "foo" }], + }); + const db = await initDb(); + const result = await db.get("vec-1"); + expect(__mocks__.get).toHaveBeenCalledWith({ ids: ["vec-1"] }); + expect(result).toEqual({ id: "vec-1", payload: { text: "foo" } }); + }); + + it("returns null when the id is not found", async () => { + __mocks__.get.mockResolvedValue({ ids: [], metadatas: [] }); + const db = await initDb(); + const result = await db.get("missing"); + expect(result).toBeNull(); + }); +}); + +describe("update", () => { + it("updates a single record with embedding and metadata", async () => { + const db = await initDb(); + await db.update("vec-1", [1, 2, 3], { text: "updated" }); + expect(__mocks__.update).toHaveBeenCalledWith({ + ids: ["vec-1"], + embeddings: [[1, 2, 3]], + metadatas: [{ text: "updated" }], + }); + }); +}); + +describe("delete", () => { + it("deletes by id", async () => { + const db = await initDb(); + await db.delete("vec-1"); + expect(__mocks__.deleteFn).toHaveBeenCalledWith({ ids: ["vec-1"] }); + }); +}); + +describe("deleteCol", () => { + it("deletes the collection and resets its cached handle", async () => { + const db = await initDb(); + await db.deleteCol(); + expect(__mocks__.deleteCollection).toHaveBeenCalledWith({ + name: "test-collection", + }); + // Cached collection promise is cleared, so the next call re-creates it. + __mocks__.getOrCreateCollection.mockClear(); + await db.search([1, 2, 3]); + expect(__mocks__.getOrCreateCollection).toHaveBeenCalledWith({ + name: "test-collection", + embeddingFunction: null, + }); + }); +}); + +describe("list", () => { + it("gets with where and limit and returns [results, count]", async () => { + __mocks__.get.mockResolvedValue({ + ids: ["a", "b"], + metadatas: [{ k: 1 }, { k: 2 }], + }); + const db = await initDb(); + const [results, count] = await db.list({ user_id: "alice" }, 50); + expect(__mocks__.get).toHaveBeenCalledWith({ + where: { user_id: { $eq: "alice" } }, + limit: 50, + }); + expect(results).toHaveLength(2); + expect(count).toBe(2); + }); +}); + +describe("getUserId", () => { + it("returns an existing user_id from the migrations collection", async () => { + __mocks__.get.mockResolvedValue({ + ids: ["mig-1"], + metadatas: [{ user_id: "u-123" }], + }); + const db = await initDb(); + const uid = await db.getUserId(); + expect(uid).toBe("u-123"); + expect(__mocks__.add).not.toHaveBeenCalled(); + }); + + it("generates and stores a new user_id when none exists", async () => { + __mocks__.get.mockResolvedValue({ ids: [], metadatas: [] }); + const db = await initDb(); + const uid = await db.getUserId(); + expect(typeof uid).toBe("string"); + expect(uid.length).toBeGreaterThan(0); + expect(__mocks__.add).toHaveBeenCalledWith( + expect.objectContaining({ + embeddings: [[0]], + metadatas: [{ user_id: uid }], + }), + ); + }); +}); + +describe("setUserId", () => { + it("upserts the user_id onto the existing migration marker", async () => { + __mocks__.get.mockResolvedValue({ + ids: ["mig-1"], + metadatas: [{ user_id: "old" }], + }); + const db = await initDb(); + await db.setUserId("u-456"); + expect(__mocks__.upsert).toHaveBeenCalledWith({ + ids: ["mig-1"], + embeddings: [[0]], + metadatas: [{ user_id: "u-456" }], + }); + }); +}); + +describe("keywordSearch", () => { + it("returns null (Chroma has no keyword search)", async () => { + const db = await initDb(); + const result = await db.keywordSearch(); + expect(result).toBeNull(); + }); +}); + +describe("generateWhereClause", () => { + it("returns undefined for undefined and empty filters", () => { + expect(ChromaDB.generateWhereClause(undefined)).toBeUndefined(); + expect(ChromaDB.generateWhereClause({})).toBeUndefined(); + }); + + it("converts a single equality filter", () => { + expect(ChromaDB.generateWhereClause({ user_id: "alice" })).toEqual({ + user_id: { $eq: "alice" }, + }); + }); + + it("combines multiple keys with $and", () => { + expect( + ChromaDB.generateWhereClause({ user_id: "alice", agent_id: "bot" }), + ).toEqual({ + $and: [{ user_id: { $eq: "alice" } }, { agent_id: { $eq: "bot" } }], + }); + }); + + it("converts an array value to $in", () => { + expect(ChromaDB.generateWhereClause({ tags: ["x", "y"] })).toEqual({ + tags: { $in: ["x", "y"] }, + }); + }); + + it("skips a wildcard filter", () => { + expect(ChromaDB.generateWhereClause({ user_id: "*" })).toBeUndefined(); + }); + + it("maps comparison operators", () => { + expect(ChromaDB.generateWhereClause({ age: { gte: 18 } })).toEqual({ + age: { $gte: 18 }, + }); + }); + + it("converts an $or block", () => { + expect( + ChromaDB.generateWhereClause({ $or: [{ a: "x" }, { b: "y" }] }), + ).toEqual({ $or: [{ a: { $eq: "x" } }, { b: { $eq: "y" } }] }); + }); + + it("collapses a single-branch $or", () => { + expect(ChromaDB.generateWhereClause({ $or: [{ a: "x" }] })).toEqual({ + a: { $eq: "x" }, + }); + }); + + it("negates a single-field $not condition", () => { + expect(ChromaDB.generateWhereClause({ $not: [{ a: "x" }] })).toEqual({ + a: { $ne: "x" }, + }); + }); + + it("applies De Morgan to a multi-field $not condition", () => { + // NOT(a=x AND b=y) is (a!=x) OR (b!=y) + expect( + ChromaDB.generateWhereClause({ $not: [{ a: "x", b: "y" }] }), + ).toEqual({ $or: [{ a: { $ne: "x" } }, { b: { $ne: "y" } }] }); + }); + + it("negates a $not comparison operator", () => { + expect( + ChromaDB.generateWhereClause({ $not: [{ age: { gt: 18 } }] }), + ).toEqual({ age: { $lte: 18 } }); + }); +}); diff --git a/mem0-ts/src/oss/src/utils/factory.ts b/mem0-ts/src/oss/src/utils/factory.ts index 2e2920834..8577f341a 100644 --- a/mem0-ts/src/oss/src/utils/factory.ts +++ b/mem0-ts/src/oss/src/utils/factory.ts @@ -18,6 +18,7 @@ import { Embedder } from "../embeddings/base"; import { LLM } from "../llms/base"; import { VectorStore } from "../vector_stores/base"; import { Qdrant } from "../vector_stores/qdrant"; +import { ChromaDB } from "../vector_stores/chroma"; import { VectorizeDB } from "../vector_stores/vectorize"; import { RedisDB } from "../vector_stores/redis"; import { ValkeyDB } from "../vector_stores/valkey"; @@ -133,6 +134,8 @@ export class VectorStoreFactory { return new MemoryVectorStore(config); case "qdrant": return new Qdrant(config as any); + case "chroma": + return new ChromaDB(config as any); case "redis": return new RedisDB(config as any); case "valkey": diff --git a/mem0-ts/src/oss/src/vector_stores/chroma.ts b/mem0-ts/src/oss/src/vector_stores/chroma.ts new file mode 100644 index 000000000..076e8fda9 --- /dev/null +++ b/mem0-ts/src/oss/src/vector_stores/chroma.ts @@ -0,0 +1,370 @@ +import { ChromaClient, CloudClient } from "chromadb"; +import { VectorStore } from "./base"; +import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types"; + +interface ChromaConfig extends VectorStoreConfig { + /** Pre-configured ChromaDB client instance. */ + client?: ChromaClient | CloudClient; + collectionName: string; + /** Host address for a ChromaDB server (defaults to the client default). */ + host?: string; + /** Port for a ChromaDB server. */ + port?: number; + /** Whether to use SSL when connecting to a ChromaDB server. */ + ssl?: boolean; + /** Path for a local ChromaDB server. */ + path?: string; + /** ChromaDB Cloud API key. */ + apiKey?: string; + /** ChromaDB Cloud tenant ID. */ + tenant?: string; + /** ChromaDB Cloud database name. */ + database?: string; +} + +const MIGRATIONS_COLLECTION = "memory_migrations"; + +/** + * ChromaDB vector store provider. + * + * Mirrors the Python SDK's `mem0.vector_stores.chroma.ChromaDB` behavior using + * the `chromadb` v3 JavaScript client. Embeddings are always supplied by mem0, + * so no embedding function is required on the collection. + */ +export class ChromaDB implements VectorStore { + private client: any; + private readonly collectionName: string; + private collectionPromise?: Promise; + private migrationsPromise?: Promise; + + constructor(config: ChromaConfig) { + if (config.client) { + this.client = config.client; + } else if (config.apiKey && config.tenant) { + this.client = new CloudClient({ + apiKey: config.apiKey, + tenant: config.tenant, + database: config.database || "mem0", + } as any); + } else { + const params: Record = {}; + if (config.host) params.host = config.host; + if (config.port) params.port = config.port; + if (config.ssl !== undefined) params.ssl = config.ssl; + if (config.path) params.path = config.path; + this.client = new ChromaClient(params as any); + } + + this.collectionName = config.collectionName; + this.initialize().catch(console.error); + } + + private async getCollection(): Promise { + if (!this.collectionPromise) { + this.collectionPromise = this.client.getOrCreateCollection({ + name: this.collectionName, + embeddingFunction: null, + }); + } + return this.collectionPromise; + } + + private async getMigrationsCollection(): Promise { + if (!this.migrationsPromise) { + this.migrationsPromise = this.client.getOrCreateCollection({ + name: MIGRATIONS_COLLECTION, + embeddingFunction: null, + }); + } + return this.migrationsPromise; + } + + private flatten(value: any): any[] { + if (Array.isArray(value) && value.length > 0 && Array.isArray(value[0])) { + return value[0]; + } + return Array.isArray(value) ? value : []; + } + + /** Parse a ChromaDB `get`/`query` response into VectorStoreResult objects. */ + private parseOutput(data: any): VectorStoreResult[] { + const ids = this.flatten(data?.ids); + const distances = this.flatten(data?.distances); + const metadatas = this.flatten(data?.metadatas); + + const length = Math.max(ids.length, metadatas.length); + const results: VectorStoreResult[] = []; + + for (let i = 0; i < length; i++) { + const rawDistance = distances[i]; + const score = + rawDistance !== undefined && rawDistance !== null + ? 1.0 / (1.0 + rawDistance) + : undefined; + + results.push({ + id: String(ids[i]), + payload: (metadatas[i] as Record) || {}, + score, + }); + } + + return results; + } + + async insert( + vectors: number[][], + ids: string[], + payloads: Record[], + ): Promise { + const collection = await this.getCollection(); + await collection.add({ + ids, + embeddings: vectors, + metadatas: payloads as any, + }); + } + + async keywordSearch(): Promise { + return null; + } + + async search( + query: number[], + topK: number = 5, + filters?: SearchFilters, + ): Promise { + const collection = await this.getCollection(); + const where = ChromaDB.generateWhereClause(filters); + const results = await collection.query({ + queryEmbeddings: [query], + nResults: topK, + where: where as any, + }); + return this.parseOutput(results); + } + + async get(vectorId: string): Promise { + const collection = await this.getCollection(); + const results = await collection.get({ ids: [vectorId] }); + const parsed = this.parseOutput(results); + return parsed.length > 0 ? parsed[0] : null; + } + + async update( + vectorId: string, + vector: number[], + payload: Record, + ): Promise { + const collection = await this.getCollection(); + await collection.update({ + ids: [vectorId], + embeddings: vector ? [vector] : undefined, + metadatas: payload ? [payload] : undefined, + } as any); + } + + async delete(vectorId: string): Promise { + const collection = await this.getCollection(); + await collection.delete({ ids: [vectorId] }); + } + + async deleteCol(): Promise { + await this.client.deleteCollection({ name: this.collectionName }); + this.collectionPromise = undefined; + } + + async list( + filters?: SearchFilters, + topK: number = 100, + ): Promise<[VectorStoreResult[], number]> { + const collection = await this.getCollection(); + const where = ChromaDB.generateWhereClause(filters); + const results = await collection.get({ where: where as any, limit: topK }); + const parsed = this.parseOutput(results); + return [parsed, parsed.length]; + } + + private generateUUID(): string { + return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace( + /[xy]/g, + function (c) { + const r = (Math.random() * 16) | 0; + const v = c === "x" ? r : (r & 0x3) | 0x8; + return v.toString(16); + }, + ); + } + + async getUserId(): Promise { + const collection = await this.getMigrationsCollection(); + const result = await collection.get({ limit: 1 }); + const ids = Array.isArray(result?.ids) ? result.ids : []; + const metadatas = Array.isArray(result?.metadatas) ? result.metadatas : []; + + if (ids.length > 0 && metadatas[0]?.user_id) { + return String(metadatas[0].user_id); + } + + const randomUserId = + Math.random().toString(36).substring(2, 15) + + Math.random().toString(36).substring(2, 15); + + await collection.add({ + ids: [this.generateUUID()], + embeddings: [[0]], + metadatas: [{ user_id: randomUserId }] as any, + }); + + return randomUserId; + } + + async setUserId(userId: string): Promise { + const collection = await this.getMigrationsCollection(); + const result = await collection.get({ limit: 1 }); + const ids = Array.isArray(result?.ids) ? result.ids : []; + const pointId = ids.length > 0 ? String(ids[0]) : this.generateUUID(); + + await collection.upsert({ + ids: [pointId], + embeddings: [[0]], + metadatas: [{ user_id: userId }] as any, + }); + } + + async initialize(): Promise { + await this.getCollection(); + await this.getMigrationsCollection(); + } + + /** Convert a single field filter into a ChromaDB where condition. */ + private static convertCondition( + key: string, + value: any, + ): Record | null { + // Wildcard - ChromaDB has no direct wildcard, so skip this filter. + if (value === "*") { + return null; + } + + if (Array.isArray(value)) { + return { [key]: { $in: value } }; + } + + if (value !== null && typeof value === "object") { + const opMap: Record = { + eq: "$eq", + ne: "$ne", + gt: "$gt", + gte: "$gte", + lt: "$lt", + lte: "$lte", + in: "$in", + nin: "$nin", + }; + const condition: Record = {}; + for (const [op, val] of Object.entries(value)) { + if (op in opMap) { + condition[key] = { [opMap[op]]: val }; + } else { + // contains/icontains and unknown operators fall back to equality. + condition[key] = { $eq: val }; + } + } + return condition; + } + + return { [key]: { $eq: value } }; + } + + /** + * Generate a properly formatted `where` clause for ChromaDB from mem0's + * universal filter format. Supports comparison operators plus $or/$not. + */ + static generateWhereClause( + filters?: SearchFilters, + ): Record | undefined { + if (!filters) { + return undefined; + } + + const negateOp: Record = { + eq: "ne", + ne: "eq", + gt: "lte", + gte: "lt", + lt: "gte", + lte: "gt", + in: "nin", + nin: "in", + }; + + const processed: any[] = []; + + for (const [key, value] of Object.entries(filters)) { + if (key === "$or" || key === "OR") { + const orConditions: any[] = []; + for (const condition of value as any[]) { + const built: Record = {}; + for (const [subKey, subValue] of Object.entries(condition)) { + const converted = ChromaDB.convertCondition(subKey, subValue); + if (converted) Object.assign(built, converted); + } + if (Object.keys(built).length > 0) orConditions.push(built); + } + if (orConditions.length > 1) { + processed.push({ $or: orConditions }); + } else if (orConditions.length === 1) { + processed.push(orConditions[0]); + } + } else if (key === "$not" || key === "NOT") { + // De Morgan: NOT(a AND b) is (NOT a) OR (NOT b), so the negated fields + // within one condition are combined with $or, and separate conditions + // are combined with $and. This mirrors the Python SDK's ChromaDB port. + const negatedPerGroup: any[] = []; + for (const condition of value as any[]) { + const negatedFields: any[] = []; + for (const [subKey, subValue] of Object.entries(condition)) { + if (subValue !== null && typeof subValue === "object") { + for (const [op, val] of Object.entries(subValue as any)) { + const neg = negateOp[op]; + if (neg) { + const converted = ChromaDB.convertCondition(subKey, { + [neg]: val, + }); + if (converted) negatedFields.push(converted); + } + } + } else { + const converted = ChromaDB.convertCondition(subKey, { + ne: subValue, + }); + if (converted) negatedFields.push(converted); + } + } + if (negatedFields.length > 1) { + negatedPerGroup.push({ $or: negatedFields }); + } else if (negatedFields.length === 1) { + negatedPerGroup.push(negatedFields[0]); + } + } + if (negatedPerGroup.length > 1) { + processed.push({ $and: negatedPerGroup }); + } else if (negatedPerGroup.length === 1) { + processed.push(negatedPerGroup[0]); + } + } else { + const converted = ChromaDB.convertCondition(key, value); + if (converted) processed.push(converted); + } + } + + if (processed.length === 0) { + return undefined; + } + if (processed.length === 1) { + return processed[0]; + } + return { $and: processed }; + } +} diff --git a/mem0-ts/tsup.config.ts b/mem0-ts/tsup.config.ts index b0f91aac6..2109c740c 100644 --- a/mem0-ts/tsup.config.ts +++ b/mem0-ts/tsup.config.ts @@ -34,6 +34,7 @@ const external = [ "mongodb", "@opensearch-project/opensearch", "@elastic/elasticsearch", + "chromadb", ]; const define = {