diff --git a/benchmarks/recall/README.md b/benchmarks/recall/README.md index 7b3cd88..6b3c9e6 100644 --- a/benchmarks/recall/README.md +++ b/benchmarks/recall/README.md @@ -167,3 +167,41 @@ time and sanitized query failures. It does not copy query text, corpus text, vectors or free-form provider errors. Credential-shaped configuration keys such as `api_key`, `password`, `secret`, `token` and `authorization` are rejected instead of being copied into an artifact. + +## Live memd producer + +The `produce` subcommand queries a running `memd` over every dataset query and +emits a `mem.recall-rankings.v1` file that the existing `run --rankings` path +consumes. Latency is measured client-side per request; the `0 ms` sentinel +warning above applies only to the offline lexical lane. + +```bash +python3 -m benchmarks.recall produce \ + --memd-url http://localhost:8080 \ + --token "$MEM_TOKEN" \ + --output /tmp/live-rankings.json \ + --dimension 1536 \ + --mode hybrid +``` + +Then score it against the lexical baseline: + +```bash +python3 -m benchmarks.recall run \ + --rankings /tmp/live-rankings.json \ + --output /tmp/live-artifact.json \ + --compare benchmarks/recall/baselines/lexical-reference.v1.json +``` + +The producer maps each API result back to a dataset `doc_id` by matching the +`path` field returned by `/v1/search` against the corpus. When multiple +documents share a path, the snippet text is used to pick the best overlap. +Query filters are translated where the API supports them: `path_prefix` becomes +`scope`, and `source_kind` becomes `type` (`image_caption` → `image`, +`text` → `text`). The `workspace` filter is not sent to the API because the +auth token determines workspace scope. + +Adjust `--dimension`, `--mode`, `--provider` and `--model` to match the +embedding configuration of the live system. The emitted `configuration` block +is populated from these flags, not hand-written, so the artifact cannot be +mistaken for the lexical reference. diff --git a/benchmarks/recall/__main__.py b/benchmarks/recall/__main__.py index 9fb888c..9df0dff 100644 --- a/benchmarks/recall/__main__.py +++ b/benchmarks/recall/__main__.py @@ -7,7 +7,9 @@ import sys import tempfile +from .dataset import load_dataset from .errors import BenchmarkError +from .live_producer import produce_rankings from .runner import ( compare_artifacts, comparison_summary, @@ -66,6 +68,24 @@ def _parser() -> argparse.ArgumentParser: type=Path, default=PACKAGE_ROOT / "fixtures" / "external-rankings.leak.v1.json", ) + + produce = subparsers.add_parser( + "produce", + help="query a live memd and emit mem.recall-rankings.v1", + ) + produce.add_argument("--memd-url", required=True, help="base URL of memd") + produce.add_argument("--token", required=True, help="bearer token for auth") + produce.add_argument("--dataset", type=Path, default=DEFAULT_DATASET) + produce.add_argument("--output", type=Path, required=True) + produce.add_argument("--limit", type=int, default=10) + produce.add_argument("--timeout", type=float, default=30.0) + produce.add_argument("--engine", default="live-memd") + produce.add_argument("--dimension", type=int, default=1536) + produce.add_argument( + "--mode", default="hybrid", choices=["lexical", "vector", "hybrid"] + ) + produce.add_argument("--provider", default="memd") + produce.add_argument("--model", default="memd-embedded") return parser @@ -100,6 +120,27 @@ def main(argv: list[str] | None = None) -> int: print(comparison_summary(comparison)) return 2 if candidate["metrics"]["overall"]["leakage_count"] else 0 + if args.command == "produce": + dataset = load_dataset(args.dataset) + rankings = produce_rankings( + dataset, + base_url=args.memd_url, + token=args.token, + limit=args.limit, + timeout=args.timeout, + engine_label=args.engine, + dimension=args.dimension, + mode=args.mode, + provider=args.provider, + model=args.model, + ) + write_json(args.output, rankings) + ok_count = sum(1 for q in rankings["queries"] if q["status"] == "ok") + err_count = sum(1 for q in rankings["queries"] if q["status"] == "error") + print(f"produced rankings: {ok_count} ok, {err_count} error") + print(f"rankings artifact: {args.output}") + return 0 + first = run_benchmark( dataset_dir=args.dataset, generated_at="2000-01-01T00:00:00+00:00", diff --git a/benchmarks/recall/live_producer.py b/benchmarks/recall/live_producer.py new file mode 100644 index 0000000..6419e2d --- /dev/null +++ b/benchmarks/recall/live_producer.py @@ -0,0 +1,283 @@ +"""Produce mem.recall-rankings.v1 from a live memd instance. + +Queries each dataset query against POST /v1/search, maps API results back to +dataset doc_ids by path, and emits the rankings JSON that the existing harness +consumes via --rankings. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import socket +import time +import unicodedata +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from .dataset import Dataset, Document, load_dataset + +PACKAGE_ROOT = Path(__file__).resolve().parent +DEFAULT_DATASET = PACKAGE_ROOT / "data" / "v1" + +_SOURCE_KIND_TO_TYPE = { + "image_caption": "image", + "text": "text", +} + + +def _build_path_index(documents: list[Document]) -> dict[str, list[Document]]: + index: dict[str, list[Document]] = {} + for doc in documents: + index.setdefault(doc.path, []).append(doc) + return index + + +def _match_doc_by_path( + api_path: str, + snippet: str, + candidates: list[Document], +) -> Document | None: + if not candidates: + return None + if len(candidates) == 1: + return candidates[0] + normalized_snippet = unicodedata.normalize("NFKC", snippet).casefold() + best: Document | None = None + best_overlap = -1 + for doc in candidates: + doc_tokens = set(unicodedata.normalize("NFKC", doc.text).casefold().split()) + overlap = sum(1 for t in normalized_snippet.split() if t in doc_tokens) + if overlap > best_overlap: + best_overlap = overlap + best = doc + return best + + +def _source_kind_to_api_type(source_kind: str) -> str | None: + return _SOURCE_KIND_TO_TYPE.get(source_kind) + + +def _coarse_host() -> str: + try: + return f"{platform.system()}/{platform.machine()}" + except Exception: + return "unknown" + + +def _query_memd( + base_url: str, + token: str, + query_text: str, + *, + scope: str = "", + type_filter: str = "", + limit: int = 10, + timeout: float = 30.0, +) -> tuple[list[dict[str, Any]], float, str | None]: + body: dict[str, Any] = {"query": query_text, "limit": limit} + if scope: + body["scope"] = scope + if type_filter: + body["type"] = type_filter + + url = base_url.rstrip("/") + "/v1/search" + data = json.dumps(body).encode("utf-8") + req = Request(url, data=data, method="POST") + req.add_header("Content-Type", "application/json") + req.add_header("Authorization", f"Bearer {token}") + + start = time.perf_counter() + try: + with urlopen(req, timeout=timeout) as resp: + payload = json.loads(resp.read().decode("utf-8")) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + results = payload.get("results", []) + return results, elapsed_ms, None + except HTTPError as exc: + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return [], elapsed_ms, f"http_{exc.code}" + except URLError: + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return [], elapsed_ms, "connection_error" + except Exception: + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return [], elapsed_ms, "unknown_error" + + +def produce_rankings( + dataset: Dataset, + *, + base_url: str, + token: str, + limit: int = 10, + timeout: float = 30.0, + engine_label: str = "live-memd", + dimension: int = 1536, + mode: str = "hybrid", + provider: str = "memd", + model: str = "memd-embedded", +) -> dict[str, Any]: + path_index = _build_path_index(list(dataset.documents)) + + query_rows: list[dict[str, Any]] = [] + for query in dataset.queries: + scope = query.filters.get("path_prefix", "") + type_filter = _source_kind_to_api_type(query.expected_source_kind) or "" + + api_results, latency_ms, error_code = _query_memd( + base_url, + token, + query.text, + scope=scope, + type_filter=type_filter, + limit=limit, + timeout=timeout, + ) + + if error_code and not api_results: + row: dict[str, Any] = { + "query_id": query.id, + "status": "error", + "latency_ms": round(latency_ms, 2), + "results": [], + "error_code": error_code, + } + query_rows.append(row) + continue + + mapped_results: list[dict[str, Any]] = [] + seen_doc_ids: set[str] = set() + for hit in api_results: + hit_path = hit.get("path", "") + snippet = hit.get("snippet", "") + candidates = path_index.get(hit_path, []) + doc = _match_doc_by_path(hit_path, snippet, candidates) + if doc is None or doc.id in seen_doc_ids: + continue + seen_doc_ids.add(doc.id) + result: dict[str, Any] = { + "doc_id": doc.id, + "citation": doc.citation, + } + score = hit.get("score") + if score is not None: + result["score"] = float(score) + mapped_results.append(result) + + status = "ok" if not error_code else "partial" + row = { + "query_id": query.id, + "status": status, + "latency_ms": round(latency_ms, 2), + "results": mapped_results, + } + if error_code: + row["error_code"] = error_code + query_rows.append(row) + + return { + "schema_version": "mem.recall-rankings.v1", + "engine": engine_label, + "configuration": { + "mode": mode, + "provider": provider, + "model": model, + "dimension": dimension, + "index": { + "kind": "pgvector", + "distance": "cosine", + }, + "search": { + "top_k": limit, + "type": "auto", + }, + }, + "hardware": { + "host": _coarse_host(), + "client": socket.gethostname(), + }, + "queries": query_rows, + } + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.recall.produce", + description="Query a live memd and emit mem.recall-rankings.v1.", + ) + parser.add_argument("--memd-url", required=True, help="base URL of memd") + parser.add_argument("--token", required=True, help="bearer token for auth") + parser.add_argument( + "--dataset", type=Path, default=DEFAULT_DATASET, help="dataset directory" + ) + parser.add_argument( + "--output", type=Path, required=True, help="output rankings JSON path" + ) + parser.add_argument( + "--limit", type=int, default=10, help="max results per query (default 10)" + ) + parser.add_argument( + "--timeout", type=float, default=30.0, help="per-query timeout in seconds" + ) + parser.add_argument( + "--engine", + default="live-memd", + help="engine label for the artifact (default: live-memd)", + ) + parser.add_argument( + "--dimension", + type=int, + default=1536, + help="embedding dimension of the live model (default: 1536)", + ) + parser.add_argument( + "--mode", + default="hybrid", + choices=["lexical", "vector", "hybrid"], + help="search mode (default: hybrid)", + ) + parser.add_argument( + "--provider", + default="memd", + help="provider label (default: memd)", + ) + parser.add_argument( + "--model", + default="memd-embedded", + help="model label (default: memd-embedded)", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + dataset = load_dataset(args.dataset) + rankings = produce_rankings( + dataset, + base_url=args.memd_url, + token=args.token, + limit=args.limit, + timeout=args.timeout, + engine_label=args.engine, + dimension=args.dimension, + mode=args.mode, + provider=args.provider, + model=args.model, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(rankings, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + ok_count = sum(1 for q in rankings["queries"] if q["status"] == "ok") + err_count = sum(1 for q in rankings["queries"] if q["status"] == "error") + print(f"wrote {args.output} ({ok_count} ok, {err_count} error)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/recall/tests/test_live_producer.py b/benchmarks/recall/tests/test_live_producer.py new file mode 100644 index 0000000..a9e3f4e --- /dev/null +++ b/benchmarks/recall/tests/test_live_producer.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +from benchmarks.recall.live_producer import ( + _build_path_index, + _match_doc_by_path, + produce_rankings, +) +from benchmarks.recall.dataset import Document, load_dataset + + +class PathIndexTest(unittest.TestCase): + def test_index_groups_by_path(self) -> None: + docs = [ + Document( + id="a", language="en", source_kind="text", workspace="alpha", + path="/notes/a.md", citation="mem://files/a", + text="alpha note", metadata={}, + ), + Document( + id="b", language="en", source_kind="text", workspace="alpha", + path="/notes/a.md", citation="mem://files/b", + text="beta note", metadata={}, + ), + ] + index = _build_path_index(docs) + self.assertEqual(len(index["/notes/a.md"]), 2) + + def test_match_single_candidate(self) -> None: + doc = Document( + id="solo", language="en", source_kind="text", workspace="alpha", + path="/notes/solo.md", citation="mem://files/solo", + text="unique content", metadata={}, + ) + result = _match_doc_by_path("/notes/solo.md", "anything", [doc]) + self.assertEqual(result, doc) + + def test_match_picks_best_snippet_overlap(self) -> None: + doc_a = Document( + id="a", language="en", source_kind="text", workspace="alpha", + path="/notes/shared.md", citation="mem://files/a", + text="saturn ring observation", metadata={}, + ) + doc_b = Document( + id="b", language="en", source_kind="text", workspace="alpha", + path="/notes/shared.md", citation="mem://files/b", + text="completely different topic", metadata={}, + ) + result = _match_doc_by_path("/notes/shared.md", "saturn ring", [doc_a, doc_b]) + self.assertEqual(result, doc_a) + + +class ProduceRankingsTest(unittest.TestCase): + def setUp(self) -> None: + self.tempdir = tempfile.TemporaryDirectory() + self.root = Path(self.tempdir.name) + self.dataset = self.root / "dataset" + self.dataset.mkdir() + (self.dataset / "dataset.json").write_text( + json.dumps({ + "schema_version": "mem.recall-dataset.v1", + "version": "unit-test-v1", + "provenance": "hand-authored synthetic data", + "license": "CC0-1.0", + "required_coverage": { + "slices": ["exact"], "languages": ["en"], "source_kinds": ["text"], + }, + }), + encoding="utf-8", + ) + (self.dataset / "corpus.jsonl").write_text( + json.dumps({ + "id": "file-en-cassini", "language": "en", "source_kind": "text", + "workspace": "alpha", "path": "/research/saturn.md", + "citation": "mem://files/file-en-cassini", + "text": "Cassini observed Saturn hexagonal storm", + "provenance": "synthetic", + }) + "\n", + encoding="utf-8", + ) + (self.dataset / "queries.jsonl").write_text( + json.dumps({ + "id": "q-en-text-exact", "text": "Cassini Saturn hexagonal storm", + "language": "en", "slice": "exact", + "filters": {"workspace": "alpha", "source_kind": "text"}, + "expected_source_kind": "text", + }) + "\n", + encoding="utf-8", + ) + (self.dataset / "qrels.json").write_text( + json.dumps({"q-en-text-exact": {"file-en-cassini": 3}}), + encoding="utf-8", + ) + + def tearDown(self) -> None: + self.tempdir.cleanup() + + @patch("benchmarks.recall.live_producer._query_memd") + def test_produce_rankings_success(self, mock_query: unittest.mock.MagicMock) -> None: + mock_query.return_value = ( + [{"path": "/research/saturn.md", "snippet": "Cassini observed Saturn hexagonal storm", "score": 0.95}], + 12.5, None, + ) + dataset = load_dataset(self.dataset) + rankings = produce_rankings( + dataset, base_url="http://localhost:8080", token="test-token", dimension=768, + ) + self.assertEqual(rankings["schema_version"], "mem.recall-rankings.v1") + self.assertEqual(rankings["engine"], "live-memd") + self.assertEqual(rankings["configuration"]["dimension"], 768) + self.assertEqual(len(rankings["queries"]), 1) + query_row = rankings["queries"][0] + self.assertEqual(query_row["query_id"], "q-en-text-exact") + self.assertEqual(query_row["status"], "ok") + self.assertGreater(query_row["latency_ms"], 0) + self.assertEqual(len(query_row["results"]), 1) + self.assertEqual(query_row["results"][0]["doc_id"], "file-en-cassini") + + @patch("benchmarks.recall.live_producer._query_memd") + def test_produce_rankings_error(self, mock_query: unittest.mock.MagicMock) -> None: + mock_query.return_value = ([], 5.0, "http_503") + dataset = load_dataset(self.dataset) + rankings = produce_rankings( + dataset, base_url="http://localhost:8080", token="test-token", + ) + query_row = rankings["queries"][0] + self.assertEqual(query_row["status"], "error") + self.assertEqual(query_row["error_code"], "http_503") + self.assertEqual(query_row["results"], []) + + +if __name__ == "__main__": + unittest.main()