feat(ts-sdk): add Chroma vector store provider (#6145)
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
sahithreddy05
co-authored by
kartik-mem0
kartik-mem0
parent
f89edb45dc
commit
22f70d50e1
@@ -6,9 +6,8 @@ description: "Use Chroma as a vector database in Mem0 for local or cloud-hosted
|
||||
|
||||
### Usage
|
||||
|
||||
#### Local Installation
|
||||
|
||||
```python
|
||||
<CodeGroup>
|
||||
```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" } });
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
<Note>
|
||||
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.
|
||||
</Note>
|
||||
|
||||
### Config
|
||||
|
||||
Here are the parameters available for configuring Chroma:
|
||||
|
||||
<Tabs>
|
||||
<Tab title="Python">
|
||||
| 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` |
|
||||
| `tenant` | ChromaDB Cloud tenant ID (for cloud usage) | `None` |
|
||||
</Tab>
|
||||
<Tab title="TypeScript">
|
||||
| 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` |
|
||||
</Tab>
|
||||
</Tabs>
|
||||
|
||||
@@ -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",
|
||||
|
||||
Generated
+63
@@ -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: {}
|
||||
|
||||
@@ -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<string, any> = {}): ChromaDB {
|
||||
return new ChromaDB({
|
||||
collectionName: "test-collection",
|
||||
host: "localhost",
|
||||
port: 8000,
|
||||
...overrides,
|
||||
} as any);
|
||||
}
|
||||
|
||||
async function initDb(overrides: Record<string, any> = {}): Promise<ChromaDB> {
|
||||
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 } });
|
||||
});
|
||||
});
|
||||
@@ -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":
|
||||
|
||||
@@ -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<any>;
|
||||
private migrationsPromise?: Promise<any>;
|
||||
|
||||
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<string, any> = {};
|
||||
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<any> {
|
||||
if (!this.collectionPromise) {
|
||||
this.collectionPromise = this.client.getOrCreateCollection({
|
||||
name: this.collectionName,
|
||||
embeddingFunction: null,
|
||||
});
|
||||
}
|
||||
return this.collectionPromise;
|
||||
}
|
||||
|
||||
private async getMigrationsCollection(): Promise<any> {
|
||||
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<string, any>) || {},
|
||||
score,
|
||||
});
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
async insert(
|
||||
vectors: number[][],
|
||||
ids: string[],
|
||||
payloads: Record<string, any>[],
|
||||
): Promise<void> {
|
||||
const collection = await this.getCollection();
|
||||
await collection.add({
|
||||
ids,
|
||||
embeddings: vectors,
|
||||
metadatas: payloads as any,
|
||||
});
|
||||
}
|
||||
|
||||
async keywordSearch(): Promise<null> {
|
||||
return null;
|
||||
}
|
||||
|
||||
async search(
|
||||
query: number[],
|
||||
topK: number = 5,
|
||||
filters?: SearchFilters,
|
||||
): Promise<VectorStoreResult[]> {
|
||||
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<VectorStoreResult | null> {
|
||||
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<string, any>,
|
||||
): Promise<void> {
|
||||
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<void> {
|
||||
const collection = await this.getCollection();
|
||||
await collection.delete({ ids: [vectorId] });
|
||||
}
|
||||
|
||||
async deleteCol(): Promise<void> {
|
||||
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<string> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
await this.getCollection();
|
||||
await this.getMigrationsCollection();
|
||||
}
|
||||
|
||||
/** Convert a single field filter into a ChromaDB where condition. */
|
||||
private static convertCondition(
|
||||
key: string,
|
||||
value: any,
|
||||
): Record<string, any> | 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<string, string> = {
|
||||
eq: "$eq",
|
||||
ne: "$ne",
|
||||
gt: "$gt",
|
||||
gte: "$gte",
|
||||
lt: "$lt",
|
||||
lte: "$lte",
|
||||
in: "$in",
|
||||
nin: "$nin",
|
||||
};
|
||||
const condition: Record<string, any> = {};
|
||||
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<string, any> | undefined {
|
||||
if (!filters) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const negateOp: Record<string, string> = {
|
||||
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<string, any> = {};
|
||||
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 };
|
||||
}
|
||||
}
|
||||
@@ -34,6 +34,7 @@ const external = [
|
||||
"mongodb",
|
||||
"@opensearch-project/opensearch",
|
||||
"@elastic/elasticsearch",
|
||||
"chromadb",
|
||||
];
|
||||
|
||||
const define = {
|
||||
|
||||
Reference in New Issue
Block a user