@batonfx/memory 0.1.0

package/test/vector-store.test.ts

@@ -0,0 +1,82 @@
+import { describe, expect, it } from "@effect/vitest"
+import { Effect } from "effect"
+import { Memory } from "@batonfx/core"
+import { VectorStore } from "../src/index"
+
+const key: Memory.Key = { agent: "agent-a", subject: "subject-a" }
+const otherSubject: Memory.Key = { agent: "agent-a", subject: "subject-b" }
+const otherAgent: Memory.Key = { agent: "agent-b", subject: "subject-a" }
+
+const document = (
+  id: string,
+  memoryKey: Memory.Key,
+  text: string,
+  embedding: ReadonlyArray<number>,
+): VectorStore.Embedded => ({
+  id,
+  key: memoryKey,
+  text,
+  embedding,
+})
+
+describe("VectorStore", () => {
+  it.effect("ranks matching-key documents by cosine score", () =>
+    Effect.gen(function* () {
+      const store = yield* VectorStore.VectorStore
+
+      yield* store.upsert([
+        document("near", key, "near", [1, 0]),
+        document("far", key, "far", [0, 1]),
+        document("also-near", key, "also-near", [0.8, 0.2]),
+      ])
+
+      const matches = yield* store.query({ key, embedding: [1, 0], limit: 2 })
+
+      expect(matches.map((match) => match.document.text)).toEqual(["near", "also-near"])
+      expect(matches[0]?.score).toBeGreaterThan(matches[1]?.score ?? 0)
+    }).pipe(Effect.provide(VectorStore.memoryLayer)),
+  )
+
+  it.effect("enforces exact agent and subject key isolation", () =>
+    Effect.gen(function* () {
+      const store = yield* VectorStore.VectorStore
+
+      yield* store.upsert([
+        document("same-id", key, "visible", [1, 0]),
+        document("same-id", otherSubject, "other subject", [1, 0]),
+        document("same-id", otherAgent, "other agent", [1, 0]),
+      ])
+
+      const matches = yield* store.query({ key, embedding: [1, 0], limit: 10 })
+
+      expect(matches.map((match) => match.document.text)).toEqual(["visible"])
+    }).pipe(Effect.provide(VectorStore.memoryLayer)),
+  )
+
+  it.effect("applies limit and minScore", () =>
+    Effect.gen(function* () {
+      const store = yield* VectorStore.VectorStore
+
+      yield* store.upsert([document("near", key, "near", [1, 0]), document("also-near", key, "also-near", [0.8, 0.2])])
+
+      const limited = yield* store.query({ key, embedding: [1, 0], limit: 1 })
+      const thresholded = yield* store.query({ key, embedding: [1, 0], limit: 10, minScore: 0.99 })
+
+      expect(limited.map((match) => match.document.text)).toEqual(["near"])
+      expect(thresholded.map((match) => match.document.text)).toEqual(["near"])
+    }).pipe(Effect.provide(VectorStore.memoryLayer)),
+  )
+
+  it.effect("fails loudly on matching-key dimension mismatch", () =>
+    Effect.gen(function* () {
+      const store = yield* VectorStore.VectorStore
+
+      yield* store.upsert([document("bad", key, "bad", [1, 0, 0])])
+
+      const failure = yield* Effect.flip(store.query({ key, embedding: [1, 0], limit: 1 }))
+
+      expect(failure._tag).toBe("@batonfx/memory/VectorStoreError")
+      expect(failure.message).toContain("dimension")
+    }).pipe(Effect.provide(VectorStore.memoryLayer)),
+  )
+})

package/test/working-memory.test.ts

@@ -0,0 +1,101 @@
+import { describe, expect, it } from "@effect/vitest"
+import { Effect, Layer, Stream } from "effect"
+import * as Ai from "effect/unstable/ai"
+import { Memory } from "@batonfx/core"
+import { WorkingMemory } from "../src/index"
+
+const key: Memory.Key = { agent: "memory-agent", subject: "subject-a" }
+const otherKey: Memory.Key = { agent: "memory-agent", subject: "subject-b" }
+
+const textPart = (text: string) => Ai.Prompt.makePart("text", { text })
+const user = (text: string) => Ai.Prompt.makeMessage("user", { content: [textPart(text)] })
+const assistant = (text: string) => Ai.Prompt.makeMessage("assistant", { content: [textPart(text)] })
+const prompt = (...messages: ReadonlyArray<Ai.Prompt.Message>) => Ai.Prompt.fromMessages(messages)
+
+const itemText = (item: Memory.Item): string =>
+  item.parts
+    .filter((part): part is Ai.Prompt.TextPart => part.type === "text")
+    .map((part) => part.text)
+    .join("")
+
+describe("WorkingMemory", () => {
+  it.effect("keeps a bounded recent tail", () =>
+    Effect.gen(function* () {
+      const memory = yield* Memory.Memory
+
+      yield* memory.remember({
+        key,
+        turn: 0,
+        terminal: true,
+        transcript: prompt(user("one"), assistant("two"), user("three")),
+      })
+
+      const recalled = yield* memory.recall({ key, turn: 0, prompt: prompt(user("current")) })
+
+      expect(recalled.map(itemText)).toEqual(["Assistant: two", "User: three"])
+    }).pipe(Effect.provide(WorkingMemory.layer({ maxMessages: 2 }))),
+  )
+
+  it.effect("summarizes overflow once and recalls summary before the recent tail", () => {
+    let summaryCalls = 0
+    let summaryPrompt = ""
+    const summaryModel = Layer.effect(
+      Ai.LanguageModel.LanguageModel,
+      Ai.LanguageModel.make({
+        generateText: (options) =>
+          Effect.sync(() => {
+            summaryCalls += 1
+            summaryPrompt = JSON.stringify(options.prompt.content)
+            return [{ type: "text", text: "summary" }]
+          }),
+        streamText: () => Stream.empty,
+      }),
+    )
+
+    return Effect.gen(function* () {
+      const memory = yield* Memory.Memory
+
+      yield* memory.remember({
+        key,
+        turn: 0,
+        terminal: true,
+        transcript: prompt(user("one"), assistant("two"), user("three"), assistant("four")),
+      })
+
+      const recalled = yield* memory.recall({ key, turn: 0, prompt: prompt(user("current")) })
+
+      expect(summaryCalls).toBe(1)
+      expect(summaryPrompt).toContain("one")
+      expect(summaryPrompt).toContain("two")
+      expect(recalled.map(itemText)).toEqual([
+        "<working-memory-summary>\nsummary\n</working-memory-summary>",
+        "User: three",
+        "Assistant: four",
+      ])
+    }).pipe(
+      Effect.provide(
+        WorkingMemory.layer({
+          maxMessages: 2,
+          summarize: { model: summaryModel },
+        }),
+      ),
+    )
+  })
+
+  it.effect("isolates state by memory key", () =>
+    Effect.gen(function* () {
+      const memory = yield* Memory.Memory
+
+      yield* memory.remember({
+        key,
+        turn: 0,
+        terminal: true,
+        transcript: prompt(user("one"), assistant("two")),
+      })
+
+      const recalled = yield* memory.recall({ key: otherKey, turn: 0, prompt: prompt(user("current")) })
+
+      expect(recalled).toEqual([])
+    }).pipe(Effect.provide(WorkingMemory.layer({ maxMessages: 2 }))),
+  )
+})

package/tsconfig.json

@@ -0,0 +1,7 @@
+{
+  "extends": "../../tsconfig.json",
+  "compilerOptions": {
+    "rootDir": "."
+  },
+  "include": ["src/**/*.ts", "test/**/*.ts"]
+}