From 516be4bb761c7686a29a0b0d060a82fc1d7b65de Mon Sep 17 00:00:00 2001 From: Drew Cain Date: Fri, 12 Jun 2026 13:33:08 -0500 Subject: [PATCH] feat: LongMemEval-S dataset, grouped-corpus runner mode, and converter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LongMemEval-S (Wu et al., ICLR 2025) gives each of its 500 questions an independent ~50-session haystack, so it cannot run as a single shared corpus. This adds: - datasets/longmemeval.py: streaming fetch from the official HF repo (~278MB) with checksum provenance, plus shape validation on load. - converters/longmemeval_to_corpus.py: one corpus per question under groups//docs plus a single queries.json whose entries carry a group field. Session ids are remapped to neutral positional ids and per-turn has_answer flags dropped — the raw dataset marks evidence sessions with an answer_ id prefix, which would leak ground truth into ingested corpora. Duplicate sessions within a haystack (15 questions) are ingested once. - runner grouped mode: when queries carry groups, each group runs as an isolated mini-benchmark — fresh provider instance, group-suffixed run id (namespacing the BM project / mem0 user), per-group corpus. A failed group is recorded in provider status and skipped; ProviderSkippedError on the first group skips the provider. - QA stage: question_date metadata is appended to the question for both answerer and judge — temporal-reasoning questions (133 of 500) are unanswerable without the reference date. - CLI: datasets fetch --dataset longmemeval-s, convert longmemeval (--max-questions for dev slices); justfile recipes incl. a 25-question dev slice; README section. Verified end-to-end against real data: converted 3 real questions (152 docs) and ran grouped retrieval with bm-local (BM 0.22.0) — recall@5 = 1.0, MRR = 1.0, evidence doc ranked first in all 3 groups, grouped metadata recorded in provider status. Known follow-up: per-group overhead is ~2.3 min with bm-local (CLI cold starts, reindex, MCP warm-up per group), ~19h extrapolated for the full 500. A warm-session/shared-config optimization across groups is the next harness task. Co-Authored-By: Claude Fable 5 Signed-off-by: Drew Cain --- .gitignore | 2 + README.md | 24 ++ justfile | 37 ++++ src/basic_memory_benchmarks/cli.py | 68 ++++-- .../converters/longmemeval_to_corpus.py | 134 ++++++++++++ .../datasets/longmemeval.py | 67 +++++- src/basic_memory_benchmarks/models.py | 3 + src/basic_memory_benchmarks/runner.py | 158 ++++++++++--- src/basic_memory_benchmarks/scoring/qa.py | 20 +- .../scoring/retrieval.py | 1 + tests/converters/__init__.py | 0 .../converters/test_longmemeval_converter.py | 132 +++++++++++ tests/test_qa_scoring.py | 22 ++ tests/test_runner_grouped.py | 207 ++++++++++++++++++ 14 files changed, 826 insertions(+), 49 deletions(-) create mode 100644 src/basic_memory_benchmarks/converters/longmemeval_to_corpus.py create mode 100644 tests/converters/__init__.py create mode 100644 tests/converters/test_longmemeval_converter.py create mode 100644 tests/test_runner_grouped.py diff --git a/.gitignore b/.gitignore index 04b8c90..0b5ce91 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,5 @@ benchmarks/datasets/locomo/locomo10.json benchmarks/datasets/locomo/locomo10.provenance.json .env benchmarks/bm-home/ +benchmarks/datasets/longmemeval/longmemeval_s.json +benchmarks/datasets/longmemeval/longmemeval_s.provenance.json diff --git a/README.md b/README.md index 52b2597..3f09430 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,30 @@ uv run bm-bench run judge --run-dir benchmarks/runs/ uv run bm-bench publish --run-dir benchmarks/runs/ ``` +## LongMemEval-S + +LongMemEval (Wu et al., ICLR 2025) gives each of its 500 questions an +independent haystack of ~50 chat sessions, so it runs in **grouped mode**: the +converter writes one corpus per question under `groups//docs`, +and the runner ingests + queries each group in isolation (fresh provider +instance, group-suffixed run id namespacing the BM project / mem0 user). + +```bash +just bench-prepare-longmemeval # fetch (~278MB) + convert all 500 +just bench-convert-longmemeval-dev # or: 25-question dev slice +just bench-run-longmemeval-dev # grouped retrieval, bm-local +``` + +Then score answers with the QA stage as usual (`run qa --run-dir ...`). The +question's ask-date is carried in query metadata and appended to the question +for both the answerer and the judge — temporal-reasoning questions are +unanswerable without it. + +Anti-leakage: the raw dataset marks evidence sessions via an `answer_` +session-id prefix and per-turn `has_answer` flags. The converter remaps all +session ids to neutral positional ids (`-s012`) and drops turn flags, so +ingested corpora carry no evidence markers. + ## Basic Memory source policy By default this project tracks Basic Memory from `main`. diff --git a/justfile b/justfile index ac99227..5c87695 100644 --- a/justfile +++ b/justfile @@ -9,6 +9,9 @@ bm_local_path_flag := if bm_local_path != "" { "--bm-local-path " + bm_local_pat locomo_dataset_path := "benchmarks/datasets/locomo/locomo10.json" locomo_output_dir := "benchmarks/generated/locomo" locomo_c1_output_dir := "benchmarks/generated/locomo-c1" +longmemeval_dataset_path := "benchmarks/datasets/longmemeval/longmemeval_s.json" +longmemeval_output_dir := "benchmarks/generated/longmemeval-s" +longmemeval_dev_output_dir := "benchmarks/generated/longmemeval-s-dev" # --- Repo maintenance --- @@ -50,6 +53,40 @@ bench-prepare-short: bench-fetch-locomo bench-convert-locomo-c1 bench-make-quick bench-prepare-long: bench-fetch-locomo bench-convert-locomo +bench-fetch-longmemeval: + uv run bm-bench datasets fetch --dataset longmemeval-s --output {{longmemeval_dataset_path}} + +bench-convert-longmemeval: + uv run bm-bench convert longmemeval --dataset-path {{longmemeval_dataset_path}} --output-dir {{longmemeval_output_dir}} + +# Dev slice: first 25 questions for fast iteration +bench-convert-longmemeval-dev: + uv run bm-bench convert longmemeval --dataset-path {{longmemeval_dataset_path}} --output-dir {{longmemeval_dev_output_dir}} --max-questions 25 + +bench-prepare-longmemeval: bench-fetch-longmemeval bench-convert-longmemeval + +# Grouped retrieval over the LongMemEval-S dev slice (bm-local only) +bench-run-longmemeval-dev: + uv run bm-bench run retrieval \ + --dataset-id longmemeval_s \ + --dataset-path {{longmemeval_dataset_path}} \ + --corpus-dir {{longmemeval_dev_output_dir}}/groups \ + --queries-path {{longmemeval_dev_output_dir}}/queries.json \ + --providers bm-local \ + {{bm_local_path_flag}} \ + --strict-providers + +# Grouped retrieval over full LongMemEval-S (slow: 500 isolated group corpora) +bench-run-longmemeval: + uv run bm-bench run retrieval \ + --dataset-id longmemeval_s \ + --dataset-path {{longmemeval_dataset_path}} \ + --corpus-dir {{longmemeval_output_dir}}/groups \ + --queries-path {{longmemeval_output_dir}}/queries.json \ + --providers bm-local,mem0-local \ + {{bm_local_path_flag}} \ + --allow-provider-skip + # --- One-command pipelines --- # Full retrieval benchmark pipeline: diff --git a/src/basic_memory_benchmarks/cli.py b/src/basic_memory_benchmarks/cli.py index 4ea96b1..042001f 100644 --- a/src/basic_memory_benchmarks/cli.py +++ b/src/basic_memory_benchmarks/cli.py @@ -10,9 +10,17 @@ from rich.console import Console from basic_memory_benchmarks.converters.locomo_to_corpus import convert_locomo_to_corpus +from basic_memory_benchmarks.converters.longmemeval_to_corpus import convert_longmemeval_to_corpus from basic_memory_benchmarks.datasets.locomo import LOCOMO_URL, fetch_locomo_dataset +from basic_memory_benchmarks.datasets.longmemeval import ( + LONGMEMEVAL_S_URL, + fetch_longmemeval_dataset, +) from basic_memory_benchmarks.models import DatasetProvenance, RunConfig -from basic_memory_benchmarks.reporting.compare import compare_provider_metric, load_retrieval_summary +from basic_memory_benchmarks.reporting.compare import ( + compare_provider_metric, + load_retrieval_summary, +) from basic_memory_benchmarks.runner import run_judge, run_qa_stage, run_retrieval from basic_memory_benchmarks.utils import sha256_file @@ -31,20 +39,29 @@ @datasets_app.command("fetch") def datasets_fetch( dataset: str = typer.Option("locomo", "--dataset"), - output: Path = typer.Option(Path("benchmarks/datasets/locomo/locomo10.json"), "--output"), - url: str = typer.Option(LOCOMO_URL, "--url"), + output: Path | None = typer.Option(None, "--output"), + url: str | None = typer.Option(None, "--url"), ) -> None: - if dataset != "locomo": - raise typer.BadParameter("Only locomo is supported in v1") - - provenance = fetch_locomo_dataset(output_path=output, url=url) - console.print(f"Downloaded {dataset} to [cyan]{output}[/cyan]") + if dataset == "locomo": + resolved_output = output or Path("benchmarks/datasets/locomo/locomo10.json") + provenance = fetch_locomo_dataset(output_path=resolved_output, url=url or LOCOMO_URL) + elif dataset == "longmemeval-s": + resolved_output = output or Path("benchmarks/datasets/longmemeval/longmemeval_s.json") + provenance = fetch_longmemeval_dataset( + output_path=resolved_output, url=url or LONGMEMEVAL_S_URL + ) + else: + raise typer.BadParameter("Supported datasets: locomo, longmemeval-s") + + console.print(f"Downloaded {dataset} to [cyan]{resolved_output}[/cyan]") console.print(f"SHA256: [green]{provenance.checksum_sha256}[/green]") @convert_app.command("locomo") def convert_locomo( - dataset_path: Path = typer.Option(Path("benchmarks/datasets/locomo/locomo10.json"), "--dataset-path"), + dataset_path: Path = typer.Option( + Path("benchmarks/datasets/locomo/locomo10.json"), "--dataset-path" + ), output_dir: Path = typer.Option(Path("benchmarks/generated/locomo"), "--output-dir"), max_conversations: int | None = typer.Option(None, "--max-conversations"), ) -> None: @@ -57,13 +74,34 @@ def convert_locomo( console.print(f"Queries: [cyan]{queries_path}[/cyan] ({query_count})") +@convert_app.command("longmemeval") +def convert_longmemeval( + dataset_path: Path = typer.Option( + Path("benchmarks/datasets/longmemeval/longmemeval_s.json"), "--dataset-path" + ), + output_dir: Path = typer.Option(Path("benchmarks/generated/longmemeval-s"), "--output-dir"), + max_questions: int | None = typer.Option(None, "--max-questions"), +) -> None: + groups_dir, queries_path, doc_count, query_count = convert_longmemeval_to_corpus( + dataset_path=dataset_path, + output_dir=output_dir, + max_questions=max_questions, + ) + console.print(f"Groups: [cyan]{groups_dir}[/cyan] ({query_count} groups, {doc_count} docs)") + console.print(f"Queries: [cyan]{queries_path}[/cyan] ({query_count})") + + @run_app.command("retrieval") def run_retrieval_command( providers: str = typer.Option("bm-local,mem0-local", "--providers"), dataset_id: str = typer.Option("locomo", "--dataset-id"), - dataset_path: Path = typer.Option(Path("benchmarks/datasets/locomo/locomo10.json"), "--dataset-path"), + dataset_path: Path = typer.Option( + Path("benchmarks/datasets/locomo/locomo10.json"), "--dataset-path" + ), corpus_dir: Path = typer.Option(Path("benchmarks/generated/locomo/docs"), "--corpus-dir"), - queries_path: Path = typer.Option(Path("benchmarks/generated/locomo/queries.json"), "--queries-path"), + queries_path: Path = typer.Option( + Path("benchmarks/generated/locomo/queries.json"), "--queries-path" + ), output_root: Path = typer.Option(Path("benchmarks/runs"), "--output-root"), run_id: str | None = typer.Option(None, "--run-id"), top_k: int = typer.Option(10, "--top-k"), @@ -148,9 +186,13 @@ def run_judge_command( def run_full_command( providers: str = typer.Option("bm-local,mem0-local", "--providers"), dataset_id: str = typer.Option("locomo", "--dataset-id"), - dataset_path: Path = typer.Option(Path("benchmarks/datasets/locomo/locomo10.json"), "--dataset-path"), + dataset_path: Path = typer.Option( + Path("benchmarks/datasets/locomo/locomo10.json"), "--dataset-path" + ), corpus_dir: Path = typer.Option(Path("benchmarks/generated/locomo/docs"), "--corpus-dir"), - queries_path: Path = typer.Option(Path("benchmarks/generated/locomo/queries.json"), "--queries-path"), + queries_path: Path = typer.Option( + Path("benchmarks/generated/locomo/queries.json"), "--queries-path" + ), output_root: Path = typer.Option(Path("benchmarks/runs"), "--output-root"), run_id: str | None = typer.Option(None, "--run-id"), top_k: int = typer.Option(10, "--top-k"), diff --git a/src/basic_memory_benchmarks/converters/longmemeval_to_corpus.py b/src/basic_memory_benchmarks/converters/longmemeval_to_corpus.py new file mode 100644 index 0000000..fe4ca35 --- /dev/null +++ b/src/basic_memory_benchmarks/converters/longmemeval_to_corpus.py @@ -0,0 +1,134 @@ +"""Convert LongMemEval-S into grouped benchmark corpora and a query manifest. + +Each LongMemEval question carries its own haystack of chat sessions, so the +output is one corpus directory per question (``groups//docs``) +plus a single ``queries.json`` whose entries name their group. The runner +ingests and queries each group in isolation, matching the official protocol. + +Anti-leakage: the raw dataset marks evidence sessions with an ``answer_`` +session-id prefix and evidence turns with ``has_answer`` flags. Session ids +are remapped to neutral positional ids and turn flags are dropped, so nothing +ingested by a provider distinguishes evidence from filler. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from basic_memory_benchmarks.datasets.longmemeval import load_longmemeval_dataset + +DATASET_ID = "longmemeval_s" + + +def _render_session_doc( + doc_id: str, + session_date: str, + turns: list[dict], +) -> str: + lines: list[str] = [ + "---", + f"title: {doc_id} ({session_date})", + "type: note", + f"source_doc_id: {doc_id}", + f"dataset_id: {DATASET_ID}", + f"session_date: {session_date}", + "---", + "", + f"# Chat session on {session_date}", + "", + "## Conversation", + ] + for turn in turns: + role = str(turn.get("role", "unknown")).capitalize() + content = str(turn.get("content", "")).strip() + if not content: + continue + # Keep each turn on one line so bullet-level chunking stays intact. + lines.append(f"- **{role}:** {' '.join(content.split())}") + return "\n".join(lines).rstrip() + "\n" + + +def convert_longmemeval_to_corpus( + dataset_path: Path, + output_dir: Path, + max_questions: int | None = None, +) -> tuple[Path, Path, int, int]: + """Convert LongMemEval-S into per-question corpora + query manifest. + + Returns: + groups_dir, queries_path, doc_count, query_count + """ + entries = load_longmemeval_dataset(dataset_path) + if max_questions is not None: + entries = entries[:max_questions] + + groups_dir = output_dir / "groups" + groups_dir.mkdir(parents=True, exist_ok=True) + + all_queries: list[dict] = [] + doc_count = 0 + + for entry in entries: + question_id = str(entry["question_id"]) + sessions = entry["haystack_sessions"] + session_ids = entry["haystack_session_ids"] + session_dates = entry["haystack_dates"] + if not (len(sessions) == len(session_ids) == len(session_dates)): + raise ValueError( + f"Question {question_id}: haystack arrays misaligned " + f"({len(sessions)} sessions, {len(session_ids)} ids, {len(session_dates)} dates)" + ) + + docs_dir = groups_dir / question_id / "docs" + docs_dir.mkdir(parents=True, exist_ok=True) + + # Neutral positional doc ids; the raw ids leak evidence via the + # "answer_" prefix. A few haystacks repeat a session id — keep the + # first occurrence so each session is ingested once. + doc_id_by_session_id: dict[str, str] = {} + for index, (session, session_id, session_date) in enumerate( + zip(sessions, session_ids, session_dates) + ): + if str(session_id) in doc_id_by_session_id: + continue + doc_id = f"{question_id}-s{index:03d}" + doc_id_by_session_id[str(session_id)] = doc_id + doc_path = docs_dir / f"{doc_id}.md" + doc_path.write_text( + _render_session_doc(doc_id, str(session_date), session), + encoding="utf-8", + ) + doc_count += 1 + + ground_truth: list[str] = [] + for answer_session_id in entry["answer_session_ids"]: + mapped = doc_id_by_session_id.get(str(answer_session_id)) + if mapped is None: + raise ValueError( + f"Question {question_id}: answer session {answer_session_id!r} " + "not present in haystack" + ) + ground_truth.append(mapped) + + is_abstention = question_id.endswith("_abs") + all_queries.append( + { + "id": question_id, + "query": str(entry["question"]).strip(), + "category": str(entry["question_type"]), + "group": question_id, + "ground_truth": sorted(ground_truth), + "expected_answer": str(entry["answer"]).strip(), + "metadata": { + "dataset_id": DATASET_ID, + "question_date": str(entry["question_date"]), + "abstention": is_abstention, + }, + } + ) + + queries_path = output_dir / "queries.json" + queries_path.write_text(json.dumps(all_queries, indent=2), encoding="utf-8") + + return groups_dir, queries_path, doc_count, len(all_queries) diff --git a/src/basic_memory_benchmarks/datasets/longmemeval.py b/src/basic_memory_benchmarks/datasets/longmemeval.py index 3f32777..60a9f2f 100644 --- a/src/basic_memory_benchmarks/datasets/longmemeval.py +++ b/src/basic_memory_benchmarks/datasets/longmemeval.py @@ -1,12 +1,69 @@ -"""LongMemEval scaffold for future implementation.""" +"""LongMemEval dataset utilities. + +LongMemEval (Wu et al., ICLR 2025) evaluates long-term conversational memory +across six question types. The -S variant gives each of the 500 questions its +own ~50-session haystack, so the benchmark runs as grouped per-question +corpora rather than one shared corpus. +""" from __future__ import annotations +import json from pathlib import Path +import httpx + +from basic_memory_benchmarks.models import DatasetProvenance +from basic_memory_benchmarks.utils import sha256_file, utc_now_iso + +LONGMEMEVAL_S_URL = ( + "https://huggingface.co/datasets/xiaowu0162/longmemeval/resolve/main/longmemeval_s" +) +LONGMEMEVAL_LICENSE_NOTE = "Dataset is owned by source authors; redistribution may be restricted." + +_REQUIRED_KEYS = { + "question_id", + "question_type", + "question", + "answer", + "question_date", + "haystack_session_ids", + "haystack_dates", + "haystack_sessions", + "answer_session_ids", +} + -def fetch_longmemeval_dataset(_: Path) -> None: - raise NotImplementedError( - "LongMemEval download is intentionally scaffolded for v1. " - "Implement in follow-up once source distribution flow is finalized." +def fetch_longmemeval_dataset(output_path: Path, url: str = LONGMEMEVAL_S_URL) -> DatasetProvenance: + output_path.parent.mkdir(parents=True, exist_ok=True) + # The -S file is ~278MB; stream to disk instead of buffering in memory. + with httpx.stream("GET", url, timeout=600, follow_redirects=True) as response: + response.raise_for_status() + with output_path.open("wb") as file: + for chunk in response.iter_bytes(): + file.write(chunk) + + checksum = sha256_file(output_path) + provenance = DatasetProvenance( + dataset_id="longmemeval_s", + source_url=url, + checksum_sha256=checksum, + license_note=LONGMEMEVAL_LICENSE_NOTE, + fetched_at_utc=utc_now_iso(), + ) + output_path.with_suffix(".provenance.json").write_text( + json.dumps(provenance.model_dump(mode="json"), indent=2), + encoding="utf-8", ) + return provenance + + +def load_longmemeval_dataset(path: Path) -> list[dict]: + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, list): + raise ValueError(f"LongMemEval payload must be a list of questions: {path}") + for index, entry in enumerate(payload): + if not isinstance(entry, dict) or not _REQUIRED_KEYS.issubset(entry): + missing = _REQUIRED_KEYS - set(entry) if isinstance(entry, dict) else _REQUIRED_KEYS + raise ValueError(f"LongMemEval entry {index} missing keys {sorted(missing)}: {path}") + return payload diff --git a/src/basic_memory_benchmarks/models.py b/src/basic_memory_benchmarks/models.py index 7c9c15c..ab78773 100644 --- a/src/basic_memory_benchmarks/models.py +++ b/src/basic_memory_benchmarks/models.py @@ -24,6 +24,8 @@ class QueryCase(BaseModel): query: str category: str category_id: int | None = None + # Grouped datasets (LongMemEval) scope each query to its own corpus group. + group: str | None = None ground_truth: list[str] = Field(default_factory=list) expected_answer: str | None = None metadata: dict[str, Any] = Field(default_factory=dict) @@ -66,6 +68,7 @@ class PerQueryRetrievalResult(BaseModel): latency_ms: float top_hit_doc_id: str | None = None retrieved_context: str = "" + metadata: dict[str, Any] = Field(default_factory=dict) class RetrievalSummary(BaseModel): diff --git a/src/basic_memory_benchmarks/runner.py b/src/basic_memory_benchmarks/runner.py index f90dfa5..c9caa3e 100644 --- a/src/basic_memory_benchmarks/runner.py +++ b/src/basic_memory_benchmarks/runner.py @@ -24,7 +24,12 @@ from basic_memory_benchmarks.reporting.artifacts import write_artifacts from basic_memory_benchmarks.scoring.judge import run_optional_judge from basic_memory_benchmarks.scoring.retrieval import evaluate_query, summarize_provider -from basic_memory_benchmarks.utils import git_sha, resolve_remote_main_sha, runtime_info, utc_now_iso +from basic_memory_benchmarks.utils import ( + git_sha, + resolve_remote_main_sha, + runtime_info, + utc_now_iso, +) def load_queries(path: Path) -> list[QueryCase]: @@ -42,6 +47,108 @@ def _resolve_bm_sha(run_config: RunConfig) -> str | None: return resolve_remote_main_sha("https://github.com/basicmachines-co/basic-memory") +def _execute_provider_flat( + *, + provider: BenchmarkProvider, + provider_name: str, + queries: list[QueryCase], + corpus_path: Path, + run_config: RunConfig, +) -> list[PerQueryRetrievalResult]: + """Classic single-corpus execution: one ingest, then every query.""" + provider_rows: list[PerQueryRetrievalResult] = [] + try: + provider.ingest(corpus_path, run_config) + for query in queries: + started = time.perf_counter() + hits = provider.search(query.query, run_config.top_k, run_config) + latency_ms = (time.perf_counter() - started) * 1000.0 + provider_rows.append( + evaluate_query( + provider=provider_name, + query=query, + hits=hits, + latency_ms=latency_ms, + ) + ) + finally: + try: + provider.cleanup(run_config) + except Exception: + # Cleanup errors should not mask run state. + pass + return provider_rows + + +def _execute_provider_grouped( + *, + provider_factory: Callable[[str], BenchmarkProvider], + provider_name: str, + queries: list[QueryCase], + corpus_path: Path, + run_config: RunConfig, +) -> tuple[list[PerQueryRetrievalResult], BenchmarkProvider, dict[str, str]]: + """Grouped execution (LongMemEval): each group is its own isolated corpus. + + Per group, a fresh provider instance ingests ``//docs`` + under a group-suffixed run id, so provider-side namespaces (BM project + name, mem0 user id) never leak content across groups. A failing group is + recorded and skipped rather than aborting the run; ProviderSkippedError on + the first group means the provider is unavailable and propagates. + """ + groups: dict[str, list[QueryCase]] = {} + for query in queries: + if query.group is None: + raise ValueError( + f"Query {query.id} has no group but the query set is grouped; " + "mixed grouped/ungrouped query files are not supported" + ) + groups.setdefault(query.group, []).append(query) + + provider_rows: list[PerQueryRetrievalResult] = [] + failed_groups: list[str] = [] + last_provider: BenchmarkProvider | None = None + for group_index, (group_id, group_queries) in enumerate(sorted(groups.items())): + group_corpus = corpus_path / group_id / "docs" + if not group_corpus.exists(): + raise FileNotFoundError(f"Missing group corpus: {group_corpus}") + group_config = run_config.model_copy(update={"run_id": f"{run_config.run_id}-{group_id}"}) + provider = provider_factory(provider_name) + try: + provider_rows.extend( + _execute_provider_flat( + provider=provider, + provider_name=provider_name, + queries=group_queries, + corpus_path=group_corpus, + run_config=group_config, + ) + ) + last_provider = provider + except ProviderSkippedError: + # Trigger: provider signals it cannot run at all (missing creds). + # Why: the first group is representative; retrying hundreds of + # groups against an unavailable provider wastes hours. + # Outcome: the provider is recorded as skipped for the whole run. + if group_index == 0: + raise + failed_groups.append(group_id) + except Exception: + failed_groups.append(group_id) + + if last_provider is None: + raise RuntimeError(f"All {len(failed_groups)} groups failed for provider {provider_name}") + + group_metadata: dict[str, str] = { + "grouped_mode": "true", + "group_count": str(len(groups)), + } + if failed_groups: + group_metadata["failed_group_count"] = str(len(failed_groups)) + group_metadata["failed_groups"] = ",".join(sorted(failed_groups)[:50]) + return provider_rows, last_provider, group_metadata + + def run_retrieval( *, run_config: RunConfig, @@ -52,6 +159,7 @@ def run_retrieval( corpus_path = Path(run_config.corpus_dir) output_root = Path(run_config.output_root) run_dir = output_root / run_config.run_id + grouped = any(query.group is not None for query in queries) retrieval_rows: list[PerQueryRetrievalResult] = [] provider_status: list[ProviderStatus] = [] @@ -60,22 +168,24 @@ def run_retrieval( rows_by_provider: dict[str, list[PerQueryRetrievalResult]] = {} for provider_name in run_config.providers: - provider = provider_factory(provider_name) - provider_rows: list[PerQueryRetrievalResult] = [] - try: - provider.ingest(corpus_path, run_config) - for query in queries: - started = time.perf_counter() - hits = provider.search(query.query, run_config.top_k, run_config) - latency_ms = (time.perf_counter() - started) * 1000.0 - provider_rows.append( - evaluate_query( - provider=provider_name, - query=query, - hits=hits, - latency_ms=latency_ms, - ) + group_metadata: dict[str, str] = {} + if grouped: + provider_rows, version_provider, group_metadata = _execute_provider_grouped( + provider_factory=provider_factory, + provider_name=provider_name, + queries=queries, + corpus_path=corpus_path, + run_config=run_config, + ) + else: + version_provider = provider_factory(provider_name) + provider_rows = _execute_provider_flat( + provider=version_provider, + provider_name=provider_name, + queries=queries, + corpus_path=corpus_path, + run_config=run_config, ) summary = summarize_provider(provider_name, provider_rows) @@ -86,7 +196,7 @@ def run_retrieval( ProviderStatus( provider=provider_name, state="ok", - metadata=provider.version_info(), + metadata={**version_provider.version_info(), **group_metadata}, ) ) except ProviderSkippedError as exc: @@ -101,12 +211,6 @@ def run_retrieval( ) if not run_config.allow_provider_skip: raise - finally: - try: - provider.cleanup(run_config) - except Exception: - # Cleanup errors should not mask run state. - pass fairness_warnings = validate_fairness(rows_by_provider) @@ -127,9 +231,7 @@ def run_retrieval( None, ), provider_versions={ - status.provider: status.metadata - for status in provider_status - if status.metadata + status.provider: status.metadata for status in provider_status if status.metadata }, dataset=dataset, runtime=RuntimeInfo(os=os_name, python_version=py_version, started_at_utc=utc_now_iso()), @@ -250,7 +352,9 @@ def run_judge( judge_summary_path = run_dir / "judge-summary.json" judge_summary_path.write_text( - json.dumps({"providers": [item.model_dump(mode="json") for item in judge_summaries]}, indent=2), + json.dumps( + {"providers": [item.model_dump(mode="json") for item in judge_summaries]}, indent=2 + ), encoding="utf-8", ) return run_dir diff --git a/src/basic_memory_benchmarks/scoring/qa.py b/src/basic_memory_benchmarks/scoring/qa.py index 39d9d3b..9db0bcb 100644 --- a/src/basic_memory_benchmarks/scoring/qa.py +++ b/src/basic_memory_benchmarks/scoring/qa.py @@ -104,18 +104,30 @@ def _is_abstention(answer: str) -> bool: return normalized == ABSTAIN_SENTINEL.strip(".").lower() +def _question_display(row: PerQueryRetrievalResult) -> str: + """Render the question with its ask-date when the dataset provides one. + + Temporal-reasoning questions ("how many weeks ago...") are unanswerable + without the reference date, and both the answerer and the judge need the + same framing. + """ + question_date = row.metadata.get("question_date") + if question_date: + return f"{row.query_text} (question asked on {question_date})" + return row.query_text + + def _score_case( row: PerQueryRetrievalResult, provider: str, answerer: LLMRunner, judge: LLMRunner, ) -> QACaseResult: + question = _question_display(row) try: - answer_result = answerer.complete( - build_answer_prompt(row.query_text, row.retrieved_context) - ) + answer_result = answerer.complete(build_answer_prompt(question, row.retrieved_context)) judge_result = judge.complete( - build_judge_prompt(row.query_text, row.expected_answer or "", answer_result.text) + build_judge_prompt(question, row.expected_answer or "", answer_result.text) ) correct, reason = parse_judge_verdict(judge_result.text) return QACaseResult( diff --git a/src/basic_memory_benchmarks/scoring/retrieval.py b/src/basic_memory_benchmarks/scoring/retrieval.py index 4febf59..b2e657e 100644 --- a/src/basic_memory_benchmarks/scoring/retrieval.py +++ b/src/basic_memory_benchmarks/scoring/retrieval.py @@ -108,6 +108,7 @@ def evaluate_query( latency_ms=latency_ms, top_hit_doc_id=top_hit_doc_id, retrieved_context=context, + metadata=query.metadata, ) diff --git a/tests/converters/__init__.py b/tests/converters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/converters/test_longmemeval_converter.py b/tests/converters/test_longmemeval_converter.py new file mode 100644 index 0000000..7c42af2 --- /dev/null +++ b/tests/converters/test_longmemeval_converter.py @@ -0,0 +1,132 @@ +"""Tests for the LongMemEval-S grouped corpus converter.""" + +from __future__ import annotations + +import json + +import pytest + +from basic_memory_benchmarks.converters.longmemeval_to_corpus import ( + convert_longmemeval_to_corpus, +) + + +def _entry( + question_id: str = "q1", + question_type: str = "single-session-user", + answer_session_ids: list[str] | None = None, + session_ids: list[str] | None = None, +) -> dict: + session_ids = session_ids or ["filler_001", "answer_abc", "filler_002"] + return { + "question_id": question_id, + "question_type": question_type, + "question": "What degree did I graduate with?", + "answer": "Business Administration", + "question_date": "2023/05/30 (Tue) 23:40", + "haystack_session_ids": session_ids, + "haystack_dates": [f"2023/05/{i + 1:02d} (Mon) 10:00" for i in range(len(session_ids))], + "haystack_sessions": [ + [ + { + "role": "user", + "content": f"hello in session number {index}", + "has_answer": sid.startswith("answer_"), + }, + {"role": "assistant", "content": "hi there"}, + ] + for index, sid in enumerate(session_ids) + ], + "answer_session_ids": answer_session_ids or ["answer_abc"], + } + + +def _write_dataset(tmp_path, entries): + path = tmp_path / "longmemeval_s.json" + path.write_text(json.dumps(entries), encoding="utf-8") + return path + + +class TestConvertLongMemEval: + def test_grouped_layout_and_queries(self, tmp_path): + dataset = _write_dataset(tmp_path, [_entry("q1"), _entry("q2")]) + out = tmp_path / "out" + + groups_dir, queries_path, doc_count, query_count = convert_longmemeval_to_corpus( + dataset_path=dataset, output_dir=out + ) + + assert doc_count == 6 + assert query_count == 2 + assert (groups_dir / "q1" / "docs").is_dir() + assert (groups_dir / "q2" / "docs").is_dir() + + queries = json.loads(queries_path.read_text()) + by_id = {q["id"]: q for q in queries} + assert by_id["q1"]["group"] == "q1" + assert by_id["q1"]["category"] == "single-session-user" + assert by_id["q1"]["expected_answer"] == "Business Administration" + assert by_id["q1"]["metadata"]["question_date"] == "2023/05/30 (Tue) 23:40" + assert by_id["q1"]["metadata"]["abstention"] is False + + def test_ground_truth_leakage_is_scrubbed(self, tmp_path): + """Evidence markers in the raw data must not survive conversion.""" + dataset = _write_dataset(tmp_path, [_entry("q1")]) + out = tmp_path / "out" + + groups_dir, queries_path, _, _ = convert_longmemeval_to_corpus( + dataset_path=dataset, output_dir=out + ) + + doc_paths = sorted((groups_dir / "q1" / "docs").glob("*.md")) + assert [p.stem for p in doc_paths] == ["q1-s000", "q1-s001", "q1-s002"] + corpus_text = "".join(p.read_text() for p in doc_paths) + assert "answer_" not in corpus_text + assert "has_answer" not in corpus_text + + queries = json.loads(queries_path.read_text()) + # Ground truth maps through the neutral ids: answer_abc was index 1. + assert queries[0]["ground_truth"] == ["q1-s001"] + + def test_duplicate_sessions_kept_once(self, tmp_path): + entry = _entry( + "q1", + session_ids=["filler_001", "answer_abc", "filler_001"], + ) + dataset = _write_dataset(tmp_path, [entry]) + + _, _, doc_count, _ = convert_longmemeval_to_corpus( + dataset_path=dataset, output_dir=tmp_path / "out" + ) + assert doc_count == 2 + + def test_abstention_flag_from_question_id(self, tmp_path): + dataset = _write_dataset(tmp_path, [_entry("q9_abs")]) + _, queries_path, _, _ = convert_longmemeval_to_corpus( + dataset_path=dataset, output_dir=tmp_path / "out" + ) + queries = json.loads(queries_path.read_text()) + assert queries[0]["metadata"]["abstention"] is True + + def test_session_date_in_doc(self, tmp_path): + dataset = _write_dataset(tmp_path, [_entry("q1")]) + groups_dir, _, _, _ = convert_longmemeval_to_corpus( + dataset_path=dataset, output_dir=tmp_path / "out" + ) + doc = (groups_dir / "q1" / "docs" / "q1-s000.md").read_text() + assert "session_date: 2023/05/01 (Mon) 10:00" in doc + assert "# Chat session on 2023/05/01 (Mon) 10:00" in doc + assert "- **User:** hello in session number 0" in doc + + def test_missing_answer_session_raises(self, tmp_path): + entry = _entry("q1", answer_session_ids=["not_in_haystack"]) + dataset = _write_dataset(tmp_path, [entry]) + with pytest.raises(ValueError, match="not present in haystack"): + convert_longmemeval_to_corpus(dataset_path=dataset, output_dir=tmp_path / "out") + + def test_max_questions(self, tmp_path): + dataset = _write_dataset(tmp_path, [_entry("q1"), _entry("q2"), _entry("q3")]) + _, _, _, query_count = convert_longmemeval_to_corpus( + dataset_path=dataset, output_dir=tmp_path / "out", max_questions=2 + ) + assert query_count == 2 diff --git a/tests/test_qa_scoring.py b/tests/test_qa_scoring.py index c9e73ec..2638496 100644 --- a/tests/test_qa_scoring.py +++ b/tests/test_qa_scoring.py @@ -217,3 +217,25 @@ def test_run_qa_stage_writes_artifacts(self, tmp_path, monkeypatch): summary = json.loads((tmp_path / "qa-summary.json").read_text()) assert summary["providers"][0]["provider"] == "bm-local" assert summary["providers"][0]["total_cases"] == 1 + + +class TestQuestionDate: + def test_question_date_reaches_answerer_and_judge(self): + row = _row("q1", "How many weeks ago did I visit the dentist?", "Three weeks ago", "ctx") + row = row.model_copy(update={"metadata": {"question_date": "2023/05/30 (Tue) 23:40"}}) + answerer = FakeRunner({}, default="Three weeks ago") + judge = FakeRunner({}, default='{"correct": true, "reason": "ok"}') + + run_qa([row], provider="bm-local", answerer=answerer, judge=judge, max_workers=1) + + assert "question asked on 2023/05/30 (Tue) 23:40" in answerer.prompts[0] + assert "question asked on 2023/05/30 (Tue) 23:40" in judge.prompts[0] + + def test_no_date_means_plain_question(self): + row = _row("q1", "Where does Joanna live?", "Austin", "ctx") + answerer = FakeRunner({}, default="Austin") + judge = FakeRunner({}, default='{"correct": true, "reason": "ok"}') + + run_qa([row], provider="bm-local", answerer=answerer, judge=judge, max_workers=1) + + assert "question asked on" not in answerer.prompts[0] diff --git a/tests/test_runner_grouped.py b/tests/test_runner_grouped.py new file mode 100644 index 0000000..88c9743 --- /dev/null +++ b/tests/test_runner_grouped.py @@ -0,0 +1,207 @@ +"""Tests for grouped (per-question corpus) retrieval execution.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from basic_memory_benchmarks.exceptions import ProviderSkippedError +from basic_memory_benchmarks.models import DatasetProvenance, RunConfig, SearchHit +from basic_memory_benchmarks.providers.base import BenchmarkProvider +from basic_memory_benchmarks.runner import run_retrieval + + +class RecordingProvider(BenchmarkProvider): + """Records lifecycle calls; class-level log shared across instances.""" + + name = "recording" + calls: list[tuple[str, str, str]] = [] # (event, corpus_or_query, run_id) + fail_groups: set[str] = set() + skip_all = False + instances = 0 + + def __init__(self) -> None: + type(self).instances += 1 + + def ingest(self, corpus_path: Path, run_config: RunConfig) -> None: + if type(self).skip_all: + raise ProviderSkippedError("creds missing") + for group_id in type(self).fail_groups: + if f"-{group_id}" in run_config.run_id: + raise RuntimeError(f"boom in {group_id}") + type(self).calls.append(("ingest", str(corpus_path), run_config.run_id)) + + def search(self, query: str, limit: int, run_config: RunConfig) -> list[SearchHit]: + type(self).calls.append(("search", query, run_config.run_id)) + return [SearchHit(source_doc_id="doc-1", text="ctx", score=1.0)] + + def cleanup(self, run_config: RunConfig) -> None: + type(self).calls.append(("cleanup", "", run_config.run_id)) + + def version_info(self) -> dict[str, str]: + return {"recording": "1.0"} + + +@pytest.fixture(autouse=True) +def _reset_recording_provider(): + RecordingProvider.calls = [] + RecordingProvider.fail_groups = set() + RecordingProvider.skip_all = False + RecordingProvider.instances = 0 + + +def _setup_grouped_corpus(tmp_path: Path, groups: list[str]) -> tuple[Path, Path]: + corpus_root = tmp_path / "groups" + queries = [] + for group_id in groups: + docs = corpus_root / group_id / "docs" + docs.mkdir(parents=True) + (docs / f"{group_id}-s000.md").write_text(f"# {group_id}\n", encoding="utf-8") + queries.append( + { + "id": group_id, + "query": f"question for {group_id}", + "category": "single-session-user", + "group": group_id, + "ground_truth": ["doc-1"], + "expected_answer": "answer", + "metadata": {"question_date": "2023/05/30"}, + } + ) + queries_path = tmp_path / "queries.json" + queries_path.write_text(json.dumps(queries), encoding="utf-8") + return corpus_root, queries_path + + +def _run_config(tmp_path: Path, corpus_root: Path, queries_path: Path) -> RunConfig: + return RunConfig( + run_id="testrun", + dataset_id="longmemeval_s", + dataset_path=str(queries_path), + corpus_dir=str(corpus_root), + queries_path=str(queries_path), + output_root=str(tmp_path / "runs"), + providers=["recording"], + ) + + +def _provenance() -> DatasetProvenance: + return DatasetProvenance( + dataset_id="longmemeval_s", + source_url="test", + checksum_sha256="0" * 64, + license_note="test", + fetched_at_utc="now", + ) + + +class TestGroupedExecution: + def test_each_group_isolated(self, tmp_path): + corpus_root, queries_path = _setup_grouped_corpus(tmp_path, ["qa", "qb"]) + config = _run_config(tmp_path, corpus_root, queries_path) + + run_dir = run_retrieval( + run_config=config, + dataset=_provenance(), + provider_factory=lambda name: RecordingProvider(), + ) + + ingests = [c for c in RecordingProvider.calls if c[0] == "ingest"] + assert len(ingests) == 2 + # Group-suffixed run ids isolate provider namespaces. + assert {run_id for _, _, run_id in ingests} == {"testrun-qa", "testrun-qb"} + # Each ingest points at its own group corpus. + assert {Path(corpus).parent.name for _, corpus, _ in ingests} == {"qa", "qb"} + # Fresh provider instance per group (plus one flat instance is never made). + assert RecordingProvider.instances == 2 + + rows = [ + json.loads(line) + for line in (run_dir / "per-query-retrieval.jsonl").read_text().splitlines() + ] + assert len(rows) == 2 + assert {row["query_id"] for row in rows} == {"qa", "qb"} + # Metadata flows into retrieval rows for the QA stage. + assert all(row["metadata"]["question_date"] == "2023/05/30" for row in rows) + + status = json.loads((run_dir / "provider-status.json").read_text()) + provider_meta = status[0]["metadata"] + assert provider_meta["grouped_mode"] == "true" + assert provider_meta["group_count"] == "2" + + def test_failed_group_recorded_and_run_continues(self, tmp_path): + corpus_root, queries_path = _setup_grouped_corpus(tmp_path, ["qa", "qb", "qc"]) + RecordingProvider.fail_groups = {"qb"} + config = _run_config(tmp_path, corpus_root, queries_path) + + run_dir = run_retrieval( + run_config=config, + dataset=_provenance(), + provider_factory=lambda name: RecordingProvider(), + ) + + rows = [ + json.loads(line) + for line in (run_dir / "per-query-retrieval.jsonl").read_text().splitlines() + ] + assert {row["query_id"] for row in rows} == {"qa", "qc"} + + status = json.loads((run_dir / "provider-status.json").read_text()) + provider_meta = status[0]["metadata"] + assert provider_meta["failed_group_count"] == "1" + assert provider_meta["failed_groups"] == "qb" + + def test_skip_on_first_group_skips_provider(self, tmp_path): + corpus_root, queries_path = _setup_grouped_corpus(tmp_path, ["qa", "qb"]) + RecordingProvider.skip_all = True + config = _run_config(tmp_path, corpus_root, queries_path) + + run_dir = run_retrieval( + run_config=config, + dataset=_provenance(), + provider_factory=lambda name: RecordingProvider(), + ) + + status = json.loads((run_dir / "provider-status.json").read_text()) + assert status[0]["state"] == "skipped" + # Only the first group was attempted. + assert RecordingProvider.instances == 1 + + def test_missing_group_corpus_raises(self, tmp_path): + corpus_root, queries_path = _setup_grouped_corpus(tmp_path, ["qa"]) + queries = json.loads(queries_path.read_text()) + queries.append({**queries[0], "id": "missing", "group": "missing"}) + queries_path.write_text(json.dumps(queries), encoding="utf-8") + config = _run_config(tmp_path, corpus_root, queries_path) + config = config.model_copy(update={"allow_provider_skip": False}) + + with pytest.raises(FileNotFoundError, match="Missing group corpus"): + run_retrieval( + run_config=config, + dataset=_provenance(), + provider_factory=lambda name: RecordingProvider(), + ) + + def test_mixed_grouped_and_ungrouped_rejected(self, tmp_path): + corpus_root, queries_path = _setup_grouped_corpus(tmp_path, ["qa"]) + queries = json.loads(queries_path.read_text()) + queries.append( + { + "id": "flat", + "query": "ungrouped question", + "category": "single_hop", + "ground_truth": [], + } + ) + queries_path.write_text(json.dumps(queries), encoding="utf-8") + config = _run_config(tmp_path, corpus_root, queries_path) + config = config.model_copy(update={"allow_provider_skip": False}) + + with pytest.raises(ValueError, match="mixed grouped/ungrouped"): + run_retrieval( + run_config=config, + dataset=_provenance(), + provider_factory=lambda name: RecordingProvider(), + )