diff --git a/benchmark/tau2/README.md b/benchmark/tau2/README.md index 7ebdb0807e..85d239dbe3 100644 --- a/benchmark/tau2/README.md +++ b/benchmark/tau2/README.md @@ -1,13 +1,16 @@ # TAU-2 Benchmark This directory contains a small OpenViking-style entry point for TAU-2 memory -evaluation. The first version is intentionally narrow: +evaluation. The scope is intentionally narrow: - fresh OpenViking Memory V2 experience-only baseline; - Memory V2 pre-write recall treatment. +- trajectory-view retrieval treatment for the refined trajectory prompt; +- experimental category-aware pre-write rerank on top of trajectory-view + memory. -Trajectory / procedure-view prompts, category rerank, and other harness-only -diagnostics are intentionally left out of this first PR. +The category-aware route is opt-in and experimental; it is meant for PR-C review +and smoke/targeted probes before any productization decision. ## Layout @@ -15,8 +18,11 @@ diagnostics are intentionally left out of this first PR. benchmark/tau2/ ├── config/ │ ├── baseline.yaml +│ ├── category_rerank.yaml +│ ├── no_memory.yaml │ ├── official.yaml -│ └── prewrite.yaml +│ ├── prewrite.yaml +│ └── trajectory.yaml ├── scripts/ │ ├── run_eval.py │ ├── setup_tau2_repo.sh @@ -24,7 +30,9 @@ benchmark/tau2/ └── run_full_eval.sh ``` -Generated artifacts are written to `benchmark/tau2/result//`. +Generated eval artifacts are written to `benchmark/tau2/result//`. +Memory corpus artifacts are cached outside the run id at +`benchmark/tau2/result/memory_corpora/` by default. ## Quick Start @@ -51,6 +59,10 @@ Plan the default benchmark without running TAU-2: python benchmark/tau2/scripts/run_eval.py --config benchmark/tau2/config/baseline.yaml --plan-only ``` +Use `config/no_memory.yaml` for same-runner no-memory baselines; it executes +through the Python wrapper so artifacts and result validation match the memory +cells. + Add `--preflight` or `--strict-preflight` when you want the runner to write a small environment/config check next to the run plan. @@ -77,6 +89,30 @@ benchmark/tau2/run_full_eval.sh \ --repeat-count 1 ``` +Plan a one-cell trajectory-view smoke: + +```bash +benchmark/tau2/run_full_eval.sh \ + --config benchmark/tau2/config/trajectory.yaml \ + --domain retail \ + --strategy-id memory_v2_trajectory_view \ + --num-tasks 1 \ + --train-num-tasks 1 \ + --repeat-count 1 +``` + +Plan a one-cell trajectory category-rerank smoke: + +```bash +benchmark/tau2/run_full_eval.sh \ + --config benchmark/tau2/config/category_rerank.yaml \ + --domain retail \ + --strategy-id memory_v2_trajectory_category_prewrite \ + --num-tasks 1 \ + --train-num-tasks 1 \ + --repeat-count 1 +``` + Run the Memory V2 8-trial matrix (`retail + airline` x 2 strategies x 8 repeats): ```bash @@ -104,20 +140,75 @@ and `OPENAI_API_BASE` for LiteLLM before running upstream TAU-2. Start the OpenViking service before executing memory cells, and verify it with `ov status`. For evidence runs, use a clean OpenViking workspace/config and set `OPENVIKING_URL` explicitly so local custom memory templates do not pollute the -Memory V2 baseline. +Memory V2 baseline. For trajectory-view evidence, start the service from this +branch and inspect generated trajectory files; changing `search_uri` alone does +not prove the new trajectory prompt was used. ## Memory Adapter -`memory_v2_experience_only` and `memory_v2_prewrite` cells run through a small -TAU-2 agent adapter in this directory: +Memory V2 cells run through a small TAU-2 agent adapter in this directory: - train by writing TAU-2 training conversations into OpenViking sessions; -- evaluate by retrieving OpenViking experience memory at the first user turn; +- evaluate by retrieving OpenViking memory at the first user turn; - for pre-write recall, retrieve again before write-like tool calls and - regenerate that step with the matched memories; + regenerate that step with the matched memories. The default benchmark + retrieves 6 pre-write candidates and injects 2, which keeps extra candidates + visible in traces without expanding the prompt budget; +- optionally run an explicit scope-prompt treatment that keeps retrieved + memories advisory and asks the agent to preserve the current task scope before + write-like tool calls; - emit artifact metadata to identify the OpenViking account, agent, corpus, retrieval mode, and simulator policy used by each cell. +The existing `train_memory_mode: experience_only` value selects the Memory V2 +session-commit path. `search_memory_type` selects which generated memory bucket +is retrieved during eval (`experiences` by default, `trajectories` for +`config/trajectory.yaml`). The runner prepares each distinct +`domain + corpus_id` once and reuses it across eval run ids when the cached +`corpus_manifest.json` is present. Different corpora may be prepared in +parallel with `benchmark.corpus_prepare_concurrency`; session commits inside one +corpus remain serial to preserve OpenViking write semantics. + +Eval cells run in parallel with `benchmark.strategy_concurrency` by default and +can be overridden with `--strategy-concurrency`. This only parallelizes read-only +TAU-2 eval cells; corpus writes inside one corpus are still serialized by the +prepare step. + +`config/category_rerank.yaml` keeps the PR-B trajectory memory route and enables +an adapter-local category-rerank probe: pre-write recall, LLM-generated category +annotation sidecars, and the same scope prompt shape used by the trajectory-view +evidence runs. The category treatment retrieves 6 candidates, keeps positive +category matches, injects at most 2 memories, skips injection when no positive +category match exists, and applies the scope/applicability prompt at the system +prompt injection point. Runtime category rerank is sidecar-only: +the runner looks up query and memory annotations from configured +`annotation_files`; if either side is missing, the cell fails instead of doing +live query-to-category mapping. Retrieval traces include +the query category, candidate memory categories, rerank reasons, selected rows, +skipped rows, scope prompt metadata, +and flat `*_category*_prompt` fields kept compatible with Harness diagnostics. +Each run summary also includes `retrieval_trace_summary`, a compact rollup of +decision nodes, category decisions, query/memory category sources, selected +category coverage, positive query-to-memory category-match coverage, +aggregate-vs-concrete memory candidate coverage, and write tool calls. Use it +as the first check that a run is using this branch's self-generated category +signal before opening the JSONL trace. Category runs whose runtime trace has +only aggregate `.overview.md` / `.abstract.md` candidates, no applied +category-rerank event, no query or memory category coverage, no positive +query-to-memory category match, no actual memory injection, no injected +concrete memory, no injected concrete positive category match, or no selected +positive category match are marked `runtime_evidence.status=diagnostic`; +`scoreboard.json` excludes those diagnostic cells from the main reward/DB +aggregates while preserving their metrics, artifacts, and +`diagnostic_reason_counts` for debugging. Corpus +manifests also include +`corpus_probe.aggregate_match_count` and `corpus_probe.concrete_match_count` so +aggregate-only corpora can be spotted before reading the eval trace; category +runs whose corpus probe is empty, or has matches but no concrete matches, are +also marked diagnostic. The corpus probe uses the category `retrieve_limit` +when category rerank is enabled, so the probe width matches the runtime +pre-write search width. + ## User Simulator Policy The runner default is the official TAU-2 user simulator if @@ -131,6 +222,14 @@ confirmation boundary to the TAU-2 user simulator guidelines; metadata such as the upstream PR link is kept in run artifacts, not in the simulator prompt. Reference: [sierra-research/tau2-bench#297](https://github.com/sierra-research/tau2-bench/pull/297). +Optional fixed-first-user fixtures keep the first simulated user turn stable +while preserving live simulator behavior after that turn: + +```bash +export TAU2_RETAIL_FIXED_FIRST_USER_FILE=/path/to/retail_fixture.json +export TAU2_AIRLINE_FIXED_FIRST_USER_FILE=/path/to/airline_fixture.json +``` + Use `config/official.yaml` with a clean TAU-2 checkout when you need an official-user-simulator parity run. If the checkout was already patched, the artifact records that boundary instead of labeling the run pure official. diff --git a/benchmark/tau2/config/baseline.yaml b/benchmark/tau2/config/baseline.yaml index 4c4a5060e7..ef692f43a1 100644 --- a/benchmark/tau2/config/baseline.yaml +++ b/benchmark/tau2/config/baseline.yaml @@ -6,7 +6,9 @@ benchmark: train_split_name: train eval_split_name: test repeat_count: 8 + strategy_concurrency: 8 task_max_concurrency: 10 + corpus_prepare_concurrency: 2 max_steps: 200 seed: 300 agent: llm_agent @@ -17,12 +19,20 @@ paths: tau2_repo: ${TAU2_REPO:-data/external_benchmarks/tau2-bench} tau2_cli: ${TAU2_CLI:-tau2} output_dir: benchmark/tau2/result + # Corpus writes are expensive and should be reused across eval run ids when + # the train split and memory prompt/config did not change. + corpus_cache_dir: benchmark/tau2/result/memory_corpora eval: # The runner default is official if this field is omitted. The OpenViking # memory benchmark config opts into a confirmation-aware TAU-2 user simulator # prompt; run_eval.py applies that small prompt patch idempotently when needed. user_simulator_policy: confirmation_aware + # Optional fixed-first-user fixtures keep the first simulated user turn stable + # while leaving later turns live. Set these env vars to fixture JSON files. + fixed_first_user_fixtures: + retail: ${TAU2_RETAIL_FIXED_FIRST_USER_FILE:-} + airline: ${TAU2_AIRLINE_FIXED_FIRST_USER_FILE:-} model: agent_llm: ${TAU2_AGENT_LLM:-openai/doubao-seed-2-0-pro-260215} @@ -33,7 +43,10 @@ openviking: url: ${OPENVIKING_URL:-http://localhost:1933} account: ${OPENVIKING_ACCOUNT:-default} agent_id: ${OPENVIKING_AGENT_ID:-tau2-openviking-agent} + reuse_corpus_across_runs: true retrieval_top_k: 4 + prewrite_retrieval_top_k: 6 + prewrite_inject_top_k: 2 replay_write_policy: read_only strategies: diff --git a/benchmark/tau2/config/category_rerank.yaml b/benchmark/tau2/config/category_rerank.yaml new file mode 100644 index 0000000000..1f169d301c --- /dev/null +++ b/benchmark/tau2/config/category_rerank.yaml @@ -0,0 +1,136 @@ +extends: trajectory.yaml + +benchmark: + name: tau2_openviking_trajectory_category_rerank + domains: + - retail + - airline + +x-trajectory-category-sidecars: &trajectory_category_sidecars + retail: + - ${TAU2_RETAIL_TRAJECTORY_FIRST_USER_QUERY_CATEGORY_ANNOTATIONS:-benchmark/tau2/result/category_annotations/tau2_pr_b_trajectory_view_retail_first_user_query_workflow_c1_20260516_merged/annotations.jsonl} + - ${TAU2_RETAIL_TRAJECTORY_QUERY_CATEGORY_ANNOTATIONS:-benchmark/tau2/result/category_annotations/tau2_pr_b_trajectory_view_retail_prewrite_query_annotations_workflow_c1_v2_20260515/annotations.jsonl} + - ${TAU2_RETAIL_TRAJECTORY_MEMORY_CATEGORY_ANNOTATIONS:-benchmark/tau2/result/category_annotations/tau2_pr_b_trajectory_view_retail_memory_workflow_c1_seed5_warm12_full_20260515_merged_memory_annotations/annotations.jsonl} + airline: + - ${TAU2_AIRLINE_TRAJECTORY_FIRST_USER_QUERY_CATEGORY_ANNOTATIONS:-benchmark/tau2/result/category_annotations/tau2_pr_b_trajectory_view_airline_first_user_query_workflow_c1_20260516_merged/annotations.jsonl} + - ${TAU2_AIRLINE_TRAJECTORY_MEMORY_CATEGORY_ANNOTATIONS:-benchmark/tau2/result/category_annotations/tau2_pr_b_trajectory_view_airline_memory_workflow_c1_warm6_full_20260516_merged_memory_annotations/annotations.jsonl} + +x-trajectory-scope-prompt: &trajectory_scope_prompt + enabled: true + injection_point: system_prompt + domain_files: + retail: benchmark/tau2/config/scope_prompts/retail_memory_scope.md + airline: benchmark/tau2/config/scope_prompts/airline_memory_scope.md + +x-prewrite-category-base: &prewrite_category_base + enabled: true + annotation_files: *trajectory_category_sidecars + apply_nodes: + - before_write_tool_call + retrieve_limit: 6 + inject_limit: 2 + positive_match_required: true + no_match_policy: skip_injection + missing_query_policy: base_rank + search_score_weight: 0.0 + +x-first-user-category-base: &first_user_category_base + enabled: true + annotation_files: *trajectory_category_sidecars + apply_nodes: + - first_user + retrieve_limit: 6 + inject_limit: 2 + positive_match_required: true + no_match_policy: skip_injection + missing_query_policy: fail_fast + search_score_weight: 0.0 + +strategies: + - id: memory_v2_trajectory_prewrite_scope + label: OpenViking Memory V2 trajectory-view pre-write recall with scope prompt + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + scope_prompt: *trajectory_scope_prompt + + - id: memory_v2_trajectory_category_prewrite_exact + label: OpenViking Memory V2 trajectory-view scope + pre-write exact-pair category rerank + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + category_rerank: + <<: *prewrite_category_base + mismatch_policy: keep_positive_match_drop_mismatch + scope_prompt: *trajectory_scope_prompt + + - id: memory_v2_trajectory_category_prewrite_priority + label: OpenViking Memory V2 trajectory-view scope + pre-write category priority fill + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + category_rerank: + <<: *prewrite_category_base + mismatch_policy: positive_priority_fill + scope_prompt: *trajectory_scope_prompt + + - id: memory_v2_trajectory_category_prewrite_strict_pair + label: OpenViking Memory V2 trajectory-view scope + pre-write strict pair-only category rerank + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + category_rerank: + <<: *prewrite_category_base + mismatch_policy: strict_pair_match_only + scope_prompt: *trajectory_scope_prompt + + - id: memory_v2_trajectory_category_first_user_exact + label: OpenViking Memory V2 trajectory-view scope + first-user exact-pair category rerank + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + category_rerank: + <<: *first_user_category_base + mismatch_policy: keep_positive_match_drop_mismatch + scope_prompt: *trajectory_scope_prompt + + - id: memory_v2_trajectory_category_first_user_priority + label: OpenViking Memory V2 trajectory-view scope + first-user category priority fill + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + category_rerank: + <<: *first_user_category_base + mismatch_policy: positive_priority_fill + scope_prompt: *trajectory_scope_prompt + + - id: memory_v2_trajectory_category_first_user_strict_pair + label: OpenViking Memory V2 trajectory-view scope + first-user strict pair-only category rerank + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + category_rerank: + <<: *first_user_category_base + mismatch_policy: strict_pair_match_only + scope_prompt: *trajectory_scope_prompt diff --git a/benchmark/tau2/config/no_memory.yaml b/benchmark/tau2/config/no_memory.yaml new file mode 100644 index 0000000000..93f35633b4 --- /dev/null +++ b/benchmark/tau2/config/no_memory.yaml @@ -0,0 +1,9 @@ +extends: baseline.yaml + +benchmark: + name: tau2_openviking_no_memory + +strategies: + - id: no_memory + label: TAU-2 no-memory baseline + memory_backend: none diff --git a/benchmark/tau2/config/scope_prompts/airline_memory_scope.md b/benchmark/tau2/config/scope_prompts/airline_memory_scope.md new file mode 100644 index 0000000000..8847796a97 --- /dev/null +++ b/benchmark/tau2/config/scope_prompts/airline_memory_scope.md @@ -0,0 +1,18 @@ + +OpenViking memories are advisory. Use them only when their trigger, preconditions, +and applicability boundary match the current airline task. + +- Do not broaden the user's requested booking, cancellation, rebooking, flight + update, passenger update, baggage update, insurance, or payment scope because a + retrieved memory describes a nearby workflow. +- Keep the current reservation scope explicit. Only use flights, passengers, + baggage entries, cabin changes, insurance choices, payment IDs, dates, and + amounts that are grounded in user input, recent tool observations, reservation + state, profile/payment state, or an explicit search/lookup result. +- Before a write tool call, verify that the selected write action matches the + user's requested operation. Do not mix cancellation, rebooking, upgrade, + downgrade, baggage, or passenger-update flows unless the current task asks for + that combined operation. +- If a memory and the current task disagree, follow the current task state and the + domain policy. + diff --git a/benchmark/tau2/config/scope_prompts/retail_memory_scope.md b/benchmark/tau2/config/scope_prompts/retail_memory_scope.md new file mode 100644 index 0000000000..65a8d61fff --- /dev/null +++ b/benchmark/tau2/config/scope_prompts/retail_memory_scope.md @@ -0,0 +1,16 @@ + +OpenViking memories are advisory. Use them only when their trigger, preconditions, +and applicability boundary match the current retail task. + +- Do not broaden the user's requested replacement, return, exchange, cancellation, + address-change, or payment scope because a retrieved memory describes a nearby + workflow. +- If the user restricts the request to the current order, same order, observed + order items, or a specific product variant, choose write arguments only from the + current tool observations or an explicitly requested catalog lookup. +- Before a write tool call, order IDs, item IDs, new item IDs, payment method IDs, + addresses, amounts, and refund/payment direction must be grounded in user input, + recent tool observations, profile/order state, or an explicit catalog lookup. +- If a memory and the current task disagree, follow the current task state and the + domain policy. + diff --git a/benchmark/tau2/config/scope_prompts/retail_same_order_variant_guard.md b/benchmark/tau2/config/scope_prompts/retail_same_order_variant_guard.md new file mode 100644 index 0000000000..2884ba0abb --- /dev/null +++ b/benchmark/tau2/config/scope_prompts/retail_same_order_variant_guard.md @@ -0,0 +1,9 @@ + +Retail exchange and modification memories are advisory. Do not broaden the user's requested replacement scope. + +- If the user says the replacement should come from the same order, the rest of that order, or an item already in that order, choose only among items visible in the current order details. +- In that case, do not call product-catalog variant lookup to find a cheaper or more available variant unless the user explicitly asks for the cheapest available variant of the product. +- If a procedure memory says to fetch all product variants but the user's wording restricts the candidate set to observed order items, follow the user's narrower scope. +- Before write tools, the new item id must be grounded in the current order observations or in the user's explicit requested catalog variant. +- Do not treat "user provided the order id" in a memory as mandatory. If the user has authenticated but does not know the order id, use current tools to retrieve the user's order list and inspect likely orders instead of stopping or repeatedly asking for the order id. + diff --git a/benchmark/tau2/config/trajectory.yaml b/benchmark/tau2/config/trajectory.yaml new file mode 100644 index 0000000000..aabded08cf --- /dev/null +++ b/benchmark/tau2/config/trajectory.yaml @@ -0,0 +1,33 @@ +extends: baseline.yaml + +benchmark: + name: tau2_openviking_trajectory_view + +strategies: + - id: memory_v2_trajectory_view + label: OpenViking Memory V2 trajectory-view first-user recall + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user + - id: memory_v2_trajectory_prewrite + label: OpenViking Memory V2 trajectory-view pre-write recall + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + - id: memory_v2_trajectory_prewrite_scope + label: OpenViking Memory V2 trajectory-view pre-write recall with scope prompt + memory_backend: openviking + train_required: true + corpus_id: memory_v2_trajectory_view + train_memory_mode: experience_only + search_memory_type: trajectories + retrieval_mode: first_user_prewrite + scope_prompt_files: + retail: benchmark/tau2/config/scope_prompts/retail_memory_scope.md + airline: benchmark/tau2/config/scope_prompts/airline_memory_scope.md diff --git a/benchmark/tau2/scripts/build_category_catalog.py b/benchmark/tau2/scripts/build_category_catalog.py new file mode 100644 index 0000000000..5c2c32c366 --- /dev/null +++ b/benchmark/tau2/scripts/build_category_catalog.py @@ -0,0 +1,167 @@ +"""Build a compact category catalog from category annotation JSONL files.""" + +from __future__ import annotations + +import argparse +import json +import re +from collections import defaultdict +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parents[3] +CATEGORY_PART_PATTERN = re.compile(r"^[a-z0-9_][a-z0-9_-]*$") +CATEGORY_PART_MAX_LENGTH = 64 + + +def _relative(path: Path) -> str: + try: + return str(path.relative_to(ROOT)) + except ValueError: + return str(path) + + +def _safe_category_part(value: Any, *, field: str) -> str: + text = str(value or "").strip() + if not text: + raise SystemExit(f"annotation category missing {field}") + if not CATEGORY_PART_PATTERN.fullmatch(text): + raise SystemExit( + f"annotation {field} must be a reusable slug id using lowercase letters, numbers, " + f"'_' or '-': {text!r}" + ) + if len(text) > CATEGORY_PART_MAX_LENGTH: + raise SystemExit( + f"annotation {field} must be a compact reusable slug id with at most " + f"{CATEGORY_PART_MAX_LENGTH} characters; put detailed boundaries in applicability: {text!r}" + ) + return text + + +def _load_annotations(paths: list[Path]) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for path in paths: + if not path.is_file(): + raise SystemExit(f"annotation file not found: {path}") + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + payload = json.loads(line) + except json.JSONDecodeError as exc: + raise SystemExit(f"invalid JSONL at {path}:{line_number}: {exc}") from exc + if not isinstance(payload, dict): + raise SystemExit(f"expected object at {path}:{line_number}") + rows.append(payload) + if not rows: + raise SystemExit("no annotations found") + return rows + + +def _dedupe(values: list[str], *, limit: int) -> list[str]: + seen: set[str] = set() + result: list[str] = [] + for value in values: + text = str(value or "").strip() + if not text or text in seen: + continue + seen.add(text) + result.append(text) + if len(result) >= limit: + break + return result + + +def _iter_texts(value: Any) -> list[str]: + if isinstance(value, list): + return [str(item).strip() for item in value if str(item).strip()] + return [] + + +def _build_catalog(rows: list[dict[str, Any]]) -> dict[str, Any]: + grouped: dict[str, list[dict[str, Any]]] = defaultdict(list) + for row in rows: + category = row.get("category") + if not isinstance(category, dict): + raise SystemExit("annotation missing category object") + category1 = _safe_category_part(category.get("category1"), field="category1") + category2 = _safe_category_part(category.get("category2"), field="category2") + grouped[f"{category1}:{category2}"].append(row) + + categories: list[dict[str, Any]] = [] + for category_id in sorted(grouped): + members = grouped[category_id] + category = members[0]["category"] + category1, category2 = category_id.split(":", 1) + category3_values = [ + str(item.get("category", {}).get("category3") or "").strip() + for item in members + if isinstance(item.get("category"), dict) + ] + positive_triggers: list[str] = [] + negative_triggers: list[str] = [] + applicability_summaries: list[str] = [] + source_annotation_ids: list[str] = [] + domains: list[str] = [] + subject_types: list[str] = [] + for item in members: + source_annotation_ids.append(str(item.get("annotation_id") or item.get("request_id") or "")) + subject = item.get("subject") if isinstance(item.get("subject"), dict) else {} + domains.append(str(subject.get("domain") or "")) + subject_types.append(str(subject.get("subject_type") or "")) + applicability = item.get("applicability") if isinstance(item.get("applicability"), dict) else {} + positive_triggers.extend(_iter_texts(applicability.get("positive_triggers"))) + negative_triggers.extend(_iter_texts(applicability.get("negative_triggers"))) + summary = str(applicability.get("applicability_summary") or "").strip() + if summary: + applicability_summaries.append(summary) + category3 = _dedupe(category3_values, limit=1) + categories.append( + { + "applicability_summaries": _dedupe(applicability_summaries, limit=2), + "category1": category1, + "category2": category2, + "category3": category3[0] if category3 else None, + "category_id": category_id, + "domains": _dedupe(domains, limit=8), + "negative_triggers": _dedupe(negative_triggers, limit=3), + "positive_triggers": _dedupe(positive_triggers, limit=3), + "source_annotation_count": len(members), + "source_annotation_ids": _dedupe(source_annotation_ids, limit=12), + "subject_types": _dedupe(subject_types, limit=8), + } + ) + + return { + "created_at": datetime.now(timezone.utc).isoformat(), + "category_count": len(categories), + "categories": categories, + "schema_version": "memory_category_catalog.v0", + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--annotations", type=Path, action="append", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--run-id", default=None) + args = parser.parse_args() + + annotation_paths = [path.resolve() for path in args.annotations] + rows = _load_annotations(annotation_paths) + catalog = _build_catalog(rows) + catalog["run_id"] = args.run_id + catalog["source_annotation_files"] = [_relative(path) for path in annotation_paths] + catalog["source_annotation_count"] = len(rows) + + output = args.output.resolve() + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(catalog, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps({"category_count": catalog["category_count"], "output": _relative(output), "status": "passed"}, ensure_ascii=False, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/tau2/scripts/build_category_requests.py b/benchmark/tau2/scripts/build_category_requests.py new file mode 100644 index 0000000000..ec773cda94 --- /dev/null +++ b/benchmark/tau2/scripts/build_category_requests.py @@ -0,0 +1,524 @@ +"""Build category annotation requests from an OpenViking TAU-2 corpus. + +The script is intentionally self-contained for the OpenViking benchmark PR: +it reads a corpus manifest, renders visible memory files into LLM prompts, and +does not import Agent Harness helpers. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parents[3] +DEFAULT_OUTPUT_ROOT = ROOT / "benchmark" / "tau2" / "result" / "category_requests" +DEFAULT_SCHEMA_REF = "memory_category_annotation.v0" + +PROMPT_TEMPLATE = """# Memory / Query Category Extraction v0 + +You are annotating a memory-related subject for retrieval and reranking. + +Return strict JSON matching `memory_category_annotation.v0`. Do not include Markdown. + +## Goal + +Assign reusable categories for both query-side and memory-side ranking: + +- `category1`: coarse workflow / outcome / skill / artifact family. +- `category2`: finer applicability class or state boundary. +- `category3`: optional narrow subcategory when clearly useful. +- `category1`, `category2`, and non-null `category3` must be compact reusable ids: + lowercase `snake_case`, using only letters, numbers, `_`, or `-`; no spaces, + commas, prose sentences, raw order/user/payment/reservation ids, or dates. +- Each category id part must be at most 64 characters. Put detailed state, + eligibility, confirmation, and boundary text in `applicability`, not in the id. + +The same schema must work for benchmark tasks, worker tasks, tool actions, skills, artifacts, trajectory segments, and memory documents. + +## Safety Boundary + +- Use only the visible input below. +- Do not infer from hidden gold answers, evaluator criteria, official expected actions, or private IDs. +- Do not preserve order IDs, emails, phone numbers, addresses, account numbers, ticket IDs, or other instance-specific identifiers as category names. +- Prefer stable reusable categories over one-off labels. +- `category1` should be broad enough to group many related subjects, but not so + broad that it mixes different workflow outcomes. Do not collapse subjects into + a generic container only because they all mention a user, account, order, + reservation, document, task, tool, or service request. +- Avoid putting product type, tool detail, exact state, or multi-step procedure + detail into `category1`; also avoid catch-all labels that only say the subject + is about management, handling, support, or processing. +- `category2` should be the compact action/state facet under that broad family. +- `category2` should still be reusable. Do not merely restate the filename or + visible title; omit product / domain object detail unless it changes + applicability. +- Category names should describe the business action, skill, artifact type, or applicability boundary. +- Category ids must be stable short slug labels, not natural-language summaries. +- Do not use recall timing or evaluation mechanics as categories unless the subject is literally about the recall/evaluation mechanism itself. +- For query-side subjects, do not encode decision-node mechanics such as + `first_user`, `pre_write`, `before_write`, `classification`, `query`, or + `retrieval` into category ids. Categorize the intended workflow or action + itself. +- For query-side subjects, ignore pipeline framing phrases such as "before + executing write-like tool calls" when naming categories. For example, a + subject with tool `cancel_pending_order` should be categorized as a pending + order cancellation workflow, not as a pre-write classification workflow. +- It is fine to mention decision-node mechanics in `evidence` or + `applicability` when they explain where the subject came from, but + `category1`, `category2`, `category3`, and `category_id` must stay about the + reusable business / skill / artifact semantics. +- `category2` must be more specific than `category1`; prefer state, precondition, or applicability boundary. +- If evidence is weak, keep the category broad and lower `confidence`. +- When no catalog is provided, still choose `category1` as if future similar + subjects will reuse it; do not create a narrow first-level category just to + fit the current item. + +## Subject Metadata + +```json +{{SUBJECT_METADATA_JSON}} +``` + +## Visible Subject Text + +```text +{{SUBJECT_TEXT}} +``` + +## Optional Category Hints + +```json +{{CATEGORY_HINTS_JSON}} +``` + +If `CATEGORY_HINTS_JSON` contains `known_category_catalog`, treat it as an existing reusable taxonomy: + +- First try to reuse an existing `category1/category2` pair from the catalog. +- Reuse a category when the visible subject clearly matches its description, examples, positive triggers, and negative boundaries. +- Reuse an existing `category1` only when the primary workflow outcome matches; + shared nouns such as user, account, order, reservation, document, task, or + tool are not sufficient evidence. +- If no existing category fits, create a new stable category and explain why in `category.new_category_reason`. +- If a close category exists but is too broad or missing a boundary, reuse `category1`, propose a more precise `category2`, and mention the relationship in `category.catalog_relation`. +- If you set `catalog_match.matched_category_id`, it must be the exact canonical `category1:category2` id from the catalog. If no exact catalog pair is reused, set it to `null`. +- Do not create or reuse catalog ids whose only difference is the runtime + decision node, retrieval stage, annotation task, or evaluation pipeline. Those + details belong in metadata/evidence, not in taxonomy. + +## Output JSON Shape + +```json +{ + "schema_version": "memory_category_annotation.v0", + "annotation_id": "...", + "producer": "llm_prompt", + "subject": { + "subject_type": "query|memory|trajectory_segment|tool_action|artifact|skill|worker_task", + "subject_id": "...", + "subject_ref": "...", + "benchmark_family": "...", + "domain": "..." + }, + "category": { + "category1": "...", + "category2": "...", + "category3": null, + "category_source": "llm_prompt", + "catalog_match": { + "matched": true, + "matched_category_id": "optional existing category id", + "decision": "reuse|refine|new", + "reason": "short reason" + }, + "confidence": 0.0, + "reason": "short reason using visible evidence only" + }, + "applicability": { + "applicability_summary": "...", + "positive_triggers": ["..."], + "negative_triggers": ["..."], + "preconditions": ["..."], + "anti_patterns": ["..."] + }, + "evidence": { + "source_fields": ["visible_subject_text"], + "evidence_spans": [ + {"text": "short visible evidence span", "role": "category_or_boundary_evidence"} + ] + }, + "ranking_features": { + "category1": "...", + "category2": "...", + "category_source": "llm_prompt", + "confidence": 0.0 + }, + "safety": { + "uses_hidden_gold": false, + "pii_or_instance_policy": "avoid_instance_ids", + "runtime_safe_inputs_only": true + } +} +``` + +`safety.pii_or_instance_policy` must be exactly one of +`avoid_instance_ids`, `redacted`, or `not_applicable`. +""" + + +@dataclass(frozen=True) +class Subject: + subject_type: str + subject_id: str + subject_ref: str + benchmark_family: str + domain: str + subject_text: str + + @property + def stable_hash(self) -> str: + payload = json.dumps( + { + "benchmark_family": self.benchmark_family, + "domain": self.domain, + "subject_id": self.subject_id, + "subject_ref": self.subject_ref, + "subject_text": self.subject_text, + "subject_type": self.subject_type, + }, + ensure_ascii=False, + sort_keys=True, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16] + + def metadata(self) -> dict[str, Any]: + return { + "benchmark_family": self.benchmark_family, + "domain": self.domain, + "subject_id": self.subject_id, + "subject_ref": self.subject_ref, + "subject_type": self.subject_type, + } + + +def _safe_key(value: str) -> str: + return "".join(ch if ch.isalnum() or ch in {"-", "_", "."} else "_" for ch in value)[:180] or "category_requests" + + +def _slug_identity(value: str) -> str: + safe = "".join(ch if ch.isalnum() else "_" for ch in value.strip().lower()) + return "_".join(part for part in safe.split("_") if part) + + +def _relative(path: Path) -> str: + try: + return str(path.relative_to(ROOT)) + except ValueError: + return str(path) + + +def _load_json(path: Path) -> dict[str, Any]: + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise SystemExit(f"expected JSON object: {path}") + return payload + + +def _compact_text(text: str, *, limit: int) -> str: + compact = "\n".join(line.rstrip() for line in text.strip().splitlines() if line.strip()) + if len(compact) <= limit: + return compact + return compact[: limit - 40].rstrip() + "\n...[truncated for category extraction]" + + +def _compact_catalog(path: Path | None, *, max_categories: int) -> dict[str, Any] | None: + if path is None: + return None + payload = _load_json(path) + categories = payload.get("categories") + if not isinstance(categories, list): + raise SystemExit(f"category catalog missing categories: {path}") + compact: list[dict[str, Any]] = [] + for category in categories[:max_categories]: + if not isinstance(category, dict): + continue + compact.append( + { + "category_id": category.get("category_id"), + "category1": category.get("category1"), + "category2": category.get("category2"), + "category3": category.get("category3"), + "positive_triggers": (category.get("positive_triggers") or [])[:1], + "negative_triggers": (category.get("negative_triggers") or [])[:1], + } + ) + return { + "catalog_ref": _relative(path.resolve()), + "category_count": len(compact), + "categories": compact, + "reuse_policy": "prefer_existing_category1_category2_pairs_before_creating_new_categories", + "schema_version": payload.get("schema_version"), + } + + +def _render_prompt(*, subject: Subject, category_hints: dict[str, Any]) -> str: + return ( + PROMPT_TEMPLATE.replace("{{SUBJECT_METADATA_JSON}}", json.dumps(subject.metadata(), ensure_ascii=False, indent=2, sort_keys=True)) + .replace("{{SUBJECT_TEXT}}", subject.subject_text.strip()) + .replace("{{CATEGORY_HINTS_JSON}}", json.dumps(category_hints, ensure_ascii=False, indent=2, sort_keys=True)) + ) + + +def _build_request(*, subject: Subject, category_hints: dict[str, Any]) -> dict[str, Any]: + request_id = f"{subject.subject_type}:{subject.subject_id}:{subject.stable_hash}" + return { + "category_hints": category_hints, + "created_at": datetime.now(timezone.utc).isoformat(), + "execution_mode": "dry_run_render_only", + "expected_output_schema_version": DEFAULT_SCHEMA_REF, + "prompt": _render_prompt(subject=subject, category_hints=category_hints), + "prompt_ref": "benchmark/tau2/scripts/build_category_requests.py::PROMPT_TEMPLATE", + "request_id": request_id, + "safety": { + "pii_or_instance_policy": "avoid_instance_ids", + "runtime_safe_inputs_only": True, + "uses_hidden_gold": False, + }, + "schema_ref": DEFAULT_SCHEMA_REF, + "schema_version": "memory_category_extraction_request.v0", + "subject": subject.metadata(), + } + + +def _memory_root(*, workspace: Path, manifest: dict[str, Any], memory_type: str) -> Path: + openviking = manifest.get("openviking") + if not isinstance(openviking, dict): + raise SystemExit("corpus manifest missing openviking block") + account = str(openviking.get("account") or "").strip() + agent_id = str(openviking.get("agent_id") or "").strip() + if not account or not agent_id: + raise SystemExit("corpus manifest missing openviking.account or openviking.agent_id") + return workspace / "viking" / account / "agent" / agent_id / "memories" / memory_type + + +def _iter_memory_files(root: Path) -> list[Path]: + if not root.is_dir(): + raise SystemExit(f"memory root not found: {root}") + files = [] + for path in sorted(root.glob("*.md")): + if path.name.startswith(".") or path.name.endswith(".abstract.md") or path.name.endswith(".overview.md"): + continue + files.append(path) + return files + + +def _subject_id(logical_uri: str) -> str: + safe = _safe_key(logical_uri) + return f"openviking_memory_{safe}" + + +def _query_subject_id(query_signature: str) -> str: + return f"tau2_query_signature_{_slug_identity(query_signature)}" + + +def _load_jsonl(path: Path) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise SystemExit(f"invalid JSONL at {path}:{line_number}: {exc}") from exc + if not isinstance(row, dict): + raise SystemExit(f"expected JSON object at {path}:{line_number}") + rows.append(row) + if not rows: + raise SystemExit(f"no query subjects found: {path}") + return rows + + +def _query_subject_from_row(row: dict[str, Any], *, default_domain: str | None, text_limit: int) -> Subject: + domain = str(row.get("domain") or default_domain or "").strip() + decision_node = str(row.get("decision_node") or "").strip() + query_signature = str(row.get("query_signature") or row.get("signature") or "").strip() + query_text = str(row.get("query_text") or row.get("query") or row.get("visible_query") or "").strip() + if not domain: + raise SystemExit(f"query subject missing domain: {row}") + if not decision_node: + raise SystemExit(f"query subject missing decision_node: {row}") + if not query_signature: + raise SystemExit(f"query subject missing query_signature: {row}") + if not query_text: + raise SystemExit(f"query subject missing query text: {row}") + tools = row.get("tools") or [] + subject_text = "\n".join( + line + for line in [ + f"Decision node: {decision_node}", + f"Query signature: {query_signature}", + f"Tools: {', '.join(str(item) for item in tools) if isinstance(tools, list) else tools}", + "", + "Visible query:", + query_text, + ] + if line != "" + ) + return Subject( + benchmark_family=str(row.get("benchmark_family") or "tau2"), + domain=domain, + subject_id=str(row.get("subject_id") or _query_subject_id(query_signature)), + subject_ref=str(row.get("subject_ref") or query_signature), + subject_text=_compact_text(subject_text, limit=text_limit), + subject_type="query", + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, default=None) + parser.add_argument("--workspace", type=Path, default=None) + parser.add_argument("--domain", default=None) + parser.add_argument("--memory-type", default="trajectories") + parser.add_argument("--query-subjects-jsonl", type=Path, default=None) + parser.add_argument("--category-catalog", type=Path, default=None) + parser.add_argument("--max-catalog-categories", type=int, default=80) + parser.add_argument("--subject-text-limit", type=int, default=3500) + parser.add_argument("--offset", type=int, default=0, help="Skip the first N sorted memory files.") + parser.add_argument("--limit", type=int, default=0, help="0 means all remaining memory files.") + parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT) + parser.add_argument("--run-id", required=True) + parser.add_argument("--hint", action="append", default=[]) + args = parser.parse_args() + if args.offset < 0: + raise SystemExit("--offset must be >= 0") + if args.limit < 0: + raise SystemExit("--limit must be >= 0") + + catalog = _compact_catalog(args.category_catalog.resolve() if args.category_catalog else None, max_categories=args.max_catalog_categories) + extra_hints: dict[str, str] = {} + for pair in args.hint: + if "=" not in pair: + raise SystemExit(f"invalid --hint {pair!r}; expected key=value") + key, value = pair.split("=", 1) + extra_hints[key.strip()] = value.strip() + + requests = [] + subjects = [] + manifest: dict[str, Any] | None = None + root: Path | None = None + all_files: list[Path] = [] + files: list[Path] = [] + if args.query_subjects_jsonl: + query_rows = _load_jsonl(args.query_subjects_jsonl.resolve()) + for row in query_rows: + subject = _query_subject_from_row(row, default_domain=args.domain, text_limit=args.subject_text_limit) + hints: dict[str, Any] = { + "benchmark": subject.benchmark_family, + "decision_node": row.get("decision_node"), + "query_signature": row.get("query_signature") or row.get("signature"), + "query_source": row.get("source") or row.get("query_source") or "query_subjects_jsonl", + "tools": row.get("tools") or [], + **extra_hints, + } + if catalog: + hints["known_category_catalog"] = catalog + requests.append(_build_request(subject=subject, category_hints=hints)) + subjects.append( + { + "query_signature": row.get("query_signature") or row.get("signature"), + "source": row.get("source") or row.get("query_source") or "query_subjects_jsonl", + "subject": subject.metadata(), + "text_sha256": hashlib.sha256(subject.subject_text.encode("utf-8")).hexdigest()[:16], + } + ) + else: + if args.manifest is None or args.workspace is None or not args.domain: + raise SystemExit("--manifest, --workspace, and --domain are required for memory request generation") + manifest = _load_json(args.manifest.resolve()) + openviking = manifest.get("openviking") if isinstance(manifest.get("openviking"), dict) else {} + agent_id = str(openviking.get("agent_id") or "").strip() + if not agent_id: + raise SystemExit("manifest openviking.agent_id is required") + root = _memory_root(workspace=args.workspace.resolve(), manifest=manifest, memory_type=args.memory_type) + all_files = _iter_memory_files(root) + files = all_files[args.offset :] + if args.limit: + files = files[: args.limit] + for path in files: + logical_uri = f"viking://agent/{agent_id}/memories/{args.memory_type}/{path.name}" + subject = Subject( + benchmark_family="tau2", + domain=args.domain, + subject_id=_subject_id(logical_uri), + subject_ref=logical_uri, + subject_text=_compact_text(path.read_text(encoding="utf-8"), limit=args.subject_text_limit), + subject_type="memory", + ) + hints: dict[str, Any] = { + "benchmark": "tau2", + "file_name": path.name, + "logical_uri": logical_uri, + "ov_uri_bucket": args.memory_type.rstrip("s"), + **extra_hints, + } + if catalog: + hints["known_category_catalog"] = catalog + requests.append(_build_request(subject=subject, category_hints=hints)) + subjects.append( + { + "bucket": args.memory_type, + "memory_path": _relative(path), + "subject": subject.metadata(), + "text_sha256": hashlib.sha256(subject.subject_text.encode("utf-8")).hexdigest()[:16], + } + ) + + if not requests: + raise SystemExit("no category requests produced") + + output_root = args.output_root if args.output_root.is_absolute() else ROOT / args.output_root + run_root = output_root / _safe_key(args.run_id) + run_root.mkdir(parents=True, exist_ok=True) + requests_path = run_root / "category_extraction_requests.jsonl" + subjects_path = run_root / "subjects.jsonl" + summary_path = run_root / "run_summary.json" + requests_path.write_text("".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in requests), encoding="utf-8") + subjects_path.write_text("".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in subjects), encoding="utf-8") + summary = { + "category_catalog_path": _relative(args.category_catalog.resolve()) if args.category_catalog else None, + "claim_boundary": "prompt_request_plan_from_openviking_manifest_visible_memory_files_no_hidden_gold", + "concrete_memory_file_count": len(files), + "domain": args.domain, + "known_category_count": catalog.get("category_count") if catalog else 0, + "limit": args.limit, + "manifest_committed_session_count": manifest.get("committed_session_count") if manifest else None, + "manifest_path": _relative(args.manifest.resolve()) if args.manifest else None, + "memory_root": _relative(root) if root else None, + "memory_type": args.memory_type, + "offset": args.offset, + "query_subjects_jsonl": _relative(args.query_subjects_jsonl.resolve()) if args.query_subjects_jsonl else None, + "request_count": len(requests), + "requests_path": _relative(requests_path), + "run_id": args.run_id, + "schema_version": "openviking_tau2_category_request_plan.v0", + "status": "passed", + "subjects_path": _relative(subjects_path), + "subject_type": "query" if args.query_subjects_jsonl else "memory", + "total_concrete_memory_file_count": len(all_files), + "workspace": _relative(args.workspace.resolve()) if args.workspace else None, + } + if manifest and manifest.get("committed_session_count") != len(files): + summary["inventory_warning"] = "committed_session_count differs from concrete memory file count" + summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/tau2/scripts/category_rerank.py b/benchmark/tau2/scripts/category_rerank.py new file mode 100644 index 0000000000..04034d8408 --- /dev/null +++ b/benchmark/tau2/scripts/category_rerank.py @@ -0,0 +1,933 @@ +from __future__ import annotations + +import json +import re +import hashlib +from itertools import combinations +from pathlib import Path +from typing import Any + + +CATEGORY_ID_PATTERN = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9_-]*:[A-Za-z0-9_][A-Za-z0-9_-]*$") + + +def _as_list(value: Any) -> list[str]: + if isinstance(value, str): + return [value] if value.strip() else [] + if isinstance(value, list): + return [str(item).strip() for item in value if str(item).strip()] + return [] + + +def _as_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +def _as_int(value: Any, default: int) -> int: + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +def _as_int_map(value: Any) -> dict[str, int]: + if not isinstance(value, dict): + return {} + parsed: dict[str, int] = {} + for key, raw_value in value.items(): + int_value = _as_int(raw_value, 0) + if int_value: + parsed[str(key)] = int_value + return parsed + + +def _as_str_map(value: Any) -> dict[str, str]: + if not isinstance(value, dict): + return {} + parsed: dict[str, str] = {} + for key, raw_value in value.items(): + text = str(raw_value or "").strip() + if text: + parsed[str(key)] = text + return parsed + + +def _public_row(row: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in row.items() if not key.startswith("_")} + + +def _score_value(value: Any) -> float: + try: + return float(value) + except (TypeError, ValueError): + return 0.0 + + +def _slug_identity(value: Any) -> str: + text = str(value or "").strip().lower() + return re.sub(r"[^a-z0-9]+", "_", text).strip("_") + + +def _category_id(category: dict[str, Any]) -> str | None: + catalog_match = category.get("catalog_match") if isinstance(category.get("catalog_match"), dict) else {} + matched_id = str(catalog_match.get("matched_category_id") or "").strip() + if matched_id: + return matched_id + category1 = str(category.get("category1") or "").strip() + category2 = str(category.get("category2") or "").strip() + if category1 and category2: + return f"{category1}:{category2}" + return category1 or None + + +def _malformed_matched_category_id(row: dict[str, Any]) -> str | None: + category = row.get("category") if isinstance(row.get("category"), dict) else {} + catalog_match = category.get("catalog_match") if isinstance(category.get("catalog_match"), dict) else {} + matched_id = str(catalog_match.get("matched_category_id") or "").strip() + if not matched_id: + return None + if not CATEGORY_ID_PATTERN.fullmatch(matched_id): + return matched_id + return None + + +def _compact_annotation(row: dict[str, Any]) -> dict[str, Any]: + category = row.get("category") if isinstance(row.get("category"), dict) else {} + ranking = row.get("ranking_features") if isinstance(row.get("ranking_features"), dict) else {} + if not category and not ranking: + return {} + category_id = _category_id(category) or _category_id(ranking) + payload = { + "matched": True, + "category_id": category_id, + "category1": category.get("category1") or ranking.get("category1"), + "category2": category.get("category2") or ranking.get("category2"), + "category3": category.get("category3") or ranking.get("category3"), + "category_source": category.get("category_source") or ranking.get("category_source") or "annotation_sidecar", + "confidence": category.get("confidence") or ranking.get("confidence"), + "annotation_id": row.get("annotation_id") or row.get("request_id"), + } + return {key: value for key, value in payload.items() if value not in (None, "", [])} + + +def _identity_variants(value: Any) -> set[str]: + text = str(value or "").strip() + if not text: + return set() + variants = {text, _slug_identity(text)} + for marker in ("/memories/", "_memories_"): + if marker not in text: + continue + suffix = text.split(marker, 1)[1].strip("/_") + if not suffix: + continue + variants.update( + { + suffix, + _slug_identity(suffix), + f"memories/{suffix}", + _slug_identity(f"memories/{suffix}"), + Path(suffix).name, + _slug_identity(Path(suffix).name), + } + ) + return {variant for variant in variants if variant} + + +def _annotation_paths(raw: Any, *, repo_root: Path) -> list[Path]: + values: list[Any] = [] + if isinstance(raw, dict): + for item in raw.values(): + if isinstance(item, list): + values.extend(item) + else: + values.append(item) + elif isinstance(raw, list): + values.extend(raw) + elif raw: + values.append(raw) + paths: list[Path] = [] + for value in values: + if not value: + continue + path = Path(str(value)).expanduser() + paths.append(path if path.is_absolute() else repo_root / path) + return paths + + +def _load_annotation_index(paths: list[Path]) -> dict[str, Any]: + index: dict[str, Any] = { + "by_key": {}, + "loaded_files": [], + "load_errors": [], + "query_count": 0, + "memory_count": 0, + } + for path in paths: + if not path.is_file(): + index["load_errors"].append({"path": str(path), "error": "file_not_found"}) + continue + loaded = 0 + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + index["load_errors"].append({"path": str(path), "line": line_number, "error": str(exc)}) + continue + if not isinstance(row, dict): + continue + malformed_category_id = _malformed_matched_category_id(row) + if malformed_category_id: + index["load_errors"].append( + { + "path": str(path), + "line": line_number, + "error": ( + "invalid matched_category_id " + f"{malformed_category_id!r}; expected ':'" + ), + } + ) + continue + subject = row.get("subject") if isinstance(row.get("subject"), dict) else {} + subject_type = str(subject.get("subject_type") or "").strip() + if subject_type == "query": + index["query_count"] += 1 + elif subject_type == "memory": + index["memory_count"] += 1 + for key in ( + row.get("annotation_id"), + row.get("request_id"), + subject.get("subject_id"), + subject.get("subject_ref"), + ): + for variant in _identity_variants(key): + index["by_key"][variant] = row + loaded += 1 + index["loaded_files"].append({"path": str(path), "rows": loaded}) + return index + + +def _lookup_annotation(index: dict[str, Any], keys: list[Any]) -> dict[str, Any] | None: + by_key = index.get("by_key") if isinstance(index.get("by_key"), dict) else {} + for key in keys: + for variant in _identity_variants(key): + row = by_key.get(variant) + if isinstance(row, dict): + return row + return None + + +def _query_lookup_keys(signature: str, *, include_hash: str | None = None) -> list[str]: + signature_slug = _slug_identity(signature) + keys = [ + signature, + signature_slug, + f"tau2_query_signature_{signature_slug}", + f"query:tau2_query_signature_{signature_slug}", + ] + if include_hash: + keys.extend([include_hash, f"query:{include_hash}"]) + return keys + + +def _write_tool_names_from_query(query: str) -> list[str]: + first_line = query.splitlines()[0] if query else "" + prefix = "Before executing write-like tool call(s):" + tool_blob = first_line.split(prefix, 1)[1] if prefix in first_line else first_line + return sorted(set(re.findall(r"\b([A-Za-z_][A-Za-z0-9_]*)\s*\(", tool_blob))) + + +def _write_tool_signature(domain: str, tools: list[str]) -> str: + return "|".join(["tau2", domain, "pre_write_action", "tools=" + ",".join(sorted(tools))]) + + +def _query_signature(domain: str, decision_node: str, query: str) -> str: + if decision_node == "before_write_tool_call": + return _write_tool_signature(domain, _write_tool_names_from_query(query)) + query_hash = hashlib.sha256(query.encode("utf-8")).hexdigest()[:16] + return "|".join(["tau2", domain, decision_node, f"query_sha256={query_hash}"]) + + +def _query_signature_candidates(domain: str, decision_node: str, query: str) -> list[str]: + primary = _query_signature(domain, decision_node, query) + if decision_node != "before_write_tool_call": + return [primary] + tools = _write_tool_names_from_query(query) + if len(tools) <= 1: + return [primary] + signatures = [primary] + for size in range(len(tools) - 1, 0, -1): + for subset in combinations(tools, size): + signature = _write_tool_signature(domain, list(subset)) + if signature not in signatures: + signatures.append(signature) + return signatures + + +class CategoryReranker: + def __init__( + self, + *, + enabled: bool, + apply_nodes: set[str], + annotation_index: dict[str, Any], + load_report: dict[str, Any], + retrieve_limit: int | None, + inject_limit: int | None, + retrieve_limits: dict[str, int] | None, + inject_limits: dict[str, int] | None, + mismatch_policy: str, + mismatch_policies: dict[str, str] | None, + positive_match_required: bool, + no_match_policy: str, + missing_query_policy: str, + search_score_weight: float, + ) -> None: + self.enabled = enabled + self.apply_nodes = apply_nodes + self.annotation_index = annotation_index + self.load_report = load_report + self.retrieve_limit = retrieve_limit + self.inject_limit = inject_limit + self.retrieve_limits = retrieve_limits or {} + self.inject_limits = inject_limits or {} + self.mismatch_policy = mismatch_policy + self.mismatch_policies = mismatch_policies or {} + self.positive_match_required = positive_match_required + self.no_match_policy = no_match_policy + self.missing_query_policy = missing_query_policy + self.search_score_weight = search_score_weight + + @classmethod + def from_payload(cls, payload: dict[str, Any] | None, *, repo_root: Path) -> "CategoryReranker": + payload = payload if isinstance(payload, dict) else {} + enabled = _as_bool(payload.get("enabled"), default=False) + apply_nodes = set(_as_list(payload.get("apply_nodes")) or ["before_write_tool_call"]) + if enabled: + annotation_files = _annotation_paths( + payload.get("annotation_files") or payload.get("category_annotation_files"), + repo_root=repo_root, + ) + if not annotation_files: + raise ValueError( + "category rerank requires LLM-generated annotation_files; " + "live query-to-category mapping is not supported" + ) + annotation_index = _load_annotation_index(annotation_files) + if annotation_index.get("load_errors"): + raise ValueError(f"category rerank sidecar failed to load: {annotation_index['load_errors']}") + if not annotation_index.get("query_count") or not annotation_index.get("memory_count"): + raise ValueError( + "category rerank sidecar must contain both query and memory annotations: " + f"query_count={annotation_index.get('query_count')} " + f"memory_count={annotation_index.get('memory_count')}" + ) + load_report = { + "loaded": True, + "source": "annotation_files", + "loaded_files": annotation_index.get("loaded_files") or [], + "query_count": annotation_index.get("query_count") or 0, + "memory_count": annotation_index.get("memory_count") or 0, + "errors": [], + } + else: + annotation_index = {"by_key": {}, "loaded_files": [], "load_errors": [], "query_count": 0, "memory_count": 0} + load_report = { + "path": None, + "loaded": False, + "domain_count": 0, + "category_count": 0, + "errors": [], + } + mismatch_policy = str(payload.get("mismatch_policy") or "").strip() + mismatch_policies = _as_str_map(payload.get("mismatch_policies")) + positive_match_policies = { + "keep_positive_match_drop_mismatch", + "positive_match_only", + "positive_priority_fill", + "strict_pair_match_only", + } + positive_match_required = _as_bool( + payload.get("positive_match_required"), + default=(mismatch_policy in positive_match_policies) + or any(policy in positive_match_policies for policy in mismatch_policies.values()), + ) + if not mismatch_policy and positive_match_required: + mismatch_policy = "keep_positive_match_drop_mismatch" + missing_query_policy = str(payload.get("missing_query_policy") or "fail_fast") + if missing_query_policy not in {"fail_fast", "base_rank", "skip_injection"}: + raise ValueError( + "category rerank missing_query_policy must be one of " + "'fail_fast', 'base_rank', or 'skip_injection'" + ) + return cls( + enabled=enabled, + apply_nodes=apply_nodes, + annotation_index=annotation_index, + load_report=load_report, + retrieve_limit=_as_int(payload.get("retrieve_limit"), 0) or None, + inject_limit=_as_int(payload.get("inject_limit"), 0) or None, + retrieve_limits=_as_int_map(payload.get("retrieve_limits")), + inject_limits=_as_int_map(payload.get("inject_limits")), + mismatch_policy=mismatch_policy or "none", + mismatch_policies=mismatch_policies, + positive_match_required=positive_match_required, + no_match_policy=str(payload.get("no_match_policy") or "skip_injection"), + missing_query_policy=missing_query_policy, + search_score_weight=float(payload.get("search_score_weight") or 0.0), + ) + + def _retrieve_limit(self, decision_node: str) -> int | None: + return self.retrieve_limits.get(decision_node) or self.retrieve_limit + + def _inject_limit(self, decision_node: str, base_limit: int) -> int: + return self.inject_limits.get(decision_node) or self.inject_limit or base_limit + + def _mismatch_policy(self, decision_node: str) -> str: + return self.mismatch_policies.get(decision_node) or self.mismatch_policy + + def search_limit(self, base_limit: int, *, decision_node: str) -> int: + if self.enabled and decision_node in self.apply_nodes: + retrieve_limit = self._retrieve_limit(decision_node) + if retrieve_limit: + return max(base_limit, retrieve_limit) + return base_limit + + def summary(self) -> dict[str, Any]: + return { + "enabled": self.enabled, + "apply_nodes": sorted(self.apply_nodes), + "retrieve_limit": self.retrieve_limit, + "inject_limit": self.inject_limit, + "retrieve_limits": dict(self.retrieve_limits), + "inject_limits": dict(self.inject_limits), + "mismatch_policy": self.mismatch_policy, + "mismatch_policies": dict(self.mismatch_policies), + "positive_match_required": self.positive_match_required, + "no_match_policy": self.no_match_policy, + "missing_query_policy": self.missing_query_policy, + "search_score_weight": self.search_score_weight, + "sidecar": self.load_report, + } + + def select( + self, + *, + domain: str, + query: str, + rows: list[dict[str, Any]], + decision_node: str, + base_limit: int, + ) -> tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]: + mismatch_policy = self._mismatch_policy(decision_node) + if not self.enabled or decision_node not in self.apply_nodes: + selected = rows[:base_limit] + trace_rows = _mark_selected(rows, selected, decision="base_rank") + diagnostics = { + "enabled": self.enabled, + "applied": False, + "decision_node": decision_node, + "decision": "node_not_enabled" if self.enabled else "disabled", + "apply_nodes": sorted(self.apply_nodes), + "raw_candidate_count": len(rows), + "selected_count": len(selected), + "retrieve_limit": self._retrieve_limit(decision_node), + "inject_limit": self._inject_limit(decision_node, base_limit), + "mismatch_policy": mismatch_policy, + "mismatch_policies": dict(self.mismatch_policies), + "no_match_policy": self.no_match_policy, + "missing_query_policy": self.missing_query_policy, + "positive_match_required": self.positive_match_required, + "selection_policy": "score_sort", + "sidecar": self.load_report, + "loaded_files": _loaded_files(self.load_report), + "load_errors": self.load_report.get("errors") or [], + } + return selected, trace_rows, diagnostics + + domain_key = str(domain).lower() + query_signatures = _query_signature_candidates(domain_key, decision_node, query) + query_signature = query_signatures[0] + query_hash = hashlib.sha256(query.encode("utf-8")).hexdigest()[:16] + query_annotations: list[dict[str, Any]] = [] + matched_query_signatures: list[str] = [] + missing_query_signatures: list[str] = [] + for index, signature in enumerate(query_signatures): + query_row = _lookup_annotation( + self.annotation_index, + _query_lookup_keys(signature, include_hash=query_hash if index == 0 else None), + ) + if not isinstance(query_row, dict): + missing_query_signatures.append(signature) + continue + query_annotation = _compact_annotation(query_row) + if not query_annotation: + raise ValueError( + "empty category query sidecar annotation: " + f"domain={domain_key} decision_node={decision_node} " + f"query_signature={signature}" + ) + query_annotation = dict(query_annotation) + query_annotation["query_signature"] = signature + query_annotations.append(query_annotation) + matched_query_signatures.append(signature) + if not query_annotations: + if self.missing_query_policy in {"base_rank", "skip_injection"}: + selected = rows[:base_limit] if self.missing_query_policy == "base_rank" else [] + decision = ( + "missing_query_sidecar_base_rank" + if self.missing_query_policy == "base_rank" + else "missing_query_sidecar_skip_injection" + ) + trace_rows = _mark_selected(rows, selected, decision=decision) + diagnostics = { + "enabled": True, + "applied": False, + "decision_node": decision_node, + "decision": decision, + "apply_nodes": sorted(self.apply_nodes), + "raw_candidate_count": len(rows), + "selected_count": len(selected), + "retrieve_limit": self._retrieve_limit(decision_node), + "inject_limit": self._inject_limit(decision_node, base_limit), + "mismatch_policy": mismatch_policy, + "mismatch_policies": dict(self.mismatch_policies), + "positive_match_required": self.positive_match_required, + "no_match_policy": self.no_match_policy, + "missing_query_policy": self.missing_query_policy, + "query_sidecar_coverage": "missing", + "query_signature": query_signature, + "query_signature_candidates": query_signatures, + "matched_query_signatures": matched_query_signatures, + "missing_query_signatures": missing_query_signatures, + "selection_policy": self.missing_query_policy, + "sidecar": self.load_report, + "loaded_files": _loaded_files(self.load_report), + "load_errors": self.load_report.get("errors") or [], + } + return selected, trace_rows, diagnostics + raise ValueError( + "missing category query sidecar annotation: " + f"domain={domain_key} decision_node={decision_node} " + f"query_signature={query_signature}" + ) + query_annotation = _merge_query_annotations(query_annotations) + query_sidecar_coverage = "covered" if matched_query_signatures and matched_query_signatures[0] == query_signature else "partial" + scored = [] + candidates = [] + for index, row in enumerate(rows): + uri = str(row.get("uri") or "") + base_uri = uri.split("#", 1)[0] + memory_row = _lookup_annotation( + self.annotation_index, + [ + uri, + base_uri, + Path(base_uri).name, + ], + ) + if not isinstance(memory_row, dict): + raise ValueError( + "missing category memory sidecar annotation: " + f"domain={domain_key} decision_node={decision_node} uri={uri}" + ) + memory_annotation = _compact_annotation(memory_row) + if not memory_annotation: + raise ValueError( + "empty category memory sidecar annotation: " + f"domain={domain_key} decision_node={decision_node} uri={uri}" + ) + best_query_annotation: dict[str, Any] | None = None + best_score = float("-inf") + best_reasons: list[str] = [] + best_match_flags: dict[str, bool] = {} + for current_query in query_annotations: + score, reasons, match_flags = _candidate_score( + current_query, + memory_annotation, + original_rank=index + 1, + original_score=_score_value(row.get("score")) * self.search_score_weight, + ) + score_key = ( + score, + 1 if match_flags.get("category2_match") else 0, + 1 if match_flags.get("category1_match") else 0, + ) + best_key = ( + best_score, + 1 if best_match_flags.get("category2_match") else 0, + 1 if best_match_flags.get("category1_match") else 0, + ) + if score_key > best_key: + best_score = score + best_reasons = reasons + best_match_flags = match_flags + best_query_annotation = current_query + candidate = { + "uri": row.get("uri"), + "raw_rank": index + 1, + "raw_score": row.get("score"), + "category_score": best_score, + "category_rerank_reasons": best_reasons, + "query_category": best_query_annotation, + "memory_category": memory_annotation, + **best_match_flags, + } + candidates.append(candidate) + scored.append((best_score, -index, row, candidate)) + + sorted_scored = sorted(scored, key=lambda item: (item[0], item[1]), reverse=True) + positive_level = "none" + if any(item[3].get("category2_match") for item in sorted_scored): + positive_level = "category2" + elif any(item[3].get("category1_match") for item in sorted_scored): + positive_level = "category1" + + decision = "soft_reranked" + filtered = sorted_scored + dropped_mismatch_count = 0 + inject_limit = self._inject_limit(decision_node, base_limit) + priority_fill: dict[str, Any] = {"applied": False} + if mismatch_policy == "strict_pair_match_only": + before_count = len(sorted_scored) + filtered = [ + item + for item in sorted_scored + if item[3].get("category1_match") and item[3].get("category2_match") + ] + if filtered: + decision = "soft_reranked_keep_strict_pair_matches" + elif self.no_match_policy == "skip_injection": + filtered = [] + decision = "no_strict_pair_category_match_skip_injection" + dropped_mismatch_count = before_count - len(filtered) + elif mismatch_policy in {"keep_positive_match_drop_mismatch", "positive_match_only"}: + before_count = len(sorted_scored) + if positive_level == "category2": + filtered = [item for item in sorted_scored if item[3].get("category2_match")] + decision = "soft_reranked_keep_category2_matches" + elif positive_level == "category1": + filtered = [item for item in sorted_scored if item[3].get("category1_match")] + decision = "soft_reranked_keep_category1_matches" + elif self.no_match_policy == "skip_injection": + filtered = [] + decision = "no_positive_category_match_skip_injection" + dropped_mismatch_count = before_count - len(filtered) + elif mismatch_policy == "positive_priority_fill": + before_count = len(sorted_scored) + if positive_level in {"category1", "category2"}: + filtered, priority_fill = _priority_fill( + query_annotation, + sorted_scored, + inject_limit=inject_limit, + ) + decision = "soft_reranked_positive_priority_fill" + elif self.no_match_policy == "skip_injection": + filtered = [] + decision = "no_positive_category_match_skip_injection" + dropped_mismatch_count = before_count - len(filtered) + elif mismatch_policy == "drop_when_match_available": + has_positive_match = positive_level in {"category1", "category2"} + if has_positive_match: + before_count = len(sorted_scored) + guarded = [ + item for item in sorted_scored if not item[3].get("category_explicit_mismatch") + ] + if guarded: + filtered = guarded + decision = "soft_reranked_with_mismatch_guard" + dropped_mismatch_count = before_count - len(filtered) + elif self.no_match_policy == "skip_injection": + dropped_mismatch_count = len(sorted_scored) + filtered = [] + decision = "no_positive_category_match_skip_injection" + elif self.no_match_policy == "skip_injection" and positive_level == "none": + dropped_mismatch_count = len(sorted_scored) + filtered = [] + decision = "no_positive_category_match_skip_injection" + + selected = [item[2] for item in filtered[:inject_limit]] + trace_rows = _mark_selected( + rows, + selected, + decision=decision, + kept_before_cap=[item[2] for item in filtered], + candidate_by_uri={ + str(candidate.get("uri") or ""): candidate for candidate in candidates + }, + query_category=query_annotation, + ) + diagnostics = { + "enabled": True, + "applied": True, + "decision_node": decision_node, + "decision": decision, + "raw_candidate_count": len(rows), + "selected_count": len(selected), + "retrieve_limit": self._retrieve_limit(decision_node), + "inject_limit": inject_limit, + "mismatch_policy": mismatch_policy, + "mismatch_policies": dict(self.mismatch_policies), + "positive_match_required": self.positive_match_required, + "positive_match_level": positive_level, + "no_match_policy": self.no_match_policy, + "missing_query_policy": self.missing_query_policy, + "selection_policy": "score_sort", + "dropped_mismatch_count": dropped_mismatch_count, + "priority_fill": priority_fill, + "kept_before_cap_ids": [str(item[2].get("uri") or "") for item in filtered], + "query_category": query_annotation, + "query_signature": query_signature, + "query_signature_candidates": query_signatures, + "matched_query_signatures": matched_query_signatures, + "missing_query_signatures": missing_query_signatures, + "query_sidecar_coverage": query_sidecar_coverage, + "candidate_count": len(candidates), + "candidates": candidates, + "sidecar": self.load_report, + "loaded_files": _loaded_files(self.load_report), + "load_errors": self.load_report.get("errors") or [], + } + return selected, trace_rows, diagnostics + + +def _loaded_files(load_report: dict[str, Any]) -> list[str]: + loaded_files = load_report.get("loaded_files") + if isinstance(loaded_files, list): + return [str(row.get("path") if isinstance(row, dict) else row) for row in loaded_files] + if load_report.get("loaded") and load_report.get("path"): + return [str(load_report["path"])] + return [] + + +def _ordered_values(payload: dict[str, Any], key: str) -> list[str]: + value = payload.get(key) + if isinstance(value, str) and value.strip(): + return [value.strip()] + if isinstance(value, list): + return list(dict.fromkeys(str(item).strip() for item in value if str(item).strip())) + return [] + + +def _merge_query_annotations(query_annotations: list[dict[str, Any]]) -> dict[str, Any]: + if not query_annotations: + return {} + merged: dict[str, Any] = { + "matched": True, + "category_source": "multi_query_sidecar" if len(query_annotations) > 1 else query_annotations[0].get("category_source"), + "annotation_id": ",".join(str(row.get("annotation_id") or "") for row in query_annotations if row.get("annotation_id")), + "query_signatures": [row.get("query_signature") for row in query_annotations if row.get("query_signature")], + } + for key in ("category_id", "category1", "category2", "category3"): + values: list[str] = [] + for row in query_annotations: + for value in _ordered_values(row, key): + if value not in values: + values.append(value) + if values: + merged[key] = values[0] if len(values) == 1 else values + confidences = [row.get("confidence") for row in query_annotations if isinstance(row.get("confidence"), (int, float))] + if confidences: + merged["confidence"] = max(confidences) + return merged + + +def _values(payload: dict[str, Any], key: str) -> set[str]: + value = payload.get(key) + if isinstance(value, str) and value.strip(): + return {value.strip()} + if isinstance(value, list): + return {str(item).strip() for item in value if str(item).strip()} + return set() + + +def _priority_fill( + query: dict[str, Any], + sorted_scored: list[tuple[float, int, dict[str, Any], dict[str, Any]]], + *, + inject_limit: int, +) -> tuple[list[tuple[float, int, dict[str, Any], dict[str, Any]]], dict[str, Any]]: + if inject_limit <= 0 or not sorted_scored: + return sorted_scored, {"applied": False, "reason": "empty_or_zero_limit"} + + selected: list[tuple[float, int, dict[str, Any], dict[str, Any]]] = [] + selected_indexes: set[int] = set() + fill_steps: list[dict[str, Any]] = [] + + def pick(level: str, category: str) -> None: + if len(selected) >= inject_limit: + return + for index, row in enumerate(sorted_scored): + if index in selected_indexes: + continue + memory_category = row[3].get("memory_category") + memory_category = memory_category if isinstance(memory_category, dict) else {} + if category not in _values(memory_category, level): + continue + selected.append(row) + selected_indexes.add(index) + fill_steps.append( + { + "level": level, + "category": category, + "uri": str(row[2].get("uri") or ""), + "category_score": row[0], + } + ) + return + + for category_id in _ordered_values(query, "category_id"): + pick("category_id", category_id) + for category1 in _ordered_values(query, "category1"): + pick("category1", category1) + + positive_remaining = [ + row + for index, row in enumerate(sorted_scored) + if index not in selected_indexes + and (row[3].get("category1_match") or row[3].get("category2_match")) + ] + other_remaining = [ + row + for index, row in enumerate(sorted_scored) + if index not in selected_indexes + and not (row[3].get("category1_match") or row[3].get("category2_match")) + ] + return selected + positive_remaining + other_remaining, { + "applied": bool(selected), + "inject_limit": inject_limit, + "selected_count": len(selected), + "fill_steps": fill_steps, + } + + +def _candidate_score( + query: dict[str, Any], + memory: dict[str, Any], + *, + original_rank: int, + original_score: float, +) -> tuple[float, list[str], dict[str, bool]]: + score = original_score - (original_rank * 0.001) + reasons = ["original_rank_tiebreak"] + if original_score: + reasons.insert(0, "openviking_score") + query_ids = _values(query, "category_id") + memory_ids = _values(memory, "category_id") + query_c1 = _values(query, "category1") + query_c2 = _values(query, "category2") + memory_c1 = _values(memory, "category1") + memory_c2 = _values(memory, "category2") + category1_match = bool(query_c1 and memory_c1 and query_c1 & memory_c1) + category_pair_match = bool(query_ids and memory_ids and query_ids & memory_ids) + category2_label_match = bool(query_c2 and memory_c2 and query_c2 & memory_c2) + category2_match = category_pair_match or bool( + not query_ids and not memory_ids and category1_match and category2_label_match + ) + if category2_match: + score += 100.0 + reasons.append("category_pair_match") + if category1_match: + score += 40.0 + reasons.append("category1_match") + if category2_label_match and not category2_match: + reasons.append("category2_label_match_without_pair") + if (query_ids and memory_ids and not category_pair_match) or ( + not query_ids and not memory_ids and query_c2 and memory_c2 and not category2_match + ): + score -= 5.0 + reasons.append("category_pair_mismatch_downrank") + if query_c1 and memory_c1 and not category1_match: + score -= 20.0 + reasons.append("category1_mismatch_downrank") + if (query_c1 or query_c2) and not (memory_c1 or memory_c2): + score -= 2.0 + reasons.append("missing_memory_category") + return ( + score, + reasons, + { + "category1_match": category1_match, + "category2_match": category2_match, + "category_pair_match": category_pair_match, + "category2_label_match": category2_label_match, + "category_explicit_mismatch": bool( + (query_c1 and memory_c1 and not category1_match) + or (query_ids and memory_ids and not category_pair_match) + or ( + not query_ids + and not memory_ids + and query_c2 + and memory_c2 + and not category2_match + ) + ), + }, + ) + + +def _row_key(row: dict[str, Any]) -> str: + return str(row.get("uri") or row.get("memory_id") or id(row)) + + +def _mark_selected( + rows: list[dict[str, Any]], + selected_rows: list[dict[str, Any]], + *, + decision: str, + kept_before_cap: list[dict[str, Any]] | None = None, + candidate_by_uri: dict[str, dict[str, Any]] | None = None, + query_category: dict[str, Any] | None = None, +) -> list[dict[str, Any]]: + selected_keys = {_row_key(row) for row in selected_rows} + kept_keys = {_row_key(row) for row in (kept_before_cap or selected_rows)} + trace_rows = [] + for index, row in enumerate(rows, start=1): + key = _row_key(row) + traced = _public_row(row) + traced["raw_rank"] = index + traced["selected_for_injection"] = key in selected_keys + traced["injected"] = bool( + traced["selected_for_injection"] and int(row.get("text_chars") or 0) > 0 + ) + if not traced["selected_for_injection"]: + traced["skipped_reason"] = ( + "category_rerank_inject_limit" if key in kept_keys else "category_rerank" + ) + if decision == "no_positive_category_match_skip_injection": + traced["skipped_reason"] = "category_rerank_no_positive_match" + candidate = (candidate_by_uri or {}).get(str(row.get("uri") or "")) + if candidate: + memory_category = candidate.get("memory_category") + memory_category = memory_category if isinstance(memory_category, dict) else {} + traced["category_rerank_score"] = candidate.get("category_score") + traced["category_rerank_reasons"] = candidate.get("category_rerank_reasons") + traced["memory_category"] = memory_category + traced["memory_category1_prompt"] = memory_category.get("category1") + traced["memory_category2_prompt"] = memory_category.get("category2") + traced["memory_category_source_prompt"] = memory_category.get("category_source") + traced["memory_category_confidence_prompt"] = memory_category.get("confidence") + candidate_query_category = candidate.get("query_category") + candidate_query_category = candidate_query_category if isinstance(candidate_query_category, dict) else query_category + if candidate_query_category: + traced["query_category1_prompt"] = candidate_query_category.get("category1") + traced["query_category2_prompt"] = candidate_query_category.get("category2") + traced["query_category_source_prompt"] = candidate_query_category.get("category_source") + traced["query_category_confidence_prompt"] = candidate_query_category.get("confidence") + traced["query_category_signature"] = candidate_query_category.get("query_signature") + traced["category1_match"] = candidate.get("category1_match") + traced["category2_match"] = candidate.get("category2_match") + traced["category_pair_match"] = candidate.get("category_pair_match") + traced["category2_label_match"] = candidate.get("category2_label_match") + traced["category_explicit_mismatch"] = candidate.get("category_explicit_mismatch") + trace_rows.append(traced) + return trace_rows diff --git a/benchmark/tau2/scripts/generate_category_annotations.py b/benchmark/tau2/scripts/generate_category_annotations.py new file mode 100755 index 0000000000..93ef2e2f90 --- /dev/null +++ b/benchmark/tau2/scripts/generate_category_annotations.py @@ -0,0 +1,613 @@ +"""Execute memory category extraction requests with an OpenAI-compatible LLM.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import time +import urllib.error +import urllib.request +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parents[3] +DEFAULT_OUTPUT_ROOT = ROOT / "outputs" / "open_benchmarks" / "memory_category_annotation_llm_v0" +CATEGORY_ID_PATTERN = re.compile(r"^[a-z0-9_][a-z0-9_-]*:[a-z0-9_][a-z0-9_-]*$") +CATEGORY_PART_PATTERN = re.compile(r"^[a-z0-9_][a-z0-9_-]*$") +CATEGORY_PART_MAX_LENGTH = 64 +CATEGORY_SOURCE_ENUM = { + "llm_prompt", + "tool_schema", + "uri_title_metadata", + "manual_taxonomy", + "outcome_policy", + "rule_fallback", + "existing_catalog", + "annotation_catalog", + "mixed", +} + + +@dataclass(frozen=True) +class WorkerResult: + status: str + raw_output: str + returncode: int | None + duration_seconds: float | None + usage: dict[str, Any] + error: dict[str, Any] | None + artifacts: dict[str, str] + + @property + def succeeded(self) -> bool: + return self.status == "succeeded" and self.returncode in (0, None) and bool(self.raw_output.strip()) + + +def _safe_key(value: str) -> str: + return "".join(ch if ch.isalnum() or ch in {"-", "_", "."} else "_" for ch in value)[:180] + + +def _iter_requests(path: Path) -> list[dict[str, Any]]: + rows = [] + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise SystemExit(f"invalid request JSONL at line {line_number}: {exc}") from exc + if not row.get("prompt"): + raise SystemExit(f"request line {line_number} missing prompt") + rows.append(row) + if not rows: + raise SystemExit(f"no requests found: {path}") + return rows + + +def _iter_jsonl_objects(path: Path) -> list[dict[str, Any]]: + if not path.exists(): + return [] + rows: list[dict[str, Any]] = [] + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + raise SystemExit(f"invalid JSONL at {path}:{line_number}: {exc}") from exc + if not isinstance(row, dict): + raise SystemExit(f"expected object at {path}:{line_number}") + rows.append(row) + return rows + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise SystemExit(f"missing required environment variable: {name}") + return value + + +def _load_backend_config() -> dict[str, str]: + return { + "api_key_env": "ARK_API_KEY", + "base_url": os.environ.get("ARK_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3").rstrip("/"), + "model": os.environ.get("DOUBAO_MODEL", "doubao-seed-2-0-pro-260215"), + "provider": os.environ.get("LLM_PROVIDER_NAME", "volcengine_ark"), + } + + +def _chat_completion( + *, + prompt: str, + backend: dict[str, str], + max_tokens: int, + timeout_seconds: int, + retry_count: int, +) -> dict[str, Any]: + payload = { + "model": backend["model"], + "messages": [{"role": "user", "content": prompt}], + "temperature": 0, + "max_tokens": max_tokens, + } + url = backend["base_url"] + "/chat/completions" + last_error: dict[str, Any] | None = None + for attempt in range(retry_count + 1): + request = urllib.request.Request( + url, + data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), + headers={ + "Authorization": "Bearer " + _required_env(backend["api_key_env"]), + "Content-Type": "application/json", + }, + method="POST", + ) + started = time.time() + try: + with urllib.request.urlopen(request, timeout=timeout_seconds) as response: + response_payload = json.loads(response.read().decode("utf-8")) + message = response_payload.get("choices", [{}])[0].get("message", {}) + return { + "content": message.get("content", ""), + "duration_seconds": round(time.time() - started, 4), + "error": None, + "returncode": 0, + "status_code": 200, + "usage": response_payload.get("usage", {}), + } + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + last_error = {"body_excerpt": body[:1000], "code": exc.code, "type": "HTTPError"} + if exc.code not in {408, 409, 429, 500, 502, 503, 504}: + break + except Exception as exc: # noqa: BLE001 - backend failures are diagnostic. + last_error = {"message": str(exc), "type": type(exc).__name__} + if attempt < retry_count: + time.sleep(min(2**attempt, 8)) + return { + "content": "", + "duration_seconds": None, + "error": last_error, + "returncode": 1, + "status_code": (last_error or {}).get("code"), + "usage": {}, + } + + +def _run_worker( + *, + prompt: str, + run_dir: Path, + backend: dict[str, str], + max_tokens: int, + timeout_seconds: int, + retry_count: int, +) -> WorkerResult: + run_dir.mkdir(parents=True, exist_ok=True) + prompt_path = run_dir / "prompt.txt" + last_message_path = run_dir / "last_message.txt" + meta_path = run_dir / "backend_response_meta.json" + prompt_path.write_text(prompt, encoding="utf-8") + response = _chat_completion( + prompt=prompt, + backend=backend, + max_tokens=max_tokens, + timeout_seconds=timeout_seconds, + retry_count=retry_count, + ) + raw_output = str(response.get("content") or "") + last_message_path.write_text(raw_output, encoding="utf-8") + provider_meta = { + "api_key_env": backend["api_key_env"], + "base_url": backend["base_url"], + "duration_seconds": response.get("duration_seconds"), + "error": response.get("error"), + "model": backend["model"], + "provider": backend["provider"], + "returncode": response.get("returncode"), + "status_code": response.get("status_code"), + "usage": response.get("usage") or {}, + } + meta_path.write_text(json.dumps(provider_meta, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return WorkerResult( + status="succeeded" if response.get("returncode") == 0 and raw_output.strip() else "failed", + raw_output=raw_output, + returncode=response.get("returncode"), + duration_seconds=response.get("duration_seconds"), + usage=response.get("usage") or {}, + error=response.get("error"), + artifacts={ + "last_message_path": str(last_message_path), + "meta_path": str(meta_path), + "prompt_path": str(prompt_path), + }, + ) + + +def _extract_json_object(text: str) -> dict[str, Any]: + stripped = text.strip() + if stripped.startswith("```"): + lines = stripped.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].startswith("```"): + lines = lines[:-1] + stripped = "\n".join(lines).strip() + try: + payload = json.loads(stripped) + if isinstance(payload, dict): + return payload + except json.JSONDecodeError: + pass + start = stripped.find("{") + end = stripped.rfind("}") + if start >= 0 and end > start: + payload = json.loads(stripped[start : end + 1]) + if isinstance(payload, dict): + return payload + raise ValueError("LLM output did not contain a JSON object") + + +def _normalize_annotation(annotation: dict[str, Any], *, request_id: str, subject: dict[str, Any] | None) -> dict[str, Any]: + annotation["schema_version"] = "memory_category_annotation.v0" + annotation["annotation_id"] = request_id + annotation["request_id"] = request_id + annotation["producer"] = "llm_prompt" + annotation["subject"] = subject + return annotation + + +def _validate_with_jsonschema(annotation: dict[str, Any], schema: dict[str, Any] | None) -> list[str] | None: + if not schema: + return None + try: + import jsonschema # type: ignore[import-not-found] + except Exception: + return None + try: + jsonschema.validate(annotation, schema) + except Exception as exc: + return [str(exc)] + return [] + + +def _basic_validate_annotation(annotation: dict[str, Any]) -> list[str]: + errors: list[str] = [] + if annotation.get("schema_version") != "memory_category_annotation.v0": + errors.append("$.schema_version must be 'memory_category_annotation.v0'") + if not isinstance(annotation.get("annotation_id"), str) or not annotation.get("annotation_id"): + errors.append("$.annotation_id must be a non-empty string") + subject = annotation.get("subject") + if not isinstance(subject, dict): + errors.append("$.subject must be an object") + else: + if not subject.get("subject_type"): + errors.append("$.subject.subject_type is required") + if not subject.get("subject_id"): + errors.append("$.subject.subject_id is required") + category = annotation.get("category") + if not isinstance(category, dict): + errors.append("$.category must be an object") + else: + for field in ("category1", "category2"): + value = category.get(field) + if not isinstance(value, str) or not value: + errors.append(f"$.category.{field} is required") + elif not CATEGORY_PART_PATTERN.fullmatch(value): + errors.append( + f"$.category.{field} must be a reusable slug id using only " + "lowercase letters, numbers, '_' or '-'" + ) + elif len(value) > CATEGORY_PART_MAX_LENGTH: + errors.append( + f"$.category.{field} must be a compact reusable slug id " + f"with at most {CATEGORY_PART_MAX_LENGTH} characters; " + "put detailed applicability boundaries in $.applicability instead" + ) + category3 = category.get("category3") + if category3 is not None and ( + not isinstance(category3, str) or not CATEGORY_PART_PATTERN.fullmatch(category3) + ): + errors.append( + "$.category.category3 must be null or a reusable slug id using only " + "lowercase letters, numbers, '_' or '-'" + ) + elif isinstance(category3, str) and len(category3) > CATEGORY_PART_MAX_LENGTH: + errors.append( + f"$.category.category3 must be a compact reusable slug id " + f"with at most {CATEGORY_PART_MAX_LENGTH} characters; " + "put detailed applicability boundaries in $.applicability instead" + ) + if category.get("category_source") not in CATEGORY_SOURCE_ENUM: + errors.append( + "$.category.category_source must be one of " + + ", ".join(sorted(CATEGORY_SOURCE_ENUM)) + ) + confidence = category.get("confidence") + if not isinstance(confidence, (int, float)) or confidence < 0 or confidence > 1: + errors.append("$.category.confidence must be a number in [0, 1]") + catalog_match = category.get("catalog_match") + if isinstance(catalog_match, dict): + matched_category_id = catalog_match.get("matched_category_id") + if matched_category_id is not None and not CATEGORY_ID_PATTERN.fullmatch(str(matched_category_id)): + errors.append( + "$.category.catalog_match.matched_category_id must be ':' " + "or null" + ) + if not isinstance(annotation.get("safety"), dict): + errors.append("$.safety must be an object") + return errors + + +def _load_schema(schema_path: Path | None) -> dict[str, Any] | None: + if not schema_path: + return None + if not schema_path.is_file(): + raise SystemExit(f"schema file not found: {schema_path}") + return json.loads(schema_path.read_text(encoding="utf-8")) + + +def _validate_annotation(annotation: dict[str, Any], schema: dict[str, Any] | None) -> list[str]: + jsonschema_errors = _validate_with_jsonschema(annotation, schema) + basic_errors = _basic_validate_annotation(annotation) + if jsonschema_errors is None: + return basic_errors + return [*jsonschema_errors, *basic_errors] + + +def _validate_annotations_file(path: Path, schema: dict[str, Any] | None) -> dict[str, Any]: + errors: list[dict[str, Any]] = [] + count = 0 + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if not line.strip(): + continue + count += 1 + try: + row = json.loads(line) + except json.JSONDecodeError as exc: + errors.append({"line_number": line_number, "message": f"invalid JSON: {exc}"}) + continue + for message in _validate_annotation(row, schema): + errors.append({"line_number": line_number, "message": message}) + return { + "status": "passed" if not errors else "failed", + "event_count": count, + "error_count": len(errors), + "errors": errors, + } + + +def _validation_retry_prompt(base_prompt: str, *, errors: list[str], previous_output: str | None) -> str: + details = "\n".join(f"- {error}" for error in errors[:12]) or "- unknown validation failure" + previous = "" + if previous_output: + previous = ( + "\n\nPrevious invalid output, for reference only. Do not copy invalid enum values or malformed ids:\n" + "```json\n" + f"{previous_output.strip()[:6000]}\n" + "```" + ) + return ( + f"{base_prompt.rstrip()}\n\n" + "Your previous answer failed the required JSON schema validation. " + "Return one corrected JSON object only, with no markdown and no explanation.\n" + "Do not invent schema fields. Do not normalize invalid enum values in prose; choose one enum value from the schema. " + "category1/category2/category3 must be compact slug ids, not prose sentences; use lowercase snake_case. " + f"Each category id part must be at most {CATEGORY_PART_MAX_LENGTH} characters; put detailed state, precondition, " + "confirmation, or eligibility boundaries in applicability fields instead of the category id. " + "For query-side subjects, category ids must describe the reusable business action, skill, artifact type, or " + "applicability boundary; do not encode decision-node mechanics such as first_user, pre_write, before_write, " + "classification, query, or retrieval into category1/category2/category3. " + "If category.catalog_match.matched_category_id is present, it must be the canonical ':' id; " + "otherwise set it to null.\n\n" + "Validation errors:\n" + f"{details}" + f"{previous}" + ) + + +def _relative(path: Path) -> str: + try: + return str(path.relative_to(ROOT)) + except ValueError: + return str(path) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--requests", type=Path, required=True) + parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT) + parser.add_argument("--run-id", default=None) + parser.add_argument("--schema-path", type=Path, default=None) + parser.add_argument("--max-tokens", type=int, default=4000) + parser.add_argument("--timeout-seconds", type=int, default=120) + parser.add_argument("--retry-count", type=int, default=1, help="Schema-validation retry count per request.") + parser.add_argument("--api-retry-count", type=int, default=1, help="Transport/backend retry count per LLM attempt.") + parser.add_argument("--limit", type=int, default=0, help="0 means all requests") + parser.add_argument( + "--resume-existing", + action="store_true", + help="Reuse existing parsed annotations/execution rows in the same run_id and only run missing requests.", + ) + args = parser.parse_args() + + requests_path = args.requests.resolve() + request_rows = _iter_requests(requests_path) + if args.limit > 0: + request_rows = request_rows[: args.limit] + request_ids = {str(row.get("request_id") or f"request_{index}") for index, row in enumerate(request_rows)} + request_id_by_subject_id = { + str(row.get("subject", {}).get("subject_id")): str(row.get("request_id") or f"request_{index}") + for index, row in enumerate(request_rows) + if isinstance(row.get("subject"), dict) and row.get("subject", {}).get("subject_id") + } + + run_id = args.run_id or datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + output_root = args.output_root if args.output_root.is_absolute() else ROOT / args.output_root + run_root = output_root / _safe_key(run_id) + run_root.mkdir(parents=True, exist_ok=True) + + schema_path = args.schema_path + schema = _load_schema(schema_path.resolve() if schema_path else None) + backend_config = _load_backend_config() + annotations_path = run_root / "annotations.jsonl" + execution_rows_path = run_root / "execution_rows.jsonl" + validation_path = run_root / "validation_report.json" + summary_path = run_root / "run_summary.json" + + loaded_annotations: list[dict[str, Any]] = _iter_jsonl_objects(annotations_path) if args.resume_existing else [] + annotations: list[dict[str, Any]] = [] + dropped_existing_annotations: list[dict[str, Any]] = [] + for line_index, row in enumerate(loaded_annotations, start=1): + validation_errors = _validate_annotation(row, schema) + if validation_errors: + dropped_existing_annotations.append( + { + "annotation_id": row.get("annotation_id"), + "line_number": line_index, + "request_id": row.get("request_id"), + "validation_errors": validation_errors, + } + ) + continue + annotations.append(row) + execution_rows: list[dict[str, Any]] = _iter_jsonl_objects(execution_rows_path) if args.resume_existing else [] + parsed_request_ids: set[str] = set() + for row in annotations: + request_id = row.get("request_id") + if isinstance(request_id, str) and request_id: + parsed_request_ids.add(request_id) + continue + annotation_id = row.get("annotation_id") + if isinstance(annotation_id, str) and annotation_id in request_ids: + parsed_request_ids.add(annotation_id) + continue + subject = row.get("subject") if isinstance(row.get("subject"), dict) else {} + subject_id = subject.get("subject_id") + if isinstance(subject_id, str) and subject_id in request_id_by_subject_id: + parsed_request_ids.add(request_id_by_subject_id[subject_id]) + if args.resume_existing and len(parsed_request_ids) > len(annotations): + raise SystemExit(f"cannot resume because parsed request ids exceed annotations: {annotations_path}") + + for index, request_row in enumerate(request_rows): + request_id = str(request_row.get("request_id") or f"request_{index}") + if args.resume_existing and request_id in parsed_request_ids: + continue + request_dir = run_root / f"{index:04d}_{_safe_key(request_id)}" + attempts: list[dict[str, Any]] = [] + parse_error: dict[str, Any] | None = None + validation_errors: list[str] = [] + previous_output: str | None = None + last_result = None + annotation: dict[str, Any] | None = None + for attempt_index in range(args.retry_count + 1): + prompt = ( + request_row["prompt"] + if attempt_index == 0 + else _validation_retry_prompt( + request_row["prompt"], + errors=validation_errors or ([parse_error["message"]] if parse_error else []), + previous_output=previous_output, + ) + ) + attempt_dir = request_dir / f"attempt_{attempt_index + 1:02d}" + result = _run_worker( + prompt=prompt, + run_dir=attempt_dir, + backend=backend_config, + max_tokens=args.max_tokens, + timeout_seconds=args.timeout_seconds, + retry_count=args.api_retry_count, + ) + last_result = result + parse_error = None + validation_errors = [] + candidate: dict[str, Any] | None = None + previous_output = result.raw_output if result.succeeded else None + if result.succeeded: + try: + candidate = _extract_json_object(result.raw_output) + candidate = _normalize_annotation( + candidate, + request_id=request_id, + subject=request_row.get("subject"), + ) + validation_errors = _validate_annotation(candidate, schema) + except Exception as exc: # noqa: BLE001 - LLM output parsing is diagnostic. + parse_error = {"type": type(exc).__name__, "message": str(exc)} + else: + validation_errors = [f"worker failed with status={result.status} returncode={result.returncode}"] + attempt_row = { + "attempt_index": attempt_index + 1, + "status": ( + "parsed" + if candidate and not validation_errors and not parse_error + else "schema_failed" + if candidate and validation_errors + else "parse_failed" + if parse_error + else "worker_failed" + ), + "worker_status": result.status, + "returncode": result.returncode, + "parse_error": parse_error, + "validation_errors": validation_errors, + "artifacts": {key: _relative(Path(value)) for key, value in result.artifacts.items()}, + "usage": result.usage, + "duration_seconds": result.duration_seconds, + } + attempts.append(attempt_row) + if candidate and not validation_errors and not parse_error: + annotation = candidate + break + row = { + "request_id": request_id, + "status": "parsed" if annotation else "failed", + "worker_status": last_result.status if last_result else "not_run", + "returncode": last_result.returncode if last_result else None, + "parse_error": parse_error, + "validation_errors": validation_errors, + "attempt_count": len(attempts), + "attempts": attempts, + "subject": request_row.get("subject"), + "artifacts": attempts[-1]["artifacts"] if attempts else {}, + "usage": attempts[-1]["usage"] if attempts else {}, + "duration_seconds": sum(float(attempt.get("duration_seconds") or 0.0) for attempt in attempts), + } + execution_rows.append(row) + if annotation: + annotations.append(annotation) + parsed_request_ids.add(request_id) + annotations_path.write_text( + "".join(json.dumps(item, ensure_ascii=False, sort_keys=True) + "\n" for item in annotations), + encoding="utf-8", + ) + execution_rows_path.write_text( + "".join(json.dumps(item, ensure_ascii=False, sort_keys=True) + "\n" for item in execution_rows), + encoding="utf-8", + ) + + validation_report = _validate_annotations_file(annotations_path, schema) if annotations else None + if validation_report: + validation_path.write_text( + json.dumps(validation_report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + summary = { + "schema_version": "memory_category_annotation_llm_summary.v0", + "run_id": run_id, + "status": "passed" if validation_report and validation_report["status"] == "passed" and len(annotations) == len(request_rows) else "failed", + "request_count": len(request_rows), + "annotation_count": len(annotations), + "failed_count": len(request_rows) - len(annotations), + "backend": "openai_compatible", + "backend_provider": backend_config.get("provider"), + "backend_model": backend_config.get("model"), + "backend_api_key_env": backend_config.get("api_key_env"), + "dropped_existing_annotation_count": len(dropped_existing_annotations), + "dropped_existing_annotations": dropped_existing_annotations[:20], + "max_tokens": args.max_tokens, + "schema_retry_count": args.retry_count, + "api_retry_count": args.api_retry_count, + "validation_retry_enabled": True, + "requests_path": _relative(requests_path), + "schema_path": _relative(schema_path.resolve()) if schema_path else None, + "annotations_path": _relative(annotations_path), + "execution_rows_path": _relative(execution_rows_path), + "validation_report_path": _relative(validation_path) if validation_report else None, + "claim_boundary": "llm_prompt_generated_category_annotation_runtime_safe_visible_context_only", + "validation_status": validation_report["status"] if validation_report else "not_run_no_annotations", + } + summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True)) + return 0 if summary["status"] == "passed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/tau2/scripts/run_category_annotation_batches.py b/benchmark/tau2/scripts/run_category_annotation_batches.py new file mode 100644 index 0000000000..a28e54015a --- /dev/null +++ b/benchmark/tau2/scripts/run_category_annotation_batches.py @@ -0,0 +1,266 @@ +"""Run category annotation in rolling catalog batches. + +This intentionally keeps OpenViking memory writes out of the loop. Each batch: +1. Builds a catalog from all prior valid annotations. +2. Renders the next batch of category requests with that catalog. +3. Runs the LLM annotation executor for the batch. +4. Adds the batch annotations to the next catalog input. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parents[3] + + +def _relative(path: Path) -> str: + try: + return str(path.relative_to(ROOT)) + except ValueError: + return str(path) + + +def _safe_key(value: str) -> str: + return "".join(ch if ch.isalnum() or ch in {"-", "_", "."} else "_" for ch in value)[:180] + + +def _run(command: list[str]) -> dict[str, Any]: + result = subprocess.run(command, cwd=ROOT, text=True, capture_output=True, check=False) + if result.returncode != 0: + raise SystemExit( + "command failed" + f"\nreturncode={result.returncode}" + f"\ncommand={' '.join(command)}" + f"\nstdout={result.stdout[-4000:]}" + f"\nstderr={result.stderr[-4000:]}" + ) + output = result.stdout.strip().splitlines()[-1] if result.stdout.strip() else "{}" + try: + payload = json.loads(output) + except json.JSONDecodeError: + payload = {"stdout": result.stdout} + return payload + + +def _batch_ranges( + *, + start_offset: int, + end_offset: int, + batch_size: int, + warmup_count: int = 0, +) -> list[tuple[int, int]]: + if start_offset < 0: + raise ValueError("start_offset must be >= 0") + if end_offset <= start_offset: + raise ValueError("end_offset must be greater than start_offset") + if batch_size <= 0: + raise ValueError("batch_size must be > 0") + if warmup_count < 0: + raise ValueError("warmup_count must be >= 0") + ranges: list[tuple[int, int]] = [] + offset = start_offset + warmup_end = min(end_offset, start_offset + warmup_count) + while offset < warmup_end: + ranges.append((offset, 1)) + offset += 1 + while offset < end_offset: + limit = min(batch_size, end_offset - offset) + ranges.append((offset, limit)) + offset += limit + return ranges + + +def _load_json(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--workspace", type=Path, required=True) + parser.add_argument("--domain", required=True) + parser.add_argument("--memory-type", default="trajectories") + parser.add_argument("--seed-annotations", type=Path, action="append", default=[]) + parser.add_argument("--start-offset", type=int, required=True) + parser.add_argument("--end-offset", type=int, required=True) + parser.add_argument("--batch-size", type=int, default=10) + parser.add_argument( + "--warmup-count", + type=int, + default=0, + help="Annotate the first N items one by one so later batches can reuse the early catalog.", + ) + parser.add_argument("--run-id-prefix", required=True) + parser.add_argument("--schema-path", type=Path, default=None) + parser.add_argument("--output-root", type=Path, default=ROOT / "benchmark" / "tau2" / "result") + parser.add_argument("--max-catalog-categories", type=int, default=80) + parser.add_argument("--subject-text-limit", type=int, default=3500) + parser.add_argument("--max-tokens", type=int, default=4000) + parser.add_argument("--timeout-seconds", type=int, default=180) + parser.add_argument("--retry-count", type=int, default=2) + parser.add_argument("--api-retry-count", type=int, default=3) + parser.add_argument("--max-category-ratio", type=float, default=0.85) + parser.add_argument("--resume-existing", action="store_true") + parser.add_argument("--hint", action="append", default=[]) + args = parser.parse_args() + + if args.max_category_ratio <= 0: + raise SystemExit("--max-category-ratio must be > 0") + + output_root = args.output_root if args.output_root.is_absolute() else ROOT / args.output_root + catalog_root = output_root / "category_catalogs" / _safe_key(args.run_id_prefix) + request_root = output_root / "category_requests" + annotation_root = output_root / "category_annotations" + summary_root = output_root / "category_batch_runs" / _safe_key(args.run_id_prefix) + catalog_root.mkdir(parents=True, exist_ok=True) + summary_root.mkdir(parents=True, exist_ok=True) + + annotation_files = [path.resolve() for path in args.seed_annotations] + batches: list[dict[str, Any]] = [] + ranges = _batch_ranges( + start_offset=args.start_offset, + end_offset=args.end_offset, + batch_size=args.batch_size, + warmup_count=args.warmup_count, + ) + + for batch_index, (offset, limit) in enumerate(ranges, start=1): + catalog_path: Path | None = None + if annotation_files: + catalog_path = catalog_root / f"catalog_before_b{batch_index:02d}_offset_{offset}.json" + command = [ + sys.executable, + str(ROOT / "benchmark/tau2/scripts/build_category_catalog.py"), + "--output", + str(catalog_path), + "--run-id", + f"{args.run_id_prefix}_catalog_before_b{batch_index:02d}", + ] + for path in annotation_files: + command.extend(["--annotations", str(path)]) + _run(command) + + request_run_id = f"{args.run_id_prefix}_requests_b{batch_index:02d}_offset_{offset}_limit_{limit}" + request_command = [ + sys.executable, + str(ROOT / "benchmark/tau2/scripts/build_category_requests.py"), + "--manifest", + str(args.manifest.resolve()), + "--workspace", + str(args.workspace.resolve()), + "--domain", + args.domain, + "--memory-type", + args.memory_type, + "--offset", + str(offset), + "--limit", + str(limit), + "--max-catalog-categories", + str(args.max_catalog_categories), + "--subject-text-limit", + str(args.subject_text_limit), + "--output-root", + str(request_root), + "--run-id", + request_run_id, + ] + if catalog_path: + request_command.extend(["--category-catalog", str(catalog_path)]) + for hint in args.hint: + request_command.extend(["--hint", hint]) + request_summary = _run(request_command) + + annotation_run_id = f"{args.run_id_prefix}_annotations_b{batch_index:02d}_offset_{offset}_limit_{limit}" + annotation_command = [ + sys.executable, + str(ROOT / "benchmark/tau2/scripts/generate_category_annotations.py"), + "--requests", + str(request_root / _safe_key(request_run_id) / "category_extraction_requests.jsonl"), + "--output-root", + str(annotation_root), + "--run-id", + annotation_run_id, + "--max-tokens", + str(args.max_tokens), + "--timeout-seconds", + str(args.timeout_seconds), + "--retry-count", + str(args.retry_count), + "--api-retry-count", + str(args.api_retry_count), + ] + if args.schema_path: + annotation_command.extend(["--schema-path", str(args.schema_path.resolve())]) + if args.resume_existing: + annotation_command.append("--resume-existing") + annotation_summary = _run(annotation_command) + + annotations_path = annotation_root / _safe_key(annotation_run_id) / "annotations.jsonl" + annotation_files.append(annotations_path.resolve()) + batches.append( + { + "annotation_run_id": annotation_run_id, + "annotations_path": _relative(annotations_path.resolve()), + "catalog_path": _relative(catalog_path.resolve()) if catalog_path else None, + "limit": limit, + "offset": offset, + "request_run_id": request_run_id, + "request_summary": request_summary, + "annotation_summary": annotation_summary, + } + ) + + final_catalog_path = catalog_root / "final_category_catalog.json" + final_command = [ + sys.executable, + str(ROOT / "benchmark/tau2/scripts/build_category_catalog.py"), + "--output", + str(final_catalog_path), + "--run-id", + f"{args.run_id_prefix}_final_catalog", + ] + for path in annotation_files: + final_command.extend(["--annotations", str(path)]) + _run(final_command) + final_catalog = _load_json(final_catalog_path) + annotation_count = int(final_catalog.get("source_annotation_count") or 0) + category_count = int(final_catalog.get("category_count") or 0) + category_ratio = category_count / annotation_count if annotation_count else 0.0 + compaction_status = "passed" if category_ratio <= args.max_category_ratio else "warning_high_category_ratio" + + summary = { + "annotation_count": annotation_count, + "batch_count": len(batches), + "batches": batches, + "category_count": category_count, + "category_count_ratio": round(category_ratio, 6), + "category_ratio_threshold": args.max_category_ratio, + "catalog_compaction_status": compaction_status, + "claim_boundary": "rolling_batch_category_annotation_no_openviking_state_mutation", + "domain": args.domain, + "end_offset": args.end_offset, + "final_catalog_path": _relative(final_catalog_path.resolve()), + "memory_type": args.memory_type, + "run_id_prefix": args.run_id_prefix, + "schema_version": "openviking_tau2_category_batch_annotation_run.v0", + "seed_annotation_files": [_relative(path) for path in args.seed_annotations], + "start_offset": args.start_offset, + "status": "passed" if compaction_status == "passed" else "warning", + "warmup_count": args.warmup_count, + } + summary_path = summary_root / "run_summary.json" + summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True)) + return 0 if summary["status"] == "passed" else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/tau2/scripts/run_eval.py b/benchmark/tau2/scripts/run_eval.py index 5458ba61ac..0d2ddade27 100755 --- a/benchmark/tau2/scripts/run_eval.py +++ b/benchmark/tau2/scripts/run_eval.py @@ -4,26 +4,51 @@ import argparse import importlib.util import json +import os import subprocess import sys +from collections import Counter +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from typing import Any -from tau2_common import ( - domains, - load_config, - output_dir, - normalize_litellm_env, - run_id, - simulator_policy_report, - split_file, - strategy_ids, - tau2_cli, - tau2_context, - tau2_repo, - user_simulator_policy, - write_json, -) +try: + from tau2_common import ( + assert_tau2_results_complete, + domains, + load_config, + normalize_litellm_env, + output_dir, + resolve_path, + run_id, + simulator_policy_report, + split_file, + strategy_ids, + tau2_context, + tau2_repo, + user_simulator_policy, + write_json, + ) +except ModuleNotFoundError: # pragma: no cover - package import path + from .tau2_common import ( + assert_tau2_results_complete, + domains, + load_config, + normalize_litellm_env, + output_dir, + resolve_path, + run_id, + simulator_policy_report, + split_file, + strategy_ids, + tau2_context, + tau2_repo, + user_simulator_policy, + write_json, + ) + + +REPO_ROOT = Path(__file__).resolve().parents[3] def _reward(sim: dict[str, Any]) -> float: @@ -46,8 +71,127 @@ def _db_match(sim: dict[str, Any]) -> bool | None: return sim.get("db_match") +def _strategy_int( + config: dict[str, Any], + strategy: dict[str, Any], + key: str, + *, + fallback_key: str | None = None, + default: int = 4, +) -> int: + openviking = config.get("openviking", {}) + value = strategy.get(key) + if value is None: + value = openviking.get(key) + if value is None and fallback_key: + value = strategy.get(fallback_key) + if value is None and fallback_key: + value = openviking.get(fallback_key) + if value is None: + value = default + return int(value) + + +def _retrieval_budget(config: dict[str, Any], strategy: dict[str, Any]) -> dict[str, int]: + retrieval_top_k = _strategy_int(config, strategy, "retrieval_top_k", default=4) + first_user_retrieval_top_k = _strategy_int( + config, + strategy, + "first_user_retrieval_top_k", + fallback_key="retrieval_top_k", + default=retrieval_top_k, + ) + first_user_inject_top_k = _strategy_int( + config, + strategy, + "first_user_inject_top_k", + fallback_key="first_user_retrieval_top_k", + default=first_user_retrieval_top_k, + ) + prewrite_retrieval_top_k = _strategy_int( + config, + strategy, + "prewrite_retrieval_top_k", + fallback_key="retrieval_top_k", + default=retrieval_top_k, + ) + prewrite_inject_top_k = _strategy_int( + config, + strategy, + "prewrite_inject_top_k", + fallback_key="prewrite_retrieval_top_k", + default=prewrite_retrieval_top_k, + ) + return { + "retrieval_top_k": retrieval_top_k, + "first_user_retrieval_top_k": first_user_retrieval_top_k, + "first_user_inject_top_k": first_user_inject_top_k, + "prewrite_retrieval_top_k": prewrite_retrieval_top_k, + "prewrite_inject_top_k": prewrite_inject_top_k, + } + + +def _memory_corpus_key_for( + *, + domain: str, + strategy: dict[str, Any], + train_num_tasks: int | None, +) -> str: + corpus_id = str(strategy.get("corpus_id") or strategy["id"]) + raw_key = strategy.get("corpus_cache_key") + if raw_key: + key = str(raw_key).format( + domain=domain, + strategy_id=strategy["id"], + corpus_id=corpus_id, + ) + else: + key = f"{domain}_{corpus_id}" + if train_num_tasks is not None: + key = f"{key}_train{train_num_tasks}" + return key + + +def _memory_corpus_dir(config: dict[str, Any], configured_run_id: str, corpus_key: str) -> Path: + raw = config.get("paths", {}).get("corpus_cache_dir") + if raw: + return resolve_path(str(raw)) / corpus_key + return output_dir(config, configured_run_id) / "memory_corpora" / corpus_key + + +def _domain_value(value: Any, domain: str) -> Any: + if isinstance(value, dict): + return value.get(domain) or value.get(str(domain).lower()) or value.get("default") + return value + + +def _search_uri_suffix(search_memory_type: str, domain: str) -> str: + if search_memory_type in {"experiences", "trajectories"}: + return search_memory_type + raise ValueError(f"Unsupported search_memory_type: {search_memory_type}") + + +def _manifest_openviking_identity(corpus_dir: Path) -> dict[str, str] | None: + manifest_path = corpus_dir / "corpus_manifest.json" + if not manifest_path.is_file(): + return None + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + return None + openviking = manifest.get("openviking") or {} + required = ("account", "user", "agent_id", "search_uri") + if not all(openviking.get(key) for key in required): + return None + values = {key: str(openviking[key]) for key in required} + if openviking.get("url"): + values["url"] = str(openviking["url"]) + return values + + def _metrics_from_tau2_results(results_path: Path) -> dict[str, Any]: data = json.loads(results_path.read_text(encoding="utf-8")) + assert_tau2_results_complete(data, context=str(results_path)) sims = data.get("simulations") or [] rewards = [_reward(sim) for sim in sims] db_values = [_db_match(sim) for sim in sims] @@ -55,7 +199,9 @@ def _metrics_from_tau2_results(results_path: Path) -> dict[str, Any]: return { "simulation_count": len(sims), "avg_reward": sum(rewards) / len(rewards) if rewards else 0.0, - "db_match_rate": (sum(1 for value in db_known if value) / len(db_known)) if db_known else None, + "db_match_rate": (sum(1 for value in db_known if value) / len(db_known)) + if db_known + else None, } @@ -80,17 +226,52 @@ def _tau2_command( if reasoning_effort: agent_llm_args = f'{{"temperature":0.0,"reasoning_effort":"{reasoning_effort}"}}' user_llm_args = f'{{"temperature":0.0,"reasoning_effort":"{reasoning_effort}"}}' + fixed_first_user_file = _fixed_first_user_file(config, domain) + scope_prompt_file = _scope_prompt_file(config, strategy, domain) - if ( - strategy.get("memory_backend") == "openviking" - and strategy.get("train_memory_mode") == "experience_only" - ): + if strategy.get("memory_backend") == "openviking": openviking = config["openviking"] corpus_id = str(strategy.get("corpus_id") or strategy["id"]) - account = f"{openviking['account']}-{configured_run_id}-{domain}-{corpus_id}" - agent_id = f"{openviking['agent_id']}-{domain}-{corpus_id}" - user = f"tau2-{domain}-{corpus_id}" - search_uri = f"viking://agent/{agent_id}/memories/experiences" + resolved_train_num_tasks = ( + train_num_tasks if train_num_tasks is not None else strategy.get("train_num_tasks") + ) + corpus_key = _memory_corpus_key_for( + domain=domain, + strategy=strategy, + train_num_tasks=resolved_train_num_tasks, + ) + corpus_dir = _memory_corpus_dir(config, configured_run_id, corpus_key) + train_memory_mode = str(strategy.get("train_memory_mode") or "") + openviking_url = str(openviking["url"]) + if train_memory_mode == "experience_only": + reuse_identity = _manifest_openviking_identity(corpus_dir) + if reuse_identity is not None: + openviking_url = reuse_identity.get("url", openviking_url) + account = reuse_identity["account"] + agent_id = reuse_identity["agent_id"] + user = reuse_identity["user"] + search_uri = reuse_identity["search_uri"] + elif openviking.get("reuse_corpus_across_runs", False): + account = f"{openviking['account']}-{corpus_key}" + agent_id = f"{openviking['agent_id']}-{domain}-{corpus_id}" + user = f"tau2-{domain}-{corpus_id}" + search_uri = "" + else: + account = f"{openviking['account']}-{configured_run_id}-{domain}-{corpus_id}" + agent_id = f"{openviking['agent_id']}-{domain}-{corpus_id}" + user = f"tau2-{domain}-{corpus_id}" + search_uri = "" + else: + return None + search_memory_type = str(strategy.get("search_memory_type", "experiences")) + search_uri_suffix = _search_uri_suffix(search_memory_type, domain) + if not search_uri: + search_uri = f"viking://agent/{agent_id}/memories/{search_uri_suffix}" + budget = _retrieval_budget(config, strategy) + category_rerank = strategy.get("category_rerank") + category_rerank = category_rerank if isinstance(category_rerank, dict) else {} + scope_prompt = strategy.get("scope_prompt") + scope_prompt = scope_prompt if isinstance(scope_prompt, dict) else {} command = [ sys.executable, str(Path(__file__).with_name("run_memory_v2_eval.py")), @@ -99,11 +280,7 @@ def _tau2_command( "--run-dir", str(output_dir(config, configured_run_id) / "memory_cells" / run_label), "--corpus-dir", - str( - output_dir(config, configured_run_id) - / "memory_corpora" - / f"{domain}_{corpus_id}" - ), + str(corpus_dir), "--run-label", run_label, "--strategy-id", @@ -127,7 +304,7 @@ def _tau2_command( "--user-llm-args", user_llm_args, "--openviking-url", - str(openviking["url"]), + openviking_url, "--openviking-account", account, "--openviking-user", @@ -137,64 +314,125 @@ def _tau2_command( "--search-uri", search_uri, "--retrieval-top-k", - str(openviking.get("retrieval_top_k", 4)), + str(budget["retrieval_top_k"]), + "--first-user-retrieval-top-k", + str(budget["first_user_retrieval_top_k"]), + "--first-user-inject-top-k", + str(budget["first_user_inject_top_k"]), + "--prewrite-retrieval-top-k", + str(budget["prewrite_retrieval_top_k"]), + "--prewrite-inject-top-k", + str(budget["prewrite_inject_top_k"]), "--retrieval-mode", str(strategy.get("retrieval_mode", "first_user")), "--seed", str(seed), ] + if fixed_first_user_file is not None: + command.extend(["--fixed-first-user-file", str(fixed_first_user_file)]) + if category_rerank.get("enabled"): + command.extend( + [ + "--category-rerank-config", + json.dumps(category_rerank, ensure_ascii=False, sort_keys=True), + ] + ) + if scope_prompt_file is not None: + command.extend(["--scope-prompt-file", str(scope_prompt_file)]) + if scope_prompt.get("enabled"): + command.extend( + [ + "--scope-prompt-config", + json.dumps(scope_prompt, ensure_ascii=False, sort_keys=True), + ] + ) if task_ids: for task_id in task_ids: command.extend(["--task-id", task_id]) elif num_tasks is not None: command.extend(["--num-tasks", str(num_tasks)]) - train_num_tasks = train_num_tasks if train_num_tasks is not None else strategy.get("train_num_tasks") - if train_num_tasks is not None: - command.extend(["--train-num-tasks", str(train_num_tasks)]) + if resolved_train_num_tasks is not None: + command.extend(["--train-num-tasks", str(resolved_train_num_tasks)]) return command if strategy.get("memory_backend") != "none": return None command = [ - tau2_cli(config), - "run", + sys.executable, + str(Path(__file__).with_name("run_memory_v2_eval.py")), + "--tau2-repo", + str(tau2_repo(config)), + "--run-dir", + str(output_dir(config, configured_run_id) / "memory_cells" / run_label), + "--run-label", + run_label, + "--strategy-id", + strategy["id"], "--domain", domain, - "--agent", - str(benchmark.get("agent", "llm_agent")), - "--user", - str(benchmark.get("user", "user_simulator")), - "--task-split-name", + "--eval-split-name", str(benchmark.get("eval_split_name", "test")), - "--num-trials", - "1", "--max-steps", str(benchmark.get("max_steps", 200)), "--max-concurrency", str(benchmark.get("task_max_concurrency", 10)), + "--base-agent", + str(benchmark.get("agent", "llm_agent")), + "--user", + str(benchmark.get("user", "user_simulator")), "--agent-llm", str(model["agent_llm"]), "--user-llm", str(model["user_llm"]), - "--save-to", - run_label, + "--agent-llm-args", + agent_llm_args, + "--user-llm-args", + user_llm_args, "--seed", str(seed), + "--no-memory", ] - - command.extend(["--agent-llm-args", agent_llm_args]) - command.extend(["--user-llm-args", user_llm_args]) + if fixed_first_user_file is not None: + command.extend(["--fixed-first-user-file", str(fixed_first_user_file)]) if task_ids: - command.append("--task-ids") - command.extend(task_ids) + for task_id in task_ids: + command.extend(["--task-id", task_id]) elif num_tasks is not None: command.extend(["--num-tasks", str(num_tasks)]) return command +def _fixed_first_user_file(config: dict[str, Any], domain: str) -> Path | None: + raw = config.get("eval", {}).get("fixed_first_user_fixture") + if raw is None: + raw = config.get("eval", {}).get("fixed_first_user_fixtures") + if isinstance(raw, dict): + raw = raw.get(domain) or raw.get("default") + if raw is None or str(raw).strip() == "": + return None + return resolve_path(str(raw)) + + +def _scope_prompt_file( + config: dict[str, Any], strategy: dict[str, Any], domain: str +) -> Path | None: + raw = strategy.get("scope_prompt_file") + if raw is None: + raw = strategy.get("scope_prompt_files") + if raw is None: + raw = config.get("openviking", {}).get("scope_prompt_file") + if raw is None: + raw = config.get("openviking", {}).get("scope_prompt_files") + if isinstance(raw, dict): + raw = raw.get(domain) or raw.get("default") + if raw is None or str(raw).strip() == "": + return None + return resolve_path(str(raw)) + + def _build_plan( config: dict[str, Any], configured_run_id: str, @@ -205,16 +443,29 @@ def _build_plan( num_tasks: int | None, train_num_tasks: int | None, repeat_count_override: int | None, + cell_concurrency_override: int | None, + strategy_concurrency_override: int | None, ) -> dict[str, Any]: repeat_count = repeat_count_override or int(config["benchmark"].get("repeat_count", 8)) base_seed = int(config["benchmark"].get("seed", 300)) + cell_timeout_seconds = int(config["benchmark"].get("cell_timeout_seconds", 0) or 0) + strategy_concurrency = strategy_concurrency_override + if strategy_concurrency is None: + strategy_concurrency = cell_concurrency_override + if strategy_concurrency is None: + strategy_concurrency = config["benchmark"].get("strategy_concurrency") + if strategy_concurrency is None: + strategy_concurrency = config["benchmark"].get("cell_concurrency", 1) + strategy_concurrency = max(1, int(strategy_concurrency or 1)) policy_report = simulator_policy_report(config) strategies = config.get("strategies") or [] if selected_strategy_ids: unknown = selected_strategy_ids - set(strategy_ids(config)) if unknown: raise ValueError(f"unknown strategy ids: {sorted(unknown)}") - strategies = [strategy for strategy in strategies if strategy["id"] in selected_strategy_ids] + strategies = [ + strategy for strategy in strategies if strategy["id"] in selected_strategy_ids + ] cells = [] plan_domains = domains(config) if selected_domains: @@ -239,6 +490,8 @@ def _build_plan( train_num_tasks=train_num_tasks, seed=seed, ) + fixed_first_user_file = _fixed_first_user_file(config, domain) + scope_prompt_file = _scope_prompt_file(config, strategy, domain) non_executable_reason = None if command is None: non_executable_reason = ( @@ -254,13 +507,46 @@ def _build_plan( "seed": seed, "run_label": run_label, "train_required": bool(strategy.get("train_required")), + "train_memory_mode": strategy.get("train_memory_mode"), "memory_backend": strategy.get("memory_backend"), "corpus_id": strategy.get("corpus_id", strategy["id"]), + "corpus_key": _memory_corpus_key_for( + domain=domain, + strategy=strategy, + train_num_tasks=( + train_num_tasks + if train_num_tasks is not None + else strategy.get("train_num_tasks") + ), + ), + "corpus_dir": str( + _memory_corpus_dir( + config, + configured_run_id, + _memory_corpus_key_for( + domain=domain, + strategy=strategy, + train_num_tasks=( + train_num_tasks + if train_num_tasks is not None + else strategy.get("train_num_tasks") + ), + ), + ) + ), "retrieval_mode": strategy.get("retrieval_mode"), + "retrieval_budget": _retrieval_budget(config, strategy), + "search_memory_type": strategy.get("search_memory_type", "experiences"), + "category_rerank": strategy.get("category_rerank") or {"enabled": False}, + "scope_prompt": strategy.get("scope_prompt") or {"enabled": False}, "adapter_status": strategy.get("adapter_status", "ready"), "executable": command is not None, "user_simulator_policy": user_simulator_policy(config), "user_simulator_policy_supported": policy_report["supported"], + "fixed_first_user_file": str(fixed_first_user_file) + if fixed_first_user_file + else None, + "scope_prompt_file": str(scope_prompt_file) if scope_prompt_file else None, "split_file": str(split_path), "command": command, "non_executable_reason": non_executable_reason, @@ -278,31 +564,33 @@ def _build_plan( "cell_count": len(cells), "executable_cell_count": executable_cell_count, "pending_cell_count": len(cells) - executable_cell_count, + "corpus_prepare_concurrency": int(config["benchmark"].get("corpus_prepare_concurrency", 1)), + "strategy_concurrency": strategy_concurrency, + "cell_concurrency": strategy_concurrency, + "cell_timeout_seconds": cell_timeout_seconds or None, "cells": cells, } def _cell_artifacts(cell: dict[str, Any], repo: Path, out: Path) -> dict[str, str]: - if cell.get("memory_backend") == "openviking": + if cell.get("memory_backend") in {"openviking", "none"}: run_dir = out / "memory_cells" / cell["run_label"] - corpus_id = str(cell.get("corpus_id") or cell["strategy_id"]) - corpus_dir = out / "memory_corpora" / f"{cell['domain']}_{corpus_id}" - return { + artifacts = { "summary": str(run_dir / f"{cell['run_label']}.summary.json"), "results": str(run_dir / f"{cell['run_label']}.json"), - "retrieval_trace": str(run_dir / f"{cell['run_label']}.retrieval_trace.jsonl"), - "corpus_manifest": str(corpus_dir / "corpus_manifest.json"), } - return { - "results": str(repo / "data" / "simulations" / f"{cell['run_label']}.json") - } + if cell.get("memory_backend") == "none": + return artifacts + corpus_dir = Path(cell["corpus_dir"]) + artifacts["retrieval_trace"] = str(run_dir / f"{cell['run_label']}.retrieval_trace.jsonl") + artifacts["corpus_manifest"] = str(corpus_dir / "corpus_manifest.json") + return artifacts + return {"results": str(repo / "data" / "simulations" / f"{cell['run_label']}.json")} def _cell_metrics(cell: dict[str, Any], artifacts: dict[str, str]) -> dict[str, Any] | None: - if cell.get("memory_backend") == "openviking": - summary_path = Path(artifacts["summary"]) - if not summary_path.is_file(): - return None + summary_path = Path(artifacts.get("summary", "")) + if summary_path.is_file(): summary = json.loads(summary_path.read_text(encoding="utf-8")) return summary.get("metrics") @@ -312,22 +600,142 @@ def _cell_metrics(cell: dict[str, Any], artifacts: dict[str, str]) -> dict[str, return _metrics_from_tau2_results(results_path) +def _cell_runtime_evidence(cell: dict[str, Any], artifacts: dict[str, str]) -> dict[str, Any]: + if cell.get("memory_backend") == "openviking": + summary_path = Path(artifacts["summary"]) + if not summary_path.is_file(): + return {"status": "missing", "reasons": ["missing_summary"]} + summary = json.loads(summary_path.read_text(encoding="utf-8")) + evidence = summary.get("runtime_evidence") + if isinstance(evidence, dict): + return { + "status": str(evidence.get("status") or "valid"), + "reasons": list(evidence.get("reasons") or []), + } + return {"status": "valid", "reasons": []} + + +def _row_is_valid_evidence(row: dict[str, Any]) -> bool: + evidence = row.get("runtime_evidence") + if not isinstance(evidence, dict): + return True + return str(evidence.get("status") or "valid") == "valid" + + +def _memory_corpus_key(cell: dict[str, Any]) -> str: + return str(cell.get("corpus_key") or f"{cell['domain']}_{cell['corpus_id']}") + + +def _tau2_subprocess_env(repo: Path) -> dict[str, str]: + env = os.environ.copy() + src = repo / "src" + pythonpath_entry = str(src if src.is_dir() else repo) + existing = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + pythonpath_entry if not existing else f"{pythonpath_entry}{os.pathsep}{existing}" + ) + return env + + +def _prepare_memory_corpus(cell: dict[str, Any], repo: Path, out: Path) -> dict[str, Any]: + key = _memory_corpus_key(cell) + manifest_path = Path(cell["corpus_dir"]) / "corpus_manifest.json" + if manifest_path.is_file(): + row = { + "domain": cell["domain"], + "strategy_id": cell["strategy_id"], + "corpus_id": str(cell.get("corpus_id") or cell["strategy_id"]), + "corpus_key": key, + "returncode": 0, + "reused": True, + "artifacts": {"corpus_manifest": str(manifest_path)}, + } + write_json(out / "corpus_prepare_results" / f"{key}.json", row) + print(f"[tau2] reusing corpus {key}", flush=True) + return row + command = list(cell["command"]) + ["--prepare-corpus-only"] + print(f"[tau2] preparing corpus {key}", flush=True) + completed = subprocess.run( + command, + cwd=repo, + env=_tau2_subprocess_env(repo), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + row = { + "domain": cell["domain"], + "strategy_id": cell["strategy_id"], + "corpus_id": str(cell.get("corpus_id") or cell["strategy_id"]), + "corpus_key": key, + "returncode": completed.returncode, + "stdout_tail": completed.stdout[-4000:], + "stderr_tail": completed.stderr[-4000:], + "artifacts": {"corpus_manifest": str(Path(cell["corpus_dir"]) / "corpus_manifest.json")}, + } + write_json(out / "corpus_prepare_results" / f"{key}.json", row) + if completed.returncode != 0: + raise RuntimeError(f"corpus prepare failed: {key} returncode={completed.returncode}") + return row + + +def _prepare_memory_corpora(plan: dict[str, Any], repo: Path, out: Path) -> list[dict[str, Any]]: + corpus_cells: dict[str, dict[str, Any]] = {} + for cell in plan["cells"]: + if cell.get("memory_backend") != "openviking" or not bool(cell.get("train_required")): + continue + corpus_cells.setdefault(_memory_corpus_key(cell), cell) + if not corpus_cells: + return [] + + worker_count = max(1, int(plan.get("corpus_prepare_concurrency") or 1)) + if worker_count == 1 or len(corpus_cells) == 1: + return [_prepare_memory_corpus(cell, repo, out) for cell in corpus_cells.values()] + + rows: list[dict[str, Any]] = [] + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = { + executor.submit(_prepare_memory_corpus, cell, repo, out): key + for key, cell in corpus_cells.items() + } + for future in as_completed(futures): + rows.append(future.result()) + return rows + + def _summarize(rows: list[dict[str, Any]]) -> dict[str, Any]: def weighted(rows_for_group: list[dict[str, Any]]) -> dict[str, Any]: metric_rows = [row for row in rows_for_group if row.get("metrics")] + valid_metric_rows = [row for row in metric_rows if _row_is_valid_evidence(row)] + diagnostic_rows = [row for row in metric_rows if not _row_is_valid_evidence(row)] + diagnostic_reason_counts: Counter[str] = Counter() + for row in diagnostic_rows: + evidence = row.get("runtime_evidence") + evidence = evidence if isinstance(evidence, dict) else {} + reasons = list(evidence.get("reasons") or []) + if not reasons: + reasons = [str(evidence.get("status") or "diagnostic")] + for reason in reasons: + diagnostic_reason_counts[str(reason)] += 1 sim_count = sum(int(row["metrics"].get("simulation_count") or 0) for row in metric_rows) + valid_sim_count = sum( + int(row["metrics"].get("simulation_count") or 0) for row in valid_metric_rows + ) reward_sum = sum( float(row["metrics"].get("avg_reward") or 0.0) * int(row["metrics"].get("simulation_count") or 0) - for row in metric_rows + for row in valid_metric_rows ) db_weighted_rows = [ row - for row in metric_rows + for row in valid_metric_rows if row["metrics"].get("db_match_rate") is not None and int(row["metrics"].get("simulation_count") or 0) > 0 ] - db_weight = sum(int(row["metrics"].get("simulation_count") or 0) for row in db_weighted_rows) + db_weight = sum( + int(row["metrics"].get("simulation_count") or 0) for row in db_weighted_rows + ) db_sum = sum( float(row["metrics"]["db_match_rate"]) * int(row["metrics"].get("simulation_count") or 0) @@ -336,8 +744,12 @@ def weighted(rows_for_group: list[dict[str, Any]]) -> dict[str, Any]: return { "cell_count": len(rows_for_group), "completed_cell_count": len(metric_rows), - "simulation_count": sim_count, - "avg_reward": reward_sum / sim_count if sim_count else None, + "valid_completed_cell_count": len(valid_metric_rows), + "diagnostic_cell_count": len(diagnostic_rows), + "diagnostic_reason_counts": dict(sorted(diagnostic_reason_counts.items())), + "diagnostic_simulation_count": sim_count - valid_sim_count, + "simulation_count": valid_sim_count, + "avg_reward": reward_sum / valid_sim_count if valid_sim_count else None, "db_match_rate": db_sum / db_weight if db_weight else None, } @@ -367,46 +779,120 @@ def weighted(rows_for_group: list[dict[str, Any]]) -> dict[str, Any]: } -def _execute_cells(plan: dict[str, Any], repo: Path, out: Path) -> list[dict[str, Any]]: - policy_report = plan.get("simulator_policy") or {} - if not policy_report.get("supported", False): - raise RuntimeError( - "configured user simulator policy is not supported by this TAU-2 checkout: " - f"{policy_report}" - ) - rows = [] - for cell in plan["cells"]: - if not cell.get("executable"): - raise RuntimeError( - f"cell is not executable yet: {cell['run_label']} " - f"(strategy_id={cell['strategy_id']}, adapter_status={cell.get('adapter_status')})" - ) - print(f"[tau2] running {cell['run_label']}") +def _execute_cell( + cell: dict[str, Any], repo: Path, out: Path, cell_timeout: int | None +) -> dict[str, Any]: + cell_result_path = out / "cell_results" / f"{cell['run_label']}.json" + if cell_result_path.is_file(): + existing_row = json.loads(cell_result_path.read_text(encoding="utf-8")) + if existing_row.get("returncode") == 0 and existing_row.get("metrics"): + print(f"[tau2] skipping completed {cell['run_label']}", flush=True) + return existing_row + + print(f"[tau2] running {cell['run_label']}", flush=True) + try: completed = subprocess.run( cell["command"], cwd=repo, + env=_tau2_subprocess_env(repo), text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False, + timeout=cell_timeout, ) + except subprocess.TimeoutExpired as exc: + stdout = exc.stdout or "" + stderr = exc.stderr or "" + if isinstance(stdout, bytes): + stdout = stdout.decode(errors="replace") + if isinstance(stderr, bytes): + stderr = stderr.decode(errors="replace") row = { "run_label": cell["run_label"], "domain": cell["domain"], "strategy_id": cell["strategy_id"], - "returncode": completed.returncode, - "stdout_tail": completed.stdout[-4000:], - "stderr_tail": completed.stderr[-4000:], + "returncode": 124, + "timed_out": True, + "timeout_seconds": cell_timeout, + "stdout_tail": stdout[-4000:], + "stderr_tail": stderr[-4000:], + "artifacts": _cell_artifacts(cell, repo, out), + "metrics": None, } - row["artifacts"] = _cell_artifacts(cell, repo, out) + write_json(cell_result_path, row) + return row + + row = { + "run_label": cell["run_label"], + "domain": cell["domain"], + "strategy_id": cell["strategy_id"], + "returncode": completed.returncode, + "stdout_tail": completed.stdout[-4000:], + "stderr_tail": completed.stderr[-4000:], + } + row["artifacts"] = _cell_artifacts(cell, repo, out) + try: row["metrics"] = _cell_metrics(cell, row["artifacts"]) - rows.append(row) - write_json(out / "cell_results" / f"{cell['run_label']}.json", row) - if completed.returncode != 0: - raise RuntimeError(f"cell failed: {cell['run_label']} returncode={completed.returncode}") + except Exception as exc: + row["returncode"] = row["returncode"] or 1 + row["metrics"] = None + row["metrics_error"] = f"{type(exc).__name__}: {exc}" + try: + row["runtime_evidence"] = _cell_runtime_evidence(cell, row["artifacts"]) + except Exception as exc: + row["runtime_evidence"] = { + "status": "invalid", + "reasons": ["runtime_evidence_error"], + "error": f"{type(exc).__name__}: {exc}", + } + write_json(cell_result_path, row) + return row + + +def _execute_cells(plan: dict[str, Any], repo: Path, out: Path) -> list[dict[str, Any]]: + policy_report = plan.get("simulator_policy") or {} + if not policy_report.get("supported", False): + raise RuntimeError( + "configured user simulator policy is not supported by this TAU-2 checkout: " + f"{policy_report}" + ) + _prepare_memory_corpora(plan, repo, out) + cells = [] + for cell in plan["cells"]: + if not cell.get("executable"): + raise RuntimeError( + f"cell is not executable yet: {cell['run_label']} " + f"(strategy_id={cell['strategy_id']}, adapter_status={cell.get('adapter_status')})" + ) + cells.append(cell) + + cell_timeout = int(plan.get("cell_timeout_seconds") or 0) or None + worker_count = max( + 1, int(plan.get("strategy_concurrency") or plan.get("cell_concurrency") or 1) + ) + if worker_count == 1 or len(cells) == 1: + return [_execute_cell(cell, repo, out, cell_timeout) for cell in cells] + + print(f"[tau2] running eval cells with concurrency={worker_count}", flush=True) + rows: list[dict[str, Any]] = [] + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = { + executor.submit(_execute_cell, cell, repo, out, cell_timeout): cell for cell in cells + } + for future in as_completed(futures): + rows.append(future.result()) return rows +def _execution_failures(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + row + for row in rows + if row.get("returncode") != 0 or row.get("timed_out") or not row.get("metrics") + ] + + def _preflight(config: dict[str, Any], out: Path, *, strict: bool) -> int: errors: list[str] = [] llm_env = normalize_litellm_env() @@ -419,7 +905,9 @@ def _preflight(config: dict[str, Any], out: Path, *, strict: bool) -> int: if strict and not llm_env["has_api_key"]: errors.append("missing LLM API key: set OPENAI_API_KEY or ARK_API_KEY") if strict and not llm_env["has_base_url"]: - errors.append("missing OpenAI-compatible base URL: set OPENAI_API_BASE, OPENAI_BASE_URL, or ARK_BASE_URL") + errors.append( + "missing OpenAI-compatible base URL: set OPENAI_API_BASE, OPENAI_BASE_URL, or ARK_BASE_URL" + ) if strict and not policy_report["supported"]: errors.append( "configured confirmation-aware user simulator policy requires a TAU-2 " @@ -440,6 +928,90 @@ def _preflight(config: dict[str, Any], out: Path, *, strict: bool) -> int: if strict and not ok: errors.append(f"missing Python module: {module}") + category_rows = [] + for strategy in config.get("strategies") or []: + category_rerank = strategy.get("category_rerank") + if not isinstance(category_rerank, dict) or not category_rerank.get("enabled"): + continue + raw_annotation_files = category_rerank.get("annotation_files") or category_rerank.get( + "category_annotation_files" + ) + annotation_values = [] + if isinstance(raw_annotation_files, dict): + for item in raw_annotation_files.values(): + annotation_values.extend(item if isinstance(item, list) else [item]) + elif isinstance(raw_annotation_files, list): + annotation_values.extend(raw_annotation_files) + elif raw_annotation_files: + annotation_values.append(raw_annotation_files) + annotation_paths = [] + for value in annotation_values: + path = Path(str(value)).expanduser() + if not path.is_absolute(): + path = REPO_ROOT / path + annotation_paths.append(path) + missing_annotation_paths = [str(path) for path in annotation_paths if not path.is_file()] + category_rows.append( + { + "strategy_id": strategy.get("id"), + "annotation_files": [str(path) for path in annotation_paths], + "exists": bool(annotation_paths) and not missing_annotation_paths, + "missing_annotation_files": missing_annotation_paths, + "apply_nodes": category_rerank.get("apply_nodes"), + "retrieve_limit": category_rerank.get("retrieve_limit"), + "inject_limit": category_rerank.get("inject_limit"), + } + ) + if strict and (not annotation_paths or missing_annotation_paths): + errors.append( + f"missing category rerank annotation sidecar for {strategy.get('id')}: " + f"{missing_annotation_paths or raw_annotation_files}" + ) + + scope_prompt_rows = [] + for strategy in config.get("strategies") or []: + scope_prompt = strategy.get("scope_prompt") + direct_scope_prompt = ( + strategy.get("scope_prompt_file") + or strategy.get("scope_prompt_files") + or config.get("openviking", {}).get("scope_prompt_file") + or config.get("openviking", {}).get("scope_prompt_files") + ) + has_scope_prompt_config = isinstance(scope_prompt, dict) and scope_prompt.get("enabled") + if not has_scope_prompt_config and not direct_scope_prompt: + continue + scope_prompt = scope_prompt if isinstance(scope_prompt, dict) else {} + domain_files = scope_prompt.get("domain_files") + domain_files = domain_files if isinstance(domain_files, dict) else {} + domain_texts = scope_prompt.get("domain_texts") + domain_texts = domain_texts if isinstance(domain_texts, dict) else {} + for domain in domains(config): + raw_prompt_path = domain_files.get(domain) + if raw_prompt_path is None: + raw_prompt_path = _domain_value(direct_scope_prompt, domain) + prompt_path = None + exists = False + if raw_prompt_path: + prompt_path = Path(str(raw_prompt_path)).expanduser() + if not prompt_path.is_absolute(): + prompt_path = REPO_ROOT / prompt_path + exists = prompt_path.is_file() + if strict and not exists: + errors.append( + f"missing scope prompt file for {strategy.get('id')} {domain}: " + f"{raw_prompt_path}" + ) + scope_prompt_rows.append( + { + "strategy_id": strategy.get("id"), + "domain": domain, + "configured": bool(raw_prompt_path or domain_texts.get(domain)), + "prompt_path": str(prompt_path) if prompt_path else None, + "exists": exists, + "injection_point": scope_prompt.get("injection_point", "system_prompt"), + } + ) + report = { "status": "failed" if errors else "ok", "strict": strict, @@ -449,6 +1021,8 @@ def _preflight(config: dict[str, Any], out: Path, *, strict: bool) -> int: "domains": domains(config), "strategies": strategy_ids(config), "imports": import_rows, + "category_rerank_sidecars": category_rows, + "scope_prompts": scope_prompt_rows, "split_files": split_rows, "errors": errors, } @@ -463,14 +1037,43 @@ def _preflight(config: dict[str, Any], out: Path, *, strict: bool) -> int: def main() -> int: parser = argparse.ArgumentParser(description="Plan or run TAU-2 benchmark cells.") - parser.add_argument("--config", type=Path, default=Path(__file__).parents[1] / "config" / "baseline.yaml") + parser.add_argument( + "--config", type=Path, default=Path(__file__).parents[1] / "config" / "baseline.yaml" + ) parser.add_argument("--run-id", default=run_id()) - parser.add_argument("--domain", action="append", help="Run only this configured domain; may be repeated.") - parser.add_argument("--repeat-count", type=int, help="Override benchmark.repeat_count for smoke runs.") - parser.add_argument("--strategy-id", action="append", help="Run only this strategy id; may be repeated.") - parser.add_argument("--task-id", action="append", help="Run only this TAU-2 task id; may be repeated.") - parser.add_argument("--num-tasks", type=int, help="Run the first N tasks from the selected split.") - parser.add_argument("--train-num-tasks", type=int, help="Train OpenViking memory on the first N train tasks.") + parser.add_argument( + "--domain", action="append", help="Run only this configured domain; may be repeated." + ) + parser.add_argument( + "--repeat-count", type=int, help="Override benchmark.repeat_count for smoke runs." + ) + parser.add_argument( + "--cell-concurrency", + type=int, + help="Deprecated alias for --strategy-concurrency.", + ) + parser.add_argument( + "--strategy-concurrency", + type=int, + help="Override benchmark.strategy_concurrency for parallel matrix cells.", + ) + parser.add_argument( + "--task-max-concurrency", + type=int, + help="Override benchmark.task_max_concurrency inside each TAU-2 cell.", + ) + parser.add_argument( + "--strategy-id", action="append", help="Run only this strategy id; may be repeated." + ) + parser.add_argument( + "--task-id", action="append", help="Run only this TAU-2 task id; may be repeated." + ) + parser.add_argument( + "--num-tasks", type=int, help="Run the first N tasks from the selected split." + ) + parser.add_argument( + "--train-num-tasks", type=int, help="Train OpenViking memory on the first N train tasks." + ) parser.add_argument( "--preflight", action="store_true", @@ -488,8 +1091,16 @@ def main() -> int: if args.plan_only and args.execute: raise SystemExit("--plan-only and --execute are mutually exclusive") + if args.cell_concurrency is not None and args.cell_concurrency < 1: + raise SystemExit("--cell-concurrency must be >= 1") + if args.strategy_concurrency is not None and args.strategy_concurrency < 1: + raise SystemExit("--strategy-concurrency must be >= 1") + if args.task_max_concurrency is not None and args.task_max_concurrency < 1: + raise SystemExit("--task-max-concurrency must be >= 1") config = load_config(args.config) + if args.task_max_concurrency is not None: + config.setdefault("benchmark", {})["task_max_concurrency"] = args.task_max_concurrency out = output_dir(config, args.run_id) out.mkdir(parents=True, exist_ok=True) if args.preflight or args.strict_preflight: @@ -506,6 +1117,8 @@ def main() -> int: num_tasks=args.num_tasks, train_num_tasks=args.train_num_tasks, repeat_count_override=args.repeat_count, + cell_concurrency_override=args.cell_concurrency, + strategy_concurrency_override=args.strategy_concurrency, ) write_json(out / "run_plan.json", plan) write_json(out / "resolved_config.json", config) @@ -514,10 +1127,15 @@ def main() -> int: if args.execute: try: rows = _execute_cells(plan, tau2_repo(config), out) - plan["status"] = "succeeded" + failures = _execution_failures(rows) + plan["status"] = "failed" if failures else "succeeded" plan["executed_cell_count"] = len(rows) + plan["failed_cell_count"] = len(failures) write_json(out / "run_plan.json", plan) write_json(out / "scoreboard.json", _summarize(rows)) + if failures: + labels = ", ".join(str(row.get("run_label")) for row in failures[:5]) + raise RuntimeError(f"{len(failures)} cell(s) failed or incomplete: {labels}") except Exception as exc: plan["status"] = "failed" plan["error"] = str(exc) diff --git a/benchmark/tau2/scripts/run_memory_v2_eval.py b/benchmark/tau2/scripts/run_memory_v2_eval.py index de5ef54411..a984783f99 100644 --- a/benchmark/tau2/scripts/run_memory_v2_eval.py +++ b/benchmark/tau2/scripts/run_memory_v2_eval.py @@ -2,15 +2,26 @@ from __future__ import annotations import argparse +import hashlib +import importlib import json import shutil import sys import time +from collections import Counter +from copy import deepcopy from pathlib import Path from typing import Any -from tau2_common import normalize_litellm_env +try: + from category_rerank import CategoryReranker +except ModuleNotFoundError: # pragma: no cover - package import path + from .category_rerank import CategoryReranker +try: + from tau2_common import assert_tau2_results_complete, normalize_litellm_env +except ModuleNotFoundError: # pragma: no cover - package import path + from .tau2_common import assert_tau2_results_complete, normalize_litellm_env AGENT_NAME = "openviking_memory_agent" REPO_ROOT = Path(__file__).resolve().parents[3] @@ -29,12 +40,21 @@ "grant_", "reboot_", ) +FIXED_FIRST_USER_NAME = "openviking_fixed_first_user_simulator" def _json(text: str) -> dict[str, Any]: return json.loads(text) if text else {} +def _as_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + def _write_json(path: Path, payload: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n") @@ -46,6 +66,27 @@ def _add_tau2_to_path(tau2_repo: Path) -> None: sys.path.insert(0, str(src if src.is_dir() else tau2_repo)) +def _patch_tau2_auxiliary_llm_defaults(llm: str, llm_args: dict[str, Any]) -> None: + # TAU-2 exposes agent/user LLMs in TextRunConfig, but NL assertion scoring + # still reads module defaults. Keep the evaluator on the same configured + # model so benchmark runs do not fall back to inaccessible upstream defaults. + patches = { + "DEFAULT_LLM_NL_ASSERTIONS": llm, + "DEFAULT_LLM_NL_ASSERTIONS_ARGS": deepcopy(llm_args), + "DEFAULT_LLM_ENV_INTERFACE": llm, + "DEFAULT_LLM_ENV_INTERFACE_ARGS": deepcopy(llm_args), + } + for module_name in ( + "tau2.config", + "tau2.evaluator.evaluator_nl_assertions", + "tau2.environment.utils.interface_agent", + ): + module = importlib.import_module(module_name) + for name, value in patches.items(): + if hasattr(module, name): + setattr(module, name, deepcopy(value)) + + def _save_to_arg(path: Path) -> str: # Some TAU-2 versions append ".json"; newer versions treat save_to as a # run directory and write results.json under it. @@ -57,6 +98,70 @@ def _compat_results_path(path: Path) -> Path: return run_dir / "results.json" +def _resolve_repo_path(raw_path: Any, *, repo_root: Path) -> Path: + path = Path(str(raw_path)).expanduser() + if not path.is_absolute(): + path = repo_root / path + return path + + +def _domain_value(mapping: Any, domain: str) -> Any: + if isinstance(mapping, dict): + return mapping.get(domain) or mapping.get(str(domain).lower()) + return None + + +def _load_scope_prompt( + payload: dict[str, Any] | None, + *, + domain: str, + repo_root: Path, +) -> tuple[str, dict[str, Any]]: + payload = payload if isinstance(payload, dict) else {} + enabled = _as_bool(payload.get("enabled"), default=False) + summary: dict[str, Any] = { + "enabled": enabled, + "domain": domain, + "injection_point": str(payload.get("injection_point") or "system_prompt"), + "loaded": False, + "loaded_files": [], + "text_chars": 0, + } + if not enabled: + summary["skipped_reason"] = "disabled" + return "", summary + + text = str(_domain_value(payload.get("domain_texts"), domain) or "").strip() + raw_path = _domain_value(payload.get("domain_files"), domain) + if raw_path: + path = _resolve_repo_path(raw_path, repo_root=repo_root) + summary["loaded_files"] = [str(path)] + if not path.is_file(): + raise FileNotFoundError(f"scope prompt file not found for {domain}: {path}") + text = path.read_text(encoding="utf-8").strip() + + if not text: + summary["skipped_reason"] = "no_domain_scope_prompt" + return "", summary + + summary["loaded"] = True + summary["text_chars"] = len(text) + return text, summary + + +def _scope_prompt_text(prompt: str) -> str: + if not prompt.strip(): + return "" + return ( + "Use this OpenViking memory applicability guard together with retrieved " + "memories. Current tool observations and the current user request remain " + "authoritative.\n\n" + "\n" + f"{prompt.strip()}\n" + "" + ) + + def _reward(sim: dict[str, Any]) -> float: info = sim.get("reward_info") or {} value = info.get("reward", sim.get("reward", 0.0)) @@ -86,7 +191,302 @@ def _metrics(results_path: Path) -> dict[str, Any]: return { "simulation_count": len(sims), "avg_reward": sum(rewards) / len(rewards) if rewards else 0.0, - "db_match_rate": (sum(1 for value in db_known if value) / len(db_known)) if db_known else None, + "db_match_rate": (sum(1 for value in db_known if value) / len(db_known)) + if db_known + else None, + } + + +def _is_aggregate_memory_uri(uri: Any) -> bool: + value = str(uri or "").split("#", 1)[0] + return value.endswith("/.overview.md") or value.endswith("/.abstract.md") + + +def _trace_category_summary(trace_path: Path) -> dict[str, Any]: + counters: Counter[str] = Counter() + decision_nodes: Counter[str] = Counter() + category_decisions: Counter[str] = Counter() + query_sidecar_coverage: Counter[str] = Counter() + query_category_sources: Counter[str] = Counter() + memory_category_sources: Counter[str] = Counter() + selected_memory_category_sources: Counter[str] = Counter() + tool_calls: Counter[str] = Counter() + trace_rows = 0 + category_event_count = 0 + + if not trace_path.is_file(): + return { + "trace_present": False, + "trace_rows": 0, + "category_event_count": 0, + } + + for line_number, line in enumerate(trace_path.read_text(encoding="utf-8").splitlines(), 1): + if not line.strip(): + continue + trace_rows += 1 + try: + row = json.loads(line) + except json.JSONDecodeError: + counters["json_decode_error_count"] += 1 + counters[f"json_decode_error_line:{line_number}"] += 1 + continue + if not isinstance(row, dict): + counters["non_object_row_count"] += 1 + continue + + decision_nodes[str(row.get("decision_node") or "unknown")] += 1 + injected_count = int(row.get("injected_count") or 0) + if str(row.get("retrieval_action_taken") or "") == "retrieve_and_inject" and ( + row.get("injected") or injected_count > 0 + ): + counters["memory_injection_event_count"] += 1 + counters["memory_injected_count"] += injected_count + for call in row.get("tool_calls") or []: + if isinstance(call, dict) and call.get("name"): + tool_calls[str(call["name"])] += 1 + + category = ( + row.get("category_rerank") if isinstance(row.get("category_rerank"), dict) else {} + ) + if category: + category_event_count += 1 + if category.get("enabled"): + counters["category_enabled_event_count"] += 1 + if category.get("applied"): + counters["category_applied_event_count"] += 1 + if category.get("decision"): + category_decisions[str(category["decision"])] += 1 + coverage = category.get("query_sidecar_coverage") + if coverage: + coverage_key = str(coverage) + query_sidecar_coverage[coverage_key] += 1 + counters[f"query_sidecar_{coverage_key}_event_count"] += 1 + query_category = ( + category.get("query_category") + if isinstance(category.get("query_category"), dict) + else {} + ) + if query_category.get("category_source"): + query_category_sources[str(query_category["category_source"])] += 1 + if query_category.get("matched"): + counters["query_category_matched_event_count"] += 1 + + for match in row.get("matches") or []: + if not isinstance(match, dict): + continue + counters["raw_match_count"] += 1 + selected = bool(match.get("selected_for_injection") or match.get("injected")) + injected = bool(match.get("injected")) + if selected: + counters["selected_match_count"] += 1 + if injected: + counters["injected_match_count"] += 1 + is_aggregate = _is_aggregate_memory_uri(match.get("uri")) + if is_aggregate: + counters["aggregate_memory_candidate_count"] += 1 + if selected: + counters["selected_aggregate_memory_count"] += 1 + else: + counters["concrete_memory_candidate_count"] += 1 + if selected: + counters["selected_concrete_memory_count"] += 1 + if injected: + counters["injected_concrete_memory_count"] += 1 + memory_source = match.get("memory_category_source_prompt") + positive_category_match = bool( + match.get("category1_match") or match.get("category2_match") + ) + if memory_source: + counters["memory_category_present_count"] += 1 + memory_category_sources[str(memory_source)] += 1 + if selected: + counters["selected_memory_category_present_count"] += 1 + selected_memory_category_sources[str(memory_source)] += 1 + if positive_category_match: + counters["memory_category_matched_count"] += 1 + if selected: + counters["selected_memory_category_matched_count"] += 1 + elif match.get("category_rerank_reasons") is not None: + counters["memory_category_missing_count"] += 1 + if positive_category_match: + counters["positive_category_match_count"] += 1 + if selected: + counters["selected_positive_category_match_count"] += 1 + if injected: + counters["injected_positive_category_match_count"] += 1 + if not is_aggregate: + counters["injected_concrete_positive_category_match_count"] += 1 + + raw_count = counters["raw_match_count"] + selected_count = counters["selected_match_count"] + injected_count = counters["injected_match_count"] + for key in [ + "aggregate_memory_candidate_count", + "concrete_memory_candidate_count", + "memory_injection_event_count", + "memory_injected_count", + "injected_match_count", + "selected_aggregate_memory_count", + "selected_concrete_memory_count", + "injected_concrete_memory_count", + "memory_category_matched_count", + "selected_memory_category_matched_count", + "injected_positive_category_match_count", + "injected_concrete_positive_category_match_count", + "query_sidecar_covered_event_count", + "query_sidecar_partial_event_count", + "query_sidecar_missing_event_count", + ]: + counters[key] += 0 + sidecar_non_missing_count = ( + counters["query_sidecar_covered_event_count"] + + counters["query_sidecar_partial_event_count"] + ) + return { + "trace_present": True, + "trace_rows": trace_rows, + "category_event_count": category_event_count, + "counts": dict(counters), + "decision_nodes": dict(decision_nodes), + "category_decisions": dict(category_decisions), + "query_sidecar_coverage": dict(query_sidecar_coverage), + "query_category_sources": dict(query_category_sources), + "memory_category_sources": dict(memory_category_sources), + "selected_memory_category_sources": dict(selected_memory_category_sources), + "tool_calls": dict(tool_calls), + "rates": { + "memory_category_candidate_coverage": ( + counters["memory_category_present_count"] / raw_count if raw_count else None + ), + "selected_memory_category_coverage": ( + counters["selected_memory_category_present_count"] / selected_count + if selected_count + else None + ), + "memory_category_match_coverage": ( + counters["memory_category_matched_count"] / raw_count if raw_count else None + ), + "selected_memory_category_match_coverage": ( + counters["selected_memory_category_matched_count"] / selected_count + if selected_count + else None + ), + "selected_positive_category_match_rate": ( + counters["selected_positive_category_match_count"] / selected_count + if selected_count + else None + ), + "injected_positive_category_match_rate": ( + counters["injected_positive_category_match_count"] / injected_count + if injected_count + else None + ), + "injected_concrete_positive_category_match_rate": ( + counters["injected_concrete_positive_category_match_count"] / injected_count + if injected_count + else None + ), + "concrete_memory_candidate_rate": ( + counters["concrete_memory_candidate_count"] / raw_count if raw_count else None + ), + "selected_concrete_memory_rate": ( + counters["selected_concrete_memory_count"] / selected_count + if selected_count + else None + ), + "injected_concrete_memory_rate": ( + counters["injected_concrete_memory_count"] / injected_count + if injected_count + else None + ), + "query_sidecar_non_missing_event_rate": ( + sidecar_non_missing_count / category_event_count if category_event_count else None + ), + "query_sidecar_full_event_rate": ( + counters["query_sidecar_covered_event_count"] / category_event_count + if category_event_count + else None + ), + }, + } + + +def _runtime_evidence_status( + *, + category_rerank: dict[str, Any], + retrieval_trace_summary: dict[str, Any], + corpus_probe: dict[str, Any] | None = None, +) -> dict[str, Any]: + reasons: list[str] = [] + if category_rerank.get("enabled"): + corpus_probe = corpus_probe if isinstance(corpus_probe, dict) else {} + probe_match_count = int(corpus_probe.get("match_count") or 0) + probe_concrete_match_count = int( + corpus_probe.get("concrete_match_count") + or corpus_probe.get("read_non_empty_count") + or 0 + ) + probe_aggregate_match_count = int(corpus_probe.get("aggregate_match_count") or 0) + if not corpus_probe: + reasons.append("missing_corpus_probe") + if corpus_probe and probe_match_count <= 0: + reasons.append("empty_corpus_probe") + if probe_match_count > 0: + if probe_concrete_match_count <= 0: + reasons.append("no_concrete_corpus_probe_matches") + if probe_aggregate_match_count == probe_match_count: + reasons.append("aggregate_only_corpus_probe") + + if not retrieval_trace_summary.get("trace_present"): + reasons.append("missing_retrieval_trace") + counts = ( + retrieval_trace_summary.get("counts") + if isinstance(retrieval_trace_summary.get("counts"), dict) + else {} + ) + rates = ( + retrieval_trace_summary.get("rates") + if isinstance(retrieval_trace_summary.get("rates"), dict) + else {} + ) + applied_count = int(counts.get("category_applied_event_count") or 0) + if retrieval_trace_summary.get("trace_present"): + if int(retrieval_trace_summary.get("category_event_count") or 0) <= 0: + reasons.append("no_category_rerank_events") + elif applied_count <= 0: + reasons.append("no_category_rerank_applied_events") + if applied_count > 0: + if int(counts.get("query_category_matched_event_count") or 0) <= 0: + reasons.append("no_query_category_coverage") + if float(rates.get("concrete_memory_candidate_rate") or 0.0) <= 0.0: + reasons.append("no_concrete_memory_candidates") + if int(counts.get("memory_category_present_count") or 0) <= 0: + reasons.append("no_memory_category_coverage") + if int(counts.get("memory_category_matched_count") or 0) <= 0: + reasons.append("no_matched_memory_categories") + if int(counts.get("memory_injection_event_count") or 0) <= 0: + reasons.append("no_memory_injection") + if ( + int(counts.get("memory_injection_event_count") or 0) > 0 + and float(rates.get("injected_concrete_memory_rate") or 0.0) <= 0.0 + ): + reasons.append("no_injected_concrete_memory") + if ( + int(counts.get("query_category_matched_event_count") or 0) > 0 + and float(rates.get("selected_positive_category_match_rate") or 0.0) <= 0.0 + ): + reasons.append("no_selected_positive_category_match") + if ( + int(counts.get("query_category_matched_event_count") or 0) > 0 + and int(counts.get("memory_injection_event_count") or 0) > 0 + and int(counts.get("injected_concrete_positive_category_match_count") or 0) <= 0 + ): + reasons.append("no_injected_concrete_positive_category_match") + + return { + "status": "diagnostic" if reasons else "valid", + "reasons": reasons, } @@ -118,12 +518,14 @@ def _tool_call_query(tool_calls: list[Any], state_messages: list[Any]) -> str: recent_user = [ str(getattr(message, "content", "") or "") for message in state_messages[-8:] - if str(getattr(message, "role", "")) == "user" and str(getattr(message, "content", "") or "").strip() + if str(getattr(message, "role", "")) == "user" + and str(getattr(message, "content", "") or "").strip() ] recent_observations = [ str(getattr(message, "content", "") or "")[:600] for message in state_messages[-12:] - if str(getattr(message, "role", "")) == "tool" and str(getattr(message, "content", "") or "").strip() + if str(getattr(message, "role", "")) == "tool" + and str(getattr(message, "content", "") or "").strip() ] parts = [ "Before executing write-like tool call(s): " + "; ".join(rendered), @@ -151,6 +553,69 @@ def _message_text(message: dict[str, Any]) -> tuple[str, str]: return "assistant", str(message.get("content") or "") +def _scenario_sha256(instructions: str) -> str: + return hashlib.sha256(instructions.encode("utf-8")).hexdigest() + + +def _load_fixed_first_user_fixture(path: Path) -> dict[str, str]: + if not path.is_file(): + raise FileNotFoundError(f"fixed-first-user fixture not found: {path}") + data = json.loads(path.read_text(encoding="utf-8")) + mapping = data.get("by_scenario_sha256") if isinstance(data, dict) else None + if not isinstance(mapping, dict) or not mapping: + raise ValueError(f"fixed-first-user fixture has no by_scenario_sha256 map: {path}") + return {str(key): str(value) for key, value in mapping.items()} + + +def _has_user_message(state: Any) -> bool: + for message in getattr(state, "messages", []) or []: + role = getattr(message, "role", None) + if str(getattr(role, "value", role)) == "user": + return True + return False + + +def _append_incoming_user_context(message: Any, state: Any) -> None: + from tau2.data_model.message import AssistantMessage, MultiToolMessage, ToolMessage + + if isinstance(message, MultiToolMessage): + state.messages.extend(message.tool_messages) + elif isinstance(message, ToolMessage): + state.messages.append(message) + elif isinstance(message, AssistantMessage) and ( + message.has_content() or message.is_tool_call() + ): + state.messages.append(message) + + +def _register_fixed_first_user(args: argparse.Namespace) -> str: + if not args.fixed_first_user_file: + return args.user + _add_tau2_to_path(args.tau2_repo) + mapping = _load_fixed_first_user_fixture(args.fixed_first_user_file) + + from tau2.data_model.message import UserMessage + from tau2.registry import registry + from tau2.user.user_simulator import UserSimulator + + class FixedFirstUserSimulator(UserSimulator): # type: ignore[misc] + def _generate_next_message(self, message: Any, state: Any) -> UserMessage: # type: ignore[override] + if not _has_user_message(state): + key = _scenario_sha256(str(self.instructions or "")) + fixed = mapping.get(key) + if fixed is None: + raise RuntimeError( + f"fixed-first-user fixture does not cover this TAU-2 scenario: sha256={key}" + ) + _append_incoming_user_context(message, state) + return UserMessage(role="user", content=fixed) + return super()._generate_next_message(message, state) + + if FIXED_FIRST_USER_NAME not in registry.get_users(): + registry.register_user(FixedFirstUserSimulator, FIXED_FIRST_USER_NAME) + return FIXED_FIRST_USER_NAME + + def _run_tau2( *, tau2_repo: Path, @@ -171,6 +636,7 @@ def _run_tau2( save_to: Path, ): _add_tau2_to_path(tau2_repo) + _patch_tau2_auxiliary_llm_defaults(agent_llm, agent_llm_args) from tau2.data_model.simulation import RunConfig, TextRunConfig from tau2.run import run_domain @@ -246,29 +712,49 @@ def _read_memory_text(client: Any, match: Any) -> tuple[str, str | None]: def _probe_corpus(args: argparse.Namespace, client: Any) -> dict[str, Any]: + probe_limit = args.retrieval_top_k + if hasattr(args, "category_reranker"): + probe_limit = args.category_reranker.search_limit( + probe_limit, + decision_node="before_write_tool_call", + ) result = client.search( query=f"{args.domain} customer service order reservation booking cancellation exchange return update", target_uri=args.search_uri, - limit=args.retrieval_top_k, + limit=probe_limit, ) memories = list(getattr(result, "memories", []) or []) reads = [] - for match in memories[: args.retrieval_top_k]: + for match in memories[:probe_limit]: uri = getattr(match, "uri", "") text, read_error = _read_memory_text(client, match) + is_aggregate = _is_aggregate_memory_uri(uri) row = { "uri": uri, "score": getattr(match, "score", None), "text_chars": len(text), "non_empty": bool(str(text).strip()), + "is_aggregate_memory": is_aggregate, + "is_concrete_memory": not is_aggregate, } if read_error: row["read_error"] = read_error reads.append(row) + aggregate_match_count = sum(1 for row in reads if row["is_aggregate_memory"]) + concrete_match_count = sum(1 for row in reads if row["is_concrete_memory"]) return { "query": f"{args.domain} customer service order reservation booking cancellation exchange return update", + "probe_limit": probe_limit, "match_count": len(memories), + "aggregate_match_count": aggregate_match_count, + "concrete_match_count": concrete_match_count, "read_non_empty_count": sum(1 for row in reads if row["non_empty"]), + "aggregate_read_non_empty_count": sum( + 1 for row in reads if row["is_aggregate_memory"] and row["non_empty"] + ), + "concrete_read_non_empty_count": sum( + 1 for row in reads if row["is_concrete_memory"] and row["non_empty"] + ), "matches": reads, } @@ -297,11 +783,14 @@ def _train(args: argparse.Namespace, train_results: Path, corpus_manifest: Path) ) data = json.loads(train_results.read_text()) + assert_tau2_results_complete(data, context=f"{args.domain} train") client = _client(args) committed = [] try: for sim in data.get("simulations") or []: - session_id = f"tau2-{args.domain}-train-{sim.get('task_id')}-trial-{sim.get('trial', 0)}" + session_id = ( + f"tau2-{args.domain}-train-{sim.get('task_id')}-trial-{sim.get('trial', 0)}" + ) created = client.create_session(session_id=session_id) sid = created.get("session_id", session_id) for msg in sim.get("messages") or []: @@ -363,20 +852,44 @@ def _register_memory_agent(args: argparse.Namespace, trace_path: Path) -> None: class OpenVikingMemoryAgent(LLMAgent): def get_init_state(self, message_history=None): state = super().get_init_state(message_history) + scope_prompt = _scope_prompt_text(args.scope_prompt_text) + if scope_prompt: + state.system_messages.append(SystemMessage(role="system", content=scope_prompt)) + self._trace( + { + "decision_node": "static_scope_prompt", + "retrieval_action_taken": "scope_prompt_static_injection", + "scope_prompt": args.scope_prompt_summary, + "injected": True, + "injected_count": 1, + } + ) if args.retrieval_mode in {"first_user", "first_user_prewrite"}: state.system_messages.append( SystemMessage(role="system", content="") ) return state - def _retrieve(self, query: str) -> tuple[str, list[dict[str, Any]]]: + def _retrieve( + self, + query: str, + *, + decision_node: str, + search_limit: int, + inject_limit: int, + ) -> tuple[str, list[dict[str, Any]], dict[str, Any]]: client = _client(args) rows: list[dict[str, Any]] = [] try: - result = client.search(query=query, target_uri=args.search_uri, limit=args.retrieval_top_k) + effective_search_limit = args.category_reranker.search_limit( + search_limit, + decision_node=decision_node, + ) + result = client.search( + query=query, target_uri=args.search_uri, limit=effective_search_limit + ) memories = list(getattr(result, "memories", []) or []) - blocks = [] - for index, match in enumerate(memories[: args.retrieval_top_k], 1): + for _index, match in enumerate(memories[:effective_search_limit], 1): uri = getattr(match, "uri", "") text, read_error = _read_memory_text(client, match) row = { @@ -384,13 +897,24 @@ def _retrieve(self, query: str) -> tuple[str, list[dict[str, Any]]]: "score": getattr(match, "score", None), "level": getattr(match, "level", None), "text_chars": len(text), + "_text": text, } if read_error: row["read_error"] = read_error rows.append(row) + selected_rows, trace_rows, category_rerank = args.category_reranker.select( + domain=args.domain, + query=query, + rows=rows, + decision_node=decision_node, + base_limit=inject_limit, + ) + blocks = [] + for index, row in enumerate(selected_rows, 1): + text = str(row.get("_text") or "") if text.strip(): - blocks.append(f"Memory {index} ({uri}):\n{text.strip()}") - return "\n\n".join(blocks), rows + blocks.append(f"Memory {index} ({row.get('uri', '')}):\n{text.strip()}") + return "\n\n".join(blocks), trace_rows, category_rerank finally: client.close() @@ -400,11 +924,18 @@ def _trace(self, event: dict[str, Any]) -> None: @staticmethod def _trace_injection_fields(block: str, matches: list[dict[str, Any]]) -> dict[str, Any]: - injected_count = sum(1 for row in matches if int(row.get("text_chars") or 0) > 0) + injected_count = sum( + 1 + for row in matches + if row.get("injected") + or (row.get("selected_for_injection", True) and int(row.get("text_chars") or 0) > 0) + ) return { "injected": bool(block.strip()), "injected_count": injected_count if block.strip() else 0, - "retrieval_action_taken": "retrieve_and_inject" if block.strip() else "retrieve_no_injection", + "retrieval_action_taken": "retrieve_and_inject" + if block.strip() + else "retrieve_no_injection", } def _generate(self, messages): @@ -476,7 +1007,8 @@ def generate_next_message(self, message, state: LLMAgentState): ( i for i, item in enumerate(state.system_messages) - if isinstance(item, SystemMessage) and item.content == "" + if isinstance(item, SystemMessage) + and item.content == "" ), None, ) @@ -484,11 +1016,16 @@ def generate_next_message(self, message, state: LLMAgentState): role_value = getattr(role, "value", role) if marker_index is not None and str(role_value) == "user": query = str(getattr(message, "content", "") or "") - block, matches = self._retrieve(query) + block, matches, category_rerank = self._retrieve( + query, + decision_node="first_user", + search_limit=args.first_user_retrieval_top_k, + inject_limit=args.first_user_inject_top_k, + ) prompt = ( "No OpenViking memory matched this user request." if not block - else "Use these OpenViking experience memories only when they match the current task:\n\n" + else "Use these OpenViking memories only when they match the current task:\n\n" + block ) state.system_messages[marker_index] = SystemMessage(role="system", content=prompt) @@ -496,8 +1033,11 @@ def generate_next_message(self, message, state: LLMAgentState): { "decision_node": "first_user", "query": query, + "search_limit": args.first_user_retrieval_top_k, + "inject_limit": args.first_user_inject_top_k, "match_count": len(matches), "matches": matches, + "category_rerank": category_rerank, **self._trace_injection_fields(block, matches), } ) @@ -508,13 +1048,21 @@ def generate_next_message(self, message, state: LLMAgentState): write_calls = [call for call in tool_calls if _is_write_tool_call(call)] if write_calls: query = _tool_call_query(write_calls, state.messages) - block, matches = self._retrieve(query) + block, matches, category_rerank = self._retrieve( + query, + decision_node="before_write_tool_call", + search_limit=args.prewrite_retrieval_top_k, + inject_limit=args.prewrite_inject_top_k, + ) self._trace( { "decision_node": "before_write_tool_call", "query": query, + "search_limit": args.prewrite_retrieval_top_k, + "inject_limit": args.prewrite_inject_top_k, "match_count": len(matches), "matches": matches, + "category_rerank": category_rerank, **self._trace_injection_fields(block, matches), "tool_calls": [ { @@ -528,8 +1076,7 @@ def generate_next_message(self, message, state: LLMAgentState): if block: prompt = ( "Before executing the pending write-like tool call, use these " - "OpenViking experience memories only when they match the current task:\n\n" - + block + "OpenViking memories only when they match the current task:\n\n" + block ) assistant_message = self._generate( state.system_messages @@ -540,6 +1087,7 @@ def generate_next_message(self, message, state: LLMAgentState): return assistant_message, state if AGENT_NAME not in registry.get_agents(): + def create_openviking_memory_agent(tools, domain_policy, **kwargs): return OpenVikingMemoryAgent( tools=tools, @@ -577,36 +1125,152 @@ def main() -> int: parser.add_argument("--user-llm", required=True) parser.add_argument("--agent-llm-args", type=_json, default={}) parser.add_argument("--user-llm-args", type=_json, default={}) - parser.add_argument("--openviking-url", required=True) - parser.add_argument("--openviking-account", required=True) - parser.add_argument("--openviking-user", required=True) - parser.add_argument("--openviking-agent-id", required=True) + parser.add_argument("--openviking-url") + parser.add_argument("--openviking-account") + parser.add_argument("--openviking-user") + parser.add_argument("--openviking-agent-id") parser.add_argument("--openviking-timeout", type=float, default=600.0) parser.add_argument("--openviking-wait-timeout", type=int, default=600) - parser.add_argument("--search-uri", required=True) + parser.add_argument("--search-uri") parser.add_argument("--retrieval-top-k", type=int, default=4) + parser.add_argument("--first-user-retrieval-top-k", type=int) + parser.add_argument("--first-user-inject-top-k", type=int) + parser.add_argument("--prewrite-retrieval-top-k", type=int) + parser.add_argument("--prewrite-inject-top-k", type=int) + parser.add_argument("--fixed-first-user-file", type=Path) + parser.add_argument("--scope-prompt-file", type=Path) parser.add_argument( "--retrieval-mode", choices=["first_user", "prewrite", "first_user_prewrite"], default="first_user", ) + parser.add_argument("--category-rerank-config", type=_json, default={}) + parser.add_argument("--scope-prompt-config", type=_json, default={}) parser.add_argument("--force-train", action="store_true") + parser.add_argument("--prepare-corpus-only", action="store_true") + parser.add_argument( + "--no-memory", + action="store_true", + help="Run the configured TAU-2 agent without OpenViking retrieval.", + ) args = parser.parse_args() normalize_litellm_env() + if not args.no_memory: + missing = [ + name + for name in ( + "openviking_url", + "openviking_account", + "openviking_user", + "openviking_agent_id", + "search_uri", + ) + if not getattr(args, name) + ] + if missing: + parser.error( + "OpenViking memory runs require: " + + ", ".join("--" + name.replace("_", "-") for name in missing) + ) + args.category_reranker = CategoryReranker.from_payload( + args.category_rerank_config, + repo_root=REPO_ROOT, + ) - args.tau2_repo = args.tau2_repo.resolve() + args.tau2_repo = args.tau2_repo.expanduser().resolve() + args.run_dir = args.run_dir.expanduser().resolve() + if args.corpus_dir: + args.corpus_dir = args.corpus_dir.expanduser().resolve() args.run_dir.mkdir(parents=True, exist_ok=True) corpus_dir = args.corpus_dir or args.run_dir corpus_dir.mkdir(parents=True, exist_ok=True) + args.first_user_retrieval_top_k = args.first_user_retrieval_top_k or args.retrieval_top_k + args.first_user_inject_top_k = args.first_user_inject_top_k or args.first_user_retrieval_top_k + args.prewrite_retrieval_top_k = args.prewrite_retrieval_top_k or args.retrieval_top_k + args.prewrite_inject_top_k = args.prewrite_inject_top_k or args.prewrite_retrieval_top_k + if args.fixed_first_user_file is not None: + args.fixed_first_user_file = args.fixed_first_user_file.expanduser().resolve() + if args.scope_prompt_file is not None: + args.scope_prompt_file = args.scope_prompt_file.expanduser().resolve() + if not args.scope_prompt_file.is_file(): + parser.error(f"--scope-prompt-file does not exist: {args.scope_prompt_file}") + if isinstance(args.scope_prompt_config, dict) and args.scope_prompt_config.get("enabled"): + parser.error( + "--scope-prompt-file and enabled --scope-prompt-config are mutually exclusive" + ) + args.scope_prompt_config = { + "enabled": True, + "domain_files": {args.domain: str(args.scope_prompt_file)}, + } + args.scope_prompt_text, args.scope_prompt_summary = _load_scope_prompt( + args.scope_prompt_config, + domain=args.domain, + repo_root=REPO_ROOT, + ) train_results = corpus_dir / "train_results.json" corpus_manifest = corpus_dir / "corpus_manifest.json" eval_results = args.run_dir / f"{args.run_label}.json" trace_path = args.run_dir / f"{args.run_label}.retrieval_trace.jsonl" summary_path = args.run_dir / f"{args.run_label}.summary.json" + if args.no_memory: + user_name = _register_fixed_first_user(args) + _run_tau2( + tau2_repo=args.tau2_repo, + domain=args.domain, + split=args.eval_split_name, + task_ids=args.task_ids, + num_tasks=args.num_tasks, + trials=1, + max_steps=args.max_steps, + max_concurrency=args.max_concurrency, + agent=args.base_agent, + user=user_name, + agent_llm=args.agent_llm, + user_llm=args.user_llm, + agent_llm_args=args.agent_llm_args, + user_llm_args=args.user_llm_args, + seed=args.seed, + save_to=eval_results, + ) + assert_tau2_results_complete( + json.loads(eval_results.read_text()), context=f"{args.domain} eval" + ) + summary = { + "run_label": args.run_label, + "domain": args.domain, + "strategy_id": args.strategy_id, + "seed": args.seed, + "fixed_first_user_file": str(args.fixed_first_user_file) + if args.fixed_first_user_file + else None, + "eval_results": str(eval_results), + "metrics": _metrics(eval_results), + } + _write_json(summary_path, summary) + print(json.dumps(summary, ensure_ascii=False, sort_keys=True)) + return 0 + corpus = _train(args, train_results, corpus_manifest) + if args.prepare_corpus_only: + print( + json.dumps( + { + "run_label": args.run_label, + "domain": args.domain, + "strategy_id": args.strategy_id, + "prepare_corpus_only": True, + "corpus": corpus, + }, + ensure_ascii=False, + sort_keys=True, + ) + ) + return 0 + trace_path.touch() _register_memory_agent(args, trace_path) + user_name = _register_fixed_first_user(args) _run_tau2( tau2_repo=args.tau2_repo, domain=args.domain, @@ -617,7 +1281,7 @@ def main() -> int: max_steps=args.max_steps, max_concurrency=args.max_concurrency, agent=AGENT_NAME, - user=args.user, + user=user_name, agent_llm=args.agent_llm, user_llm=args.user_llm, agent_llm_args=args.agent_llm_args, @@ -625,15 +1289,38 @@ def main() -> int: seed=args.seed, save_to=eval_results, ) + assert_tau2_results_complete( + json.loads(eval_results.read_text()), context=f"{args.domain} eval" + ) + category_summary = args.category_reranker.summary() + trace_summary = _trace_category_summary(trace_path) summary = { "run_label": args.run_label, "domain": args.domain, "strategy_id": args.strategy_id, "retrieval_mode": args.retrieval_mode, + "retrieval": { + "first_user_retrieval_top_k": args.first_user_retrieval_top_k, + "first_user_inject_top_k": args.first_user_inject_top_k, + "prewrite_retrieval_top_k": args.prewrite_retrieval_top_k, + "prewrite_inject_top_k": args.prewrite_inject_top_k, + }, "seed": args.seed, + "fixed_first_user_file": str(args.fixed_first_user_file) + if args.fixed_first_user_file + else None, + "scope_prompt_file": str(args.scope_prompt_file) if args.scope_prompt_file else None, "corpus": corpus, + "category_rerank": category_summary, + "scope_prompt": args.scope_prompt_summary, "eval_results": str(eval_results), "retrieval_trace": str(trace_path), + "retrieval_trace_summary": trace_summary, + "runtime_evidence": _runtime_evidence_status( + category_rerank=category_summary, + retrieval_trace_summary=trace_summary, + corpus_probe=corpus.get("corpus_probe") if isinstance(corpus, dict) else None, + ), "metrics": _metrics(eval_results), } _write_json(summary_path, summary) diff --git a/benchmark/tau2/scripts/tau2_common.py b/benchmark/tau2/scripts/tau2_common.py index a8b5ce2013..4f4505bd87 100755 --- a/benchmark/tau2/scripts/tau2_common.py +++ b/benchmark/tau2/scripts/tau2_common.py @@ -5,13 +5,13 @@ import re import shutil import subprocess +from collections import Counter from datetime import datetime, timezone from pathlib import Path from typing import Any import yaml - TAU2_DIR = Path(__file__).resolve().parents[1] REPO_ROOT = TAU2_DIR.parents[1] CONFIRMATION_AWARE_UPSTREAM_PR = "https://github.com/sierra-research/tau2-bench/pull/297" @@ -63,6 +63,7 @@ def normalize_litellm_env() -> dict[str, Any]: def render_env(value: Any) -> Any: if isinstance(value, str): + def replace(match: re.Match[str]) -> str: name = match.group(1) default = match.group(2) or "" @@ -79,11 +80,7 @@ def replace(match: re.Match[str]) -> str: def deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: merged = dict(base) for key, value in override.items(): - if ( - key in merged - and isinstance(merged[key], dict) - and isinstance(value, dict) - ): + if key in merged and isinstance(merged[key], dict) and isinstance(value, dict): merged[key] = deep_merge(merged[key], value) else: merged[key] = value @@ -125,6 +122,66 @@ def write_json(path: Path, payload: Any) -> None: ) +def tau2_result_failures(data: dict[str, Any], *, expected_trials: int = 1) -> list[str]: + tasks = data.get("tasks") or [] + simulations = data.get("simulations") or [] + failures: list[str] = [] + if tasks: + expected_task_ids = { + str(task.get("id", task.get("task_id"))) for task in tasks if isinstance(task, dict) + } + observed_task_ids = {str(sim.get("task_id")) for sim in simulations} + expected = len(tasks) * expected_trials + if len(simulations) != expected: + failures.append(f"expected {expected} simulations, found {len(simulations)}") + if observed_task_ids != expected_task_ids: + missing = sorted(expected_task_ids - observed_task_ids) + extra = sorted(observed_task_ids - expected_task_ids) + failures.append( + f"simulation task ids do not match tasks: missing={missing[:10]} extra={extra[:10]}" + ) + expected_pairs = { + (task_id, trial) for task_id in expected_task_ids for trial in range(expected_trials) + } + observed_pairs = [ + (str(sim.get("task_id")), int(sim.get("trial", 0))) for sim in simulations + ] + duplicate_pairs = sorted( + pair for pair, count in Counter(observed_pairs).items() if count != 1 + ) + missing_pairs = sorted(expected_pairs - set(observed_pairs)) + if duplicate_pairs or missing_pairs: + failures.append( + "simulation task/trial coverage mismatch: " + f"missing={missing_pairs[:10]} duplicate={duplicate_pairs[:10]}" + ) + + for sim in simulations: + info = sim.get("info") or {} + termination_reason = str(sim.get("termination_reason") or "") + if info.get("failed_after_attempts") or "infrastructure_error" in termination_reason: + failures.append( + "task=" + f"{sim.get('task_id')} trial={sim.get('trial', 0)} " + f"termination={termination_reason} error={info.get('error') or info.get('error_type')}" + ) + elif not sim.get("messages"): + failures.append( + f"task={sim.get('task_id')} trial={sim.get('trial', 0)} has no messages" + ) + return failures + + +def assert_tau2_results_complete( + data: dict[str, Any], *, context: str, expected_trials: int = 1 +) -> None: + failures = tau2_result_failures(data, expected_trials=expected_trials) + if failures: + preview = "; ".join(failures[:5]) + more = f"; ... {len(failures) - 5} more" if len(failures) > 5 else "" + raise RuntimeError(f"{context} produced invalid TAU-2 results: {preview}{more}") + + def strategy_ids(config: dict[str, Any]) -> list[str]: strategies = config.get("strategies") or [] if not isinstance(strategies, list): @@ -219,9 +276,7 @@ def user_simulator_policy(config: dict[str, Any]) -> str: policy = config.get("eval", {}).get("user_simulator_policy", "official") policy = str(policy) if policy not in {"official", "confirmation_aware"}: - raise ValueError( - "eval.user_simulator_policy must be 'official' or 'confirmation_aware'" - ) + raise ValueError("eval.user_simulator_policy must be 'official' or 'confirmation_aware'") return policy diff --git a/openviking/prompts/templates/memory/trajectories.yaml b/openviking/prompts/templates/memory/trajectories.yaml index bd894dfefa..b467b17b39 100644 --- a/openviking/prompts/templates/memory/trajectories.yaml +++ b/openviking/prompts/templates/memory/trajectories.yaml @@ -1,6 +1,6 @@ memory_type: trajectories description: | - A record of agent execution in one business domain within a conversation. + A compact, reusable view of how the agent handled one task in a conversation. Extract when the agent worked through identifiable tasks involving decisions, tool calls, or multi-step actions. Skip pure chitchat or simple Q&A with no execution trace. @@ -32,31 +32,41 @@ fields: - name: content type: string description: | - Execution trace in EXACTLY this format: - - Goal: - - Trajectory: - 1. . - Actions: . - Progress: . - 2. <...>. - Actions: <...>. - Progress: <...>. - 3. <... continue sequentially to capture the entire execution> - - Result: - - Fail reason: - - Rules for 'Trajectory' section: - - STANDALONE COMPLETENESS: The trajectory must be comprehensive enough that a reader can fully understand exactly how the agent performed the task, what attempts it made, and what tools it used, entirely without the raw logs. - - EXHAUSTIVE TRACKING: Record every logical interaction, tool use, and system response in chronological order as a numbered list. - - SUMMARIZE LONG TEXTS: Do NOT record exact tool responses, raw JSON payloads, or verbatim user/agent messages, especially when they are very long. You MUST use a concise, summarized version that captures the core meaning without losing any key information, constraints, or data points critical to the task. - - CAPTURE ALL MISTAKES: You MUST explicitly detail any errors, false attempts, wrong function calls, or agent misunderstandings. Do not gloss over failures, dead ends, or retry loops. - - INTENT & PROGRESS TRACKING: Every step must clearly state the current agent intention, summarize the 'Actions/Events' taken, and conclude with 'Progress' (evaluating what was achieved or blocked). - - TOOL TRACKING (under 'Actions'): - * Format (Success): "Called -> Response: ." - * Format (Error/Mistake): "Called -> Error/Issue: | Context: ." + Procedure-like trajectory view in EXACTLY this format: + + # + - Domain: + - Trigger: + - Preconditions: + 1. + 2. <...> + - Procedure: + 1. + 2. + 3. + - Anti-patterns: + - + - Applicability Boundary: + - + - + - Result: + - Evidence: + + Rules: + - Write for future agent execution, not for human audit. Prefer clear instructions over chronological narration. + - Preserve the successful or best-known path: critical reads, required policy checks, confirmation steps, write-tool ordering, and final user-facing completion. + - Keep negative lessons in Anti-patterns, not mixed into Procedure. + - Keep the memory grounded in this session, but abstract away user-specific names, raw IDs, exact payloads, and long tool responses. + - Do not include raw user/order/reservation/payment/card IDs, user names, email/phone/address values, exact dates, case-specific amounts, exact budgets, route pairs, airport pairs, passenger/bag counts, card suffixes, flight numbers, order numbers, or raw tool payloads. Replace them with semantic descriptions such as "a delivered order", "an ineligible basic-economy reservation", "the saved payment method", or "a policy-ineligible cancellation". Stable policy constants may be kept when they are needed for future execution. + - The Evidence line is not an exception: it may name the source status and lesson type, but must not carry case-specific values such as exact amounts, routes, dates, counts, product names, or customer-specific state. + - Bad Evidence: "a mechanical keyboard exceeded the $200 threshold". Good Evidence: "a product variant exceeded the user-specified threshold". + - Bad Evidence: "the ORD-to-LAX reservation on May 12 used card ending 1234". Good Evidence: "a matching reservation required the saved payment method and policy checks before the write action". + - Mention tool names when they are part of the reusable path, but summarize observations instead of copying raw JSON. + - If the session failed or was partial, still write the best reusable lesson: put the corrected approach in Procedure and the failure cause in Anti-patterns / Evidence. + - Avoid broad SOPs. The Trigger and Applicability Boundary should make this record narrower than a whole domain workflow. General Rules: - - Use exactly the 4 labels (Goal, Trajectory, Result, Fail reason) in this exact order. - - Goal and Result are ONE sentence each. + - Use exactly the title, Domain line, and 7 labels above in this exact order. + - Trigger, Result, and Evidence are ONE sentence each. - No extra headings, free paragraphs, or closing remarks. merge_op: patch diff --git a/openviking/session/memory/agent_trajectory_context_provider.py b/openviking/session/memory/agent_trajectory_context_provider.py index 2316052f77..1f9f3d25c0 100644 --- a/openviking/session/memory/agent_trajectory_context_provider.py +++ b/openviking/session/memory/agent_trajectory_context_provider.py @@ -14,7 +14,6 @@ from openviking.session.memory.session_extract_context_provider import ( SessionExtractContextProvider, ) -from openviking.storage.viking_fs import VikingFS from openviking_cli.utils import get_logger logger = get_logger(__name__) @@ -24,14 +23,20 @@ class AgentTrajectoryContextProvider(SessionExtractContextProvider): - """Phase 1 provider: extract trajectory summaries from conversation.""" + """Phase 1 provider: extract reusable trajectory-view memories.""" def instruction(self) -> str: output_language = self._output_language - return f"""You are a memory extraction agent. Summarize this agent session as a trajectory record. - -One session = one trajectory. Always output exactly one, no exceptions. -Sub-tasks, pivots, errors, and follow-ups are numbered steps inside that one record — not separate trajectories. + return f"""You are a memory extraction agent. Convert this agent session into a reusable trajectory-view memory. + +One session = one trajectory-view record. Always output exactly one record. +Write the record as a compact procedure-like view of the useful execution pattern, +not as a raw transcript. Keep the future agent's decision points, tool path, +confirmation/write boundary, failure corrections, and applicability boundary. +Generalize case evidence; do not copy raw user names, identifiers, dates, amounts, +exact budgets, routes, counts, payment details, or tool payloads into the reusable memory. +Sub-tasks, pivots, errors, and follow-ups are folded into that one record as steps, +guardrails, or evidence — not separate trajectories. Output a JSON object with a `trajectories` array containing exactly one item. Follow field descriptions in the schema. JSON only, no explanation. diff --git a/tests/benchmark/test_tau2_category_rerank.py b/tests/benchmark/test_tau2_category_rerank.py new file mode 100644 index 0000000000..955d1ed607 --- /dev/null +++ b/tests/benchmark/test_tau2_category_rerank.py @@ -0,0 +1,1253 @@ +import json +import hashlib +from pathlib import Path +from types import SimpleNamespace + +from benchmark.tau2.scripts.category_rerank import CategoryReranker +from benchmark.tau2.scripts.run_eval import ( + _cell_artifacts, + _cell_metrics, + _summarize, + _tau2_command, +) +from benchmark.tau2.scripts.tau2_common import load_config +from benchmark.tau2.scripts.run_memory_v2_eval import ( + _load_scope_prompt, + _probe_corpus, + _runtime_evidence_status, + _trace_category_summary, +) +from benchmark.tau2.scripts.build_category_catalog import _build_catalog +from benchmark.tau2.scripts.generate_category_annotations import _basic_validate_annotation +from benchmark.tau2.scripts.run_category_annotation_batches import _batch_ranges + + +def _annotation( + *, + subject_type: str, + subject_id: str, + subject_ref: str, + category1: str, + category2: str, +) -> dict: + return { + "schema_version": "memory_category_annotation.v0", + "annotation_id": f"{subject_type}:{subject_id}", + "request_id": f"{subject_type}:{subject_id}", + "producer": "llm_prompt", + "subject": { + "benchmark_family": "tau2", + "domain": "retail", + "subject_type": subject_type, + "subject_id": subject_id, + "subject_ref": subject_ref, + }, + "category": { + "category1": category1, + "category2": category2, + "category_source": "existing_catalog", + "confidence": 1.0, + "catalog_match": { + "matched": True, + "decision": "reuse", + "matched_category_id": f"{category1}:{category2}", + }, + }, + "ranking_features": { + "category1": category1, + "category2": category2, + "category_source": "existing_catalog", + "confidence": 1.0, + }, + } + + +def _write_sidecar(tmp_path: Path, *, first_user_query: str | None = None) -> Path: + rows = [ + _annotation( + subject_type="query", + subject_id=( + "tau2_query_signature_" + "tau2_retail_pre_write_action_tools_exchange_delivered_order_items" + ), + subject_ref="unit#query", + category1="retail_order_post_shipment_service_request", + category2="delivered_order_exchange", + ), + _annotation( + subject_type="memory", + subject_id="delivered_return.md", + subject_ref="viking://agent/demo/memories/trajectories/delivered_return.md", + category1="retail_order_post_shipment_service_request", + category2="delivered_order_return", + ), + _annotation( + subject_type="memory", + subject_id="delivered_exchange.md", + subject_ref="viking://agent/demo/memories/trajectories/delivered_exchange.md", + category1="retail_order_post_shipment_service_request", + category2="delivered_order_exchange", + ), + _annotation( + subject_type="memory", + subject_id="pending_cancel.md", + subject_ref="viking://agent/demo/memories/trajectories/pending_cancel.md", + category1="retail_order_cancellation", + category2="pending_order_cancel", + ), + ] + if first_user_query is not None: + query_hash = hashlib.sha256(first_user_query.encode("utf-8")).hexdigest()[:16] + rows.append( + _annotation( + subject_type="query", + subject_id=f"tau2_query_signature_tau2_retail_first_user_query_sha256_{query_hash}", + subject_ref="unit#first_user", + category1="retail_order_post_shipment_service_request", + category2="delivered_order_exchange", + ) + ) + path = tmp_path / "annotations.jsonl" + path.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n") + return path + + +def _reranker(tmp_path: Path) -> CategoryReranker: + sidecar = _write_sidecar(tmp_path) + return CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["before_write_tool_call"], + "retrieve_limit": 6, + "inject_limit": 2, + "mismatch_policy": "keep_positive_match_drop_mismatch", + "positive_match_required": True, + "no_match_policy": "skip_injection", + "search_score_weight": 0.0, + }, + repo_root=Path(__file__).resolve().parents[2], + ) + + +def _has_key_fragment(value: object, fragment: str) -> bool: + if isinstance(value, dict): + return any( + fragment in str(key).lower() or _has_key_fragment(item, fragment) + for key, item in value.items() + ) + if isinstance(value, list): + return any(_has_key_fragment(item, fragment) for item in value) + return False + + +def test_category_rerank_config_matches_s89_alignment_shape() -> None: + repo_root = Path(__file__).resolve().parents[2] + config = load_config(repo_root / "benchmark/tau2/config/category_rerank.yaml") + strategies = {row["id"]: row for row in config["strategies"]} + category_strategy = strategies["memory_v2_trajectory_category_prewrite_exact"] + + assert config["benchmark"]["reasoning_effort"] == "high" + assert config["benchmark"]["domains"] == ["retail", "airline"] + assert config["openviking"]["retrieval_top_k"] == 4 + assert category_strategy["memory_backend"] == "openviking" + assert category_strategy["train_memory_mode"] == "experience_only" + assert category_strategy["search_memory_type"] == "trajectories" + assert category_strategy["retrieval_mode"] == "first_user_prewrite" + assert category_strategy["corpus_id"] == "memory_v2_trajectory_view" + + category_rerank = category_strategy["category_rerank"] + assert category_rerank["enabled"] is True + assert category_rerank["apply_nodes"] == ["before_write_tool_call"] + assert category_rerank["retrieve_limit"] == 6 + assert category_rerank["inject_limit"] == 2 + assert category_rerank["mismatch_policy"] == "keep_positive_match_drop_mismatch" + assert category_rerank["positive_match_required"] is True + assert category_rerank["no_match_policy"] == "skip_injection" + assert category_rerank["missing_query_policy"] == "base_rank" + assert category_rerank["search_score_weight"] == 0.0 + assert "annotation_files" in category_rerank + + scope_prompt = category_strategy["scope_prompt"] + assert scope_prompt["enabled"] is True + assert scope_prompt["injection_point"] == "system_prompt" + assert scope_prompt["domain_files"] == { + "retail": "benchmark/tau2/config/scope_prompts/retail_memory_scope.md", + "airline": "benchmark/tau2/config/scope_prompts/airline_memory_scope.md", + } + + assert "memory_v2_trajectory_prewrite_scope" in strategies + assert "memory_v2_trajectory_category_prewrite_priority" in strategies + assert "memory_v2_trajectory_category_prewrite_strict_pair" in strategies + first_user_strategy = strategies["memory_v2_trajectory_category_first_user_exact"] + first_user_rerank = first_user_strategy["category_rerank"] + assert first_user_rerank["apply_nodes"] == ["first_user"] + assert first_user_rerank["retrieve_limit"] == 6 + assert first_user_rerank["inject_limit"] == 2 + assert first_user_rerank["missing_query_policy"] == "fail_fast" + assert "memory_v2_trajectory_category_first_user_priority" in strategies + assert "memory_v2_trajectory_category_first_user_strict_pair" in strategies + assert _has_key_fragment(category_strategy, "annotation") + + +def test_category_rerank_keeps_positive_category_match(tmp_path: Path) -> None: + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/delivered_return.md", + "score": 0.99, + "_text": "return memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "score": 0.25, + "_text": "exchange memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/pending_cancel.md", + "score": 0.95, + "_text": "cancel memory body", + }, + ] + + selected, trace_rows, diagnostics = _reranker(tmp_path).select( + domain="retail", + query="Before executing write-like tool call(s): exchange_delivered_order_items({})", + rows=rows, + decision_node="before_write_tool_call", + base_limit=4, + ) + + assert diagnostics["applied"] is True + assert diagnostics["decision"] == "soft_reranked_keep_category2_matches" + assert diagnostics["mismatch_policy"] == "keep_positive_match_drop_mismatch" + assert diagnostics["positive_match_level"] == "category2" + assert diagnostics["inject_limit"] == 2 + assert diagnostics["query_category"]["category_id"] == "retail_order_post_shipment_service_request:delivered_order_exchange" + assert [row["uri"] for row in selected] == [ + "viking://agent/demo/memories/trajectories/delivered_exchange.md" + ] + assert trace_rows[0]["selected_for_injection"] is False + assert trace_rows[0]["category1_match"] is True + assert trace_rows[0]["category2_match"] is False + assert trace_rows[0]["category2_label_match"] is False + assert trace_rows[1]["selected_for_injection"] is True + assert trace_rows[1]["category_pair_match"] is True + assert trace_rows[1]["query_category1_prompt"] == "retail_order_post_shipment_service_request" + assert trace_rows[1]["memory_category1_prompt"] == "retail_order_post_shipment_service_request" + assert trace_rows[2]["selected_for_injection"] is False + assert trace_rows[2]["skipped_reason"] == "category_rerank" + + +def test_build_category_catalog_groups_category_pairs() -> None: + rows = [ + _annotation( + subject_type="memory", + subject_id="delivered_exchange_a.md", + subject_ref="viking://agent/demo/memories/trajectories/delivered_exchange_a.md", + category1="retail_order_post_shipment_service_request", + category2="delivered_order_exchange", + ), + _annotation( + subject_type="memory", + subject_id="delivered_exchange_b.md", + subject_ref="viking://agent/demo/memories/trajectories/delivered_exchange_b.md", + category1="retail_order_post_shipment_service_request", + category2="delivered_order_exchange", + ), + _annotation( + subject_type="memory", + subject_id="pending_cancel.md", + subject_ref="viking://agent/demo/memories/trajectories/pending_cancel.md", + category1="retail_order_cancellation", + category2="pending_order_cancel", + ), + ] + + catalog = _build_catalog(rows) + + assert catalog["schema_version"] == "memory_category_catalog.v0" + assert catalog["category_count"] == 2 + grouped = {row["category_id"]: row for row in catalog["categories"]} + assert ( + grouped[ + "retail_order_post_shipment_service_request:delivered_order_exchange" + ]["source_annotation_count"] + == 2 + ) + assert grouped["retail_order_cancellation:pending_order_cancel"][ + "source_annotation_count" + ] == 1 + + +def test_category_annotation_batch_ranges() -> None: + assert _batch_ranges(start_offset=20, end_offset=53, batch_size=10) == [ + (20, 10), + (30, 10), + (40, 10), + (50, 3), + ] + assert _batch_ranges(start_offset=20, end_offset=53, batch_size=10, warmup_count=3) == [ + (20, 1), + (21, 1), + (22, 1), + (23, 10), + (33, 10), + (43, 10), + ] + + +def test_category_rerank_rejects_malformed_matched_category_id(tmp_path: Path) -> None: + sidecar = _write_sidecar(tmp_path) + rows = [json.loads(line) for line in sidecar.read_text().splitlines()] + rows[0]["category"]["catalog_match"]["matched_category_id"] = ( + "retail_order_post_shipment_service_request:" + ) + sidecar.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n") + + try: + CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["before_write_tool_call"], + }, + repo_root=Path(__file__).resolve().parents[2], + ) + except ValueError as exc: + assert "invalid matched_category_id" in str(exc) + assert "expected ':'" in str(exc) + else: + raise AssertionError("malformed category id should fail fast") + + +def test_category_annotation_validation_rejects_prose_category_ids() -> None: + row = _annotation( + subject_type="memory", + subject_id="partial_cancel.md", + subject_ref="viking://agent/demo/memories/trajectories/partial_cancel.md", + category1="retail_order_cancellation_related_service_request", + category2=( + "partial pending unshipped order item cancellation handling with " + "identity verification" + ), + ) + + errors = _basic_validate_annotation(row) + + assert any("$.category.category2 must be a reusable slug id" in error for error in errors) + + +def test_category_annotation_validation_rejects_overlong_category_ids() -> None: + row = _annotation( + subject_type="memory", + subject_id="partial_cancel.md", + subject_ref="viking://agent/demo/memories/trajectories/partial_cancel.md", + category1="retail_order_cancellation", + category2=( + "pending_order_partial_cancellation_with_identity_verification_" + "human_transfer_and_user_confirmation" + ), + ) + + errors = _basic_validate_annotation(row) + + assert any("$.category.category2 must be a compact reusable slug id" in error for error in errors) + + +def test_category_catalog_rejects_prose_category_ids() -> None: + row = _annotation( + subject_type="memory", + subject_id="partial_cancel.md", + subject_ref="viking://agent/demo/memories/trajectories/partial_cancel.md", + category1="retail_order_cancellation_related_service_request", + category2="partial pending cancellation with human transfer", + ) + + try: + _build_catalog([row]) + except SystemExit as exc: + assert "category2 must be a reusable slug id" in str(exc) + else: + raise AssertionError("catalog builder should reject prose category ids") + + +def test_category_catalog_rejects_overlong_category_ids() -> None: + row = _annotation( + subject_type="memory", + subject_id="partial_cancel.md", + subject_ref="viking://agent/demo/memories/trajectories/partial_cancel.md", + category1="retail_order_cancellation", + category2=( + "pending_order_partial_cancellation_with_identity_verification_" + "human_transfer_and_user_confirmation" + ), + ) + + try: + _build_catalog([row]) + except SystemExit as exc: + assert "category2 must be a compact reusable slug id" in str(exc) + else: + raise AssertionError("catalog builder should reject overlong category ids") + + +def test_category_rerank_missing_query_can_fall_back_to_base_rank(tmp_path: Path) -> None: + sidecar = _write_sidecar(tmp_path) + reranker = CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["before_write_tool_call"], + "retrieve_limit": 6, + "inject_limit": 2, + "mismatch_policy": "keep_positive_match_drop_mismatch", + "positive_match_required": True, + "no_match_policy": "skip_injection", + "missing_query_policy": "base_rank", + "search_score_weight": 0.0, + }, + repo_root=Path(__file__).resolve().parents[2], + ) + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/pending_cancel.md", + "score": 0.99, + "_text": "cancel memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "score": 0.25, + "_text": "exchange memory body", + }, + ] + + selected, trace_rows, diagnostics = reranker.select( + domain="retail", + query="Before executing write-like tool call(s): update_unknown_order({})", + rows=rows, + decision_node="before_write_tool_call", + base_limit=1, + ) + + assert diagnostics["applied"] is False + assert diagnostics["decision"] == "missing_query_sidecar_base_rank" + assert diagnostics["query_sidecar_coverage"] == "missing" + assert diagnostics["missing_query_policy"] == "base_rank" + assert [row["uri"] for row in selected] == [ + "viking://agent/demo/memories/trajectories/pending_cancel.md" + ] + assert trace_rows[0]["selected_for_injection"] is True + assert trace_rows[1]["selected_for_injection"] is False + + +def test_category_rerank_combo_query_uses_covered_tool_subquery(tmp_path: Path) -> None: + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/delivered_return.md", + "score": 0.99, + "_text": "return memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "score": 0.25, + "_text": "exchange memory body", + }, + ] + + selected, trace_rows, diagnostics = _reranker(tmp_path).select( + domain="retail", + query=( + "Before executing write-like tool call(s): " + "return_delivered_order_items({}); exchange_delivered_order_items({})" + ), + rows=rows, + decision_node="before_write_tool_call", + base_limit=4, + ) + + assert diagnostics["applied"] is True + assert diagnostics["query_sidecar_coverage"] == "partial" + assert diagnostics["decision"] == "soft_reranked_keep_category2_matches" + assert any( + signature.endswith("tools=exchange_delivered_order_items") + for signature in diagnostics["matched_query_signatures"] + ) + assert any( + signature.endswith("tools=exchange_delivered_order_items,return_delivered_order_items") + for signature in diagnostics["missing_query_signatures"] + ) + assert [row["uri"] for row in selected] == [ + "viking://agent/demo/memories/trajectories/delivered_exchange.md" + ] + assert trace_rows[1]["query_category_signature"].endswith("tools=exchange_delivered_order_items") + + +def test_category_rerank_priority_fills_category1_matches(tmp_path: Path) -> None: + first_user_query = "I need to exchange an item in a delivered order for a replacement." + sidecar = _write_sidecar(tmp_path, first_user_query=first_user_query) + reranker = CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["first_user"], + "retrieve_limit": 6, + "inject_limit": 2, + "retrieve_limits": {"first_user": 6}, + "inject_limits": {"first_user": 4}, + "mismatch_policy": "positive_priority_fill", + "positive_match_required": True, + "no_match_policy": "skip_injection", + "search_score_weight": 0.0, + }, + repo_root=Path(__file__).resolve().parents[2], + ) + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/delivered_return.md", + "score": 0.9, + "_text": "return memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "score": 0.1, + "_text": "exchange memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/pending_cancel.md", + "score": 0.8, + "_text": "cancel memory body", + }, + ] + + selected, trace_rows, diagnostics = reranker.select( + domain="retail", + query=first_user_query, + rows=rows, + decision_node="first_user", + base_limit=4, + ) + + assert diagnostics["decision"] == "soft_reranked_positive_priority_fill" + assert diagnostics["positive_match_level"] == "category2" + assert diagnostics["retrieve_limit"] == 6 + assert diagnostics["inject_limit"] == 4 + assert [row["uri"] for row in selected] == [ + "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "viking://agent/demo/memories/trajectories/delivered_return.md", + "viking://agent/demo/memories/trajectories/pending_cancel.md", + ] + assert reranker.search_limit(4, decision_node="first_user") == 6 + assert trace_rows[0]["selected_for_injection"] is True + assert trace_rows[1]["selected_for_injection"] is True + assert trace_rows[2]["selected_for_injection"] is True + + +def test_category_rerank_strict_pair_keeps_only_category1_and_category2_match(tmp_path: Path) -> None: + first_user_query = "I need to exchange an item in a delivered order for a replacement." + sidecar = _write_sidecar(tmp_path, first_user_query=first_user_query) + reranker = CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["first_user"], + "retrieve_limit": 6, + "inject_limit": 2, + "retrieve_limits": {"first_user": 6}, + "inject_limits": {"first_user": 4}, + "mismatch_policy": "strict_pair_match_only", + "positive_match_required": True, + "no_match_policy": "skip_injection", + "search_score_weight": 0.0, + }, + repo_root=Path(__file__).resolve().parents[2], + ) + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/delivered_return.md", + "score": 0.9, + "_text": "return memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "score": 0.1, + "_text": "exchange memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/pending_cancel.md", + "score": 0.8, + "_text": "cancel memory body", + }, + ] + + selected, trace_rows, diagnostics = reranker.select( + domain="retail", + query=first_user_query, + rows=rows, + decision_node="first_user", + base_limit=4, + ) + + assert diagnostics["decision"] == "soft_reranked_keep_strict_pair_matches" + assert diagnostics["positive_match_level"] == "category2" + assert diagnostics["dropped_mismatch_count"] == 2 + assert [row["uri"] for row in selected] == [ + "viking://agent/demo/memories/trajectories/delivered_exchange.md" + ] + assert trace_rows[0]["category1_match"] is True + assert trace_rows[0]["category2_match"] is False + assert trace_rows[0]["selected_for_injection"] is False + assert trace_rows[1]["category1_match"] is True + assert trace_rows[1]["category2_match"] is True + assert trace_rows[1]["selected_for_injection"] is True + assert trace_rows[2]["selected_for_injection"] is False + + +def test_category_rerank_strict_pair_skips_category1_only_matches(tmp_path: Path) -> None: + first_user_query = "I need to exchange an item in a delivered order for a replacement." + sidecar = _write_sidecar(tmp_path, first_user_query=first_user_query) + reranker = CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["first_user"], + "retrieve_limits": {"first_user": 6}, + "inject_limits": {"first_user": 4}, + "mismatch_policy": "strict_pair_match_only", + "positive_match_required": True, + "no_match_policy": "skip_injection", + "search_score_weight": 0.0, + }, + repo_root=Path(__file__).resolve().parents[2], + ) + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/delivered_return.md", + "score": 0.9, + "_text": "return memory body", + }, + ] + + selected, trace_rows, diagnostics = reranker.select( + domain="retail", + query=first_user_query, + rows=rows, + decision_node="first_user", + base_limit=4, + ) + + assert selected == [] + assert diagnostics["decision"] == "no_strict_pair_category_match_skip_injection" + assert diagnostics["positive_match_level"] == "category1" + assert diagnostics["dropped_mismatch_count"] == 1 + assert trace_rows[0]["category1_match"] is True + assert trace_rows[0]["category2_match"] is False + assert trace_rows[0]["selected_for_injection"] is False + + +def test_category_rerank_supports_node_specific_mismatch_policy(tmp_path: Path) -> None: + first_user_query = "I need to exchange an item in a delivered order for a replacement." + sidecar = _write_sidecar(tmp_path, first_user_query=first_user_query) + reranker = CategoryReranker.from_payload( + { + "enabled": True, + "annotation_files": [str(sidecar)], + "apply_nodes": ["first_user", "before_write_tool_call"], + "retrieve_limits": {"first_user": 6, "before_write_tool_call": 6}, + "inject_limits": {"first_user": 4, "before_write_tool_call": 2}, + "mismatch_policy": "keep_positive_match_drop_mismatch", + "mismatch_policies": { + "first_user": "positive_priority_fill", + "before_write_tool_call": "keep_positive_match_drop_mismatch", + }, + "positive_match_required": True, + "no_match_policy": "skip_injection", + "search_score_weight": 0.0, + }, + repo_root=Path(__file__).resolve().parents[2], + ) + rows = [ + { + "uri": "viking://agent/demo/memories/trajectories/delivered_return.md", + "score": 0.9, + "_text": "return memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "score": 0.1, + "_text": "exchange memory body", + }, + { + "uri": "viking://agent/demo/memories/trajectories/pending_cancel.md", + "score": 0.8, + "_text": "cancel memory body", + }, + ] + + first_user_selected, _, first_user_diagnostics = reranker.select( + domain="retail", + query=first_user_query, + rows=rows, + decision_node="first_user", + base_limit=4, + ) + prewrite_selected, _, prewrite_diagnostics = reranker.select( + domain="retail", + query="Before executing write-like tool call(s): exchange_delivered_order_items({})", + rows=rows, + decision_node="before_write_tool_call", + base_limit=4, + ) + + assert first_user_diagnostics["mismatch_policy"] == "positive_priority_fill" + assert first_user_diagnostics["decision"] == "soft_reranked_positive_priority_fill" + assert [row["uri"] for row in first_user_selected] == [ + "viking://agent/demo/memories/trajectories/delivered_exchange.md", + "viking://agent/demo/memories/trajectories/delivered_return.md", + "viking://agent/demo/memories/trajectories/pending_cancel.md", + ] + assert prewrite_diagnostics["mismatch_policy"] == "keep_positive_match_drop_mismatch" + assert prewrite_diagnostics["decision"] == "soft_reranked_keep_category2_matches" + assert [row["uri"] for row in prewrite_selected] == [ + "viking://agent/demo/memories/trajectories/delivered_exchange.md" + ] + + +def test_category_rerank_skips_non_target_node(tmp_path: Path) -> None: + rows = [ + {"uri": "viking://agent/demo/memories/trajectories/one.md", "score": 0.2}, + {"uri": "viking://agent/demo/memories/trajectories/two.md", "score": 0.1}, + ] + + selected, trace_rows, diagnostics = _reranker(tmp_path).select( + domain="retail", + query="exchange_delivered_order_items", + rows=rows, + decision_node="first_user", + base_limit=1, + ) + + assert diagnostics["applied"] is False + assert diagnostics["decision"] == "node_not_enabled" + assert [row["uri"] for row in selected] == [ + "viking://agent/demo/memories/trajectories/one.md" + ] + assert trace_rows[0]["selected_for_injection"] is True + assert trace_rows[1]["selected_for_injection"] is False + + +def test_scope_prompt_loads_domain_file(tmp_path: Path) -> None: + prompt = tmp_path / "retail_scope.md" + prompt.write_text("same order") + + text, summary = _load_scope_prompt( + {"enabled": True, "domain_files": {"retail": str(prompt)}}, + domain="retail", + repo_root=Path(__file__).resolve().parents[2], + ) + + assert "same order" in text + assert summary["enabled"] is True + assert summary["loaded"] is True + assert summary["loaded_files"] == [str(prompt)] + + +def test_scope_prompt_skips_unconfigured_domain(tmp_path: Path) -> None: + prompt = tmp_path / "retail_scope.md" + prompt.write_text("retail only") + + text, summary = _load_scope_prompt( + {"enabled": True, "domain_files": {"retail": str(prompt)}}, + domain="airline", + repo_root=Path(__file__).resolve().parents[2], + ) + + assert text == "" + assert summary["loaded"] is False + assert summary["skipped_reason"] == "no_domain_scope_prompt" + + +def test_trace_category_summary_counts_runtime_sources(tmp_path: Path) -> None: + trace = tmp_path / "retrieval_trace.jsonl" + rows = [ + { + "decision_node": "static_scope_prompt", + "retrieval_action_taken": "scope_prompt_static_injection", + "injected": True, + }, + { + "decision_node": "before_write_tool_call", + "retrieval_action_taken": "retrieve_and_inject", + "injected": True, + "injected_count": 1, + "tool_calls": [{"name": "exchange_delivered_order_items"}], + "category_rerank": { + "enabled": True, + "applied": True, + "decision": "soft_reranked_keep_category2_matches", + "query_sidecar_coverage": "covered", + "query_category": { + "matched": True, + "category_source": "existing_catalog", + }, + }, + "matches": [ + { + "uri": "viking://agent/example/memories/trajectories/delivered_exchange.md", + "selected_for_injection": True, + "injected": True, + "memory_category_source_prompt": "existing_catalog", + "memory_category1_prompt": ["retail_order_post_shipment_service_request"], + "memory_category2_prompt": ["delivered_order_exchange"], + "category2_match": True, + }, + { + "uri": "viking://agent/example/memories/trajectories/.overview.md", + "selected_for_injection": False, + "category_rerank_reasons": ["missing_memory_category"], + }, + { + "uri": "viking://agent/example/memories/trajectories/pending_cancel.md", + "selected_for_injection": False, + "memory_category_source_prompt": "existing_catalog", + "memory_category1_prompt": ["retail_order_cancellation"], + "memory_category2_prompt": ["pending_order_cancel"], + "category1_match": False, + "category2_match": False, + }, + ], + }, + ] + trace.write_text( + "\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n" + ) + + summary = _trace_category_summary(trace) + + assert summary["trace_present"] is True + assert summary["decision_nodes"]["before_write_tool_call"] == 1 + assert summary["category_decisions"]["soft_reranked_keep_category2_matches"] == 1 + assert summary["query_sidecar_coverage"]["covered"] == 1 + assert summary["query_category_sources"]["existing_catalog"] == 1 + assert summary["selected_memory_category_sources"]["existing_catalog"] == 1 + assert summary["tool_calls"]["exchange_delivered_order_items"] == 1 + assert summary["rates"]["memory_category_candidate_coverage"] == 2 / 3 + assert summary["rates"]["selected_memory_category_coverage"] == 1.0 + assert summary["rates"]["memory_category_match_coverage"] == 1 / 3 + assert summary["rates"]["selected_memory_category_match_coverage"] == 1.0 + assert summary["counts"]["aggregate_memory_candidate_count"] == 1 + assert summary["counts"]["concrete_memory_candidate_count"] == 2 + assert summary["counts"]["memory_injection_event_count"] == 1 + assert summary["counts"]["memory_injected_count"] == 1 + assert summary["counts"]["injected_match_count"] == 1 + assert summary["counts"]["injected_concrete_memory_count"] == 1 + assert summary["counts"]["injected_positive_category_match_count"] == 1 + assert summary["counts"]["injected_concrete_positive_category_match_count"] == 1 + assert summary["counts"]["query_sidecar_covered_event_count"] == 1 + assert summary["counts"]["query_sidecar_partial_event_count"] == 0 + assert summary["counts"]["query_sidecar_missing_event_count"] == 0 + assert summary["counts"]["memory_category_present_count"] == 2 + assert summary["counts"]["memory_category_matched_count"] == 1 + assert summary["rates"]["concrete_memory_candidate_rate"] == 2 / 3 + assert summary["rates"]["selected_concrete_memory_rate"] == 1.0 + assert summary["rates"]["injected_positive_category_match_rate"] == 1.0 + assert summary["rates"]["injected_concrete_positive_category_match_rate"] == 1.0 + assert summary["rates"]["injected_concrete_memory_rate"] == 1.0 + assert summary["rates"]["query_sidecar_non_missing_event_rate"] == 1.0 + assert summary["rates"]["query_sidecar_full_event_rate"] == 1.0 + + +def test_runtime_evidence_marks_aggregate_only_category_diagnostic() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 1, + "aggregate_match_count": 1, + "concrete_match_count": 0, + }, + retrieval_trace_summary={ + "trace_present": True, + "counts": { + "category_applied_event_count": 1, + "query_category_matched_event_count": 1, + "memory_category_present_count": 1, + "memory_category_matched_count": 0, + }, + "rates": { + "concrete_memory_candidate_rate": 0.0, + "selected_positive_category_match_rate": 0.0, + }, + }, + ) + + assert evidence["status"] == "diagnostic" + assert "aggregate_only_corpus_probe" in evidence["reasons"] + assert "no_concrete_corpus_probe_matches" in evidence["reasons"] + assert "no_concrete_memory_candidates" in evidence["reasons"] + assert "no_matched_memory_categories" in evidence["reasons"] + assert "no_memory_injection" in evidence["reasons"] + assert "no_selected_positive_category_match" in evidence["reasons"] + + +def test_runtime_evidence_requires_applied_category_events() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 1, + "aggregate_match_count": 0, + "concrete_match_count": 1, + }, + retrieval_trace_summary={ + "trace_present": True, + "category_event_count": 1, + "counts": { + "category_enabled_event_count": 1, + "category_applied_event_count": 0, + }, + "rates": { + "concrete_memory_candidate_rate": 1.0, + "selected_positive_category_match_rate": 1.0, + }, + }, + ) + + assert evidence["status"] == "diagnostic" + assert "no_category_rerank_applied_events" in evidence["reasons"] + + +def test_runtime_evidence_requires_query_category_coverage() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 1, + "aggregate_match_count": 0, + "concrete_match_count": 1, + }, + retrieval_trace_summary={ + "trace_present": True, + "category_event_count": 1, + "counts": { + "category_applied_event_count": 1, + "query_category_matched_event_count": 0, + "memory_category_present_count": 1, + "memory_category_matched_count": 1, + }, + "rates": { + "concrete_memory_candidate_rate": 1.0, + "selected_positive_category_match_rate": 1.0, + }, + }, + ) + + assert evidence["status"] == "diagnostic" + assert "no_query_category_coverage" in evidence["reasons"] + + +def test_runtime_evidence_accepts_valid_category_runtime_coverage() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 2, + "aggregate_match_count": 0, + "concrete_match_count": 2, + }, + retrieval_trace_summary={ + "trace_present": True, + "category_event_count": 1, + "counts": { + "category_applied_event_count": 1, + "query_category_matched_event_count": 1, + "memory_category_present_count": 2, + "memory_category_matched_count": 1, + "memory_injection_event_count": 1, + "injected_concrete_memory_count": 1, + "injected_positive_category_match_count": 1, + "injected_concrete_positive_category_match_count": 1, + }, + "rates": { + "concrete_memory_candidate_rate": 1.0, + "selected_positive_category_match_rate": 1.0, + "injected_concrete_memory_rate": 1.0, + "injected_positive_category_match_rate": 1.0, + }, + }, + ) + + assert evidence == {"status": "valid", "reasons": []} + + +def test_runtime_evidence_requires_actual_memory_injection() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 1, + "aggregate_match_count": 0, + "concrete_match_count": 1, + }, + retrieval_trace_summary={ + "trace_present": True, + "category_event_count": 1, + "counts": { + "category_applied_event_count": 1, + "query_category_matched_event_count": 1, + "memory_category_present_count": 1, + "memory_category_matched_count": 1, + "memory_injection_event_count": 0, + }, + "rates": { + "concrete_memory_candidate_rate": 1.0, + "selected_positive_category_match_rate": 1.0, + }, + }, + ) + + assert evidence["status"] == "diagnostic" + assert "no_memory_injection" in evidence["reasons"] + + +def test_runtime_evidence_requires_injected_concrete_memory() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 2, + "aggregate_match_count": 1, + "concrete_match_count": 1, + }, + retrieval_trace_summary={ + "trace_present": True, + "category_event_count": 1, + "counts": { + "category_applied_event_count": 1, + "query_category_matched_event_count": 1, + "memory_category_present_count": 1, + "memory_category_matched_count": 1, + "memory_injection_event_count": 1, + "injected_concrete_memory_count": 0, + "injected_positive_category_match_count": 1, + }, + "rates": { + "concrete_memory_candidate_rate": 0.5, + "selected_positive_category_match_rate": 1.0, + "injected_concrete_memory_rate": 0.0, + "injected_positive_category_match_rate": 1.0, + }, + }, + ) + + assert evidence["status"] == "diagnostic" + assert "no_injected_concrete_memory" in evidence["reasons"] + + +def test_runtime_evidence_requires_injected_concrete_positive_category_match() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={ + "match_count": 2, + "aggregate_match_count": 1, + "concrete_match_count": 1, + }, + retrieval_trace_summary={ + "trace_present": True, + "category_event_count": 1, + "counts": { + "category_applied_event_count": 1, + "query_category_matched_event_count": 1, + "memory_category_present_count": 1, + "memory_category_matched_count": 1, + "memory_injection_event_count": 2, + "injected_concrete_memory_count": 1, + "injected_positive_category_match_count": 1, + "injected_concrete_positive_category_match_count": 0, + }, + "rates": { + "concrete_memory_candidate_rate": 0.5, + "selected_positive_category_match_rate": 1.0, + "injected_concrete_memory_rate": 0.5, + "injected_positive_category_match_rate": 0.5, + "injected_concrete_positive_category_match_rate": 0.0, + }, + }, + ) + + assert evidence["status"] == "diagnostic" + assert "no_injected_concrete_positive_category_match" in evidence["reasons"] + + +def test_scoreboard_excludes_diagnostic_runtime_evidence() -> None: + scoreboard = _summarize( + [ + { + "domain": "airline", + "strategy_id": "memory_v2_trajectory_category_prewrite", + "metrics": { + "simulation_count": 1, + "avg_reward": 1.0, + "db_match_rate": 1.0, + }, + "runtime_evidence": { + "status": "diagnostic", + "reasons": [ + "no_concrete_memory_candidates", + "no_query_category_coverage", + ], + }, + }, + { + "domain": "airline", + "strategy_id": "memory_v2_trajectory_category_prewrite", + "metrics": { + "simulation_count": 1, + "avg_reward": 0.5, + "db_match_rate": 0.0, + }, + "runtime_evidence": {"status": "valid", "reasons": []}, + }, + ] + ) + + domain = scoreboard["strategies"]["memory_v2_trajectory_category_prewrite"][ + "domains" + ]["airline"] + assert domain["completed_cell_count"] == 2 + assert domain["valid_completed_cell_count"] == 1 + assert domain["diagnostic_cell_count"] == 1 + assert domain["diagnostic_reason_counts"] == { + "no_concrete_memory_candidates": 1, + "no_query_category_coverage": 1, + } + assert domain["diagnostic_simulation_count"] == 1 + assert domain["simulation_count"] == 1 + assert domain["avg_reward"] == 0.5 + assert domain["db_match_rate"] == 0.0 + + +def test_no_memory_strategy_uses_wrapper_command(tmp_path: Path) -> None: + config = { + "benchmark": { + "eval_split_name": "test", + "max_steps": 7, + "task_max_concurrency": 2, + "agent": "llm_agent", + "user": "user_simulator", + "reasoning_effort": "high", + }, + "model": { + "agent_llm": "agent-model", + "user_llm": "user-model", + }, + "paths": { + "tau2_repo": str(tmp_path / "tau2-bench"), + "output_dir": str(tmp_path / "result"), + }, + } + + command = _tau2_command( + config, + domain="airline", + strategy={"id": "no_memory", "memory_backend": "none"}, + configured_run_id="baseline_run", + run_label="baseline_run_airline_no_memory_r1", + task_ids=["18"], + num_tasks=None, + train_num_tasks=None, + seed=300, + ) + + assert command is not None + assert command[1].endswith("run_memory_v2_eval.py") + assert "--no-memory" in command + assert "--openviking-url" not in command + assert command[command.index("--strategy-id") + 1] == "no_memory" + assert command[command.index("--base-agent") + 1] == "llm_agent" + assert command[command.index("--task-id") + 1] == "18" + assert command[command.index("--run-dir") + 1].endswith( + "result/baseline_run/memory_cells/baseline_run_airline_no_memory_r1" + ) + + +def test_no_memory_artifacts_read_wrapper_summary_metrics(tmp_path: Path) -> None: + out = tmp_path / "out" + cell = { + "memory_backend": "none", + "domain": "airline", + "strategy_id": "no_memory", + "run_label": "baseline_run_airline_no_memory_r1", + } + + artifacts = _cell_artifacts(cell, repo=tmp_path / "tau2-bench", out=out) + assert set(artifacts) == {"summary", "results"} + + summary_path = Path(artifacts["summary"]) + summary_path.parent.mkdir(parents=True) + summary_path.write_text( + json.dumps( + { + "metrics": { + "simulation_count": 1, + "avg_reward": 0.0, + "db_match_rate": 0.0, + } + } + ) + ) + + assert _cell_metrics(cell, artifacts) == { + "simulation_count": 1, + "avg_reward": 0.0, + "db_match_rate": 0.0, + } + + +def test_runtime_evidence_marks_empty_corpus_probe_diagnostic() -> None: + evidence = _runtime_evidence_status( + category_rerank={"enabled": True}, + corpus_probe={"match_count": 0}, + retrieval_trace_summary={"trace_present": True, "counts": {}, "rates": {}}, + ) + + assert evidence["status"] == "diagnostic" + assert "empty_corpus_probe" in evidence["reasons"] + + +def test_probe_corpus_counts_aggregate_and_concrete_matches(tmp_path: Path) -> None: + class FakeClient: + def __init__(self) -> None: + self.limit: int | None = None + + def search(self, **kwargs: object) -> SimpleNamespace: + self.limit = int(kwargs["limit"]) + return SimpleNamespace( + memories=[ + SimpleNamespace( + uri="viking://agent/a/memories/trajectories/.overview.md", + score=0.2, + ), + SimpleNamespace( + uri="viking://agent/a/memories/trajectories/concrete.md#chunk_0001", + score=0.1, + ), + ] + ) + + def read(self, uri: str) -> str: + return f"body for {uri}" + + client = FakeClient() + probe = _probe_corpus( + SimpleNamespace( + category_reranker=_reranker(tmp_path), + domain="airline", + search_uri="viking://agent/a/memories/trajectories", + retrieval_top_k=4, + ), + client, + ) + + assert client.limit == 6 + assert probe["probe_limit"] == 6 + assert probe["match_count"] == 2 + assert probe["aggregate_match_count"] == 1 + assert probe["concrete_match_count"] == 1 + assert probe["aggregate_read_non_empty_count"] == 1 + assert probe["concrete_read_non_empty_count"] == 1 + assert probe["matches"][0]["is_aggregate_memory"] is True + assert probe["matches"][1]["is_concrete_memory"] is True