feat(reedy): Appendix A · Phase 3.2 — MemoryConsolidator (#4304)

2026-05-26 13:46:11 +08:00
parent 568b8c0a88
commit 49664ecb75
2 changed files with 450 additions and 0 deletions
@@ -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<typeof import('ai')>();
+  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<ChatModel['getLanguageModel']>,
+  };
+}
+
+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<MemoryConsolidatorOptions> = {}) {
+    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.');
+    });
+  });
+});
@@ -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<ConsolidatedMemory[]> {
+    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<string> {
+    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<z.infer<typeof memoryRowSchema>>;
+}
+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;
+}