diff --git a/catalog.json b/catalog.json index 03fac76..f916bd5 100644 --- a/catalog.json +++ b/catalog.json @@ -71,6 +71,13 @@ } }, "providers": { + "automem": { + "key": "automem", + "description": "AutoMem graph+vector memory (FalkorDB + Qdrant) self-spun via Docker with FastEmbed (BAAI/bge-base-en-v1.5, 768d, no embedding API keys); recall traverses graph relations.", + "kind": "local", + "link": "https://automem.ai", + "logo": "https://www.google.com/s2/favicons?sz=32&domain=automem.ai" + }, "bm25": { "key": "bm25", "description": "Keyword search baseline. No embeddings \u2014 splits docs into chunks and uses BM25 ranking.", diff --git a/outputs/beam/automem/rag/100k.json.gz b/outputs/beam/automem/rag/100k.json.gz new file mode 100644 index 0000000..e82ac3e Binary files /dev/null and b/outputs/beam/automem/rag/100k.json.gz differ diff --git a/outputs/beam/automem/rag/10m.json.gz b/outputs/beam/automem/rag/10m.json.gz new file mode 100644 index 0000000..4215487 Binary files /dev/null and b/outputs/beam/automem/rag/10m.json.gz differ diff --git a/outputs/beam/automem/rag/1m.json.gz b/outputs/beam/automem/rag/1m.json.gz new file mode 100644 index 0000000..989b3ec Binary files /dev/null and b/outputs/beam/automem/rag/1m.json.gz differ diff --git a/outputs/beam/automem/rag/500k.json.gz b/outputs/beam/automem/rag/500k.json.gz new file mode 100644 index 0000000..c7a956a Binary files /dev/null and b/outputs/beam/automem/rag/500k.json.gz differ diff --git a/outputs/locomo/automem/rag/locomo10.json.gz b/outputs/locomo/automem/rag/locomo10.json.gz new file mode 100644 index 0000000..9359fdf Binary files /dev/null and b/outputs/locomo/automem/rag/locomo10.json.gz differ diff --git a/outputs/longmemeval/automem/rag/s.json.gz b/outputs/longmemeval/automem/rag/s.json.gz new file mode 100644 index 0000000..6299c88 Binary files /dev/null and b/outputs/longmemeval/automem/rag/s.json.gz differ diff --git a/outputs/personamem/automem/rag/32k.json.gz b/outputs/personamem/automem/rag/32k.json.gz new file mode 100644 index 0000000..e4fb070 Binary files /dev/null and b/outputs/personamem/automem/rag/32k.json.gz differ diff --git a/src/memory_bench/memory/AUTOMEM_REPRODUCE.md b/src/memory_bench/memory/AUTOMEM_REPRODUCE.md new file mode 100644 index 0000000..69d9fec --- /dev/null +++ b/src/memory_bench/memory/AUTOMEM_REPRODUCE.md @@ -0,0 +1,68 @@ +# Reproducing the AutoMem results + +AutoMem is a graph + vector memory service (FalkorDB + Qdrant behind a Flask API). The +provider is **self-spinning**: it brings the whole stack up via Docker, runs the benchmark, +and tears it down. You need **Docker** and a **GEMINI_API_KEY** (the shared answer/judge +key every provider uses) — **no embedding API keys**. + +## One command per split + +```bash +GEMINI_API_KEY=... \ +OMB_ANSWER_LLM=gemini OMB_ANSWER_MODEL=gemini-3.1-pro-preview \ +OMB_JUDGE_LLM=gemini OMB_JUDGE_MODEL=gemini-2.5-flash-lite \ +uv run omb run --memory automem --dataset locomo --split locomo10 +``` + +Swap `--dataset/--split` for the others (`longmemeval/s`, `personamem/32k`, +`beam/{100k,500k,1m,10m}`); the env block stays identical. `initialize()` runs +`docker compose up` on `automem_compose.yml` (AutoMem + FalkorDB + Qdrant), waits for +`/health`, and an atexit-registered `cleanup()` runs `docker compose down -v`. Ports are +chosen per run, so concurrent/repeat runs don't collide; a crash also tears the stack down. + +## Pinned configuration + +| Knob | Value | +|---|---| +| AutoMem image | `ghcr.io/verygoodplugins/automem:amb-v1` (override with `AUTOMEM_IMAGE`) | +| FalkorDB | `falkordb/falkordb:v4.18.3` (pinned) | +| Qdrant | `qdrant/qdrant:v1.11.3` (pinned) | +| Embeddings | FastEmbed local, `BAAI/bge-base-en-v1.5`, 768-dim (`EMBEDDING_PROVIDER=local`, no API key) | +| Answer LLM | `gemini-3.1-pro-preview` | +| Judge LLM | `gemini-2.5-flash-lite` | +| Mode | `rag` | + +## How the provider uses AutoMem + +- **Ingest** → `POST /memory/batch`, one memory per document, chunked at `AUTOMEM_MAX_CHARS` + (default 1800) on sentence/paragraph boundaries with timestamps backdated to the source. + After ingest the provider waits for AutoMem's enrichment queue to settle so the graph it + queries is fully built. +- **Retrieve** → `GET /recall`, scoped to the run's tags, with graph relation expansion on + (`expand_relations` + `expand_respect_tags`). Content is read from + `result["memory"]["content"]`. + +## Reported metrics + +Each run writes per-query `retrieve_time_ms`, `context_tokens`, and `correct`/`score` to +`outputs/{dataset}/{name}/rag/{split}.json`. Recall latency is wall-clock around +`memory.retrieve()` on local hardware (FastEmbed in-process, single-query/RAG mode) and is +environment-relative — summarize with the median (P50), and treat it as a per-run figure, +not a cross-system axis. + +## Tuning knobs (env) + +| Env | Default | Purpose | +|---|---|---| +| `AUTOMEM_IMAGE` | `ghcr.io/verygoodplugins/automem:amb-v1` | AutoMem image tag | +| `AUTOMEM_MAX_CHARS` | `1800` | chunk size (under AutoMem's 2000 hard limit) | +| `AUTOMEM_RECALL_K` | (harness `k`) | override retrieval depth | +| `AUTOMEM_ENRICH_SETTLE_SECONDS` | `120` | max wait for enrichment to drain after ingest | +| `AUTOMEM_HOST` / `AUTOMEM_*_PORT` | localhost / free | override where the provider reaches the stack | + +## Tests + +```bash +uv run --with pytest pytest tests/test_automem_provider.py +``` +(The provider's HTTP contract and chunking/extraction helpers are unit-tested without Docker.) diff --git a/src/memory_bench/memory/__init__.py b/src/memory_bench/memory/__init__.py index fe30a2e..97f0f5e 100644 --- a/src/memory_bench/memory/__init__.py +++ b/src/memory_bench/memory/__init__.py @@ -1,4 +1,5 @@ from .base import MemoryProvider +from .automem import AutoMemMemoryProvider from .bm25 import BM25MemoryProvider from .cognee import CogneeMemoryProvider from .hindsight import HindsightCloudMemoryProvider, HindsightHTTPMemoryProvider, HindsightMemoryProvider @@ -11,6 +12,7 @@ from .supermemory import SupermemoryMemoryProvider REGISTRY: dict[str, type[MemoryProvider]] = { + "automem": AutoMemMemoryProvider, "bm25": BM25MemoryProvider, "cognee": CogneeMemoryProvider, "hindsight": HindsightMemoryProvider, diff --git a/src/memory_bench/memory/automem.py b/src/memory_bench/memory/automem.py new file mode 100644 index 0000000..8bb9005 --- /dev/null +++ b/src/memory_bench/memory/automem.py @@ -0,0 +1,219 @@ +"""AutoMem provider for the Agent Memory Benchmark. + +Self-spinning: initialize() brings up AutoMem (GHCR image) + FalkorDB + Qdrant via +docker compose with FastEmbed local embeddings (no API keys); cleanup() tears it down. +ingest() POSTs /memory (chunked, backdated); retrieve() GETs /recall and extracts +content from item["memory"]["content"] (top-level "content" is always empty). +""" +from __future__ import annotations +import atexit, http.client, json, os, re, socket, subprocess, time, urllib.error, urllib.parse, urllib.request, uuid +from pathlib import Path +from ..models import Document +from .base import MemoryProvider + +_COMPOSE = Path(__file__).parent / "automem_compose.yml" +_DEFAULT_IMAGE = "ghcr.io/verygoodplugins/automem:amb-v1" + + +def _slug(value: str) -> str: + return re.sub(r"[^a-z0-9]+", "-", str(value).lower()).strip("-")[:48] or "x" + + +def _chunk_text(text: str, max_chars: int) -> list: + text = (text or "").strip() + if len(text) <= max_chars: + return [text] if text else [] + chunks, buf = [], "" + for seg in re.split(r"(?<=[.!?])\s+|\n+", text): + seg = seg.strip() + if not seg: + continue + while len(seg) > max_chars: + if buf: + chunks.append(buf); buf = "" + chunks.append(seg[:max_chars]); seg = seg[max_chars:] + if not seg: + continue + if len(buf) + 1 + len(seg) <= max_chars: + buf = f"{buf} {seg}".strip() + else: + if buf: + chunks.append(buf) + buf = seg + if buf: + chunks.append(buf) + return chunks + + +def _extract_content(item: dict) -> str: + mem = item.get("memory") if isinstance(item.get("memory"), dict) else {} + return mem.get("content") or item.get("content") or mem.get("summary") or "" + + +def _free_port() -> int: + s = socket.socket() + s.bind(("", 0)) + port = s.getsockname()[1] + s.close() + return port + + +class AutoMemMemoryProvider(MemoryProvider): + name = "automem" + description = "AutoMem graph+vector memory (FalkorDB+Qdrant) self-spun via Docker with FastEmbed; relations traversed on recall." + kind = "local" + provider = "automem" + variant = "docker" + link = "https://github.com/verygoodplugins/automem" + concurrency = 4 + + def __init__(self): + self._image = os.environ.get("AUTOMEM_IMAGE", _DEFAULT_IMAGE) + self._token = os.environ.get("AUTOMEM_TOKEN", "benchmark-token") + self._max_chars = int(os.environ.get("AUTOMEM_MAX_CHARS", "1800")) + self._batch_size = int(os.environ.get("AUTOMEM_BATCH_SIZE", "50")) # /memory/batch, max 500 + self._k_override = os.environ.get("AUTOMEM_RECALL_K") + self._enrich_settle_s = int(os.environ.get("AUTOMEM_ENRICH_SETTLE_SECONDS", "120")) + # Backpressure during ingest so the async enrichment queue can't grow unbounded + # and OOM-kill the container (set max_pending<=0 to disable). + self._enrich_drain_every = int(os.environ.get("AUTOMEM_ENRICH_DRAIN_EVERY", "100")) + self._enrich_max_pending = int(os.environ.get("AUTOMEM_ENRICH_MAX_PENDING", "150")) + self._project = f"automem_amb_{uuid.uuid4().hex[:8]}" + self._endpoint = None + self._run_tag = f"ambrun-{uuid.uuid4().hex[:8]}" + self._compose_env = None + + def initialize(self) -> None: + # Honor pre-set ports if a caller pins them (e.g. to reach the stack from + # another network namespace); otherwise pick free ones for concurrent runs. + api = int(os.environ.get("AUTOMEM_API_PORT") or _free_port()) + falk = int(os.environ.get("AUTOMEM_FALKOR_PORT") or _free_port()) + qdr = int(os.environ.get("AUTOMEM_QDRANT_PORT") or _free_port()) + self._compose_env = {**os.environ, "AUTOMEM_IMAGE": self._image, + "AUTOMEM_API_PORT": str(api), "AUTOMEM_FALKOR_PORT": str(falk), + "AUTOMEM_QDRANT_PORT": str(qdr)} + subprocess.run(["docker", "compose", "-p", self._project, "-f", str(_COMPOSE), "up", "-d"], + env=self._compose_env, check=True) + # Ensure the stack is torn down even if the harness crashes before calling cleanup(). + atexit.register(self.cleanup) + # Default localhost; override AUTOMEM_HOST to reach the stack on another host. + self._endpoint = f"http://{os.environ.get('AUTOMEM_HOST', 'localhost')}:{api}" + for _ in range(60): + try: + body = self._req("GET", "/health") + if body.get("status") in {"ok", "healthy", "degraded"} and body.get("falkordb") == "connected": + return + except Exception: + pass + time.sleep(2) + raise RuntimeError("AutoMem stack did not become healthy") + + def cleanup(self) -> None: + if self._compose_env is not None: + subprocess.run(["docker", "compose", "-p", self._project, "-f", str(_COMPOSE), "down", "-v"], + env=self._compose_env, check=False) + + def _req(self, method, path, *, params=None, body=None, _retries=5): + url = f"{self._endpoint}{path}" + if params: + url = f"{url}?{urllib.parse.urlencode(params, doseq=True)}" + data = json.dumps(body).encode() if body is not None else None + headers = {"X-Api-Key": self._token} + if data is not None: + headers["Content-Type"] = "application/json" + req = urllib.request.Request(url, data=data, headers=headers, method=method) + for attempt in range(_retries): + try: + with urllib.request.urlopen(req, timeout=120) as r: + raw = r.read() + return json.loads(raw) if raw else {} + except urllib.error.HTTPError: + # A real HTTP response (e.g. 400 for over-limit content) — never retry; + # callers handle these (ingest skips 400s). + raise + except (urllib.error.URLError, http.client.RemoteDisconnected, + http.client.IncompleteRead, ConnectionError, TimeoutError, OSError): + # AutoMem drops connections under ingest/enrichment load; retry with backoff + # so a transient blip doesn't kill a multi-thousand-query run. + if attempt == _retries - 1: + raise + time.sleep(min(2 ** attempt, 20)) + + def _scoped_tags(self, user_id): + tags = [self._run_tag] + if user_id: + tags.append(f"ambuser-{_slug(user_id)}") + return tags + + def ingest(self, documents) -> None: + # Flatten docs into chunked memory items, then POST /memory/batch (which + # batch-embeds + UNWINDs graph writes). One-at-a-time POST /memory is ~0.8s/doc + # on CPU FastEmbed; batching is the difference between hours and days. + items = [] + for doc in documents: + tags = self._scoped_tags(doc.user_id) + pieces = _chunk_text(doc.content, self._max_chars) + for i, piece in enumerate(pieces): + meta = {"amb_doc_id": doc.id, "amb_user_id": doc.user_id, "amb_run": self._run_tag} + if len(pieces) > 1: + meta["amb_chunk"] = i + item = {"content": piece, "tags": tags, "importance": 0.6, "metadata": meta} + if doc.timestamp: + item["timestamp"] = doc.timestamp + items.append(item) + posted = 0 + for start in range(0, len(items), self._batch_size): + batch = items[start:start + self._batch_size] + try: + self._req("POST", "/memory/batch", body={"memories": batch}) + posted += len(batch) + except urllib.error.HTTPError: + # A bad item (e.g. over-limit) fails the whole batch; fall back to + # per-item posts for this batch so one bad chunk doesn't drop the rest. + for item in batch: + try: + self._req("POST", "/memory", body=item) + posted += 1 + except urllib.error.HTTPError as exc: + if exc.code == 400: + continue + raise + # Bound the async enrichment backlog (no-op when ENRICHMENT_ENABLED=false). + if posted % self._enrich_drain_every < self._batch_size: + self._bound_enrichment_queue() + self._settle_enrichment() + + def _bound_enrichment_queue(self) -> None: + if self._enrich_max_pending <= 0: + return + for _ in range(120): # up to ~4 min of backpressure + try: + pending = self._req("GET", "/health").get("enrichment", {}).get("pending", 0) + except Exception: + return + if pending <= self._enrich_max_pending: + return + time.sleep(2) + + def _settle_enrichment(self) -> None: + if self._enrich_settle_s <= 0: + return + waited = 0 + while waited < self._enrich_settle_s: + try: + pending = self._req("GET", "/health").get("enrichment", {}).get("pending", 0) + except Exception: + pending = 0 + if not pending: + return + time.sleep(3); waited += 3 + + def retrieve(self, query, k=10, user_id=None, query_timestamp=None): + k_eff = int(self._k_override) if self._k_override else k + params = {"query": query, "limit": k_eff, "tags": self._scoped_tags(user_id), + "tag_mode": "all", "tag_match": "exact", "recency_bias": "auto", + "expand_relations": "true", "expand_respect_tags": "true"} + resp = self._req("GET", "/recall", params=params) + results = resp.get("results", []) if isinstance(resp, dict) else [] + docs = [Document(id=str(r.get("id") or ""), content=_extract_content(r)) for r in results] + return docs, {"results": results, "count": len(results)} diff --git a/src/memory_bench/memory/automem_compose.yml b/src/memory_bench/memory/automem_compose.yml new file mode 100644 index 0000000..cd002ad --- /dev/null +++ b/src/memory_bench/memory/automem_compose.yml @@ -0,0 +1,23 @@ +# Spun by the AutoMem provider. Ports + project name are overridden per run. +services: + falkordb: + image: falkordb/falkordb:v4.18.3 # pinned (was :latest) for reproducibility — graph module 41803, the version the submission ran + ports: ["${AUTOMEM_FALKOR_PORT:-6379}:6379"] + qdrant: + image: qdrant/qdrant:v1.11.3 + ports: ["${AUTOMEM_QDRANT_PORT:-6333}:6333"] + automem: + image: ${AUTOMEM_IMAGE:-ghcr.io/verygoodplugins/automem:amb-v1} + depends_on: [falkordb, qdrant] + ports: ["${AUTOMEM_API_PORT:-8001}:8001"] + environment: + PORT: "8001" + FALKORDB_HOST: falkordb + FALKORDB_PORT: "6379" + QDRANT_URL: http://qdrant:6333 + EMBEDDING_PROVIDER: local + VECTOR_SIZE: "768" + AUTOMEM_API_TOKEN: benchmark-token + JIT_ENRICHMENT_ENABLED: "true" + ENRICHMENT_ENABLE_SUMMARIES: "false" # no LLM summaries (lean, no external deps) + ENRICHMENT_ENABLED: "false" # async enrichment worker off (fast + low-mem); JIT inline recall enrichment stays on diff --git a/tests/test_automem_provider.py b/tests/test_automem_provider.py new file mode 100644 index 0000000..09a8e35 --- /dev/null +++ b/tests/test_automem_provider.py @@ -0,0 +1,105 @@ +import io, json +from contextlib import contextmanager +from memory_bench.memory.automem import ( + _chunk_text, _extract_content, _free_port, AutoMemMemoryProvider, +) +from memory_bench.models import Document +import memory_bench.memory.automem as m + + +def test_chunk_long_text_under_limit(): + text = " ".join(f"sentence {i}." for i in range(600)) + assert len(text) > 1800 + chunks = _chunk_text(text, 1800) + assert len(chunks) > 1 + assert all(len(c) <= 1800 for c in chunks) + +def test_chunk_short_text_single(): + assert _chunk_text("hello world", 1800) == ["hello world"] + +def test_extract_content_from_nested_memory(): + assert _extract_content({"id": "m1", "memory": {"content": "answer", "summary": "s"}}) == "answer" + +def test_extract_content_falls_back_to_summary(): + assert _extract_content({"id": "m1", "memory": {"content": "", "summary": "fallback"}}) == "fallback" + +def test_extract_content_empty_when_missing(): + assert _extract_content({"id": "m1"}) == "" + +def test_free_port_returns_open_port(): + p = _free_port() + assert 1024 < p < 65536 + + +class _FakeHTTP: + def __init__(self, responses): + self.responses = responses; self.calls = [] + @contextmanager + def urlopen(self, req, timeout=None): + body = req.data.decode() if req.data else None + self.calls.append((req.get_method(), req.full_url, dict(req.headers), body)) + payload = self.responses.pop(0) if self.responses else {} + yield io.BytesIO(json.dumps(payload).encode()) + + +def test_ingest_batches_chunks(monkeypatch): + fake = _FakeHTTP([{"stored": 99} for _ in range(10)]) + monkeypatch.setattr(m.urllib.request, "urlopen", fake.urlopen) + p = AutoMemMemoryProvider() + p._endpoint = "http://x:8001"; p._token = "t"; p._run_tag = "ambrun-test" + p._enrich_settle_s = 0; p._enrich_max_pending = 0 + long_doc = Document(id="d", content=" ".join(f"s{i}." for i in range(600)), user_id="u1") + p.ingest([long_doc]) + batches = [c for c in fake.calls if c[1].endswith("/memory/batch")] + assert len(batches) >= 1 # chunks go out as a batch, not one-by-one + items = json.loads(batches[0][3])["memories"] + assert len(items) > 1 + for it in items: + assert len(it["content"]) <= 1800 + assert "ambrun-test" in it["tags"] + +def test_retrieve_extracts_nested_content(monkeypatch): + fake = _FakeHTTP([{"results": [{"id": "m9", "memory": {"content": "answer"}}]}]) + monkeypatch.setattr(m.urllib.request, "urlopen", fake.urlopen) + p = AutoMemMemoryProvider() + p._endpoint = "http://x:8001"; p._token = "t"; p._run_tag = "ambrun-test" + docs, _ = p.retrieve("q", k=5, user_id="u1") + assert docs[0].content == "answer" + assert "expand_relations=true" in fake.calls[0][1] + + +def test_req_retries_remote_disconnected(monkeypatch): + import http.client + calls = {"n": 0} + class _OK: + def __enter__(self): return self + def __exit__(self, *a): return False + def read(self): return b'{"ok": 1}' + def flaky_urlopen(req, timeout=None): + calls["n"] += 1 + if calls["n"] == 1: + raise http.client.RemoteDisconnected("Remote end closed connection") + return _OK() + monkeypatch.setattr(m.urllib.request, "urlopen", flaky_urlopen) + monkeypatch.setattr(m.time, "sleep", lambda *_: None) + p = AutoMemMemoryProvider() + p._endpoint = "http://x:8001"; p._token = "t" + assert p._req("GET", "/health") == {"ok": 1} + assert calls["n"] == 2 # retried past the disconnect + + +def test_req_does_not_retry_http_error(monkeypatch): + import urllib.error + calls = {"n": 0} + def boom(req, timeout=None): + calls["n"] += 1 + raise urllib.error.HTTPError(req.full_url, 400, "bad", {}, None) + monkeypatch.setattr(m.urllib.request, "urlopen", boom) + p = AutoMemMemoryProvider() + p._endpoint = "http://x:8001"; p._token = "t" + try: + p._req("POST", "/memory", body={"content": "x"}) + raised = False + except urllib.error.HTTPError: + raised = True + assert raised and calls["n"] == 1 # 400 raised immediately, not retried