@dex-ai/memory 0.3.3

package/src/episodic.ts

@@ -0,0 +1,144 @@
+/**
+ * Episodic memory — summarized past turns/tasks.
+ *
+ * Writes:
+ *   - record(summary): inserts into episodic + embeds into episodic_vec.
+ *
+ * Reads:
+ *   - recall(userId, queryEmbedding, opts): returns the union of
+ *     most-recent N and top-K similar episodes (de-duped by id).
+ *
+ * The caller supplies the embedding vector — this module doesn't own the
+ * embedder (keeps the module testable without Transformers.js).
+ */
+import type { Database } from 'bun:sqlite';
+
+
+export interface EpisodeRow {
+  id: string;
+  userId: string;
+  summary: string;
+  metadata: Record<string, unknown>;
+  createdAt: number;
+}
+
+interface EpisodeSqlRow {
+  id: string;
+  user_id: string;
+  summary: string;
+  metadata: string;
+  created_at: number;
+}
+
+function rowToEpisode(r: EpisodeSqlRow): EpisodeRow {
+  return {
+    id: r.id,
+    userId: r.user_id,
+    summary: r.summary,
+    metadata: JSON.parse(r.metadata) as Record<string, unknown>,
+    createdAt: r.created_at,
+  };
+}
+
+/** Packs a number[] into Float32Array bytes for sqlite-vec. */
+function toVecBytes(embedding: number[]): Uint8Array {
+  const f = new Float32Array(embedding);
+  return new Uint8Array(f.buffer, f.byteOffset, f.byteLength);
+}
+
+export interface RecallOptions {
+  /** Most-recent episodes to include. Default 3. */
+  recentLimit?: number;
+  /** Most-similar episodes to include via vector search. Default 3. */
+  similarLimit?: number;
+}
+
+export class EpisodicStore {
+  readonly #db: Database;
+  readonly #insert;
+  readonly #insertVec;
+  readonly #listRecent;
+  readonly #similar;
+
+  constructor(db: Database) {
+    this.#db = db;
+    this.#insert = db.prepare(
+      'INSERT INTO episodic(id, user_id, summary, metadata, created_at) VALUES(?, ?, ?, ?, ?)',
+    );
+    this.#insertVec = db.prepare('INSERT INTO episodic_vec(id, embedding) VALUES(?, ?)');
+    this.#listRecent = db.prepare(
+      'SELECT * FROM episodic WHERE user_id = ? ORDER BY created_at DESC LIMIT ?',
+    );
+    // sqlite-vec vector search: returns ids by distance ascending (closer = more similar).
+    // We join back onto episodic to filter by user_id and return full rows.
+    this.#similar = db.prepare(`
+      SELECT e.*, v.distance as distance
+      FROM episodic e
+      JOIN (
+        SELECT id, distance
+        FROM episodic_vec
+        WHERE embedding MATCH ?
+        ORDER BY distance
+        LIMIT ?
+      ) v ON v.id = e.id
+      WHERE e.user_id = ?
+      ORDER BY v.distance
+    `);
+  }
+
+  async record(input: {
+    userId: string;
+    summary: string;
+    embedding: number[];
+    metadata?: Record<string, unknown>;
+  }): Promise<EpisodeRow> {
+    const now = Date.now();
+    const id = crypto.randomUUID();
+    const metaJson = JSON.stringify(input.metadata ?? {});
+    this.#db.transaction(() => {
+      this.#insert.run(id, input.userId, input.summary, metaJson, now);
+      this.#insertVec.run(id, toVecBytes(input.embedding));
+    })();
+    return {
+      id,
+      userId: input.userId,
+      summary: input.summary,
+      metadata: input.metadata ?? {},
+      createdAt: now,
+    };
+  }
+
+  /**
+   * Recall: union of recent N and top-K similar, de-duped by id and sorted by time desc.
+   * When the query embedding is omitted, only recent episodes are returned.
+   */
+  async recall(
+    userId: string,
+    queryEmbedding: number[] | undefined,
+    opts: RecallOptions = {},
+  ): Promise<EpisodeRow[]> {
+    const recentLimit = opts.recentLimit ?? 3;
+    const similarLimit = opts.similarLimit ?? 3;
+
+    const recent = this.#listRecent.all(userId, recentLimit) as EpisodeSqlRow[];
+    const byId = new Map<string, EpisodeSqlRow>();
+    for (const r of recent) byId.set(r.id, r);
+
+    if (queryEmbedding !== undefined && similarLimit > 0) {
+      // Fetch more than similarLimit to allow user_id filtering in the join
+      // dropping rows from other users.
+      const similar = this.#similar.all(
+        toVecBytes(queryEmbedding),
+        similarLimit * 4,
+        userId,
+      ) as Array<EpisodeSqlRow & { distance: number }>;
+      for (const s of similar.slice(0, similarLimit)) {
+        if (!byId.has(s.id)) byId.set(s.id, s);
+      }
+    }
+
+    const merged = Array.from(byId.values()).map(rowToEpisode);
+    merged.sort((a, b) => b.createdAt - a.createdAt);
+    return merged;
+  }
+}

package/src/extension.test.ts

@@ -0,0 +1,266 @@
+import { describe, expect, test } from 'bun:test';
+import type { Message } from '@dex-ai/sdk';
+import { Agent } from '@dex-ai/sdk';
+import { memoryExtension } from './extension';
+import { fakeEmbedder, scriptedProviderExtension, FAKE_PROVIDER, FAKE_MODEL } from './_fakes';
+
+const userMsg = (text: string): Message => ({
+  role: 'user',
+  content: [{ type: 'text', text }],
+});
+
+function findToolResult(
+  messages: ReadonlyArray<Message>,
+  toolName: string,
+): unknown {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const msg = messages[i]!;
+    if (msg.role !== 'tool') continue;
+    for (const c of msg.content) {
+      if (c.type === 'tool-result' && c.toolName === toolName) {
+        return c.output;
+      }
+    }
+  }
+  return null;
+}
+
+describe('memoryExtension', () => {
+  test('stores a user-supplied fact via remember_fact tool and recalls it next turn', async () => {
+    const providerExt = scriptedProviderExtension({
+      steps: [
+        {
+          kind: 'tool-call',
+          toolName: 'memory_remember_fact',
+          input: { subject: 'user', predicate: 'prefers', object: 'TypeScript' },
+        },
+        { kind: 'text', text: 'noted' },
+      ],
+      generateReplies: ['turn summary', '{"facts":[]}'],
+    });
+
+    const agent = await Agent.create({
+      provider: FAKE_PROVIDER, model: FAKE_MODEL,
+      extensions: [providerExt,
+        memoryExtension({
+          path: ':memory:',
+          userId: 'alice',
+          embed: fakeEmbedder(),
+          autoWrite: false, // off for this test so we're not also hitting extractFacts
+        }),
+      ],
+    });
+
+    const s = agent.generate({ input: [userMsg('remember I like TS')] });
+    await s.result;
+    for await (const _ of s) { /* drain */ }
+    await agent.dispose();
+
+    // Shutting down already happened; test DB is closed. Re-verify state by
+    // opening a fresh agent pointed at the same file — but :memory: doesn't
+    // persist across connections. So instead assert from the tool-result
+    // that was committed into history during the turn.
+    const out = findToolResult(agent.context.messages, 'memory_remember_fact');
+    expect(out).toBeDefined();
+    expect((out as { type: string }).type).toBe('json');
+  });
+
+  test('onRequest injects facts as a synthetic system message that never persists', async () => {
+    const providerExt = scriptedProviderExtension({
+      steps: [
+        // turn 1: remember a fact
+        {
+          kind: 'tool-call',
+          toolName: 'memory_remember_fact',
+          input: { subject: 'user', predicate: 'uses', object: 'Bun' },
+        },
+        { kind: 'text', text: 'ok' },
+        // turn 2: plain text reply
+        { kind: 'text', text: 'hi again' },
+      ],
+      generateReplies: [],
+    });
+
+    const agent = await Agent.create({
+      provider: FAKE_PROVIDER, model: FAKE_MODEL,
+      extensions: [providerExt,
+        memoryExtension({
+          path: ':memory:',
+          userId: 'alice',
+          embed: fakeEmbedder(),
+          autoWrite: false,
+        }),
+      ],
+    });
+
+    // Turn 1 — write the fact.
+    const s1 = agent.generate({ input: [userMsg('use bun')] });
+    await s1.result;
+    for await (const _ of s1) { /* drain */ }
+
+    // Turn 2 — the synthetic system message should be prepended but NOT appear
+    // in actx.messages (history stays clean).
+    const beforeLen = agent.context.messages.length;
+    const s2 = agent.generate({ input: [userMsg('what stack?')] });
+    await s2.result;
+    for await (const _ of s2) { /* drain */ }
+    const afterLen = agent.context.messages.length;
+
+    // +1 user, +1 assistant (text-only turn) = +2 messages total.
+    expect(afterLen - beforeLen).toBe(2);
+    // No injected 'system' message got into history either.
+    for (const m of agent.context.messages) {
+      if (m.role === 'system') {
+        for (const c of m.content) {
+          if (c.type === 'text') {
+            expect(c.text).not.toMatch(/Known facts:/);
+          }
+        }
+      }
+    }
+
+    await agent.dispose();
+  });
+});
+
+describe('memoryExtension — auto-write', () => {
+  test('fires summarize + extractFacts at iteration stop (fire-and-forget) and awaits at dispose', async () => {
+    const providerExt = scriptedProviderExtension({
+      steps: [{ kind: 'text', text: 'hey there' }],
+      generateReplies: [
+        'user greeted the assistant',                                      // summarize call
+        '{"facts":[{"subject":"user","predicate":"greeted","object":"at start"}]}', // extract call
+      ],
+    });
+
+    const agent = await Agent.create({
+      provider: FAKE_PROVIDER, model: FAKE_MODEL,
+      extensions: [providerExt,
+        memoryExtension({
+          path: ':memory:',
+          userId: 'alice',
+          embed: fakeEmbedder(),
+          autoWrite: true,
+        }),
+      ],
+    });
+
+    const s = agent.generate({ input: [userMsg('hello')] });
+    await s.result;
+    for await (const _ of s) { /* drain */ }
+    // onAgentStop awaits pending background writes.
+    await agent.dispose();
+
+    // We can't inspect :memory: after close. Instead, keep the DB alive by using
+    // a persistent file path — do the same flow against a temp file.
+    expect(true).toBe(true); // sanity; real check in the next test below
+  });
+
+  test('auto-write persists summary + extracted fact (file-backed)', async () => {
+    const { mkdtemp, rm } = await import('node:fs/promises');
+    const { tmpdir } = await import('node:os');
+    const { join } = await import('node:path');
+    const { MemoryDb } = await import('./db');
+    const { EpisodicStore } = await import('./episodic');
+    const { SemanticStore } = await import('./semantic');
+
+    const dir = await mkdtemp(join(tmpdir(), 'dex-mem-ext-'));
+    const path = join(dir, 'mem.db');
+    try {
+      const providerExt = scriptedProviderExtension({
+        steps: [{ kind: 'text', text: 'reply' }],
+        generateReplies: [
+          'user greeted the assistant',
+          '{"facts":[{"subject":"user","predicate":"greeted","object":"at start"}]}',
+        ],
+      });
+
+      const agent = await Agent.create({
+        provider: FAKE_PROVIDER, model: FAKE_MODEL,
+        extensions: [providerExt,
+          memoryExtension({
+            path,
+            userId: 'alice',
+            embed: fakeEmbedder(),
+            autoWrite: true,
+          }),
+        ],
+      });
+
+      const s = agent.generate({ input: [userMsg('hello')] });
+      await s.result;
+      for await (const _ of s) { /* drain */ }
+      await agent.dispose();
+
+      // Re-open to inspect.
+      const db = new MemoryDb({ path });
+      const ep = new EpisodicStore(db.db);
+      const sem = new SemanticStore(db.db);
+
+      const episodes = await ep.recall('alice', undefined, { recentLimit: 10, similarLimit: 0 });
+      expect(episodes.length).toBe(1);
+      expect(episodes[0]!.summary).toBe('user greeted the assistant');
+
+      const facts = await sem.list('alice');
+      expect(facts.length).toBe(1);
+      expect(facts[0]!.subject).toBe('user');
+      expect(facts[0]!.object).toBe('at start');
+      expect(facts[0]!.source).toBe('extracted');
+
+      db.close();
+    } finally {
+      await rm(dir, { recursive: true, force: true });
+    }
+  });
+});
+
+describe('memoryExtension — procedural tools', () => {
+  test('store_procedure then get_procedure round-trips', async () => {
+    const providerExt = scriptedProviderExtension({
+      steps: [
+        {
+          kind: 'tool-call',
+          toolName: 'memory_store_procedure',
+          input: {
+            title: 'deploy-dex',
+            body: '1. bun run typecheck\n2. bun run test\n3. git tag',
+            tags: ['deploy', 'release'],
+          },
+        },
+        {
+          kind: 'tool-call',
+          toolName: 'memory_get_procedure',
+          input: { title: 'deploy-dex' },
+        },
+        { kind: 'text', text: 'done' },
+      ],
+    });
+
+    const agent = await Agent.create({
+      provider: FAKE_PROVIDER, model: FAKE_MODEL,
+      extensions: [providerExt,
+        memoryExtension({
+          path: ':memory:',
+          userId: 'alice',
+          embed: fakeEmbedder(),
+          autoWrite: false,
+        }),
+      ],
+    });
+
+    const s = agent.generate({ input: [userMsg('save and fetch')] });
+    await s.result;
+    for await (const _ of s) { /* drain */ }
+    await agent.dispose();
+
+    const stored = findToolResult(agent.context.messages, 'memory_store_procedure');
+    expect((stored as { type: string }).type).toBe('json');
+
+    const fetched = findToolResult(agent.context.messages, 'memory_get_procedure');
+    expect((fetched as { type: string }).type).toBe('json');
+    const body = (fetched as { value: { title: string; body: string } | null }).value;
+    expect(body).not.toBeNull();
+    expect(body!.title).toBe('deploy-dex');
+    expect(body!.body).toMatch(/bun run typecheck/);
+  });
+});