From cf15ebe27d696cd9abba680c78bb125b631efb11 Mon Sep 17 00:00:00 2001 From: Kyle D Date: Mon, 13 Oct 2025 14:40:50 -0700 Subject: [PATCH 1/2] feat: Add ChromaDB Memory Provider --- packages/memory/package.json | 5 + packages/memory/src/providers/chroma.ts | 184 ++++++++++++++++++++++++ 2 files changed, 189 insertions(+) create mode 100644 packages/memory/src/providers/chroma.ts diff --git a/packages/memory/package.json b/packages/memory/package.json index 6f1b0ca..fa1cb14 100644 --- a/packages/memory/package.json +++ b/packages/memory/package.json @@ -22,18 +22,23 @@ }, "peerDependencies": { "@upstash/redis": "^1.34.3", + "chromadb": "^3.0.17", "drizzle-orm": "^0.36.0" }, "peerDependenciesMeta": { "@upstash/redis": { "optional": true }, + "chromadb": { + "optional": true + }, "drizzle-orm": { "optional": true } }, "devDependencies": { "@upstash/redis": "^1.35.5", + "chromadb": "^3.0.17", "drizzle-orm": "^0.44.6", "tsup": "^8.5.0", "typescript": "^5.9.3" diff --git a/packages/memory/src/providers/chroma.ts b/packages/memory/src/providers/chroma.ts new file mode 100644 index 0000000..305c74d --- /dev/null +++ b/packages/memory/src/providers/chroma.ts @@ -0,0 +1,184 @@ +import type { ChromaClient, Collection, EmbeddingFunction } from "chromadb"; +import type { + ConversationMessage, + MemoryProvider, + MemoryScope, + WorkingMemory, +} from "../types.js"; + +type WorkingMemoryMetadata = { + scope: MemoryScope; + chatId?: string; + userId?: string; + updatedAt: number; +}; + +type MessageMetadata = { + chatId: string; + userId?: string; + role: "user" | "assistant" | "system"; + timestamp: number; +}; + +export interface ChromaMemoryProviderConfig { + wmCollectionName: string; + messagesPrefix: string; +} + +/** + * Chroma provider for memory storage + * + * Architecture: + * - All working memory stored in ONE collection (e.g., "working_memory") + * - Messages sharded by chatId - each chat gets its own collection (e.g., "messages_chat-123") + * + * @example + * ```ts + * import { CloudClient } from 'chromadb'; + * import { ChromaMemoryProvider } from '@ai-sdk-tools/memory'; + * + * const client = new CloudClient(); + * const provider = new ChromaMemoryProvider(client); + * ``` + */ +export class ChromaMemoryProvider implements MemoryProvider { + private wmCollectionName: string; + + constructor( + private client: ChromaClient, + private ef: EmbeddingFunction, + options: ChromaMemoryProviderConfig = { + wmCollectionName: "working_memory", + messagesPrefix: "messages", + }, + ) { + this.wmCollectionName = options.wmCollectionName; + this.messagesPrefix = options.messagesPrefix; + } + + private messagesPrefix: string; + + async getWorkingMemory(params: { + chatId?: string; + userId?: string; + scope: MemoryScope; + }): Promise { + const collection = await this.client.getOrCreateCollection({ + name: this.wmCollectionName, + embeddingFunction: this.ef, + }); + const id = this.getWMId(params.scope, params.chatId, params.userId); + + const result = await collection.get({ + ids: [id], + }); + + if (!result.documents.length || !result.documents[0]) return null; + + const metadata = result.metadatas[0] as WorkingMemoryMetadata | null; + return { + content: result.documents[0] as string, + updatedAt: new Date(metadata?.updatedAt || Date.now()), + }; + } + + async updateWorkingMemory(params: { + chatId?: string; + userId?: string; + scope: MemoryScope; + content: string; + }): Promise { + const collection = await this.client.getOrCreateCollection({ + name: this.wmCollectionName, + embeddingFunction: this.ef, + }); + const id = this.getWMId(params.scope, params.chatId, params.userId); + + const metadata: WorkingMemoryMetadata = { + scope: params.scope, + chatId: params.chatId, + userId: params.userId, + updatedAt: Date.now(), + }; + + await collection.upsert({ + ids: [id], + documents: [params.content], + metadatas: [metadata], + }); + } + + async saveMessage(message: ConversationMessage): Promise { + const collection = await this.getOrCreateMessagesCollection( + message.chatId, + ); + + const messageId = `${message.timestamp.getTime()}-${Math.random().toString(36).substring(7)}`; + + const metadata: MessageMetadata = { + chatId: message.chatId, + userId: message.userId, + role: message.role, + timestamp: message.timestamp.getTime(), + }; + + await collection.add({ + ids: [messageId], + documents: [message.content], + metadatas: [metadata], + }); + } + + async getMessages(params: { + chatId: string; + limit?: number; + }): Promise { + const collection = await this.getOrCreateMessagesCollection( + params.chatId, + ); + + const result = await collection.get({ + where: { chatId: params.chatId }, + }); + + if (!result.documents.length) return []; + + const messages = result.documents + .map((doc, i) => { + const metadata = result.metadatas[i] as MessageMetadata | null; + if (!metadata || !doc) return null; + + const message: ConversationMessage = { + chatId: metadata.chatId, + userId: metadata.userId, + role: metadata.role, + content: doc as string, + timestamp: new Date(metadata.timestamp), + }; + return message; + }) + .filter((msg): msg is ConversationMessage => msg !== null) + .sort((a, b) => a.timestamp.getTime() - b.timestamp.getTime()); + + return params.limit ? messages.slice(-params.limit) : messages; + } + + private async getOrCreateMessagesCollection( + chatId: string, + ): Promise { + const collectionName = `${this.messagesPrefix}_${chatId}`; + return await this.client.getOrCreateCollection({ + name: collectionName, + embeddingFunction: this.ef, + }); + } + + private getWMId( + scope: MemoryScope, + chatId?: string, + userId?: string, + ): string { + const id = scope === "chat" ? chatId : userId; + return `${scope}:${id}`; + } +} \ No newline at end of file From de0a3575822a9d4b9a7a451b8828b39adc5907a6 Mon Sep 17 00:00:00 2001 From: Kyle D Date: Mon, 13 Oct 2025 16:16:51 -0700 Subject: [PATCH 2/2] export ChromaMemoryProvider from module --- packages/memory/src/index.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/memory/src/index.ts b/packages/memory/src/index.ts index 45a996a..417a9bc 100644 --- a/packages/memory/src/index.ts +++ b/packages/memory/src/index.ts @@ -20,6 +20,7 @@ export { createSqliteWorkingMemorySchema, SQL_SCHEMAS, } from "./providers/drizzle-schema.js"; +export { ChromaMemoryProvider } from "./providers/chroma.js"; export { InMemoryProvider } from "./providers/in-memory.js"; export { UpstashProvider } from "./providers/upstash.js";