ref(memory)!: remove MemoryStore

packages/junior-evals/evals/memory/actors.eval.ts

@@ -1,19 +1,14 @@
 import { expect } from "vitest";
 import { describeEval } from "vitest-evals";
-import { getDb } from "@/chat/db";
 import { readActorIdentity } from "@/chat/plugins/viewer";
-import type { MemoryDb } from "@sentry/junior-memory";
-import {
-  juniorMemoryEmbeddings,
-  juniorMemoryMemories,
-} from "../../../junior-memory/src/db/schema";
 import {
   mention,
   rubric,
   slackEvals,
   steer,
   threadMessage,
 } from "../../src/helpers";
+import { clearMemories, readMemories, type MemoryThread } from "./helpers";
 
 /**
  * Passive memory learning when a run has more than one Actor.
@@ -46,34 +41,6 @@ const CAROL = {
   full_name: "Carol Example",
 };
 
-interface MemoryThread {
-  channel_type?: "channel" | "group" | "im" | "mpim";
-  channel_id: string;
-  id: string;
-  thread_ts: string;
-}
-
-function memoryDb(): MemoryDb {
-  return getDb() as unknown as MemoryDb;
-}
-
-function memorySourceKey(thread: MemoryThread): string {
-  return `slack:${memoryTeamId}:${thread.channel_id}:${thread.thread_ts}`;
-}
-
-async function readMemories(thread: MemoryThread) {
-  const rows = await memoryDb()
-    .select()
-    .from(juniorMemoryMemories)
-    .orderBy(juniorMemoryMemories.createdAtMs, juniorMemoryMemories.id);
-  return rows.filter((memory) => memory.sourceKey === memorySourceKey(thread));
-}
-
-async function clearMemories() {
-  await memoryDb().delete(juniorMemoryEmbeddings);
-  await memoryDb().delete(juniorMemoryMemories);
-}
-
 async function memoriesForActor(
   rows: Awaited<ReturnType<typeof readMemories>>,
   slackUserId: string,

packages/junior-evals/evals/memory/helpers.ts

@@ -4,12 +4,14 @@ import { getDb, getSqlExecutor } from "@/chat/db";
 import { upsertIdentity } from "@/chat/identities/sql";
 import { completeText, resolveGatewayModel } from "@/chat/pi/client";
 import { createPluginEmbedder } from "@/chat/plugins/model";
-import { createMemoryStore, type MemoryDb } from "@sentry/junior-memory";
-import { createSlackSource } from "@sentry/junior-plugin-api";
 import {
-  juniorMemoryEmbeddings,
-  juniorMemoryMemories,
-} from "../../../junior-memory/src/db/schema";
+  clearAll,
+  countEmbeddings,
+  createMemory,
+  listBySource,
+  type MemoryDb,
+} from "@sentry/junior-memory/testing";
+import { createSlackSource } from "@sentry/junior-plugin-api";
 import { TEST_USER_ID } from "@junior-tests/fixtures/slack/factories/ids";
 
 export const memoryPluginOverrides = {
@@ -47,37 +49,34 @@ export async function seedMemory(args: {
   if (!identity.userId) {
     throw new Error("Eval memory Actor did not resolve to a User");
   }
-  const store = createMemoryStore(
-    memoryDb(),
-    {
-      conversationId: `slack:${args.thread.channel_id}:${args.thread.thread_ts}`,
-      actor: {
-        platform: "slack",
-        teamId: memoryTeamId,
-        userId: actorUserId,
-      },
-      source: createSlackSource({
-        channelId: args.thread.channel_id,
-        messageTs: args.thread.thread_ts,
-        teamId: memoryTeamId,
-        threadTs: args.thread.thread_ts,
-        visibility:
-          args.thread.channel_type === "channel" ? "public" : "private",
-      }),
-      userId: identity.userId,
+  const context = {
+    conversationId: `slack:${args.thread.channel_id}:${args.thread.thread_ts}`,
+    actor: {
+      platform: "slack" as const,
+      teamId: memoryTeamId,
+      userId: actorUserId,
     },
-    { embedder: evalMemoryEmbedder },
-  );
+    source: createSlackSource({
+      channelId: args.thread.channel_id,
+      messageTs: args.thread.thread_ts,
+      teamId: memoryTeamId,
+      threadTs: args.thread.thread_ts,
+      visibility: args.thread.channel_type === "channel" ? "public" : "private",
+    }),
+    userId: identity.userId,
+  };
   const input = {
     content: args.content,
     idempotencyKey: args.idempotencyKey,
     kind: args.kind ?? "preference",
   };
-  if (args.subject === "conversation") {
-    await store.createConversationMemory(input);
-    return;
-  }
-  await store.createMemory(input);
+  await createMemory({
+    context,
+    db: memoryDb(),
+    embedder: evalMemoryEmbedder,
+    input,
+    subjectType: args.subject === "conversation" ? "conversation" : "user",
+  });
 }
 
 function memoryDb(): MemoryDb {
@@ -89,24 +88,16 @@ function memorySourceKey(thread: MemoryThread): string {
 }
 
 export async function readMemories(thread: MemoryThread) {
-  const rows = await memoryDb()
-    .select()
-    .from(juniorMemoryMemories)
-    .orderBy(juniorMemoryMemories.createdAtMs, juniorMemoryMemories.id);
-  return rows.filter((memory) => memory.sourceKey === memorySourceKey(thread));
+  return listBySource(memoryDb(), memorySourceKey(thread));
 }
 
-/** Count vector rows for memories seeded in one eval thread. */
+/** Count embeddings for memories from one eval thread. */
 export async function countMemoryEmbeddings(thread: MemoryThread) {
   const memories = await readMemories(thread);
-  if (memories.length === 0) {
-    return 0;
-  }
-  const memoryIds = new Set(memories.map((memory) => memory.id));
-  const rows = await memoryDb()
-    .select({ memoryId: juniorMemoryEmbeddings.memoryId })
-    .from(juniorMemoryEmbeddings);
-  return rows.filter((row) => memoryIds.has(row.memoryId)).length;
+  return countEmbeddings(
+    memoryDb(),
+    memories.map((memory) => memory.id),
+  );
 }
 
 /** Read the durable memories currently eligible for recall in one eval thread. */
@@ -124,8 +115,7 @@ export async function readActiveMemories(
 }
 
 export async function clearMemories() {
-  await memoryDb().delete(juniorMemoryEmbeddings);
-  await memoryDb().delete(juniorMemoryMemories);
+  await clearAll(memoryDb());
 }
 
 export function visibleAssistantText(result: {

packages/junior-evals/vitest.evals.behavioral.config.ts

@@ -47,6 +47,10 @@ export default defineConfig({
   resolve: {
     alias: {
       "@": path.resolve(juniorPackageRoot, "src"),
+      "@sentry/junior-memory/testing": path.resolve(
+        memoryPackageRoot,
+        "src/testing.ts",
+      ),
       "@sentry/junior-memory": path.resolve(memoryPackageRoot, "src/index.ts"),
       "@sentry/junior-plugin-api": path.resolve(
         pluginApiPackageRoot,

packages/junior-memory/package.json

@@ -15,6 +15,10 @@
     ".": {
       "types": "./src/index.ts",
       "default": "./dist/index.js"
+    },
+    "./testing": {
+      "types": "./src/testing.ts",
+      "default": "./dist/testing.js"
     }
   },
   "files": [

packages/junior-memory/src/agent.ts

@@ -9,7 +9,7 @@ import {
   memorySupersessionInputSchema,
   type MemorySupersessionDecision,
   type MemorySupersessionInput,
-} from "./store";
+} from "./create";
 import {
   MEMORY_KINDS,
   memoryRuntimeContextSchema,

packages/junior-memory/src/api.ts

@@ -10,7 +10,7 @@ import {
   type PluginRouteApp,
   type User,
 } from "@sentry/junior-plugin-api";
-import type { MemoryDb } from "./store";
+import type { MemoryDb } from "./memories";
 import {
   archiveMemory,
   getMemory,

packages/junior-memory/src/cli/search.ts

@@ -5,7 +5,7 @@ import type {
   PluginCliHost,
 } from "@sentry/junior-plugin-api";
 import { juniorMemoryMemories } from "../db/schema";
-import type { MemoryDb } from "../store";
+import type { MemoryDb } from "../memories";
 import { MEMORY_SCOPES, type MemoryScope } from "../types";
 import { formatMemory } from "./format";
 

packages/junior-memory/src/cli/show.ts

@@ -5,7 +5,7 @@ import type {
 } from "@sentry/junior-plugin-api";
 import { eq } from "drizzle-orm";
 import { juniorMemoryMemories } from "../db/schema";
-import type { MemoryDb } from "../store";
+import type { MemoryDb } from "../memories";
 import { formatMemory } from "./format";
 
 async function runShow(