MCPMem

import { v4 as uuidv4 } from 'uuid'; import type { Memory, MemoryInput, SearchOptions, SearchResult } from '../types/index.js'; import { DatabaseConnection } from './connection.js'; export class MemoryRepository { private db; constructor(dbPathOrConnection?: string | DatabaseConnection) { if (typeof dbPathOrConnection === 'string') { this.db = DatabaseConnection.getInstance(dbPathOrConnection).getDb(); } else { this.db = (dbPathOrConnection || DatabaseConnection.getInstance()).getDb(); } } async create(input: MemoryInput, embedding?: number[]): Promise<Memory> { const id = uuidv4(); const now = new Date(); const insertMemory = this.db.prepare(` INSERT INTO memories (id, content, metadata, created_at, updated_at) VALUES (?, ?, ?, ?, ?) `); insertMemory.run( id, input.content, input.metadata ? JSON.stringify(input.metadata) : null, now.toISOString(), now.toISOString(), ); if (embedding) { const insertEmbedding = this.db.prepare(` INSERT INTO memory_embeddings (memory_id, embedding) VALUES (?, ?) `); const vectorBuffer = Buffer.from(new Float32Array(embedding).buffer); insertEmbedding.run(id, vectorBuffer); } return this.findById(id) as Memory; } findById(id: string): Memory | null { const stmt = this.db.prepare(` SELECT id, content, metadata, created_at, updated_at FROM memories WHERE id = ? `); const row = stmt.get(id) as | { id: string; content: string; metadata: string | null; created_at: string; updated_at: string; } | undefined; if (!row) return null; return this.rowToMemory(row); } findAll(limit = 100): Memory[] { const stmt = this.db.prepare(` SELECT id, content, metadata, created_at, updated_at FROM memories ORDER BY created_at DESC LIMIT ? `); const rows = stmt.all(limit) as { id: string; content: string; metadata: string | null; created_at: string; updated_at: string; }[]; return rows.map(row => this.rowToMemory(row)); } async update( id: string, input: Partial<MemoryInput>, embedding?: number[], ): Promise<Memory | null> { const existing = this.findById(id); if (!existing) return null; const updates: string[] = []; const values: (string | number)[] = []; if (input.content !== undefined) { updates.push('content = ?'); values.push(input.content); } if (input.metadata !== undefined) { updates.push('metadata = ?'); values.push(input.metadata ? JSON.stringify(input.metadata) : ''); } if (updates.length > 0) { values.push(id); const stmt = this.db.prepare(` UPDATE memories SET ${updates.join(', ')} WHERE id = ? `); stmt.run(...values); } if (embedding) { const updateEmbedding = this.db.prepare(` INSERT OR REPLACE INTO memory_embeddings (memory_id, embedding) VALUES (?, ?) `); const vectorBuffer = Buffer.from(new Float32Array(embedding).buffer); updateEmbedding.run(id, vectorBuffer); } return this.findById(id); } delete(id: string): boolean { const deleteEmbedding = this.db.prepare( 'DELETE FROM memory_embeddings WHERE memory_id = ?', ); const deleteMemory = this.db.prepare('DELETE FROM memories WHERE id = ?'); deleteEmbedding.run(id); const result = deleteMemory.run(id); return result.changes > 0; } clearAll(): number { const deleteEmbeddings = this.db.prepare('DELETE FROM memory_embeddings'); const deleteMemories = this.db.prepare('DELETE FROM memories'); deleteEmbeddings.run(); const result = deleteMemories.run(); return result.changes; } searchSimilar(queryEmbedding: number[], options: SearchOptions = {}): SearchResult[] { const { limit = 10, threshold } = options; const stmt = this.db.prepare(` SELECT m.id, m.content, m.metadata, m.created_at, m.updated_at, distance FROM memory_embeddings me JOIN memories m ON m.id = me.memory_id WHERE me.embedding MATCH ? AND k = ? ORDER BY distance ASC `); const queryBuffer = Buffer.from(new Float32Array(queryEmbedding).buffer); const rows = stmt.all(queryBuffer, limit) as { id: string; content: string; metadata: string | null; created_at: string; updated_at: string; distance: number; }[]; const results: SearchResult[] = rows.map(row => { const distance: number = row.distance; const similarity = Math.max(0, 1 - distance / 2); return { memory: this.rowToMemory(row), similarity, }; }); if (typeof threshold === 'number') { return results.filter(r => r.similarity >= threshold); } return results; } count(): number { const stmt = this.db.prepare('SELECT COUNT(*) as count FROM memories'); const result = stmt.get() as { count: number }; return result.count; } findWithoutEmbeddings(limit = 100): { id: string; content: string }[] { const stmt = this.db.prepare(` SELECT m.id as id, m.content as content FROM memories m LEFT JOIN memory_embeddings me ON me.memory_id = m.id WHERE me.memory_id IS NULL ORDER BY m.created_at ASC LIMIT ? `); const rows = stmt.all(limit) as { id: string; content: string }[]; return rows; } upsertEmbedding(memoryId: string, embedding: number[]): void { const stmt = this.db.prepare(` INSERT OR REPLACE INTO memory_embeddings (memory_id, embedding) VALUES (?, ?) `); const vectorBuffer = Buffer.from(new Float32Array(embedding).buffer); stmt.run(memoryId, vectorBuffer); } countWithoutEmbeddings(): number { const stmt = this.db.prepare(` SELECT COUNT(1) as count FROM memories m LEFT JOIN memory_embeddings me ON me.memory_id = m.id WHERE me.memory_id IS NULL `); const row = stmt.get() as { count: number }; return row.count; } private rowToMemory(row: { id: string; content: string; metadata: string | null; created_at: string; updated_at: string; }): Memory { return { id: row.id, content: row.content, metadata: row.metadata ? JSON.parse(row.metadata) : undefined, created_at: new Date(row.created_at), updated_at: new Date(row.updated_at), }; } }