From d062da113a964e43ae40885233ab5a1d582b8d0c Mon Sep 17 00:00:00 2001 From: Nina Doinjashvili Date: Sat, 8 Aug 2026 13:46:49 +0300 Subject: [PATCH] Fix correctness gaps and make the evaluation honest (v0.3.0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MCP server — the headline feature — did not start on a fresh install, and the published results overstated what the benchmark showed. Both are fixed, along with the durability and data-loss issues underneath them. Correctness - MCP server: `mcp` 2.0 renamed FastMCP to mcp.server.mcpserver.MCPServer, so `pip install -e ".[mcp]"` produced a server that exited claiming the package was missing. Load either class; both majors are now tested in CI. - memory_boot ignored its own budget: `remaining or None` turned a fully consumed budget into "no cap", returning 6.6x the requested tokens. The budget is now shared across handoff and recall by MemoryStore.boot. - Concurrent agents silently lost memories: the store rewrote the whole file from a snapshot taken at startup. Writes now take a cross-process lock, merge in anything another agent appended, and land atomically via os.replace. Reads reload when the file changes, so a running server sees another agent's writes. - Fixed a TOCTOU bug in that reload: load() stamped the file after reading it, so a store replaced mid-read was recorded as current and never reloaded again. Two agents then chose the same id and a write vanished. Stamp is now taken before the read; regression test included. - Switching embedders bricked a store (512-d query against 384-d vectors). Mismatched dim or backend now re-embeds from the stored text. - Ids are derived from the ids present, surviving explicit ids and deletions. Memory you can correct - Added forget/update plus memory_list, memory_update, memory_forget and the matching CLI commands. An append-only store keeps recalling stale state. - write_with_status reports when a near-duplicate was dropped; the tools no longer claim a save that did not happen. - Added a relevance floor so an off-topic task returns nothing instead of k unrelated memories the agent cannot tell are noise. Calibrated per embedder from a published sweep. Evaluation - Added a random-k control. It shows the same 78% token saving at 0.07 recall, which is the point: the saving is arithmetic (k of n), and only the 0.86 recall gap is evidence that retrieval works. - Added a paraphrase arm. The default embedder is lexical, and recall falls 0.93 -> 0.43 when queries stop sharing vocabulary with the memories. This is now published rather than implied by the word "semantic". - Added a scaling experiment with distractor memories, so "cost stays fixed as the store grows" is measured on a store that actually grows. - Dropped the precision-vs-baseline comparison; the baseline's 0.10 is just |relevant|/|store|. Also: base64 float16 embeddings (~5x smaller on disk), ingest_markdown chunks oversized sections and counts only real writes, tiktoken in dev so the numbers are exact, CI fails when committed results go stale, and tests went 17 -> 55. Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 42 ++++ .gitignore | 1 + README.md | 105 +++++++--- eval/dataset.json | 58 +++++- eval/results.json | 105 ++++++++-- eval/results.md | 51 ++++- eval/run_eval.py | 339 +++++++++++++++++++++++++++------ examples/quickstart.py | 32 +++- pyproject.toml | 8 +- scripts/ingest_markdown.py | 63 ++++-- src/agent_memory/__init__.py | 15 +- src/agent_memory/cli.py | 95 +++++++-- src/agent_memory/embeddings.py | 63 +++++- src/agent_memory/mcp_server.py | 108 +++++++++-- src/agent_memory/store.py | 329 ++++++++++++++++++++++++++++++-- tests/test_eval.py | 84 ++++++-- tests/test_mcp_server.py | 139 ++++++++++++++ tests/test_persistence.py | 232 ++++++++++++++++++++++ tests/test_store.py | 83 ++++++++ 19 files changed, 1741 insertions(+), 211 deletions(-) create mode 100644 tests/test_mcp_server.py create mode 100644 tests/test_persistence.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8631c63..fe55fe0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,6 +9,7 @@ jobs: test: runs-on: ubuntu-latest strategy: + fail-fast: false matrix: python-version: ["3.10", "3.12"] env: @@ -20,8 +21,49 @@ jobs: with: python-version: ${{ matrix.python-version }} - name: Install + # `dev` includes mcp and tiktoken: the MCP server is the headline + # feature and must be exercised, and the pinned tokenizer is what makes + # the published evaluation numbers reproducible. run: pip install -e ".[dev]" - name: Run tests run: pytest -q - name: Run evaluation run: python eval/run_eval.py + - name: Fail if the published results are stale + # The README quotes these numbers. If a change moves them, the results + # files must be regenerated in the same commit. + run: git diff --exit-code -- eval/results.md eval/results.json + - name: Run the quickstart example + run: python examples/quickstart.py + + mcp-versions: + # The MCP server broke once because `mcp` 2.0 renamed the server class and + # nothing in CI imported it. Both majors are tested from now on. + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + mcp-version: ["mcp>=1.0,<2", "mcp>=2.0"] + env: + AGENT_MEMORY_EMBEDDER: hashing + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install with ${{ matrix.mcp-version }} + run: | + pip install -e ".[dev]" + pip install "${{ matrix.mcp-version }}" + - name: Show resolved version + run: pip show mcp | head -2 + - name: Test the MCP server + run: pytest -q tests/test_mcp_server.py + - name: Start the server entry point + # Smoke-test the console script an agent config actually launches: + # it must reach the stdio loop rather than exit on an import error. + run: | + timeout 10s agent-memory-mcp < /dev/null; status=$? + if [ $status -ne 0 ] && [ $status -ne 124 ]; then + echo "agent-memory-mcp failed to start (exit $status)"; exit 1 + fi diff --git a/.gitignore b/.gitignore index 9c8069a..a713478 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ __pycache__/ build/ dist/ .venv/ +.venv*/ venv/ .env .DS_Store diff --git a/README.md b/README.md index 6d07b97..44a5622 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Agent Memory Engine -**One shared memory for all your coding agents — Claude Code, Codex CLI, Cursor — over MCP, with token-budgeted recall so it never floods a context window.** +**One shared memory for your coding agents — Claude Code, Codex CLI, Cursor — over MCP, with token-budgeted recall so it never floods a context window.** Every coding agent forgets everything between sessions, and none of them can read another's notes: what Claude Code learned about your codebase is invisible to Codex, and vice versa. The common fix — a pile of Markdown memory files loaded into every prompt — has the opposite problem: it burns more tokens every week as the pile grows. @@ -10,16 +10,43 @@ This engine is the small piece in between. Agents write atomic facts, decisions ## Results -On a labeled benchmark of an agent working across many sessions on one codebase (14 durable memories, 7 fresh-session tasks, top-k = 3): +On a hand-labeled benchmark of an agent working across many sessions on one codebase (14 durable memories, 7 fresh-session tasks, top-k = 3): | Arm | Avg context tokens | Recall | Precision | Tokens saved vs baseline | | --- | ---: | ---: | ---: | ---: | | No memory (control) | 0 | 0.00 | 0.00 | 100% | -| Full context (baseline — load every memory) | 550 | 1.00 | 0.10 | 0% | -| **Semantic recall (this engine)** | **120** | **0.93** | **0.43** | **78%** | -| **Budget recall (≤ 120 tokens, hard cap)** | **113** | **0.93** | **0.43** | **80%** | +| Full context (baseline — load every memory) | 424 | 1.00 | 0.10 | 0% | +| *Random k (control)* | *93* | *0.07* | *0.05* | *78%* | +| **Targeted retrieval (this engine)** | **93** | **0.93** | **0.43** | **78%** | +| **Budget recall (≤ 120 tokens, hard cap)** | **93** | **0.93** | **0.43** | **78%** | -**Semantic recall retrieves 93% of the relevant memories while loading 78% fewer context tokens** than dumping every memory file — and under a **hard 120-token cap recall is unchanged**, so the cost of memory stays fixed as the store grows. Numbers are reproduced by `python eval/run_eval.py` and regenerated in CI with the deterministic offline embedder, so they are byte-stable. +**Read the random row first.** It loads three memories picked at random, so it reports the *same* 78% token saving with 0.07 recall. The saving is arithmetic — you loaded 3 of 14 memories — and proves nothing on its own. The claim worth making is the **0.86 recall gap between random and retrieval at identical token cost**. + +Reproduce with `python eval/run_eval.py`; the full report is in [eval/results.md](eval/results.md). + +### What these numbers don't show + +- **The default embedder is lexical, not semantic.** It is feature hashing over word and character n-grams, so it matches shared wording. Re-run the same tasks with paraphrased queries that avoid the memories' vocabulary and recall drops from **0.93 to 0.43**. That gap is published in the results, not tuned away, and it is the reason the optional `sentence-transformers` backend exists. + + | Query phrasing | Word overlap with gold | Recall | + | --- | ---: | ---: | + | Developer phrasing (as labelled) | 40% | 0.93 | + | Outsider paraphrase | 3% | 0.43 | + +- **The precision column is nearly meaningless for the baseline.** Full context scores 0.10 because that is `|relevant| / |store|` — an artefact of loading everything. +- **7 tasks, 10 gold labels, written by the same person who wrote the retriever.** One retrieval either way moves recall by ~0.07. This is an engineering check, not a production-scale claim. + +### Does the cost stay flat as the store grows? + +Adding unrelated memories from the same project, re-measuring against the same labels: + +| Memories in store | Full context tokens | Budget recall tokens | Budget recall | +| ---: | ---: | ---: | ---: | +| 14 | 424 | 93 | 0.93 | +| 34 | 805 | 87 | 0.93 | +| 54 | 1159 | 82 | 0.93 | + +The baseline grows with the store. The budgeted arm does not, and recall holds. ## The cross-agent loop @@ -39,10 +66,12 @@ sequenceDiagram Three habits, three tools: - **`memory_boot(task, budget_tokens)`** — call once at session start: the latest handoff from the *previous agent, whichever tool it was*, plus the memories most relevant to the task, packed under one token budget. -- **`memory_write(text, type)`** — persist a durable fact, decision or issue the moment it's learned. Near-duplicates are dropped automatically. +- **`memory_write(text, type)`** — persist a durable fact, decision or issue the moment it's learned. Near-duplicates are dropped, and the tool says so rather than reporting a save that didn't happen. - **`memory_handoff(done, next_steps, warnings)`** — call at session end so the next agent continues instead of rediscovering. -Tool outputs are deliberately compact — no scores, no timestamps — because everything a memory tool returns is paid for again in the calling agent's context window. +Plus `memory_list`, `memory_update` and `memory_forget`, because memory you cannot correct is worse than no memory: a stale `state` entry keeps being recalled and quietly misleads every later session. + +Tool outputs are deliberately compact — no scores, no timestamps — because everything a memory tool returns is paid for again in the calling agent's context window. The flip side is that the agent can't judge relevance itself, so weak matches are filtered out before they're returned (see [Relevance floor](#relevance-floor)). ## Hook it up to your agents @@ -50,7 +79,7 @@ Tool outputs are deliberately compact — no scores, no timestamps — because e pip install -e ".[mcp]" ``` -Point every agent at the **same store path**, and give each its own name: +Works with both `mcp` 1.x and 2.x. Point every agent at the **same store path**, and give each its own name: **Claude Code** @@ -83,24 +112,42 @@ env = { AGENT_MEMORY_PATH = "~/.agent_memory/store.json", AGENT_MEMORY_AGENT = " } ``` -Each memory now carries its origin (`[decision · claude-code]`, `[handoff · codex]`), and a handoff written in one tool is picked up by the next via `memory_boot`. +Each memory now carries its origin (`[decision · claude-code]`, `[handoff · codex]`), and a handoff written in one tool is picked up by the next via `memory_boot`. Agents may run at the same time: writes take a lock, merge, and land atomically, and an already-running server picks up another agent's writes on its next read. ## How it works -- **Pluggable embedder** ([embeddings.py](src/agent_memory/embeddings.py)) — a small `Embedder` protocol with two backends: a dependency-free, deterministic `HashingEmbedder` (default; runs offline and in CI) and a `SentenceTransformerEmbedder` for real semantic embeddings. Same API, swap by installing one package. +- **Pluggable embedder** ([embeddings.py](src/agent_memory/embeddings.py)) — a small `Embedder` protocol with two backends: a dependency-free, deterministic `HashingEmbedder` (default; runs offline and in CI) and a `SentenceTransformerEmbedder` for real semantic embeddings. Same API, swap by installing one package. Each carries its own calibrated relevance floor, because the two score on different scales. - **Vector store** ([store.py](src/agent_memory/store.py)) — exact brute-force cosine over a numpy matrix, JSON persistence, near-duplicate suppression on write, and greedy **budget packing**: recall walks results in relevance order and skips anything that would overflow the caller's token budget. An agent's memory is hundreds of entries, not millions, so this is instant and exact; FAISS/Chroma can slot behind the same API later. -- **MCP server** ([mcp_server.py](src/agent_memory/mcp_server.py)) — exposes `memory_boot`, `memory_recall`, `memory_write`, `memory_handoff`, `memory_stats` as agent tools. -- **Evaluation harness** ([run_eval.py](eval/run_eval.py)) — measures token cost and retrieval precision/recall across four arms on a labeled dataset, including the hard-budget arm. +- **Durable writes** — every write takes a cross-process lock, re-reads anything another agent appended, and replaces the file atomically. A crash mid-write cannot truncate the store, and two agents writing at once cannot silently overwrite each other. Embeddings are stored as base64 float16, which keeps the file about 5× smaller than JSON float lists. +- **MCP server** ([mcp_server.py](src/agent_memory/mcp_server.py)) — exposes `memory_boot`, `memory_recall`, `memory_write`, `memory_handoff`, `memory_list`, `memory_update`, `memory_forget` and `memory_stats` as agent tools. +- **Evaluation harness** ([run_eval.py](eval/run_eval.py)) — measures token cost and retrieval precision/recall across five arms, including the random control and the hard-budget arm, plus the paraphrase gap, a store-growth scaling test and the relevance-floor sweep. + +### Relevance floor + +Without a floor, a query about something the store knows nothing about still returns *k* memories — and since tool output hides the scores, the calling agent has no way to tell noise from signal. `min_score` drops weak matches. The shipped default comes from a published sweep: + +| min_score | Labelled recall | Off-topic memories returned | +| ---: | ---: | ---: | +| 0.00 | 0.93 | 12/12 | +| 0.10 | 0.93 | 10/12 | +| **0.15** | **0.93** | **6/12** | +| 0.20 | 0.93 | 2/12 | + +The floor lives on the embedder (`HashingEmbedder.recommended_min_score = 0.15`), overridable with `AGENT_MEMORY_MIN_SCORE` or the `--min-score` flag. A higher floor cuts more noise on this benchmark, but the lowest-scoring genuinely-relevant memory in it sits at 0.13, so the default stays close to that observed edge. The `sentence-transformers` value is **not** calibrated here — measure it on your own store. ## Quickstart (60 seconds) ```bash -pip install -e . # core engine + CLI (numpy only) -python eval/run_eval.py # reproduce the results table above +pip install -e ".[dev]" # engine, CLI, MCP server, exact tokenizer +python eval/run_eval.py # reproduce the results above +python examples/quickstart.py agent-memory --agent claude-code write "Bookings are stored in UTC." --type decision agent-memory --agent claude-code handoff --done "Fixed double emails." --next "Add rate limiting." agent-memory boot "continue the rate limiting work" --budget 200 +agent-memory list # ids, so you can fix mistakes +agent-memory update mem_0003 "Bookings are stored in UTC; the UI converts." +agent-memory forget mem_0007 ``` In Python: @@ -111,7 +158,9 @@ from agent_memory import MemoryStore store = MemoryStore(path="~/.agent_memory/store.json") store.write("The chatbot uses Google Gemini in server/gemini-chat.ts.", type="decision", agent="claude-code") -for hit in store.recall("where is the chatbot configured", k=3, budget_tokens=150): + +handoff, hits = store.boot("where is the chatbot configured", budget_tokens=150) +for hit in hits: print(hit.score, hit.entry.text) ``` @@ -121,14 +170,25 @@ for hit in store.recall("where is the chatbot configured", k=3, budget_tokens=15 pip install -e ".[real]" # sentence-transformers + tiktoken ``` -`default_embedder()` automatically prefers the real model when installed. Force the offline embedder anywhere with `AGENT_MEMORY_EMBEDDER=hashing` (CI does this for stable numbers). +`default_embedder()` prefers the real model when installed, and warns on stderr if it falls back rather than switching silently. Force the offline embedder with `AGENT_MEMORY_EMBEDDER=hashing` (CI does this for stable numbers), or demand the real one with `AGENT_MEMORY_EMBEDDER=sentence-transformers` to turn a missing dependency into an error instead of a downgrade. + +Switching backends on an existing store is safe: vectors from a different embedder aren't comparable, so the store detects the mismatch on load and re-embeds from the stored text instead of returning meaningless scores. ## Migrate an existing Markdown scaffold ```bash -python scripts/ingest_markdown.py path/to/agent-memory/ # one memory per file/section +python scripts/ingest_markdown.py path/to/agent-memory/ --path ~/.agent_memory/store.json ``` +One memory per `##` section, with long sections split so that every ingested memory stays small enough to be recalled under a budget. + +## Limits and trust boundary + +- **Memories are injected into agent context verbatim.** Anything an agent writes to the store — including text it read from a webpage, an issue tracker or a dependency — will be replayed into a *different* agent's context later, where it reads as trusted project knowledge. Don't point a shared store at untrusted input, and skim `agent-memory list` occasionally. +- **The store is a local file.** No auth, no encryption, no server. It belongs next to your code, not on a shared host. +- **Retrieval is lexical by default.** See the paraphrase gap above. +- **The benchmark is small and self-authored.** It is a regression check for the engine, not evidence about your codebase. + ## Develop ```bash @@ -136,18 +196,13 @@ pip install -e ".[dev]" pytest -q ``` -## Design notes - -- **Small on purpose.** A memory layer that itself eats context defeats its point. Recall output is compact, budget-capped, and score-free; the whole engine is a few hundred lines with one required dependency (numpy). -- **Offline-safe by default.** No model download or network needed to run, test, or evaluate — the heavy backends are optional extras. This keeps CI deterministic and the repo cloneable-and-runnable in one step. -- **Honest evaluation.** The benchmark is small and labeled by hand; the one missed memory (`task_1`) is reported rather than tuned away. The harness includes a no-memory control and a full-context baseline so the comparison is fair. -- **Boring where it counts.** Brute-force exact search instead of an ANN index, JSON instead of a database — chosen deliberately for the actual scale and swappable behind the same API. +CI runs the suite on Python 3.10 and 3.12, runs the evaluation and fails if the committed results are stale, and separately tests the MCP server against both `mcp` 1.x and 2.x. ## Roadmap - LLM-based compaction: summarize/dedup `state` and `worklog` memories, extract durable facts from a raw session transcript. - Recency- and type-aware ranking (decay old `state`, never drop `decision`). -- `sentence-transformers` numbers published alongside the hashing baseline in CI. +- `sentence-transformers` numbers published alongside the hashing baseline in CI, including a calibrated relevance floor. - Optional FAISS backend for large stores. ## License diff --git a/eval/dataset.json b/eval/dataset.json index 3750f1e..f59ec03 100644 --- a/eval/dataset.json +++ b/eval/dataset.json @@ -1,6 +1,8 @@ { "name": "mr-studio-multisession", "description": "A coding agent has worked across many sessions on the MR Studio web app and accumulated durable memories. Each task is a fresh session: the agent is about to do something and needs the relevant prior memories pulled into context. `relevant_ids` is the human-labeled gold set of memories that actually matter for that task.", + "query_note": "Each task carries two phrasings of the same request. `query` is written the way a developer who knows the codebase would type it, and shares vocabulary with the gold memories. `paraphrase` describes the same task from the outside — a bug report or a product ask — deliberately avoiding the words used in the memories. The gap between the two is the honest measure of how much the default lexical retriever depends on shared wording.", + "distractor_note": "`distractors` are additional plausible memories from the same project that are not relevant to any task. They exist to grow the store for the scaling experiment, which tests whether recall cost stays flat as memory accumulates.", "memories": [ {"id": "mem_0001", "type": "project", "text": "The MR Studio web app is a React + Vite frontend with an Express backend, written in TypeScript. Run the dev server with `npm run dev` from the project root; it listens on port 5002."}, {"id": "mem_0002", "type": "project", "text": "The database is PostgreSQL accessed through the Drizzle ORM. The schema lives in shared/schema.ts and migrations are managed with drizzle-kit."}, @@ -18,12 +20,54 @@ {"id": "mem_0014", "type": "decision", "text": "Calendar sync writes confirmed bookings to Google Calendar via server/google-calendar.ts using a service account; a sync failure must not block the booking from being saved."} ], "tasks": [ - {"id": "task_1", "query": "Fix the bug where customers receive two confirmation emails for a single booking.", "relevant_ids": ["mem_0005", "mem_0013"]}, - {"id": "task_2", "query": "Add rate limiting to the chat endpoint so the Gemini API key cannot be abused.", "relevant_ids": ["mem_0011", "mem_0004"]}, - {"id": "task_3", "query": "A booking shows the wrong time to the customer. Investigate how timezones are handled for bookings.", "relevant_ids": ["mem_0003"]}, - {"id": "task_4", "query": "Update the list of services and prices the chatbot tells customers about.", "relevant_ids": ["mem_0004", "mem_0012"]}, - {"id": "task_5", "query": "Make the gallery images publicly viewable while keeping all other uploads private.", "relevant_ids": ["mem_0007"]}, - {"id": "task_6", "query": "Add a new admin-only settings page and make sure it is protected like the other admin routes.", "relevant_ids": ["mem_0006"]}, - {"id": "task_7", "query": "The chat widget sometimes crashes when Gemini replies. Make the chat resilient to bad responses.", "relevant_ids": ["mem_0008"]} + {"id": "task_1", "query": "Fix the bug where customers receive two confirmation emails for a single booking.", "paraphrase": "Customers complain they get duplicate messages in their inbox after reserving a slot.", "relevant_ids": ["mem_0005", "mem_0013"]}, + {"id": "task_2", "query": "Add rate limiting to the chat endpoint so the Gemini API key cannot be abused.", "paraphrase": "Stop people from hammering the assistant endpoint and running up our LLM bill.", "relevant_ids": ["mem_0011", "mem_0004"]}, + {"id": "task_3", "query": "A booking shows the wrong time to the customer. Investigate how timezones are handled for bookings.", "paraphrase": "A client saw the wrong hour for their appointment. Check how clock offsets are handled.", "relevant_ids": ["mem_0003"]}, + {"id": "task_4", "query": "Update the list of services and prices the chatbot tells customers about.", "paraphrase": "Change what the assistant quotes for treatments and their cost.", "relevant_ids": ["mem_0004", "mem_0012"]}, + {"id": "task_5", "query": "Make the gallery images publicly viewable while keeping all other uploads private.", "paraphrase": "Let visitors see the photo showcase while other uploads stay locked down.", "relevant_ids": ["mem_0007"]}, + {"id": "task_6", "query": "Add a new admin-only settings page and make sure it is protected like the other admin routes.", "paraphrase": "Build a settings screen only staff can open, protected like the rest of the back office.", "relevant_ids": ["mem_0006"]}, + {"id": "task_7", "query": "The chat widget sometimes crashes when Gemini replies. Make the chat resilient to bad responses.", "paraphrase": "The assistant window dies when the model returns nothing. Harden it.", "relevant_ids": ["mem_0008"]} + ], + "distractors": [ + {"type": "project", "text": "CI runs on GitHub Actions from .github/workflows/ci.yml on every pull request and on pushes to main."}, + {"type": "decision", "text": "Production deploys are manual: build the frontend, then restart the Node process behind the reverse proxy."}, + {"type": "project", "text": "Static marketing pages are pre-rendered at build time and served from the CDN edge cache."}, + {"type": "decision", "text": "We use Vitest for unit tests and Playwright for the two end-to-end smoke flows."}, + {"type": "worklog", "text": "Bumped the Vite version to 5.4 and removed the legacy rollup workaround from the config."}, + {"type": "issue", "text": "The Playwright smoke test is flaky on cold starts because the dev server needs about four seconds to boot."}, + {"type": "project", "text": "Analytics events are sent to a self-hosted Plausible instance; no third-party trackers are used anywhere."}, + {"type": "decision", "text": "The favicon and social preview images are generated from a single source SVG at build time."}, + {"type": "worklog", "text": "Removed 14 unused npm dependencies and cut the production bundle by roughly 180 kilobytes."}, + {"type": "project", "text": "Error monitoring uses a self-hosted Sentry; the DSN is read from the SENTRY_DSN environment variable."}, + {"type": "decision", "text": "Log output is structured JSON in production and human-readable pretty print in development."}, + {"type": "issue", "text": "Safari 16 renders the sticky footer one pixel off; the workaround uses translateZ(0) on the container."}, + {"type": "worklog", "text": "Migrated the contact form markup to semantic fieldset and legend elements for screen readers."}, + {"type": "project", "text": "The repository uses pnpm workspaces with the client, server and shared packages in separate folders."}, + {"type": "decision", "text": "Commit messages follow Conventional Commits so the changelog can be generated automatically."}, + {"type": "worklog", "text": "Added a robots.txt and an XML sitemap generated from the public route manifest."}, + {"type": "issue", "text": "The staging database occasionally holds a stale materialized view; refreshing it manually resolves the drift."}, + {"type": "project", "text": "Feature flags are plain environment variables read once at server startup, not a third-party service."}, + {"type": "decision", "text": "Dates in the admin tables are displayed in ISO format to avoid ambiguity for staff in different locales."}, + {"type": "worklog", "text": "Replaced the custom scroll spy on the landing page with an IntersectionObserver."}, + {"type": "project", "text": "The design tokens were exported from Figma and checked in as JSON under design/tokens.json."}, + {"type": "decision", "text": "Long-running report generation happens in a background worker, not inside a request handler."}, + {"type": "issue", "text": "Hot module replacement drops the websocket connection when the laptop wakes from sleep; a page refresh fixes it."}, + {"type": "worklog", "text": "Split the giant utils module into four focused modules and removed the circular import."}, + {"type": "project", "text": "Fonts are self-hosted as woff2 subsets; nothing is loaded from Google Fonts at runtime."}, + {"type": "decision", "text": "Form validation logic is shared between client and server through the schemas in the shared package."}, + {"type": "worklog", "text": "Added skeleton loading states to the three slowest pages in the admin dashboard."}, + {"type": "issue", "text": "The pagination component computes the wrong page count when the result set is exactly one page long."}, + {"type": "project", "text": "Backups run nightly with pg_dump to encrypted object storage and are retained for thirty days."}, + {"type": "decision", "text": "Secrets never enter the repository; local development reads them from an untracked .env file."}, + {"type": "worklog", "text": "Converted the remaining class components in the admin area to function components with hooks."}, + {"type": "project", "text": "The health check endpoint at /healthz returns build metadata and the database connectivity status."}, + {"type": "decision", "text": "Database migrations run as a separate deploy step so a failed migration never leaves a half-started server."}, + {"type": "issue", "text": "Very large uploads time out behind the proxy because the default body size limit was never raised."}, + {"type": "worklog", "text": "Documented the local setup steps in CONTRIBUTING so new contributors can start without asking."}, + {"type": "project", "text": "Lint and format are enforced by ESLint and Prettier through a pre-commit hook."}, + {"type": "decision", "text": "The API returns problem+json error bodies so the client can render messages without string matching."}, + {"type": "worklog", "text": "Added an npm script that seeds the local database with a small realistic fixture set."}, + {"type": "issue", "text": "Two integration tests share a fixture and fail when run in parallel; they are pinned to one worker for now."}, + {"type": "project", "text": "Client-side routing uses wouter rather than react-router to keep the bundle small."} ] } diff --git a/eval/results.json b/eval/results.json index 0e49497..564d4a9 100644 --- a/eval/results.json +++ b/eval/results.json @@ -1,7 +1,8 @@ { "k": 3, "budget": 120, - "exact_tokenizer": false, + "seed": 0, + "exact_tokenizer": true, "embedder": "HashingEmbedder", "n_memories": 14, "n_tasks": 7, @@ -13,22 +14,28 @@ "token_reduction_vs_baseline": 100.0 }, "full_context": { - "avg_context_tokens": 550.0, + "avg_context_tokens": 424.0, "precision": 0.102, "recall": 1.0, "token_reduction_vs_baseline": 0.0 }, - "semantic_recall": { - "avg_context_tokens": 120.1, + "random_k": { + "avg_context_tokens": 93.4, + "precision": 0.048, + "recall": 0.071, + "token_reduction_vs_baseline": 78.0 + }, + "retrieval": { + "avg_context_tokens": 93.3, "precision": 0.429, "recall": 0.929, - "token_reduction_vs_baseline": 78.2 + "token_reduction_vs_baseline": 78.0 }, "budget_recall": { - "avg_context_tokens": 113.0, + "avg_context_tokens": 93.3, "precision": 0.429, "recall": 0.929, - "token_reduction_vs_baseline": 79.5 + "token_reduction_vs_baseline": 78.0 } }, "per_task": [ @@ -45,7 +52,7 @@ "mem_0003" ], "recall": 0.5, - "tokens": 116 + "tokens": 92 }, { "task": "task_2", @@ -60,7 +67,7 @@ "mem_0010" ], "recall": 1.0, - "tokens": 131 + "tokens": 101 }, { "task": "task_3", @@ -74,7 +81,7 @@ "mem_0014" ], "recall": 1.0, - "tokens": 116 + "tokens": 92 }, { "task": "task_4", @@ -89,7 +96,7 @@ "mem_0011" ], "recall": 1.0, - "tokens": 114 + "tokens": 86 }, { "task": "task_5", @@ -103,7 +110,7 @@ "mem_0014" ], "recall": 1.0, - "tokens": 128 + "tokens": 99 }, { "task": "task_6", @@ -117,7 +124,7 @@ "mem_0012" ], "recall": 1.0, - "tokens": 103 + "tokens": 80 }, { "task": "task_7", @@ -131,7 +138,75 @@ "mem_0010" ], "recall": 1.0, - "tokens": 133 + "tokens": 103 + } + ], + "phrasing": { + "query": { + "recall": 0.929, + "avg_query_word_overlap_with_gold": 0.403 + }, + "paraphrase": { + "recall": 0.429, + "avg_query_word_overlap_with_gold": 0.032 + } + }, + "scaling": [ + { + "n_memories": 14, + "full_context_tokens": 424, + "budget_recall_tokens": 93.3, + "budget_recall": 0.929 + }, + { + "n_memories": 34, + "full_context_tokens": 805, + "budget_recall_tokens": 87.4, + "budget_recall": 0.929 + }, + { + "n_memories": 54, + "full_context_tokens": 1159, + "budget_recall_tokens": 82.1, + "budget_recall": 0.929 + } + ], + "min_score_sweep": [ + { + "min_score": 0.0, + "labelled_recall": 0.929, + "off_topic_memories_returned": 12, + "off_topic_max": 12 + }, + { + "min_score": 0.05, + "labelled_recall": 0.929, + "off_topic_memories_returned": 12, + "off_topic_max": 12 + }, + { + "min_score": 0.1, + "labelled_recall": 0.929, + "off_topic_memories_returned": 10, + "off_topic_max": 12 + }, + { + "min_score": 0.15, + "labelled_recall": 0.929, + "off_topic_memories_returned": 6, + "off_topic_max": 12 + }, + { + "min_score": 0.2, + "labelled_recall": 0.929, + "off_topic_memories_returned": 2, + "off_topic_max": 12 + }, + { + "min_score": 0.25, + "labelled_recall": 0.929, + "off_topic_memories_returned": 2, + "off_topic_max": 12 } ] -} \ No newline at end of file +} diff --git a/eval/results.md b/eval/results.md index a8c5f4b..bb6c09c 100644 --- a/eval/results.md +++ b/eval/results.md @@ -1,15 +1,54 @@ # Evaluation results - Benchmark: **14 memories**, **7 tasks**, top-k = **3**, budget = **120 tokens** -- Embedder: `HashingEmbedder` · exact tokenizer: `False` +- Embedder: `HashingEmbedder` · exact tokenizer: `True` | Arm | Avg context tokens | Recall | Precision | Tokens saved vs baseline | | --- | ---: | ---: | ---: | ---: | | No memory (control) | 0 | 0.00 | 0.00 | 100% | -| Full context (baseline) | 550 | 1.00 | 0.10 | 0% | -| Semantic recall (this engine) | 120 | 0.93 | 0.43 | 78% | -| Budget recall (≤ 120 tokens) | 113 | 0.93 | 0.43 | 80% | +| Full context (baseline — load every memory) | 424 | 1.00 | 0.10 | 0% | +| Random k (control — the saving without the retrieval) | 93 | 0.07 | 0.05 | 78% | +| Targeted retrieval (this engine) | 93 | 0.93 | 0.43 | 78% | +| Budget recall (≤ 120 tokens) | 93 | 0.93 | 0.43 | 78% | -**Headline:** semantic recall retrieves **93%** of the relevant memories while loading **78% fewer context tokens** than dumping every memory file, and at **43%** precision vs **10%** for the baseline. Under a hard **120-token cap** recall is still **93%** — the cost of memory stays fixed as the store grows. +**Read the random arm first.** It loads the same number of memories as the engine, so it reports the same ~78% token saving — at **0.07** recall against the engine's **0.93**. The saving is arithmetic (k memories out of n); only the recall gap is evidence that retrieval does anything. -_Numbers above use the offline hashing embedder so they are byte-stable in CI. Installing `sentence-transformers` (the `real` extra) raises recall further on paraphrased queries._ +Precision is reported for completeness, but the baseline's **0.10** is just `|relevant| / |store|` — an artefact of loading everything, not a meaningful comparison. + +## Does it survive a rephrase? + +| Query phrasing | Word overlap with gold memories | Recall | +| --- | ---: | ---: | +| Developer phrasing (as labelled) | 40% | 0.93 | +| Outsider paraphrase (vocabulary avoided) | 3% | 0.43 | + +The default embedder is feature hashing — lexical, not semantic. When the query stops sharing words with the memory, recall drops by **0.50**. That gap is the honest limit of the offline default, and the reason the optional `sentence-transformers` backend exists. + +## Does the cost stay flat as the store grows? + +| Memories in store | Full context tokens | Budget recall tokens | Budget recall | +| ---: | ---: | ---: | ---: | +| 14 | 424 | 93 | 0.93 | +| 34 | 805 | 87 | 0.93 | +| 54 | 1159 | 82 | 0.93 | + +Distractor memories are added from the same project. The baseline grows linearly; the budgeted arm does not. Recall is measured against the same labels throughout, so any drop is real interference from the added memories. + +## Relevance floor (`min_score`) + +| min_score | Labelled recall | Off-topic memories returned | +| ---: | ---: | ---: | +| 0.00 | 0.93 | 12/12 | +| 0.05 | 0.93 | 12/12 | +| 0.10 | 0.93 | 10/12 | +| 0.15 | 0.93 | 6/12 | +| 0.20 | 0.93 | 2/12 | +| 0.25 | 0.93 | 2/12 | + +Four deliberately off-topic queries stand in for a task the store knows nothing about. Without a floor the engine returns k memories anyway, and the MCP layer strips scores, so the agent cannot tell. `HashingEmbedder.recommended_min_score` is set from this sweep. + +## Honest limits + +- 7 tasks and 10 gold labels, all hand-written by the author. This is an engineering check, not a production-scale claim; one retrieval either way moves recall by ~0.07. +- The labels, the queries and the retriever all come from the same person, which is exactly the setup that flatters a retriever. The paraphrase arm exists to push back on that. +- Numbers use the offline hashing embedder so they are stable in CI. The `sentence-transformers` backend is not measured here. diff --git a/eval/run_eval.py b/eval/run_eval.py index ffdb581..073de99 100644 --- a/eval/run_eval.py +++ b/eval/run_eval.py @@ -1,17 +1,36 @@ -"""Evaluate semantic recall against full-context and no-memory baselines. +"""Evaluate targeted retrieval against full-context, random and no-memory arms. The claim under test: *retrieving only the top-k relevant memories loads far fewer context tokens than dumping every memory file, without dropping the facts -the agent needs.* This harness measures exactly that on a labeled benchmark and -prints a results table. +the agent needs.* -Four arms per task: +Two things make that claim easy to overstate, so both are measured here: + + * **The token saving is arithmetic, not skill.** Loading 3 of 14 memories + costs ~21% of the tokens no matter which 3 you pick. The `random_k` arm + picks k memories at random and therefore reports the *same* saving with + useless recall. Only the gap in recall between `random_k` and the retriever + is evidence that retrieval works. + * **The default embedder is lexical, not semantic.** It matches shared + wording, so a benchmark whose queries reuse the vocabulary of the answers + flatters it. Every task carries a `paraphrase` that avoids that vocabulary, + and both are reported side by side. + +Arms per task: * no_memory — control. The agent gets nothing (0 tokens, 0 recall). * full_context — the Markdown-scaffold baseline: load every memory. - * semantic_recall — this engine: load only the top-k by vector similarity. + * random_k — control. k memories chosen at random: the token saving + without the retrieval. + * retrieval — this engine: load only the top-k by vector similarity. * budget_recall — this engine under a hard token cap (recall never loads more than --budget tokens, whatever the store size). +Also reported: + * the same retrieval arm re-run on the paraphrased queries; + * a scaling experiment that grows the store with distractor memories, to test + whether recall cost really stays flat as memory accumulates; + * a sweep of the `min_score` relevance floor used by the CLI and MCP server. + Run: python eval/run_eval.py (writes eval/results.md + results.json) """ @@ -19,6 +38,7 @@ import argparse import json +import random import sys from pathlib import Path @@ -26,9 +46,17 @@ ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT / "src")) -from agent_memory import MemoryStore, count_tokens # noqa: E402 +from agent_memory import HashingEmbedder, MemoryStore, count_tokens # noqa: E402 from agent_memory.tokens import using_exact_tokenizer # noqa: E402 +ARM_LABELS = { + "no_memory": "No memory (control)", + "full_context": "Full context (baseline — load every memory)", + "random_k": "Random k (control — the saving without the retrieval)", + "retrieval": "Targeted retrieval (this engine)", + "budget_recall": "Budget recall (hard token cap)", +} + def load_dataset(path: Path) -> dict: return json.loads(path.read_text()) @@ -43,54 +71,57 @@ def _prf(retrieved_ids: list[str], gold_ids: set[str]) -> tuple[float, float]: return precision, recall -def evaluate(dataset: dict, k: int = 3, budget: int = 120) -> dict: - store = MemoryStore() - for m in dataset["memories"]: +def mean(xs: list[float]) -> float: + return sum(xs) / len(xs) if xs else 0.0 + + +def build_store(memories: list[dict]) -> MemoryStore: + # Explicit embedder: the published numbers are the offline one's, and we do + # not want them to change silently on a machine with sentence-transformers. + store = MemoryStore(embedder=HashingEmbedder()) + for m in memories: store.write(m["text"], type=m["type"], id=m["id"]) + return store - all_ids = [m["id"] for m in dataset["memories"]] - full_tokens = sum(count_tokens(m["text"]) for m in dataset["memories"]) - arms = { - "no_memory": {"tokens": [], "precision": [], "recall": []}, - "full_context": {"tokens": [], "precision": [], "recall": []}, - "semantic_recall": {"tokens": [], "precision": [], "recall": []}, - "budget_recall": {"tokens": [], "precision": [], "recall": []}, - } +def evaluate(dataset: dict, k: int = 3, budget: int = 120, seed: int = 0) -> dict: + memories = dataset["memories"] + store = build_store(memories) + rng = random.Random(seed) + + all_ids = [m["id"] for m in memories] + full_tokens = sum(count_tokens(m["text"]) for m in memories) + by_id = {m["id"]: m for m in memories} + + arms = {name: {"tokens": [], "precision": [], "recall": []} for name in ARM_LABELS} per_task = [] for task in dataset["tasks"]: gold = set(task["relevant_ids"]) - # no_memory - arms["no_memory"]["tokens"].append(0) - arms["no_memory"]["precision"].append(0.0) - arms["no_memory"]["recall"].append(0.0) + def record(name: str, ids: list[str], tokens: int) -> tuple[float, float]: + p, r = _prf(ids, gold) + arms[name]["tokens"].append(tokens) + arms[name]["precision"].append(p) + arms[name]["recall"].append(r) + return p, r - # full_context - p, r = _prf(all_ids, gold) - arms["full_context"]["tokens"].append(full_tokens) - arms["full_context"]["precision"].append(p) - arms["full_context"]["recall"].append(r) + record("no_memory", [], 0) + record("full_context", all_ids, full_tokens) + + # Control: the token saving with none of the retrieval. + picked = rng.sample(all_ids, min(k, len(all_ids))) + record("random_k", picked, sum(count_tokens(by_id[i]["text"]) for i in picked)) - # semantic_recall hits = store.recall(task["query"], k=k) retrieved_ids = [h.entry.id for h in hits] tokens = sum(h.entry.tokens for h in hits) - p, r = _prf(retrieved_ids, gold) - arms["semantic_recall"]["tokens"].append(tokens) - arms["semantic_recall"]["precision"].append(p) - arms["semantic_recall"]["recall"].append(r) + _, r = record("retrieval", retrieved_ids, tokens) - # budget_recall — same engine, hard token cap bhits = store.recall(task["query"], k=k, budget_tokens=budget) - bids = [h.entry.id for h in bhits] btokens = sum(h.entry.tokens for h in bhits) assert btokens <= budget, "budget packing must never overflow" - bp, br = _prf(bids, gold) - arms["budget_recall"]["tokens"].append(btokens) - arms["budget_recall"]["precision"].append(bp) - arms["budget_recall"]["recall"].append(br) + record("budget_recall", [h.entry.id for h in bhits], btokens) per_task.append( { @@ -103,9 +134,6 @@ def evaluate(dataset: dict, k: int = 3, budget: int = 120) -> dict: } ) - def mean(xs: list[float]) -> float: - return sum(xs) / len(xs) if xs else 0.0 - summary = {} baseline_tokens = mean(arms["full_context"]["tokens"]) for name, m in arms.items(): @@ -124,23 +152,129 @@ def mean(xs: list[float]) -> float: return { "k": k, "budget": budget, + "seed": seed, "exact_tokenizer": using_exact_tokenizer(), "embedder": store.stats()["embedder"], - "n_memories": len(dataset["memories"]), + "n_memories": len(memories), "n_tasks": len(dataset["tasks"]), "summary": summary, "per_task": per_task, + "phrasing": phrasing_gap(dataset, k=k), + "scaling": scaling(dataset, k=k, budget=budget), + "min_score_sweep": min_score_sweep(dataset, k=k), } +def phrasing_gap(dataset: dict, k: int = 3) -> dict: + """Recall on the developer phrasing vs an outsider's paraphrase. + + The default embedder matches shared wording. This is the number that says + how much of the headline recall comes from the benchmark's phrasing. + """ + store = build_store(dataset["memories"]) + out = {} + for field in ("query", "paraphrase"): + recalls, overlaps = [], [] + for task in dataset["tasks"]: + if field not in task: + continue + gold = set(task["relevant_ids"]) + hits = store.recall(task[field], k=k) + _, r = _prf([h.entry.id for h in hits], gold) + recalls.append(r) + overlaps.append(_word_overlap(task[field], dataset["memories"], gold)) + out[field] = { + "recall": round(mean(recalls), 3), + "avg_query_word_overlap_with_gold": round(mean(overlaps), 3), + } + return out + + +def _word_overlap(query: str, memories: list[dict], gold: set[str]) -> float: + """Share of the query's content words that also appear in its gold memories.""" + import re + + stop = set( + "a an the to is are of in on for and it that this with be as at by from " + "how do we make sure so they their there these it's its can should".split() + ) + words = {w for w in re.findall(r"[a-z0-9]+", query.lower()) if w not in stop} + if not words: + return 0.0 + gold_text = " ".join(m["text"] for m in memories if m["id"] in gold).lower() + gold_words = set(re.findall(r"[a-z0-9]+", gold_text)) + return len(words & gold_words) / len(words) + + +def scaling(dataset: dict, k: int = 3, budget: int = 120) -> list[dict]: + """Grow the store with distractors and re-measure. + + "The cost of memory stays fixed as the store grows" is a claim about a store + that grows, so the store has to actually grow for it to mean anything. + """ + distractors = dataset.get("distractors", []) + rows = [] + for extra in (0, len(distractors) // 2, len(distractors)): + memories = list(dataset["memories"]) + memories += [ + {"id": f"dis_{i:04d}", "type": d["type"], "text": d["text"]} + for i, d in enumerate(distractors[:extra]) + ] + store = build_store(memories) + full_tokens = sum(count_tokens(m["text"]) for m in memories) + recalls, tokens = [], [] + for task in dataset["tasks"]: + gold = set(task["relevant_ids"]) + hits = store.recall(task["query"], k=k, budget_tokens=budget) + _, r = _prf([h.entry.id for h in hits], gold) + recalls.append(r) + tokens.append(sum(h.entry.tokens for h in hits)) + rows.append( + { + "n_memories": len(memories), + "full_context_tokens": full_tokens, + "budget_recall_tokens": round(mean(tokens), 1), + "budget_recall": round(mean(recalls), 3), + } + ) + return rows + + +def min_score_sweep(dataset: dict, k: int = 3) -> list[dict]: + """Calibrate the relevance floor used by the CLI and MCP server. + + Labelled recall on real queries must not drop, while queries about things + the store knows nothing about should return nothing at all. + """ + store = build_store(dataset["memories"]) + off_topic = [ + "how do I bake sourdough bread at home", + "what is the capital of Peru", + "my cat will not eat her food", + "best hiking boots for winter walking", + ] + rows = [] + for threshold in (0.0, 0.05, 0.10, 0.15, 0.20, 0.25): + recalls = [] + for task in dataset["tasks"]: + gold = set(task["relevant_ids"]) + hits = store.recall(task["query"], k=k, min_score=threshold) + _, r = _prf([h.entry.id for h in hits], gold) + recalls.append(r) + junk = sum(len(store.recall(q, k=k, min_score=threshold)) for q in off_topic) + rows.append( + { + "min_score": threshold, + "labelled_recall": round(mean(recalls), 3), + "off_topic_memories_returned": junk, + "off_topic_max": k * len(off_topic), + } + ) + return rows + + def render_markdown(results: dict) -> str: s = results["summary"] - rows = [ - ("No memory (control)", s["no_memory"]), - ("Full context (baseline)", s["full_context"]), - ("Semantic recall (this engine)", s["semantic_recall"]), - (f"Budget recall (≤ {results['budget']} tokens)", s["budget_recall"]), - ] lines = [ "# Evaluation results", "", @@ -153,26 +287,98 @@ def render_markdown(results: dict) -> str: "| Arm | Avg context tokens | Recall | Precision | Tokens saved vs baseline |", "| --- | ---: | ---: | ---: | ---: |", ] - for label, m in rows: + for name, label in ARM_LABELS.items(): + m = s[name] + if name == "budget_recall": + label = f"Budget recall (≤ {results['budget']} tokens)" lines.append( f"| {label} | {m['avg_context_tokens']:.0f} | {m['recall']:.2f} " f"| {m['precision']:.2f} | {m['token_reduction_vs_baseline']:.0f}% |" ) - sem = s["semantic_recall"] - bud = s["budget_recall"] + + ret, rnd = s["retrieval"], s["random_k"] + lines += [ + "", + "**Read the random arm first.** It loads the same number of memories as " + f"the engine, so it reports the same ~{rnd['token_reduction_vs_baseline']:.0f}% " + "token saving — at " + f"**{rnd['recall']:.2f}** recall against the engine's **{ret['recall']:.2f}**. " + "The saving is arithmetic (k memories out of n); only the recall gap is " + "evidence that retrieval does anything.", + "", + "Precision is reported for completeness, but the baseline's " + f"**{s['full_context']['precision']:.2f}** is just " + "`|relevant| / |store|` — an artefact of loading everything, not a " + "meaningful comparison.", + "", + "## Does it survive a rephrase?", + "", + "| Query phrasing | Word overlap with gold memories | Recall |", + "| --- | ---: | ---: |", + ] + ph = results["phrasing"] + for field, label in ( + ("query", "Developer phrasing (as labelled)"), + ("paraphrase", "Outsider paraphrase (vocabulary avoided)"), + ): + row = ph.get(field) + if row: + lines.append( + f"| {label} | {row['avg_query_word_overlap_with_gold']:.0%} | " + f"{row['recall']:.2f} |" + ) lines += [ "", - f"**Headline:** semantic recall retrieves **{sem['recall']:.0%}** of the " - f"relevant memories while loading **{sem['token_reduction_vs_baseline']:.0f}% " - f"fewer context tokens** than dumping every memory file, and at " - f"**{sem['precision']:.0%}** precision vs " - f"**{s['full_context']['precision']:.0%}** for the baseline. Under a hard " - f"**{results['budget']}-token cap** recall is still **{bud['recall']:.0%}** — " - f"the cost of memory stays fixed as the store grows.", + "The default embedder is feature hashing — lexical, not semantic. When the " + "query stops sharing words with the memory, recall drops by " + f"**{(ph['query']['recall'] - ph['paraphrase']['recall']):.2f}**. That gap is " + "the honest limit of the offline default, and the reason the optional " + "`sentence-transformers` backend exists.", "", - "_Numbers above use the offline hashing embedder so they are byte-stable " - "in CI. Installing `sentence-transformers` (the `real` extra) raises recall " - "further on paraphrased queries._", + "## Does the cost stay flat as the store grows?", + "", + "| Memories in store | Full context tokens | Budget recall tokens | Budget recall |", + "| ---: | ---: | ---: | ---: |", + ] + for row in results["scaling"]: + lines.append( + f"| {row['n_memories']} | {row['full_context_tokens']} | " + f"{row['budget_recall_tokens']:.0f} | {row['budget_recall']:.2f} |" + ) + lines += [ + "", + "Distractor memories are added from the same project. The baseline grows " + "linearly; the budgeted arm does not. Recall is measured against the same " + "labels throughout, so any drop is real interference from the added memories.", + "", + "## Relevance floor (`min_score`)", + "", + "| min_score | Labelled recall | Off-topic memories returned |", + "| ---: | ---: | ---: |", + ] + for row in results["min_score_sweep"]: + lines.append( + f"| {row['min_score']:.2f} | {row['labelled_recall']:.2f} | " + f"{row['off_topic_memories_returned']}/{row['off_topic_max']} |" + ) + lines += [ + "", + "Four deliberately off-topic queries stand in for a task the store knows " + "nothing about. Without a floor the engine returns k memories anyway, and " + "the MCP layer strips scores, so the agent cannot tell. `HashingEmbedder." + "recommended_min_score` is set from this sweep.", + "", + "## Honest limits", + "", + f"- {results['n_tasks']} tasks and " + f"{sum(len(t['gold']) for t in results['per_task'])} gold labels, all " + "hand-written by the author. This is an engineering check, not a " + "production-scale claim; one retrieval either way moves recall by ~0.07.", + "- The labels, the queries and the retriever all come from the same person, " + "which is exactly the setup that flatters a retriever. The paraphrase arm " + "exists to push back on that.", + "- Numbers use the offline hashing embedder so they are stable in CI. The " + "`sentence-transformers` backend is not measured here.", ] return "\n".join(lines) @@ -183,6 +389,9 @@ def main() -> None: parser.add_argument( "--budget", type=int, default=120, help="token cap for the budget arm" ) + parser.add_argument( + "--seed", type=int, default=0, help="seed for the random-k control arm" + ) parser.add_argument( "--dataset", type=Path, default=Path(__file__).parent / "dataset.json" ) @@ -190,12 +399,18 @@ def main() -> None: args = parser.parse_args() dataset = load_dataset(args.dataset) - results = evaluate(dataset, k=args.k, budget=args.budget) + results = evaluate(dataset, k=args.k, budget=args.budget, seed=args.seed) md = render_markdown(results) (args.out_dir / "results.md").write_text(md + "\n") - (args.out_dir / "results.json").write_text(json.dumps(results, indent=2)) + (args.out_dir / "results.json").write_text(json.dumps(results, indent=2) + "\n") print(md) + if not results["exact_tokenizer"]: + print( + "\n[warning] tiktoken is not installed, so token counts are approximate " + 'and will not match the published table. Install it with: pip install -e ".[dev]"', + file=sys.stderr, + ) if __name__ == "__main__": diff --git a/examples/quickstart.py b/examples/quickstart.py index c09cdd2..8ec07f6 100644 --- a/examples/quickstart.py +++ b/examples/quickstart.py @@ -9,9 +9,10 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) -from agent_memory import MemoryStore # noqa: E402 +from agent_memory import MemoryStore, default_min_score # noqa: E402 store = MemoryStore() # in-memory; pass path=... to persist +floor = default_min_score(store.embedder) # --- session 1: Claude Code learns things and hands off ----------------- store.write("Bookings are stored in UTC and shown in Asia/Tbilisi.", @@ -20,16 +21,37 @@ type="decision", agent="claude-code") store.write("Admin routes are guarded by requireAdmin in server/auth.ts.", type="decision", agent="claude-code") +store.write("The public chat endpoint /api/chat has no rate limiter; the Gemini " + "API key is exposed to abuse until one is added.", + type="issue", agent="claude-code") store.write("Done: fixed double confirmation emails. " "Next: add rate limiting to the public chat endpoint.", type="handoff", agent="claude-code") +# Writing the same thing twice does not duplicate it, and says so. +_, stored = store.write_with_status( + "Admin routes are guarded by requireAdmin in server/auth.ts.", + type="decision", agent="claude-code") +print(f"Re-writing a known fact stored a new memory? {stored}\n") + # --- session 2: a different agent (Codex) boots from the same store ----- task = "add rate limiting to the chat endpoint" print(f"Task: {task}\n") -handoff = store.latest("handoff") -print(f"Last handoff [{handoff.agent}]: {handoff.text}\n") - -for hit in store.recall(task, k=3, budget_tokens=100): +handoff, hits = store.boot(task, k=3, budget_tokens=100, min_score=floor) +if handoff: + print(f"Last handoff [{handoff.agent}]: {handoff.text}\n") +for hit in hits: print(f" {hit.score:.2f} [{hit.entry.type} · {hit.entry.agent}] {hit.entry.text}") + +# --- a task the store knows nothing about -------------------------------- +_, unrelated = store.boot("choose a color palette for the print brochure", + k=3, budget_tokens=100, min_score=floor) +print(f"\nMemories recalled for an unrelated task: {len(unrelated)} " + f"(the relevance floor is {floor}; without it this would return 3)") + +# --- correcting memory ---------------------------------------------------- +stale = store.write("Deploys are triggered from the old Jenkins box.", type="state") +store.update(stale.id, text="Deploys are triggered from GitHub Actions.") +store.forget(stale.id) +print(f"Memories after update + forget: {store.stats()['count']}") diff --git a/pyproject.toml b/pyproject.toml index 4fc88e7..7c4d04f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "agent-memory-engine" -version = "0.2.0" +version = "0.3.0" description = "One shared memory for coding agents (Claude Code, Codex, Cursor) over MCP, with token-budgeted recall and an evaluation harness." readme = "README.md" requires-python = ">=3.10" @@ -17,8 +17,12 @@ dependencies = ["numpy>=1.24"] # Real semantic embeddings + exact token accounting. Drop-in, no code change. real = ["sentence-transformers>=2.2", "tiktoken>=0.5"] # Model Context Protocol server so agents can call the engine as tools. +# Both the 1.x (FastMCP) and 2.x (MCPServer) APIs are supported. mcp = ["mcp>=1.0"] -dev = ["pytest>=7.0"] +# `dev` pulls in mcp and tiktoken on purpose: the MCP server is the headline +# feature and must be exercised in CI, and pinning the tokenizer is what makes +# the published evaluation numbers reproducible rather than approximate. +dev = ["pytest>=7.0", "mcp>=1.0", "tiktoken>=0.5"] [project.urls] Repository = "https://github.com/Ninadnj/ai-agent-memory-scaffold" diff --git a/scripts/ingest_markdown.py b/scripts/ingest_markdown.py index 753f0fa..15a357b 100644 --- a/scripts/ingest_markdown.py +++ b/scripts/ingest_markdown.py @@ -16,7 +16,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) -from agent_memory import MEMORY_TYPES, MemoryStore # noqa: E402 +from agent_memory import MEMORY_TYPES, MemoryStore, count_tokens # noqa: E402 # Map common scaffold filenames to memory types. _FILE_TYPE = { @@ -34,29 +34,70 @@ def type_for(filename: str) -> str: return _FILE_TYPE.get(stem, "fact") -def sections(text: str) -> list[str]: - """Split a Markdown file into `##` sections, dropping headings and blanks.""" - chunks = re.split(r"^##\s+.*$", text, flags=re.MULTILINE) - return [c.strip() for c in chunks if c.strip()] +def sections(text: str, max_tokens: int = 120) -> list[str]: + """Split a Markdown file into `##` sections, dropping headings and blanks. + + Long sections are split further on blank lines. A memory bigger than a + typical recall budget can never be packed into one, so it would sit in the + store permanently unreachable — chunking keeps every ingested memory + retrievable. + """ + chunks = [c.strip() for c in re.split(r"^##\s+.*$", text, flags=re.MULTILINE)] + out: list[str] = [] + for chunk in chunks: + if not chunk: + continue + if count_tokens(chunk) <= max_tokens: + out.append(chunk) + continue + buffer: list[str] = [] + for para in re.split(r"\n\s*\n", chunk): + para = para.strip() + if not para: + continue + candidate = "\n\n".join(buffer + [para]) + if buffer and count_tokens(candidate) > max_tokens: + out.append("\n\n".join(buffer)) + buffer = [para] + else: + buffer.append(para) + if buffer: + out.append("\n\n".join(buffer)) + return out def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("source", type=Path, help="directory of .md memory files") parser.add_argument("--path", type=Path, default=None, help="output store.json") + parser.add_argument( + "--max-tokens", + type=int, + default=120, + help="split sections larger than this so they stay recallable", + ) args = parser.parse_args() + if args.path is None: + parser.error( + "--path is required: without it the ingested memories would be built " + "in memory and thrown away. Example: --path ~/.agent_memory/store.json" + ) + store = MemoryStore(path=args.path) - written = 0 + written = skipped = 0 for md in sorted(args.source.glob("*.md")): mem_type = type_for(md.name) if mem_type not in MEMORY_TYPES: mem_type = "fact" - for body in sections(md.read_text()): - store.write(body, type=mem_type, metadata={"source": md.name}) - written += 1 - dest = args.path or "(in-memory, not saved — pass --path)" - print(f"Ingested {written} memories from {args.source} -> {dest}") + for body in sections(md.read_text(), max_tokens=args.max_tokens): + _, stored = store.write_with_status( + body, type=mem_type, metadata={"source": md.name} + ) + written += stored + skipped += not stored + note = f" ({skipped} skipped as near-duplicates)" if skipped else "" + print(f"Ingested {written} memories from {args.source} -> {args.path}{note}") if __name__ == "__main__": diff --git a/src/agent_memory/__init__.py b/src/agent_memory/__init__.py index de74b14..21358fa 100644 --- a/src/agent_memory/__init__.py +++ b/src/agent_memory/__init__.py @@ -1,24 +1,33 @@ -"""Agent Memory Engine — semantic memory + handoffs for AI coding agents.""" +"""Agent Memory Engine — durable memory and handoffs for AI coding agents.""" from .embeddings import ( Embedder, HashingEmbedder, SentenceTransformerEmbedder, default_embedder, + default_min_score, +) +from .store import ( + MEMORY_TYPES, + STORE_FORMAT, + MemoryEntry, + MemoryStore, + RecallHit, ) -from .store import MEMORY_TYPES, MemoryEntry, MemoryStore, RecallHit from .tokens import count_tokens -__version__ = "0.1.0" +__version__ = "0.3.0" __all__ = [ "MemoryStore", "MemoryEntry", "RecallHit", "MEMORY_TYPES", + "STORE_FORMAT", "Embedder", "HashingEmbedder", "SentenceTransformerEmbedder", "default_embedder", + "default_min_score", "count_tokens", ] diff --git a/src/agent_memory/cli.py b/src/agent_memory/cli.py index 9be4934..6c4a34b 100644 --- a/src/agent_memory/cli.py +++ b/src/agent_memory/cli.py @@ -4,6 +4,9 @@ agent-memory recall "how are timezones handled" -k 3 --budget 200 agent-memory handoff --done "Fixed double emails." --next "Add rate limiting." agent-memory boot "continue the booking bug fix" + agent-memory list --type decision + agent-memory update mem_0003 "Bookings are stored in UTC; UI converts." + agent-memory forget mem_0007 agent-memory stats Uses a JSON store at $AGENT_MEMORY_PATH (default ~/.agent_memory/store.json). @@ -17,6 +20,7 @@ import os from pathlib import Path +from .embeddings import default_min_score from .store import MEMORY_TYPES, MemoryStore DEFAULT_STORE = Path( @@ -24,24 +28,44 @@ ).expanduser() DEFAULT_AGENT = os.environ.get("AGENT_MEMORY_AGENT", "") +# Sentinel: resolved per-embedder once the store is open (see default_min_score). +AUTO_MIN_SCORE = -1.0 + def _store(args) -> MemoryStore: return MemoryStore(path=args.path) +def _min_score(args, store: MemoryStore) -> float: + if getattr(args, "min_score", AUTO_MIN_SCORE) != AUTO_MIN_SCORE: + return args.min_score + return default_min_score(store.embedder) + + def _tag(entry) -> str: return f"{entry.type} · {entry.agent}" if entry.agent else entry.type def cmd_write(args) -> None: - entry = _store(args).write(args.text, type=args.type, agent=args.agent) - print(f"Saved {entry.id} ({entry.type}).") + entry, stored = _store(args).write_with_status( + args.text, type=args.type, agent=args.agent + ) + if stored: + print(f"Saved {entry.id} ({entry.type}).") + else: + print(f"Not saved — near-duplicate of {entry.id}: {entry.text}") def cmd_recall(args) -> None: - hits = _store(args).recall(args.query, k=args.k, budget_tokens=args.budget) + store = _store(args) + hits = store.recall( + args.query, + k=args.k, + budget_tokens=args.budget, + min_score=_min_score(args, store), + ) if not hits: - print("No memories yet.") + print("No relevant memories.") return for h in hits: print(f"{h.score:.2f} [{_tag(h.entry)}] {h.entry.text}") @@ -51,25 +75,45 @@ def cmd_handoff(args) -> None: text = f"Done: {args.done} Next: {args.next}" if args.watch_out: text += f" Watch out: {args.watch_out}" - entry = _store(args).write(text, type="handoff", agent=args.agent) - print(f"Handoff saved ({entry.id}).") + entry, stored = _store(args).write_with_status( + text, type="handoff", agent=args.agent + ) + if stored: + print(f"Handoff saved ({entry.id}).") + else: + print(f"Identical handoff already stored as {entry.id}; nothing written.") def cmd_boot(args) -> None: store = _store(args) - handoff = store.latest("handoff") - remaining = args.budget + handoff, hits = store.boot( + args.task, k=5, budget_tokens=args.budget, min_score=_min_score(args, store) + ) if handoff is not None: print(f"Last handoff [{_tag(handoff)}]: {handoff.text}") - if remaining is not None: - remaining = max(0, remaining - handoff.tokens) - hits = store.recall(args.task, k=5, budget_tokens=remaining) for h in hits: - if handoff is not None and h.entry.id == handoff.id: - continue print(f"- [{_tag(h.entry)}] {h.entry.text}") +def cmd_list(args) -> None: + entries = [e for e in reversed(_store(args).all()) if not args.type or e.type == args.type] + if not entries: + print("No memories stored.") + return + for e in entries[: args.limit]: + print(f"{e.id} [{_tag(e)}] {e.text}") + + +def cmd_update(args) -> None: + entry = _store(args).update(args.id, text=args.text) + print(f"Updated {entry.id}." if entry else f"No memory with id {args.id}.") + + +def cmd_forget(args) -> None: + ok = _store(args).forget(args.id) + print(f"Forgot {args.id}." if ok else f"No memory with id {args.id}.") + + def cmd_stats(args) -> None: s = _store(args).stats() print(f"{s['count']} memories | {s['total_tokens']} tokens | {s['embedder']}") @@ -92,6 +136,14 @@ def build_parser() -> argparse.ArgumentParser: r.add_argument("query") r.add_argument("-k", type=int, default=5) r.add_argument("--budget", type=int, default=None, help="max context tokens") + r.add_argument( + "--min-score", + type=float, + default=AUTO_MIN_SCORE, + dest="min_score", + help="drop matches weaker than this cosine score (0 disables; " + "default is calibrated per embedder)", + ) r.set_defaults(func=cmd_recall) h = sub.add_parser("handoff", help="save a handoff for the next agent") @@ -103,8 +155,25 @@ def build_parser() -> argparse.ArgumentParser: b = sub.add_parser("boot", help="latest handoff + relevant memories") b.add_argument("task") b.add_argument("--budget", type=int, default=300, help="max context tokens") + b.add_argument( + "--min-score", type=float, default=AUTO_MIN_SCORE, dest="min_score" + ) b.set_defaults(func=cmd_boot) + ls = sub.add_parser("list", help="list memories with their ids") + ls.add_argument("--type", default="", choices=[""] + sorted(MEMORY_TYPES)) + ls.add_argument("--limit", type=int, default=20) + ls.set_defaults(func=cmd_list) + + u = sub.add_parser("update", help="replace the text of a memory") + u.add_argument("id") + u.add_argument("text") + u.set_defaults(func=cmd_update) + + f = sub.add_parser("forget", help="delete a memory that is wrong or stale") + f.add_argument("id") + f.set_defaults(func=cmd_forget) + s = sub.add_parser("stats", help="show store stats") s.set_defaults(func=cmd_stats) return parser diff --git a/src/agent_memory/embeddings.py b/src/agent_memory/embeddings.py index 8b1910f..89219de 100644 --- a/src/agent_memory/embeddings.py +++ b/src/agent_memory/embeddings.py @@ -20,6 +20,7 @@ import hashlib import os import re +import sys from typing import Protocol, runtime_checkable import numpy as np @@ -41,6 +42,10 @@ class Embedder(Protocol): """Anything that turns text into unit-norm vectors of a fixed dimension.""" dim: int + # Cosine score below which a match should be treated as unrelated. Every + # embedder has its own scale, so this travels with the backend rather than + # being a single constant in the retrieval code. + recommended_min_score: float def embed(self, texts: list[str]) -> np.ndarray: # (n, dim) float32 ... @@ -64,6 +69,14 @@ def _features(text: str) -> list[str]: class HashingEmbedder: """Deterministic feature-hashing embedder. No dependencies, no network.""" + # Calibrated on eval/dataset.json: see the `min_score` sweep in + # eval/results.md. At 0.15 labelled recall is untouched while half the + # matches for off-topic queries are dropped. A higher floor cuts more noise + # on that benchmark, but the lowest-scoring gold memory in it sits at 0.13, + # so 0.15 is deliberately close to the observed floor of "genuinely + # relevant" rather than as aggressive as the sweep alone would allow. + recommended_min_score = 0.15 + def __init__(self, dim: int = 512) -> None: self.dim = dim @@ -88,6 +101,12 @@ def embed(self, texts: list[str]) -> np.ndarray: class SentenceTransformerEmbedder: """Real semantic embeddings. Optional: needs ``sentence-transformers``.""" + # NOT calibrated in this repo — the offline benchmark runs on the hashing + # embedder, and MiniLM cosines sit on a higher scale (unrelated pairs often + # score 0.1-0.3). Treat this as a conservative starting point and measure on + # your own store before relying on it. + recommended_min_score = 0.25 + def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None: from sentence_transformers import SentenceTransformer # lazy import @@ -102,15 +121,45 @@ def embed(self, texts: list[str]) -> np.ndarray: return vecs.astype(np.float32) +def default_min_score(embedder: Embedder) -> float: + """Relevance floor to use for an agent-facing call. + + ``AGENT_MEMORY_MIN_SCORE`` overrides; otherwise the backend's own + calibrated value is used. Library callers of ``MemoryStore.recall`` get no + floor unless they ask for one — this is applied at the CLI/MCP boundary, + where the caller is an agent that cannot see the scores. + """ + override = os.environ.get("AGENT_MEMORY_MIN_SCORE") + if override is not None: + return float(override) + return float(getattr(embedder, "recommended_min_score", 0.0)) + + def default_embedder() -> Embedder: """Prefer the real model when available; fall back to hashing. - Force the offline embedder with ``AGENT_MEMORY_EMBEDDER=hashing`` (CI does - this so results are byte-stable). + ``AGENT_MEMORY_EMBEDDER`` selects explicitly: ``hashing`` forces the offline + embedder (CI does this so results are byte-stable), ``sentence-transformers`` + demands the real one and raises if it cannot be loaded. The default is + ``auto``, which tries the real model and warns — loudly, on stderr — before + falling back, because the two produce incompatible vectors and a silent + switch is how a store ends up half-embedded by each. """ - if os.environ.get("AGENT_MEMORY_EMBEDDER", "").lower() != "hashing": - try: - return SentenceTransformerEmbedder() - except Exception: - pass + choice = os.environ.get("AGENT_MEMORY_EMBEDDER", "auto").lower() + if choice == "hashing": + return HashingEmbedder() + try: + return SentenceTransformerEmbedder() + except Exception as exc: + if choice in {"sentence-transformers", "sentence_transformers", "real"}: + raise RuntimeError( + f"AGENT_MEMORY_EMBEDDER={choice} but sentence-transformers could " + f'not be loaded: {exc}. Install it with: pip install "agent-memory-engine[real]"' + ) from exc + print( + f"[agent-memory] sentence-transformers unavailable ({exc.__class__.__name__}); " + "using the offline HashingEmbedder. Set AGENT_MEMORY_EMBEDDER=hashing to " + "silence this.", + file=sys.stderr, + ) return HashingEmbedder() diff --git a/src/agent_memory/mcp_server.py b/src/agent_memory/mcp_server.py index 8116365..32e1b13 100644 --- a/src/agent_memory/mcp_server.py +++ b/src/agent_memory/mcp_server.py @@ -7,7 +7,9 @@ The tool outputs are deliberately compact — recall is token-budgeted and scores/timestamps are omitted — because everything a memory tool returns is -paid for again in the calling agent's context window. +paid for again in the calling agent's context window. The flip side is that the +agent cannot judge relevance itself, so weak matches are filtered out here +rather than passed along unlabelled. Requires the optional `mcp` dependency (`pip install "agent-memory-engine[mcp]"`). The imports are deferred so the rest of the package works without it. @@ -20,7 +22,9 @@ import os from pathlib import Path +from typing import Optional +from .embeddings import default_min_score from .store import MEMORY_TYPES, MemoryStore DEFAULT_STORE = Path( @@ -31,6 +35,31 @@ DEFAULT_AGENT = os.environ.get("AGENT_MEMORY_AGENT", "") +def _load_server_class(): + """Return the FastMCP-style server class across `mcp` major versions. + + `mcp` 1.x exposes it as `mcp.server.fastmcp.FastMCP`; 2.x renamed the module + and the class to `mcp.server.mcpserver.MCPServer`. Both take a name and + provide `.tool()` and `.run()`, so the rest of this module is unchanged. + """ + try: + from mcp.server.fastmcp import FastMCP # mcp 1.x + + return FastMCP + except ImportError: + pass + try: + from mcp.server.mcpserver import MCPServer # mcp 2.x + + return MCPServer + except ImportError as exc: + raise SystemExit( + "Could not load an MCP server class from the installed 'mcp' package " + f"({exc}). Install a supported version with: " + 'pip install "agent-memory-engine[mcp]"' + ) from exc + + def _tag(entry) -> str: return f"{entry.type} · {entry.agent}" if entry.agent else entry.type @@ -39,19 +68,19 @@ def _render(hits) -> str: return "\n".join(f"- [{_tag(h.entry)}] {h.entry.text}" for h in hits) -def build_server(store_path: Path = DEFAULT_STORE, agent: str = DEFAULT_AGENT): - """Construct the FastMCP server. Imports `mcp` lazily so importing this - module never hard-fails when the optional dependency is absent.""" - try: - from mcp.server.fastmcp import FastMCP - except ImportError as exc: # pragma: no cover - depends on optional extra - raise SystemExit( - "The MCP server needs the 'mcp' package. " - 'Install it with: pip install "agent-memory-engine[mcp]"' - ) from exc +def build_server( + store_path: Path = DEFAULT_STORE, + agent: str = DEFAULT_AGENT, + min_score: Optional[float] = None, +): + """Construct the MCP server. Imports `mcp` lazily so importing this module + never hard-fails when the optional dependency is absent.""" + server_class = _load_server_class() store = MemoryStore(path=store_path) - server = FastMCP("agent-memory") + if min_score is None: + min_score = default_min_score(store.embedder) + server = server_class("agent-memory") @server.tool() def memory_write(text: str, type: str = "fact") -> str: @@ -59,29 +88,34 @@ def memory_write(text: str, type: str = "fact") -> str: state, handoff, worklog, fact. Near-duplicates are skipped.""" if type not in MEMORY_TYPES: return f"Error: type must be one of {sorted(MEMORY_TYPES)}." - entry = store.write(text, type=type, agent=agent) + entry, stored = store.write_with_status(text, type=type, agent=agent) + if not stored: + return ( + f"Not saved — near-duplicate of {entry.id}: {entry.text!r} " + "Use memory_update to revise it if this supersedes it." + ) return f"Saved {entry.id} ({entry.type})." @server.tool() def memory_recall(query: str, k: int = 5, budget_tokens: int = 300) -> str: """Recall the most relevant memories for `query`, never exceeding `budget_tokens` of context. Set budget_tokens=0 for no cap.""" - hits = store.recall(query, k=k, budget_tokens=budget_tokens or None) + hits = store.recall( + query, k=k, budget_tokens=budget_tokens or None, min_score=min_score + ) return _render(hits) if hits else "No relevant memories." @server.tool() def memory_boot(task: str, budget_tokens: int = 300) -> str: """Call once at the start of a session: returns the latest handoff from the previous agent plus the memories most relevant to `task`, packed - under one token budget.""" + under one memory-content token budget.""" parts: list[str] = [] - remaining = budget_tokens - handoff = store.latest("handoff") + handoff, hits = store.boot( + task, k=5, budget_tokens=budget_tokens, min_score=min_score + ) if handoff is not None: parts.append(f"Last handoff [{_tag(handoff)}]: {handoff.text}") - remaining = max(0, remaining - handoff.tokens) - hits = store.recall(task, k=5, budget_tokens=remaining or None) - hits = [h for h in hits if handoff is None or h.entry.id != handoff.id] if hits: parts.append(_render(hits)) return "\n".join(parts) if parts else "Empty store — start fresh." @@ -93,9 +127,41 @@ def memory_handoff(done: str, next_steps: str, warnings: str = "") -> str: text = f"Done: {done} Next: {next_steps}" if warnings: text += f" Watch out: {warnings}" - entry = store.write(text, type="handoff", agent=agent) + entry, stored = store.write_with_status(text, type="handoff", agent=agent) + if not stored: + return f"Identical handoff already stored as {entry.id}; nothing written." return f"Handoff saved ({entry.id}). The next agent gets it via memory_boot." + @server.tool() + def memory_update(id: str, text: str) -> str: + """Replace the text of an existing memory. Use this when a fact changes, + instead of writing a second memory that contradicts the first.""" + entry = store.update(id, text=text) + if entry is None: + return f"No memory with id {id}." + return f"Updated {entry.id} ({entry.type})." + + @server.tool() + def memory_forget(id: str) -> str: + """Delete a memory that is wrong or has gone stale. Find ids with + memory_list.""" + return ( + f"Forgot {id}." if store.forget(id) else f"No memory with id {id}." + ) + + @server.tool() + def memory_list(type: str = "", limit: int = 20) -> str: + """List stored memories with their ids, newest first, so they can be + updated or forgotten. Optionally filter by `type`.""" + entries = [e for e in reversed(store.all()) if not type or e.type == type] + if not entries: + return "No memories stored." + shown = entries[:limit] + lines = [f"- {e.id} [{_tag(e)}] {e.text}" for e in shown] + if len(entries) > len(shown): + lines.append(f"... and {len(entries) - len(shown)} more.") + return "\n".join(lines) + @server.tool() def memory_stats() -> str: """Summarize what is in the memory store.""" diff --git a/src/agent_memory/store.py b/src/agent_memory/store.py index b009590..1896b6f 100644 --- a/src/agent_memory/store.py +++ b/src/agent_memory/store.py @@ -5,15 +5,24 @@ memories here and recalls the top-k relevant ones for the task at hand via vector similarity. That is what cuts context tokens while keeping the facts the agent actually needs. + +Persistence is a single JSON file, but writes are careful: they take a lock, +re-read anything another process appended, and land atomically. Two agents +pointed at one store append to it instead of overwriting each other. """ from __future__ import annotations +import base64 import json +import os +import re +import time +from contextlib import contextmanager from dataclasses import asdict, dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import Optional +from typing import Iterator, Optional import numpy as np @@ -32,6 +41,12 @@ "fact", } +# Bumped when the on-disk layout changes. v2 stores embeddings as base64 +# float16 instead of JSON float lists (~5x smaller, same ranking). +STORE_FORMAT = 2 + +_ID_RE = re.compile(r"^mem_(\d+)$") + def _now_iso() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") @@ -59,8 +74,57 @@ class RecallHit: score: float +@contextmanager +def _file_lock(target: Path, timeout: float = 10.0, stale_after: float = 60.0) -> Iterator[None]: + """Cross-process advisory lock for one store file. + + An exclusive-create lock file is portable (POSIX and Windows) and needs no + extra dependency. A lock older than `stale_after` is assumed to belong to a + crashed process and is broken, so a dead agent can't wedge the store. + """ + lock = target.with_name(target.name + ".lock") + lock.parent.mkdir(parents=True, exist_ok=True) + deadline = time.monotonic() + timeout + while True: + try: + fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + break + except FileExistsError: + try: + if time.time() - lock.stat().st_mtime > stale_after: + lock.unlink(missing_ok=True) + continue + except FileNotFoundError: + continue # released while we looked; retry immediately + if time.monotonic() > deadline: + raise TimeoutError( + f"could not lock {target} after {timeout}s; " + f"remove {lock} if no agent is running" + ) + time.sleep(0.02) + try: + os.close(fd) + yield + finally: + lock.unlink(missing_ok=True) + + +def _encode_vector(vec: np.ndarray) -> str: + """float16 + base64. Precision loss is ~1e-3 — far below what ranking needs.""" + return base64.b64encode(np.asarray(vec, dtype=np.float16).tobytes()).decode("ascii") + + +def _decode_vector(raw: str | list[float]) -> np.ndarray: + if isinstance(raw, str): + vec = np.frombuffer(base64.b64decode(raw), dtype=np.float16).astype(np.float32) + else: # v1 stores kept a plain JSON list of floats + vec = np.asarray(raw, dtype=np.float32) + norm = float(np.linalg.norm(vec)) + return vec / norm if norm else vec # re-normalise after the float16 round-trip + + class MemoryStore: - """In-memory vector store with JSON persistence. + """Vector store over a JSON file. Small by design: an agent's durable memory for one project is hundreds of entries, not millions, so a brute-force cosine search over a numpy matrix @@ -71,10 +135,11 @@ class MemoryStore: def __init__( self, path: Optional[str | Path] = None, embedder: Optional[Embedder] = None ) -> None: - self.path = Path(path) if path else None + self.path = Path(path).expanduser() if path else None self.embedder = embedder or default_embedder() self._entries: list[MemoryEntry] = [] self._matrix = np.zeros((0, self.embedder.dim), dtype=np.float32) + self._stamp: Optional[tuple[int, int]] = None if self.path and self.path.exists(): self.load() @@ -88,8 +153,58 @@ def write( dedup_threshold: float = 0.97, agent: str = "", ) -> MemoryEntry: + """Save one memory. Returns the entry — the existing one if this text + near-duplicates something already stored.""" + entry, _ = self.write_with_status( + text, + type=type, + metadata=metadata, + id=id, + dedup_threshold=dedup_threshold, + agent=agent, + ) + return entry + + def write_with_status( + self, + text: str, + type: str = "fact", + metadata: Optional[dict] = None, + id: Optional[str] = None, + dedup_threshold: float = 0.97, + agent: str = "", + ) -> tuple[MemoryEntry, bool]: + """Like `write`, but also reports whether the text was actually stored. + + Returns `(entry, stored)`. `stored=False` means the write was dropped as + a near-duplicate and `entry` is the memory already on file — callers + that report back to an agent must not claim a save happened. + """ if type not in MEMORY_TYPES: raise ValueError(f"unknown memory type {type!r}; use one of {MEMORY_TYPES}") + + if self.path is None: + return self._append(text, type, metadata, id, dedup_threshold, agent) + + # Under the lock: pick up anything another agent appended, then write. + with _file_lock(self.path): + self._reload_if_changed() + entry, stored = self._append( + text, type, metadata, id, dedup_threshold, agent + ) + if stored: + self._save_unlocked() + return entry, stored + + def _append( + self, + text: str, + type: str, + metadata: Optional[dict], + id: Optional[str], + dedup_threshold: float, + agent: str, + ) -> tuple[MemoryEntry, bool]: vec = self.embedder.embed([text])[0] # Skip near-duplicates so repeated handoffs don't bloat the store. @@ -97,10 +212,10 @@ def write( sims = self._matrix @ vec best = int(np.argmax(sims)) if sims[best] >= dedup_threshold: - return self._entries[best] + return self._entries[best], False entry = MemoryEntry( - id=id or f"mem_{len(self._entries) + 1:04d}", + id=id or self._next_id(), type=type, text=text, metadata=metadata or {}, @@ -108,9 +223,80 @@ def write( ) self._entries.append(entry) self._matrix = np.vstack([self._matrix, vec[None, :]]) - if self.path: - self.save() - return entry + return entry, True + + def _next_id(self) -> str: + """Smallest unused `mem_NNNN`. Derived from the ids actually present, so + it survives explicit ids, deletions and concurrent appends.""" + used = {e.id for e in self._entries} + highest = 0 + for entry_id in used: + match = _ID_RE.match(entry_id) + if match: + highest = max(highest, int(match.group(1))) + candidate = highest + 1 + while f"mem_{candidate:04d}" in used: + candidate += 1 + return f"mem_{candidate:04d}" + + def forget(self, entry_id: str) -> bool: + """Delete one memory. Returns False if that id isn't in the store. + + Memory that can't be corrected is worse than no memory: a stale `state` + entry keeps being recalled and quietly misleads every later session. + """ + if self.path is None: + return self._remove(entry_id) + with _file_lock(self.path): + self._reload_if_changed() + removed = self._remove(entry_id) + if removed: + self._save_unlocked() + return removed + + def _remove(self, entry_id: str) -> bool: + for i, entry in enumerate(self._entries): + if entry.id == entry_id: + del self._entries[i] + self._matrix = np.delete(self._matrix, i, axis=0) + return True + return False + + def update( + self, + entry_id: str, + text: Optional[str] = None, + type: Optional[str] = None, + ) -> Optional[MemoryEntry]: + """Revise a memory in place, re-embedding when the text changes. + + Use this when a fact changes rather than writing a second, contradictory + memory — both would otherwise be recalled together. + """ + if type is not None and type not in MEMORY_TYPES: + raise ValueError(f"unknown memory type {type!r}; use one of {MEMORY_TYPES}") + if self.path is None: + return self._revise(entry_id, text, type) + with _file_lock(self.path): + self._reload_if_changed() + entry = self._revise(entry_id, text, type) + if entry is not None: + self._save_unlocked() + return entry + + def _revise( + self, entry_id: str, text: Optional[str], type: Optional[str] + ) -> Optional[MemoryEntry]: + for i, entry in enumerate(self._entries): + if entry.id != entry_id: + continue + if text is not None and text != entry.text: + entry.text = text + self._matrix[i] = self.embedder.embed([text])[0] + if type is not None: + entry.type = type + return entry + return None # ---- reading ------------------------------------------------------- def recall( @@ -119,6 +305,8 @@ def recall( k: int = 5, type_filter: Optional[str] = None, budget_tokens: Optional[int] = None, + exclude_ids: Optional[set[str]] = None, + min_score: float = 0.0, ) -> list[RecallHit]: """Top-k most relevant memories, optionally under a hard token budget. @@ -126,7 +314,12 @@ def recall( order: an entry that would overflow the remaining budget is skipped and the next-best one is tried. The result never costs more than the budget — the caller controls exactly how much context this loads. + + `min_score` drops weak matches entirely. Without it a query unrelated to + anything in the store still returns k memories, and the agent reading + them has no way to tell they are noise. """ + self._reload_if_changed() if not self._entries: return [] qvec = self.embedder.embed([query])[0] @@ -135,7 +328,12 @@ def recall( hits: list[RecallHit] = [] remaining = budget_tokens for idx in order: + score = float(sims[idx]) + if score < min_score: + break # sorted by score, so nothing further can qualify entry = self._entries[idx] + if exclude_ids and entry.id in exclude_ids: + continue if type_filter and entry.type != type_filter: continue if remaining is not None: @@ -143,22 +341,59 @@ def recall( if cost > remaining: continue # doesn't fit; a smaller lower-ranked one may remaining -= cost - hits.append(RecallHit(entry=entry, score=float(sims[idx]))) + hits.append(RecallHit(entry=entry, score=score)) if len(hits) >= k: break return hits + def boot( + self, + task: str, + k: int = 5, + budget_tokens: Optional[int] = 300, + min_score: float = 0.0, + ) -> tuple[Optional[MemoryEntry], list[RecallHit]]: + """Return the latest handoff plus relevant memories for a new session. + + The budget applies to memory content across both parts. If the latest + handoff is too large to fit, it is skipped and the full budget remains + available for relevant memories. + """ + remaining = budget_tokens + latest_handoff = self.latest("handoff") + included_handoff: Optional[MemoryEntry] = None + excluded_ids: set[str] = set() + + if latest_handoff is not None: + excluded_ids.add(latest_handoff.id) + if remaining is None or latest_handoff.tokens <= remaining: + included_handoff = latest_handoff + if remaining is not None: + remaining -= latest_handoff.tokens + + hits = self.recall( + task, + k=k, + budget_tokens=remaining, + exclude_ids=excluded_ids, + min_score=min_score, + ) + return included_handoff, hits + def latest(self, type: str) -> Optional[MemoryEntry]: """Most recently written entry of a type (e.g. the last handoff).""" + self._reload_if_changed() for entry in reversed(self._entries): if entry.type == type: return entry return None def all(self) -> list[MemoryEntry]: + self._reload_if_changed() return list(self._entries) def stats(self) -> dict: + self._reload_if_changed() by_type: dict[str, int] = {} for e in self._entries: by_type[e.type] = by_type.get(e.type, 0) + 1 @@ -172,35 +407,91 @@ def stats(self) -> dict: # ---- persistence --------------------------------------------------- def save(self, path: Optional[str | Path] = None) -> None: - target = Path(path) if path else self.path + target = Path(path).expanduser() if path else self.path if target is None: raise ValueError("no path set for this store") + with _file_lock(target): + self._save_unlocked(target) + + def _save_unlocked(self, path: Optional[Path] = None) -> None: + """Serialise atomically: a crash mid-write must not truncate the store.""" + target = path or self.path + assert target is not None target.parent.mkdir(parents=True, exist_ok=True) payload = { + "format": STORE_FORMAT, "embedder": type(self.embedder).__name__, "dim": self.embedder.dim, "entries": [ - {**asdict(e), "embedding": self._matrix[i].tolist()} + {**asdict(e), "embedding": _encode_vector(self._matrix[i])} for i, e in enumerate(self._entries) ], } - target.write_text(json.dumps(payload, indent=2, ensure_ascii=False)) + tmp = target.with_name(f"{target.name}.{os.getpid()}.tmp") + tmp.write_text(json.dumps(payload, indent=2, ensure_ascii=False)) + os.replace(tmp, target) # atomic on POSIX and Windows + if target == self.path: + self._stamp = self._read_stamp() def load(self, path: Optional[str | Path] = None) -> None: - target = Path(path) if path else self.path + target = Path(path).expanduser() if path else self.path if target is None or not target.exists(): return + # Stamp BEFORE reading. Reads are not locked (a running server reloads on + # every recall), so another agent can replace the file mid-read. Stamping + # afterwards would pair the new stamp with the content we already read, + # and every later freshness check would wrongly conclude we were current. + # Stamping first can only cause a redundant reload, never a skipped one. + stamp = self._read_stamp() if target == self.path else None payload = json.loads(target.read_text()) - self._entries = [] - vectors: list[list[float]] = [] + + # A store written by a different embedder holds vectors that are not + # comparable with ours — different dimension (a hard crash on the first + # matmul) or, worse, the same dimension from a different model (silently + # meaningless scores). Re-embed from the text instead. + stored_dim = payload.get("dim") + stored_embedder = payload.get("embedder") + reembed = ( + stored_dim != self.embedder.dim + or stored_embedder != type(self.embedder).__name__ + ) + + entries: list[MemoryEntry] = [] + vectors: list[Optional[np.ndarray]] = [] + known = {f.name for f in MemoryEntry.__dataclass_fields__.values()} for raw in payload.get("entries", []): embedding = raw.pop("embedding", None) - self._entries.append(MemoryEntry(**raw)) - if embedding is None: # re-embed if the file predates stored vectors - embedding = self.embedder.embed([raw["text"]])[0].tolist() - vectors.append(embedding) + entries.append(MemoryEntry(**{k: v for k, v in raw.items() if k in known})) + if embedding is None or reembed: + vectors.append(None) # filled in below, in one batch + else: + vectors.append(_decode_vector(embedding)) + + missing = [i for i, v in enumerate(vectors) if v is None] + if missing: + fresh = self.embedder.embed([entries[i].text for i in missing]) + for slot, i in enumerate(missing): + vectors[i] = fresh[slot] + + self._entries = entries self._matrix = ( np.array(vectors, dtype=np.float32) if vectors else np.zeros((0, self.embedder.dim), dtype=np.float32) ) + if target == self.path: + self._stamp = stamp + + def _read_stamp(self) -> Optional[tuple[int, int]]: + try: + st = self.path.stat() # type: ignore[union-attr] + except (OSError, AttributeError): + return None + return (st.st_mtime_ns, st.st_size) + + def _reload_if_changed(self) -> None: + """Pick up writes made by another process since we last read the file.""" + if self.path is None or not self.path.exists(): + return + if self._read_stamp() != self._stamp: + self.load() diff --git a/tests/test_eval.py b/tests/test_eval.py index 59d9c37..67995e2 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -12,28 +12,82 @@ def _dataset(): return json.loads((ROOT / "eval" / "dataset.json").read_text()) -def test_semantic_recall_uses_fewer_tokens_than_baseline(): - r = evaluate(_dataset(), k=3) - sem = r["summary"]["semantic_recall"] - base = r["summary"]["full_context"] - assert sem["avg_context_tokens"] < base["avg_context_tokens"] - assert sem["token_reduction_vs_baseline"] >= 50 +def _results(**kw): + return evaluate(_dataset(), **kw) -def test_recall_is_high_enough_to_be_useful(): - r = evaluate(_dataset(), k=3) - assert r["summary"]["semantic_recall"]["recall"] >= 0.8 +def test_retrieval_uses_fewer_tokens_than_the_full_context_baseline(): + r = _results(k=3) + assert ( + r["summary"]["retrieval"]["avg_context_tokens"] + < r["summary"]["full_context"]["avg_context_tokens"] + ) + + +def test_the_token_saving_is_not_evidence_on_its_own(): + """The random control must show the same saving as the engine. + If this ever fails, the token-reduction number has started measuring + something other than "we loaded k of n memories" and the README's framing + needs revisiting. + """ + r = _results(k=3) + engine = r["summary"]["retrieval"]["token_reduction_vs_baseline"] + random_arm = r["summary"]["random_k"]["token_reduction_vs_baseline"] + assert abs(engine - random_arm) < 15 -def test_semantic_recall_beats_baseline_precision(): - r = evaluate(_dataset(), k=3) - sem = r["summary"]["semantic_recall"] - base = r["summary"]["full_context"] - assert sem["precision"] > base["precision"] + +def test_retrieval_beats_the_random_control_by_a_wide_margin(): + """This — not the token saving — is the claim worth making.""" + r = _results(k=3) + assert r["summary"]["retrieval"]["recall"] - r["summary"]["random_k"]["recall"] > 0.5 + + +def test_recall_is_high_enough_to_be_useful(): + assert _results(k=3)["summary"]["retrieval"]["recall"] >= 0.8 def test_budget_arm_respects_cap_and_stays_useful(): - r = evaluate(_dataset(), k=3, budget=120) + r = _results(k=3, budget=120) bud = r["summary"]["budget_recall"] assert bud["avg_context_tokens"] <= 120 assert bud["recall"] >= 0.7 + + +def test_paraphrase_arm_is_reported_and_is_the_weaker_number(): + """The lexical default loses recall when wording is not shared. + + The point of the assertion is that the gap stays *measured and published*, + not that it stays small. + """ + ph = _results(k=3)["phrasing"] + assert ph["paraphrase"]["recall"] <= ph["query"]["recall"] + assert ph["paraphrase"]["avg_query_word_overlap_with_gold"] < ( + ph["query"]["avg_query_word_overlap_with_gold"] + ) + + +def test_budget_cost_stays_flat_while_the_store_grows(): + rows = _results(k=3, budget=120)["scaling"] + assert len(rows) >= 3 + assert rows[-1]["n_memories"] > rows[0]["n_memories"] * 2 + # The baseline grows with the store; the budgeted arm does not. + assert rows[-1]["full_context_tokens"] > 2 * rows[0]["full_context_tokens"] + assert rows[-1]["budget_recall_tokens"] <= 120 + assert rows[-1]["budget_recall"] >= 0.7 + + +def test_min_score_sweep_supports_the_configured_floor(): + from agent_memory import HashingEmbedder + + rows = {row["min_score"]: row for row in _results(k=3)["min_score_sweep"]} + floor = HashingEmbedder.recommended_min_score + assert floor in rows, "the shipped floor must appear in the published sweep" + assert rows[floor]["labelled_recall"] == rows[0.0]["labelled_recall"] + assert rows[floor]["off_topic_memories_returned"] < rows[0.0]["off_topic_memories_returned"] + + +def test_random_arm_is_deterministic_for_a_given_seed(): + assert _results(k=3, seed=7)["summary"]["random_k"] == ( + _results(k=3, seed=7)["summary"]["random_k"] + ) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 0000000..476bd50 --- /dev/null +++ b/tests/test_mcp_server.py @@ -0,0 +1,139 @@ +"""Tests for the MCP server. + +This is the headline feature and it used to have no coverage at all, which is +how it came to be broken against `mcp` 2.x: the package renamed the server +class, the import failed, and nothing noticed. These tests exercise the tools +through the real server object on whatever `mcp` version is installed. +""" + +import asyncio + +import pytest + +mcp = pytest.importorskip("mcp", reason="needs the optional 'mcp' extra") + +from agent_memory.mcp_server import build_server # noqa: E402 + + +def call(server, name: str, **args) -> str: + """Invoke a tool and return its text, across mcp 1.x and 2.x result shapes.""" + result = asyncio.run(server.call_tool(name, args)) + content = getattr(result, "content", None) # mcp 2.x: CallToolResult + if content is None: # mcp 1.x: list, or (list, dict) + content = result[0] if isinstance(result, tuple) else result + return content[0].text + + +@pytest.fixture +def server(tmp_path, monkeypatch): + monkeypatch.setenv("AGENT_MEMORY_EMBEDDER", "hashing") + return build_server(store_path=tmp_path / "store.json", agent="claude-code") + + +def test_server_builds_and_exposes_the_documented_tools(server): + names = {t.name for t in asyncio.run(server.list_tools())} + assert { + "memory_write", + "memory_recall", + "memory_boot", + "memory_handoff", + "memory_update", + "memory_forget", + "memory_list", + "memory_stats", + } <= names + + +def test_write_then_recall_roundtrip(server): + assert "Saved" in call(server, "memory_write", text="Bookings are stored in UTC.", type="decision") + out = call(server, "memory_recall", query="how are timezones handled for bookings") + assert "UTC" in out + + +def test_write_reports_duplicates_instead_of_claiming_a_save(server): + text = "Admin routes are guarded by requireAdmin in server/auth.ts." + first = call(server, "memory_write", text=text, type="decision") + second = call(server, "memory_write", text=text, type="decision") + assert first.startswith("Saved") + assert "Not saved" in second and "near-duplicate" in second + + +def test_write_rejects_unknown_type(server): + assert "Error" in call(server, "memory_write", text="something", type="nope") + + +def test_recall_filters_out_unrelated_memories(server): + call(server, "memory_write", text="Bookings are stored in UTC.", type="decision") + call(server, "memory_write", text="The chatbot uses Google Gemini.", type="decision") + assert call(server, "memory_recall", query="how do I bake sourdough bread") == ( + "No relevant memories." + ) + + +def test_boot_never_exceeds_its_budget(server): + call( + server, + "memory_handoff", + done="Finished the booking refactor across the router, the email helper " + "and the calendar sync, and re-ran the integration suite twice.", + next_steps="Rate limit the chat endpoint before deploying to production.", + ) + for text in [ + "Bookings are stored in UTC and converted in the UI layer only.", + "The chatbot uses Google Gemini and its prompt lives in gemini-chat.ts.", + "Admin routes are guarded by requireAdmin in server/auth.ts.", + "Uploaded images go to object storage with ACLs in objectAcl.ts.", + ]: + call(server, "memory_write", text=text, type="decision") + + from agent_memory.tokens import count_tokens + + for budget in (30, 60, 120, 300): + out = call(server, "memory_boot", task="continue the booking work", budget_tokens=budget) + # Strip the rendering scaffolding; only memory text is charged to the budget. + body = out.replace("Last handoff [handoff · claude-code]: ", "") + body = "\n".join(line.lstrip("- ") for line in body.splitlines()) + content = "".join( + part.split("] ", 1)[-1] if "] " in part else part for part in body.splitlines() + ) + assert count_tokens(content) <= budget, f"budget {budget} exceeded: {out!r}" + + +def test_boot_on_an_empty_store(server): + assert "Empty store" in call(server, "memory_boot", task="anything") + + +def test_update_and_forget(server): + call(server, "memory_write", text="The API listens on port 5002.", type="project") + listed = call(server, "memory_list") + entry_id = listed.split()[1] + + assert "Updated" in call(server, "memory_update", id=entry_id, text="The API listens on port 8080.") + assert "8080" in call(server, "memory_recall", query="which port does the API listen on") + + assert "Forgot" in call(server, "memory_forget", id=entry_id) + assert "No memory with id" in call(server, "memory_forget", id=entry_id) + assert call(server, "memory_list") == "No memories stored." + + +def test_update_and_forget_report_missing_ids(server): + assert "No memory with id" in call(server, "memory_update", id="mem_9999", text="x") + assert "No memory with id" in call(server, "memory_forget", id="mem_9999") + + +def test_handoff_is_picked_up_by_a_second_agent_on_the_same_store(tmp_path, monkeypatch): + monkeypatch.setenv("AGENT_MEMORY_EMBEDDER", "hashing") + path = tmp_path / "shared.json" + claude = build_server(store_path=path, agent="claude-code") + call(claude, "memory_handoff", done="Fixed double emails.", next_steps="Add rate limiting.") + + codex = build_server(store_path=path, agent="codex") + out = call(codex, "memory_boot", task="continue the rate limiting work") + assert "rate limiting" in out.lower() + assert "claude-code" in out # provenance survives the handoff + + +def test_stats_reports_the_store(server): + call(server, "memory_write", text="Bookings are stored in UTC.", type="decision") + out = call(server, "memory_stats") + assert "1 memories" in out and "decision=1" in out diff --git a/tests/test_persistence.py b/tests/test_persistence.py new file mode 100644 index 0000000..7f9a63b --- /dev/null +++ b/tests/test_persistence.py @@ -0,0 +1,232 @@ +"""Persistence, concurrency and recovery. + +The engine's pitch is one store shared by several agents, so the failure that +matters most is two agents writing at once and one of them silently winning. +These tests cover that, plus the recovery paths around switching embedders and +crash-safe writes. +""" + +import json +import subprocess +import sys +import textwrap + +import numpy as np +import pytest + +from agent_memory import HashingEmbedder, MemoryStore + +SRC = str((__import__("pathlib").Path(__file__).resolve().parent.parent / "src")) + + +def store_at(path): + return MemoryStore(path=path, embedder=HashingEmbedder()) + + +def test_two_open_stores_both_keep_their_writes(tmp_path): + path = tmp_path / "shared.json" + seed = store_at(path) + seed.write("Shared baseline fact about the project.", type="project") + + # Both agents opened the store before either wrote — the classic lost-update + # setup: each holds a snapshot from before the other's write. + claude = store_at(path) + codex = store_at(path) + claude.write("Claude Code: bookings are stored in UTC.", type="decision", agent="claude-code") + codex.write("Codex: auth uses signed session cookies.", type="decision", agent="codex") + + final = store_at(path) + assert final.stats()["count"] == 3 + agents = {e.agent for e in final.all()} + assert {"claude-code", "codex"} <= agents + + +def test_parallel_processes_do_not_lose_writes(tmp_path): + path = tmp_path / "parallel.json" + worker = textwrap.dedent( + f""" + import sys + sys.path.insert(0, {SRC!r}) + from agent_memory import MemoryStore, HashingEmbedder + store = MemoryStore(path={str(path)!r}, embedder=HashingEmbedder()) + who = sys.argv[1] + for i in range(5): + store.write( + f"Agent {{who}} learned distinct fact number {{i}} about subsystem {{who}}{{i}}.", + type="fact", agent=who, + ) + """ + ) + script = tmp_path / "worker.py" + script.write_text(worker) + + procs = [subprocess.Popen([sys.executable, str(script), f"agent{i}"]) for i in range(6)] + assert all(p.wait() == 0 for p in procs) + + final = store_at(path) + ids = [e.id for e in final.all()] + assert len(ids) == 30, f"lost {30 - len(ids)} writes" + assert len(set(ids)) == 30, "ids must stay unique across processes" + + +def test_an_open_store_sees_another_agents_writes(tmp_path): + """A long-running MCP server must not serve a snapshot from startup.""" + path = tmp_path / "live.json" + reader = store_at(path) # e.g. Claude Code's server, already running + reader.write("Bookings are stored in UTC.", type="decision", agent="claude-code") + + writer = store_at(path) # e.g. Codex, in another process + writer.write( + "Done: added the rate limiter. Next: re-run the integration suite.", + type="handoff", + agent="codex", + ) + + handoff = reader.latest("handoff") + assert handoff is not None and handoff.agent == "codex" + assert reader.stats()["count"] == 2 + assert any("rate limiter" in h.entry.text for h in reader.recall("rate limiting", k=3)) + + +def test_stamp_is_taken_before_the_read_not_after(tmp_path): + """Regression: a file replaced mid-read must not be recorded as 'current'. + + Reads are unlocked, so another agent can replace the store between our stat + and our read. Stamping after the read pairs the *new* stamp with the *old* + content, and every later freshness check then wrongly concludes we are up to + date — which showed up as two agents choosing the same id and one write + disappearing. + """ + path = tmp_path / "toctou.json" + writer = store_at(path) + writer.write("First fact, present at read time.", type="fact") + + reader = MemoryStore(path=None, embedder=HashingEmbedder()) + reader.path = path + + original_read_text = type(path).read_text + + def read_then_mutate(self, *args, **kwargs): + content = original_read_text(self, *args, **kwargs) + if self == path: # simulate another agent landing a write mid-read + writer.write("Second fact, written during the read.", type="fact") + return content + + monkeypatch = pytest.MonkeyPatch() + try: + monkeypatch.setattr(type(path), "read_text", read_then_mutate) + reader.load() + finally: + monkeypatch.undo() + + assert reader.stats()["count"] == 2, "the store must notice it read stale content" + + +def test_ids_survive_explicit_ids_and_deletions(tmp_path): + store = store_at(tmp_path / "ids.json") + store.write("first entry", id="mem_0001") + store.write("a second entry, wholly different", id="mem_0009") + store.write("a third distinct entry about deployment") + store.forget("mem_0009") + store.write("a fourth separate entry about analytics") + ids = [e.id for e in store.all()] + assert len(set(ids)) == len(ids) + assert "mem_0001" in ids + + +def test_switching_embedder_reembeds_instead_of_crashing(tmp_path): + class OtherEmbedder: # a different dimension, like sentence-transformers + dim = 384 + recommended_min_score = 0.25 + + def embed(self, texts): + rng = np.random.default_rng(len(texts[0])) + vecs = rng.normal(size=(len(texts), 384)).astype(np.float32) + return vecs / np.linalg.norm(vecs, axis=1, keepdims=True) + + path = tmp_path / "switch.json" + first = MemoryStore(path=path, embedder=OtherEmbedder()) + first.write("Admin routes are guarded by requireAdmin in server/auth.ts.", type="decision") + first.write("Bookings are stored in UTC.", type="decision") + + reopened = store_at(path) # 512-d hashing embedder over a 384-d store + assert reopened.stats()["count"] == 2 + hits = reopened.recall("how do we protect admin pages", k=1) + assert hits and "Admin routes" in hits[0].entry.text + + +def test_save_is_atomic_and_leaves_no_partial_file(tmp_path): + path = tmp_path / "atomic.json" + store = store_at(path) + for i in range(10): + store.write(f"Fact number {i} about a distinct subsystem {i}.", type="fact") + assert json.loads(path.read_text())["entries"] + assert not list(tmp_path.glob("*.tmp")), "temp files must be renamed away" + assert not list(tmp_path.glob("*.lock")), "locks must be released" + + +def test_legacy_float_list_stores_still_load(tmp_path): + """v1 wrote embeddings as JSON float lists; those stores must keep working.""" + path = tmp_path / "legacy.json" + store = store_at(path) + store.write("Bookings are stored in UTC.", type="decision") + + payload = json.loads(path.read_text()) + vec = HashingEmbedder().embed(["Bookings are stored in UTC."])[0] + payload.pop("format", None) + payload["entries"][0]["embedding"] = vec.tolist() + path.write_text(json.dumps(payload)) + + reopened = store_at(path) + assert reopened.stats()["count"] == 1 + assert reopened.recall("timezone handling for bookings", k=1) + + +def test_unknown_fields_in_a_newer_store_do_not_break_load(tmp_path): + path = tmp_path / "future.json" + store = store_at(path) + store.write("Bookings are stored in UTC.", type="decision") + + payload = json.loads(path.read_text()) + payload["entries"][0]["confidence"] = 0.9 # written by a future version + path.write_text(json.dumps(payload)) + + assert store_at(path).stats()["count"] == 1 + + +def test_forget_and_update_persist(tmp_path): + path = tmp_path / "edit.json" + store = store_at(path) + keep = store.write("The API listens on port 5002.", type="project") + drop = store.write("A stale note about the old deployment box.", type="state") + + assert store.update(keep.id, text="The API listens on port 8080.") is not None + assert store.forget(drop.id) is True + assert store.forget(drop.id) is False + + reopened = store_at(path) + assert [e.text for e in reopened.all()] == ["The API listens on port 8080."] + hits = reopened.recall("which port does the API listen on", k=1) + assert "8080" in hits[0].entry.text + + +def test_update_reembeds_so_recall_follows_the_new_text(tmp_path): + store = store_at(tmp_path / "reembed.json") + entry = store.write("The chatbot uses Google Gemini.", type="decision") + store.update(entry.id, text="The chatbot uses Anthropic Claude via the Messages API.") + hits = store.recall("which model does the chatbot use", k=1) + assert "Claude" in hits[0].entry.text + + +def test_update_rejects_unknown_type(tmp_path): + store = store_at(tmp_path / "badtype.json") + entry = store.write("something", type="fact") + with pytest.raises(ValueError): + store.update(entry.id, type="not_a_type") + + +def test_paths_with_a_tilde_are_expanded(tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + store = MemoryStore(path="~/nested/store.json", embedder=HashingEmbedder()) + store.write("Bookings are stored in UTC.", type="decision") + assert (tmp_path / "nested" / "store.json").exists() diff --git a/tests/test_store.py b/tests/test_store.py index 861c235..57df390 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -68,6 +68,89 @@ def test_budget_recall_skips_oversized_and_packs_smaller(store): assert all(len(t) < 500 for t in texts), "oversized entry must be skipped" +def test_recall_can_exclude_an_entry(store): + excluded = store.all()[0] + hits = store.recall("bookings UTC", k=3, exclude_ids={excluded.id}) + assert all(hit.entry.id != excluded.id for hit in hits) + + +def test_boot_shares_one_budget_between_handoff_and_recall(store): + handoff = store.write( + "Done: fixed email retries. Next: protect admin routes.", + type="handoff", + agent="claude-code", + ) + budget = handoff.tokens + 20 + + included_handoff, hits = store.boot( + "protect admin routes", + k=5, + budget_tokens=budget, + ) + + assert included_handoff == handoff + assert handoff.tokens + sum(hit.entry.tokens for hit in hits) <= budget + assert all(hit.entry.id != handoff.id for hit in hits) + + +def test_boot_skips_handoff_that_cannot_fit(store): + oversized = store.write( + ("long handoff " * 100) + "next deploy safely", + type="handoff", + agent="claude-code", + ) + budget = 25 + + included_handoff, hits = store.boot( + "protect admin routes", + k=5, + budget_tokens=budget, + ) + + assert oversized.tokens > budget + assert included_handoff is None + assert sum(hit.entry.tokens for hit in hits) <= budget + assert all(hit.entry.id != oversized.id for hit in hits) + + +def test_min_score_drops_unrelated_matches(store): + """Without a floor, a query about nothing in the store still returns k + memories — and the MCP layer hides the scores, so the agent can't tell.""" + assert store.recall("how do I bake sourdough bread at home", k=3) != [] + assert store.recall("how do I bake sourdough bread at home", k=3, min_score=0.15) == [] + + +def test_min_score_keeps_genuine_matches(store): + hits = store.recall("how do we protect admin pages", k=3, min_score=0.15) + assert hits and hits[0].entry.text.startswith("Admin routes") + + +def test_write_with_status_reports_whether_it_stored(store): + text = "Deploys go out through GitHub Actions on every push to main." + entry, stored = store.write_with_status(text, type="decision") + assert stored is True + same, stored_again = store.write_with_status(text, type="decision") + assert stored_again is False and same.id == entry.id + + +def test_forget_removes_the_entry_from_recall(store): + entry = store.write("A stale note about the retired staging box.", type="state") + assert store.forget(entry.id) is True + assert all(h.entry.id != entry.id for h in store.recall("staging box", k=5)) + assert store.forget(entry.id) is False + + +def test_update_changes_text_and_ranking(store): + entry = store.write("The API listens on port 5002.", type="project") + updated = store.update(entry.id, text="The API listens on port 8080.") + assert updated is not None and "8080" in updated.text + assert "8080" in store.recall("which port does the API listen on", k=1)[0].entry.text + + +def test_update_returns_none_for_unknown_id(store): + assert store.update("mem_9999", text="whatever") is None + + def test_agent_attribution_roundtrip(tmp_path): from agent_memory import HashingEmbedder, MemoryStore