diff --git a/packages/memory/src/index.ts b/packages/memory/src/index.ts index a516ebc..863a0a5 100644 --- a/packages/memory/src/index.ts +++ b/packages/memory/src/index.ts @@ -1 +1,2 @@ export * from "./buffer.js"; +export * from "./ledger.js"; diff --git a/packages/memory/src/ledger.ts b/packages/memory/src/ledger.ts new file mode 100644 index 0000000..69e6d20 --- /dev/null +++ b/packages/memory/src/ledger.ts @@ -0,0 +1,228 @@ +import Database from "better-sqlite3"; + +export interface LedgerEntry { + id: string; + ownerId: string; + timestamp: string; + locationId: string | null; + involvedEntityIds: string[]; + + content: string; + quotes: string[]; + importance: number; + embedding: number[]; +} + +export class LedgerRepository { + constructor(private db: Database.Database) { + // Enable foreign keys for cascading deletes + this.db.exec("PRAGMA foreign_keys = ON;"); + this.initializeSchema(); + } + + private initializeSchema(): void { + this.db.exec(` + CREATE TABLE IF NOT EXISTS ledger_entries ( + id TEXT PRIMARY KEY, + owner_id TEXT NOT NULL, + timestamp TEXT NOT NULL, + location_id TEXT, + content TEXT NOT NULL, + quotes_json TEXT, + importance INTEGER NOT NULL, + embedding BLOB, + FOREIGN KEY (owner_id) REFERENCES objects(id) ON DELETE CASCADE + ); + + CREATE TABLE IF NOT EXISTS ledger_involved_entities ( + entry_id TEXT NOT NULL, + entity_id TEXT NOT NULL, + PRIMARY KEY (entry_id, entity_id), + FOREIGN KEY (entry_id) REFERENCES ledger_entries(id) ON DELETE CASCADE + ); + `); + } + + save(entry: LedgerEntry): void { + const insertEntry = this.db.prepare(` + INSERT INTO ledger_entries (id, owner_id, timestamp, location_id, content, quotes_json, importance, embedding) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + owner_id = excluded.owner_id, + timestamp = excluded.timestamp, + location_id = excluded.location_id, + content = excluded.content, + quotes_json = excluded.quotes_json, + importance = excluded.importance, + embedding = excluded.embedding + `); + + const insertEntity = this.db.prepare(` + INSERT OR IGNORE INTO ledger_involved_entities (entry_id, entity_id) + VALUES (?, ?) + `); + + const deleteEntities = this.db.prepare(` + DELETE FROM ledger_involved_entities WHERE entry_id = ? + `); + + this.db.transaction(() => { + insertEntry.run( + entry.id, + entry.ownerId, + entry.timestamp, + entry.locationId, + entry.content, + JSON.stringify(entry.quotes), + entry.importance, + entry.embedding.length > 0 + ? Buffer.from(new Float32Array(entry.embedding).buffer) + : null + ); + + deleteEntities.run(entry.id); + for (const entityId of entry.involvedEntityIds) { + insertEntity.run(entry.id, entityId); + } + })(); + } + + load(id: string): LedgerEntry | null { + const row = this.db + .prepare( + ` + SELECT id, owner_id, timestamp, location_id, content, quotes_json, importance, embedding + FROM ledger_entries + WHERE id = ? + ` + ) + .get(id) as any; + + if (!row) return null; + + const entitiesRows = this.db + .prepare( + ` + SELECT entity_id FROM ledger_involved_entities WHERE entry_id = ? + ` + ) + .all(id) as { entity_id: string }[]; + + let embedding: number[] = []; + if (row.embedding) { + const buffer = row.embedding as Buffer; + const floatArray = new Float32Array( + buffer.buffer, + buffer.byteOffset, + buffer.byteLength / Float32Array.BYTES_PER_ELEMENT + ); + embedding = Array.from(floatArray); + } + + return { + id: row.id, + ownerId: row.owner_id, + timestamp: row.timestamp, + locationId: row.location_id, + involvedEntityIds: entitiesRows.map((er) => er.entity_id), + content: row.content, + quotes: JSON.parse(row.quotes_json), + importance: row.importance, + embedding: embedding, + }; + } + + /** + * Retrieves relevant ledger entries using Phase 1: Deterministic Heuristic Filtering + * Filters by: + * 1. locationId matches current location + * 2. involvedEntityIds overlaps with current involved entities + * 3. importance >= 8 (high salience) + */ + getRelevant( + ownerId: string, + currentLocationId: string | null, + currentInvolvedEntityIds: string[], + limit: number = 20 + ): LedgerEntry[] { + let query = ` + SELECT DISTINCT le.id, le.owner_id, le.timestamp, le.location_id, le.content, le.quotes_json, le.importance, le.embedding + FROM ledger_entries le + LEFT JOIN ledger_involved_entities lie ON le.id = lie.entry_id + WHERE le.owner_id = ? + AND ( + le.importance >= 8 + `; + + const params: any[] = [ownerId]; + + if (currentLocationId) { + query += ` OR le.location_id = ?`; + params.push(currentLocationId); + } + + if (currentInvolvedEntityIds.length > 0) { + const placeholders = currentInvolvedEntityIds.map(() => "?").join(","); + query += ` OR lie.entity_id IN (${placeholders})`; + params.push(...currentInvolvedEntityIds); + } + + query += ` + ) + ORDER BY le.timestamp DESC + LIMIT ? + `; + params.push(limit); + + const rows = this.db.prepare(query).all(...params) as any[]; + + if (rows.length === 0) return []; + + const entryIds = rows.map((r) => r.id); + const placeholders = entryIds.map(() => "?").join(","); + const entitiesRows = this.db + .prepare( + ` + SELECT entry_id, entity_id FROM ledger_involved_entities + WHERE entry_id IN (${placeholders}) + ` + ) + .all(...entryIds) as { entry_id: string; entity_id: string }[]; + + const entitiesMap = new Map(); + for (const er of entitiesRows) { + if (!entitiesMap.has(er.entry_id)) { + entitiesMap.set(er.entry_id, []); + } + entitiesMap.get(er.entry_id)!.push(er.entity_id); + } + + return rows.map((row) => { + let embedding: number[] = []; + if (row.embedding) { + const buffer = row.embedding as Buffer; + const floatArray = new Float32Array( + buffer.buffer, + buffer.byteOffset, + buffer.byteLength / Float32Array.BYTES_PER_ELEMENT + ); + embedding = Array.from(floatArray); + } + return { + id: row.id, + ownerId: row.owner_id, + timestamp: row.timestamp, + locationId: row.location_id, + involvedEntityIds: entitiesMap.get(row.id) || [], + content: row.content, + quotes: JSON.parse(row.quotes_json), + importance: row.importance, + embedding: embedding, + }; + }); + } + + delete(id: string): void { + this.db.prepare(`DELETE FROM ledger_entries WHERE id = ?`).run(id); + } +} diff --git a/packages/memory/tests/ledger.test.ts b/packages/memory/tests/ledger.test.ts new file mode 100644 index 0000000..8f28b31 --- /dev/null +++ b/packages/memory/tests/ledger.test.ts @@ -0,0 +1,124 @@ +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import Database from "better-sqlite3"; +import { LedgerRepository, LedgerEntry } from "../src/ledger"; + +describe("LedgerRepository", () => { + let db: Database.Database; + let repo: LedgerRepository; + + beforeEach(() => { + db = new Database(":memory:"); + + // We need to create a dummy objects table to satisfy foreign keys + db.exec(` + CREATE TABLE objects ( + id TEXT PRIMARY KEY + ); + `); + + db.exec(` + INSERT INTO objects (id) VALUES ('alice'), ('bob'), ('charlie'); + `); + + repo = new LedgerRepository(db); + }); + + afterEach(() => { + db.close(); + }); + + it("should save and load a ledger entry", () => { + const entry: LedgerEntry = { + id: "mem1", + ownerId: "alice", + timestamp: new Date().toISOString(), + locationId: "loc1", + involvedEntityIds: ["bob", "charlie"], + content: "Alice met Bob and Charlie at the market.", + quotes: ["Hi guys!"], + importance: 5, + embedding: [0.1, 0.2, 0.3], + }; + + repo.save(entry); + + const loaded = repo.load("mem1"); + expect(loaded).toBeDefined(); + expect(loaded?.id).toBe("mem1"); + expect(loaded?.ownerId).toBe("alice"); + expect(loaded?.locationId).toBe("loc1"); + expect(loaded?.involvedEntityIds.sort()).toEqual(["bob", "charlie"].sort()); + expect(loaded?.content).toBe(entry.content); + expect(loaded?.quotes).toEqual(entry.quotes); + expect(loaded?.importance).toBe(5); + + // Check float precision + expect(loaded?.embedding[0]).toBeCloseTo(0.1); + expect(loaded?.embedding[1]).toBeCloseTo(0.2); + expect(loaded?.embedding[2]).toBeCloseTo(0.3); + }); + + it("should return null for non-existent entry", () => { + const loaded = repo.load("missing"); + expect(loaded).toBeNull(); + }); + + it("should retrieve relevant memories based on Phase 1 heuristics", () => { + repo.save({ + id: "mem_high_salience", + ownerId: "alice", + timestamp: "2024-01-01T10:00:00.000Z", + locationId: "loc2", + involvedEntityIds: [], + content: "Alice found a magical sword.", + quotes: [], + importance: 9, // high salience + embedding: [], + }); + + repo.save({ + id: "mem_location", + ownerId: "alice", + timestamp: "2024-01-02T10:00:00.000Z", + locationId: "loc1", // matches query + involvedEntityIds: [], + content: "Alice sat on a bench.", + quotes: [], + importance: 2, + embedding: [], + }); + + repo.save({ + id: "mem_social", + ownerId: "alice", + timestamp: "2024-01-03T10:00:00.000Z", + locationId: "loc2", + involvedEntityIds: ["bob"], // matches query + content: "Alice waved at Bob.", + quotes: [], + importance: 3, + embedding: [], + }); + + repo.save({ + id: "mem_irrelevant", + ownerId: "alice", + timestamp: "2024-01-04T10:00:00.000Z", + locationId: "loc3", + involvedEntityIds: ["charlie"], + content: "Alice sneezed.", + quotes: [], + importance: 2, + embedding: [], + }); + + const relevant = repo.getRelevant("alice", "loc1", ["bob"]); + + expect(relevant).toHaveLength(3); + const ids = relevant.map((r) => r.id); + expect(ids).toContain("mem_high_salience"); // due to importance >= 8 + expect(ids).toContain("mem_location"); // due to locationId + expect(ids).toContain("mem_social"); // due to involvedEntityIds + expect(ids).not.toContain("mem_irrelevant"); + }); +});