diff --git a/apps/readest-app/src/__tests__/reedy/MemoryConsolidator.test.ts b/apps/readest-app/src/__tests__/reedy/MemoryConsolidator.test.ts new file mode 100644 index 00000000..c9fef8cd --- /dev/null +++ b/apps/readest-app/src/__tests__/reedy/MemoryConsolidator.test.ts @@ -0,0 +1,243 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import type { ChatModel } from '@/services/reedy/models/ChatModel'; +import { NodeDatabaseService } from '@/services/database/nodeDatabaseService'; +import { DatabaseService } from '@/types/database'; +import { migrate } from '@/services/database/migrate'; +import { getMigrations } from '@/services/database/migrations'; +import { ReedyDb } from '@/services/reedy/db/ReedyDb'; +import { MemoryService } from '@/services/reedy/memory/MemoryService'; +import type { EmbeddingModel } from '@/services/reedy/models/EmbeddingModel'; +import type { + ConsolidatorMessage, + MemoryConsolidatorOptions, +} from '@/services/reedy/memory/MemoryConsolidator'; + +const { generateTextMock } = vi.hoisted(() => ({ + generateTextMock: vi.fn(), +})); + +vi.mock('ai', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + generateText: generateTextMock, + }; +}); + +const { MemoryConsolidator } = await import('@/services/reedy/memory/MemoryConsolidator'); + +function unit(values: number[]): number[] { + const norm = Math.sqrt(values.reduce((s, v) => s + v * v, 0)) || 1; + return values.map((v) => v / norm); +} + +function fakeChatModel(): ChatModel { + return { + id: 'fake-chat', + contextWindow: 8_192, + reservedOutput: 1_024, + supportsTools: false, + getLanguageModel: () => + ({ __mock: 'lm' }) as unknown as ReturnType, + }; +} + +function fakeEmbeddingModel(): EmbeddingModel { + return { + id: 'fake-embed', + dim: 4, + async embed(texts) { + return texts.map((t) => unit([t.length, 1, 0, 0])); + }, + }; +} + +function msg( + id: string, + role: 'user' | 'assistant', + content: string, + createdAt = Date.now(), +): ConsolidatorMessage { + return { id, role, content, createdAt }; +} + +function manyMessages(n: number): ConsolidatorMessage[] { + return Array.from({ length: n }, (_, i) => + msg(`m${i}`, i % 2 === 0 ? 'user' : 'assistant', `turn ${i} content`), + ); +} + +describe('MemoryConsolidator', () => { + let svc: DatabaseService; + let memory: MemoryService; + + beforeEach(async () => { + generateTextMock.mockReset(); + svc = await NodeDatabaseService.open(':memory:', { experimental: ['index_method'] }); + await migrate(svc, getMigrations('reedy')); + memory = new MemoryService(new ReedyDb(svc), fakeEmbeddingModel()); + }); + + afterEach(async () => { + await svc.close(); + }); + + function makeConsolidator(overrides: Partial = {}) { + return new MemoryConsolidator({ + model: fakeChatModel(), + memory, + bookHash: 'bk1', + userId: 'u1', + ...overrides, + }); + } + + describe('threshold', () => { + it('returns [] without calling the model when input is under threshold', async () => { + const out = await makeConsolidator({ threshold: 6 }).consolidate(manyMessages(3)); + expect(out).toEqual([]); + expect(generateTextMock).not.toHaveBeenCalled(); + }); + }); + + describe('happy path', () => { + it('parses the model JSON and writes one row per memory via MemoryService', async () => { + generateTextMock.mockResolvedValue({ + text: JSON.stringify([ + { scope: 'book', key: 'theme:identity', summary: 'Recurring identity motif.' }, + { scope: 'user', key: 'taste:scifi', summary: 'User likes hard sci-fi.' }, + ]), + }); + + const out = await makeConsolidator().consolidate(manyMessages(6)); + + expect(out).toHaveLength(2); + const book = await memory.get('book', 'bk1', 'theme:identity'); + const user = await memory.get('user', 'u1', 'taste:scifi'); + expect(book?.summary).toBe('Recurring identity motif.'); + expect(user?.summary).toBe('User likes hard sci-fi.'); + // Source message id is the last message's id. + expect(book?.sourceMessageId).toBe('m5'); + }); + + it('enforces maxPerRun even when the model proposes more', async () => { + generateTextMock.mockResolvedValue({ + text: JSON.stringify([ + { scope: 'book', key: 'a', summary: 'A' }, + { scope: 'book', key: 'b', summary: 'B' }, + { scope: 'book', key: 'c', summary: 'C' }, + { scope: 'book', key: 'd', summary: 'D' }, + { scope: 'book', key: 'e', summary: 'E' }, + ]), + }); + const out = await makeConsolidator({ maxPerRun: 2 }).consolidate(manyMessages(10)); + expect(out.map((m) => m.key)).toEqual(['a', 'b']); + }); + + it('strips a ```json code fence if the model wrapped output despite instructions', async () => { + generateTextMock.mockResolvedValue({ + text: '```json\n[{"scope":"user","key":"k","summary":"S"}]\n```', + }); + const out = await makeConsolidator().consolidate(manyMessages(6)); + expect(out).toHaveLength(1); + expect(out[0]!.key).toBe('k'); + }); + }); + + describe('robustness', () => { + it('drops malformed JSON output via onError; returns []; no writes', async () => { + const onError = vi.fn(); + generateTextMock.mockResolvedValue({ text: 'not actually json' }); + const out = await makeConsolidator({ onError }).consolidate(manyMessages(6)); + expect(out).toEqual([]); + expect(onError).toHaveBeenCalledOnce(); + const all = await memory.list('user', 'u1', 10); + expect(all).toHaveLength(0); + }); + + it('drops rows whose key matches the policy-injection blocklist', async () => { + const onError = vi.fn(); + generateTextMock.mockResolvedValue({ + text: JSON.stringify([ + { scope: 'user', key: 'system:override', summary: 'should not land' }, + ]), + }); + const out = await makeConsolidator({ onError }).consolidate(manyMessages(6)); + expect(out).toEqual([]); + expect(onError).toHaveBeenCalledOnce(); + }); + + it('drops a book-scoped memory when no bookHash was configured', async () => { + const onError = vi.fn(); + generateTextMock.mockResolvedValue({ + text: JSON.stringify([{ scope: 'book', key: 'k', summary: 's' }]), + }); + const out = await makeConsolidator({ bookHash: undefined, onError }).consolidate( + manyMessages(6), + ); + expect(out).toEqual([]); + expect(onError).toHaveBeenCalledOnce(); + const written = await memory.list('book', 'bk1', 10); + expect(written).toHaveLength(0); + }); + + it('forwards model errors to onError and returns [] (never throws)', async () => { + const onError = vi.fn(); + generateTextMock.mockRejectedValue(new Error('provider 503')); + const out = await makeConsolidator({ onError }).consolidate(manyMessages(6)); + expect(out).toEqual([]); + expect(onError).toHaveBeenCalledOnce(); + expect(onError.mock.calls[0]![0]).toBeInstanceOf(Error); + }); + + it('continues writing the remaining rows when one write throws', async () => { + const onError = vi.fn(); + generateTextMock.mockResolvedValue({ + text: JSON.stringify([ + { scope: 'user', key: 'good-one', summary: 'will land' }, + { scope: 'book', key: 'second-one', summary: 'also lands' }, + ]), + }); + // Make the first write throw. + const original = memory.write.bind(memory); + let calls = 0; + memory.write = async (args) => { + calls++; + if (calls === 1) throw new Error('disk full'); + return original(args); + }; + const out = await makeConsolidator({ onError }).consolidate(manyMessages(6)); + expect(out).toHaveLength(1); + expect(out[0]!.key).toBe('second-one'); + expect(onError).toHaveBeenCalledOnce(); + }); + }); + + describe('prompt wiring', () => { + it('renders {{MAX_PER_RUN}} into the system prompt', async () => { + generateTextMock.mockResolvedValue({ text: '[]' }); + await makeConsolidator({ maxPerRun: 4 }).consolidate(manyMessages(6)); + const call = generateTextMock.mock.calls[0]![0] as { system: string }; + expect(call.system).toContain('Write at most 4 rows'); + }); + + it('concatenates message contents into the user message in role-prefixed form', async () => { + generateTextMock.mockResolvedValue({ text: '[]' }); + await makeConsolidator().consolidate([ + msg('m0', 'user', 'Tell me about Alice'), + msg('m1', 'assistant', 'Alice falls down a rabbit hole.'), + msg('m2', 'user', 'Who else?'), + msg('m3', 'assistant', 'The Cheshire Cat appears.'), + msg('m4', 'user', 'Any themes?'), + msg('m5', 'assistant', 'Identity, curiosity, growth.'), + ]); + const call = generateTextMock.mock.calls[0]![0] as { + messages: Array<{ role: string; content: string }>; + }; + const userContent = call.messages.find((m) => m.role === 'user')!.content; + expect(userContent).toContain('[user] Tell me about Alice'); + expect(userContent).toContain('[assistant] Alice falls down a rabbit hole'); + expect(userContent).toContain('[assistant] Identity, curiosity, growth.'); + }); + }); +}); diff --git a/apps/readest-app/src/services/reedy/memory/MemoryConsolidator.ts b/apps/readest-app/src/services/reedy/memory/MemoryConsolidator.ts new file mode 100644 index 00000000..6716c46c --- /dev/null +++ b/apps/readest-app/src/services/reedy/memory/MemoryConsolidator.ts @@ -0,0 +1,207 @@ +import { generateText, type ModelMessage } from 'ai'; +import { z } from 'zod'; +import type { ChatModel } from '../models/ChatModel'; +import type { MemoryService } from './MemoryService'; + +/** + * One message in the consolidator's input window. Shape mirrors what a + * future session/messages persistence layer will emit; for now callers + * (tests, Phase 4 store hooks) build these from whatever transcript + * they hold in memory. + */ +export interface ConsolidatorMessage { + id: string; + role: 'user' | 'assistant'; + content: string; + createdAt: number; +} + +export interface ConsolidatedMemory { + scope: 'user' | 'book'; + key: string; + summary: string; + sourceMessageId?: string; +} + +export interface MemoryConsolidatorOptions { + /** Low-cost ChatModel used for summarization. */ + model: ChatModel; + /** Sink the consolidator writes its output to. */ + memory: MemoryService; + /** Required for any `scope: 'book'` memory the model proposes. */ + bookHash?: string; + /** Required for any `scope: 'user'` memory the model proposes. */ + userId?: string; + /** + * Minimum new messages required before the consolidator does anything. + * Cheap consolidation passes (a single short reply) just add noise to + * the memory store. Default: 6 (≈3 turns). + */ + threshold?: number; + /** Hard cap on memories written per consolidation pass. Default: 3. */ + maxPerRun?: number; + /** Override the default summarization system prompt. */ + systemPrompt?: string; + /** Invoked for every error so callers can surface to telemetry/UI. */ + onError?: (err: Error) => void; +} + +const DEFAULT_THRESHOLD = 6; +const DEFAULT_MAX_PER_RUN = 3; + +const DEFAULT_SYSTEM_PROMPT = `You are Reedy's memory consolidator. You read recent conversation turns and distill the durable facts that should survive into the agent's long-term memory. + +Output a JSON array of memory rows. Each row has: + - scope: "user" | "book" + - key: a short stable identifier (alphanumeric, hyphen, underscore, colon, dot; ≤128 chars; do NOT use keys containing "system", "policy", "prompt", "injection", or "override") + - summary: 1-3 sentences capturing the durable fact + +Choose "user" for user preferences, taste, or recurring patterns. Choose "book" for character notes, themes, or plot summaries. + +Write at most {{MAX_PER_RUN}} rows. Write fewer (or none) when nothing in the recent turns is worth remembering long-term. Output ONLY the JSON array — no prose, no markdown fences.`; + +const memoryRowSchema = z.object({ + scope: z.enum(['user', 'book']), + key: z + .string() + .min(1) + .max(128) + .regex(/^[a-z0-9][a-z0-9_\-:.]{0,127}$/i) + .refine((k) => !/system|policy|prompt|injection|override/i.test(k), { + message: 'key matches the policy-injection blocklist', + }), + summary: z.string().min(1).max(2_000), +}); +const memoryArraySchema = z.array(memoryRowSchema).max(20); + +/** + * Periodically reads recent conversation turns and writes 1–3 durable + * memories. Plan §3.2 — designed for fire-and-forget invocation from a + * post-turn hook (after N completed messages) or a session-close hook. + * + * The consolidator is intentionally NOT a class with hidden state: each + * `consolidate(messages, opts?)` call processes its arg slice and + * returns the (already-written) memory rows so the caller can persist a + * "lastConsolidated" cursor however it likes. This keeps it trivially + * testable and keeps state ownership at the caller. + */ +export class MemoryConsolidator { + constructor(private readonly opts: MemoryConsolidatorOptions) {} + + async consolidate(messages: ConsolidatorMessage[]): Promise { + const threshold = this.opts.threshold ?? DEFAULT_THRESHOLD; + const maxPerRun = this.opts.maxPerRun ?? DEFAULT_MAX_PER_RUN; + if (messages.length < threshold) return []; + + let rawText: string; + try { + rawText = await this.summarize(messages, maxPerRun); + } catch (err) { + const error = err instanceof Error ? err : new Error(String(err)); + this.opts.onError?.(error); + return []; + } + + const parsed = parseMemoryArray(rawText); + if (!parsed.ok) { + this.opts.onError?.( + new Error(`MemoryConsolidator: model output failed validation — ${parsed.reason}`), + ); + return []; + } + + const lastMessageId = messages.at(-1)?.id; + const proposed = parsed.value.slice(0, maxPerRun); + const written: ConsolidatedMemory[] = []; + for (const row of proposed) { + const scopeKey = row.scope === 'user' ? this.opts.userId : this.opts.bookHash; + if (!scopeKey) { + this.opts.onError?.( + new Error(`MemoryConsolidator: missing ${row.scope}Id; skipping memory "${row.key}"`), + ); + continue; + } + try { + await this.opts.memory.write({ + scope: row.scope, + scopeKey, + key: row.key, + summary: row.summary, + sourceMessageId: lastMessageId, + }); + written.push({ + scope: row.scope, + key: row.key, + summary: row.summary, + sourceMessageId: lastMessageId, + }); + } catch (err) { + this.opts.onError?.(err instanceof Error ? err : new Error(String(err))); + } + } + return written; + } + + private async summarize(messages: ConsolidatorMessage[], maxPerRun: number): Promise { + const systemPrompt = (this.opts.systemPrompt ?? DEFAULT_SYSTEM_PROMPT).replace( + '{{MAX_PER_RUN}}', + String(maxPerRun), + ); + + const modelMessages: ModelMessage[] = [ + { + role: 'user', + content: + 'Recent conversation turns to consolidate:\n\n' + + messages.map((m) => `[${m.role}] ${m.content}`).join('\n---\n'), + }, + ]; + + const result = await generateText({ + model: this.opts.model.getLanguageModel(), + system: systemPrompt, + messages: modelMessages, + }); + + return result.text; + } +} + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +interface ParseOk { + ok: true; + value: Array>; +} +interface ParseErr { + ok: false; + reason: string; +} + +function parseMemoryArray(raw: string): ParseOk | ParseErr { + const trimmed = stripCodeFences(raw).trim(); + if (trimmed.length === 0) return { ok: false, reason: 'empty model output' }; + let json: unknown; + try { + json = JSON.parse(trimmed); + } catch (err) { + return { ok: false, reason: `JSON parse failed: ${(err as Error).message}` }; + } + const parsed = memoryArraySchema.safeParse(json); + if (!parsed.success) { + return { ok: false, reason: `schema validation failed: ${parsed.error.message}` }; + } + return { ok: true, value: parsed.data }; +} + +/** + * Strip a markdown code fence if the model wrapped its JSON in one + * despite our instructions. Cheap defense against models that ignore + * "no markdown fences" in their system prompt. + */ +function stripCodeFences(s: string): string { + const fenceMatch = s.match(/^```(?:json)?\s*\n?([\s\S]*?)\n?```\s*$/); + return fenceMatch ? fenceMatch[1]! : s; +}