#!/usr/bin/env python3 """Evaluate DWS tool retrieval methods against reviewed Catalog metadata. The primary intent benchmark deliberately excludes ``use_when``, ``avoid_when``, and examples from the indexed corpus. The reviewed ``use_when`` entries become queries, so a method cannot win by retrieving the exact sentence it indexed. The script has no required third-party dependency. Pass ``--dense-model`` to add a SentenceTransformers dense retriever and BM25F+dense RRF. For example:: uv run --with fastembed==0.7.3 \ scripts/dev/eval_tool_search_ranking.py \ --dense-backend fastembed \ --dense-model BAAI/bge-small-zh-v1.5 \ --output /tmp/dws-tool-search-eval.json The output is aggregate evidence, not a generated production baseline. Release gates need a separately reviewed, versioned query/qrels dataset. """ from __future__ import annotations import argparse import json import math import random import re import statistics import time import unicodedata from collections import Counter, defaultdict from dataclasses import dataclass from pathlib import Path from typing import Any, Iterable, Protocol REPO_ROOT = Path(__file__).resolve().parents[2] DEFAULT_CATALOG = REPO_ROOT / ".worktrees/policy-tmp/tool-search-schema-catalog" DEFAULT_WORKFLOWS = REPO_ROOT / "scripts/testdata/tool_search_workflows.json" FIELD_WEIGHTS = { "identity": 8.0, "summary": 5.0, "description": 2.0, "parameters": 2.0, } DEFAULT_PROJECTION = "proxy_v1" PROJECTION_SPECS = { "identity_only": { "identity": True, }, "identity_summary": { "identity": True, "summary": True, }, "identity_summary_description": { "identity": True, "summary": True, "description": True, }, "proxy_v1": { "identity": True, "summary": True, "description": True, "parameter_names": True, }, "with_parameter_descriptions": { "identity": True, "summary": True, "description": True, "parameter_names": True, "parameter_descriptions": True, }, "production_v1_candidate": { "identity": True, "summary": True, "description": True, "parameter_names": True, "parameter_descriptions": True, "parameter_types": True, "aliases": True, }, "production_with_use_when_leakage_upper_bound": { "identity": True, "summary": True, "description": True, "parameter_names": True, "parameter_descriptions": True, "parameter_types": True, "aliases": True, "use_when": True, }, } CAMEL_BOUNDARY = re.compile(r"(?<=[a-z0-9])(?=[A-Z])") IDENTIFIER = re.compile(r"[A-Za-z0-9]+(?:[._-][A-Za-z0-9]+)+") TEXT_CHUNK = re.compile(r"[a-z0-9]+|[\u3400-\u9fff]+") @dataclass(frozen=True) class ToolDocument: canonical: str product: str cli_path: str identities: tuple[str, ...] fields: dict[str, str] use_when: tuple[str, ...] avoid_when: tuple[str, ...] examples: tuple[str, ...] @property def dense_text(self) -> str: return "\n".join( value for value in self.fields.values() if value.strip() ) @dataclass(frozen=True) class SingleCase: query: str gold: str product: str kind: str = "" @dataclass(frozen=True) class WorkflowCase: case_id: str query: str required: tuple[str, ...] subqueries: tuple[str, ...] class Ranker(Protocol): name: str def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: ... def _strings(value: Any) -> list[str]: if isinstance(value, str): return [value] if isinstance(value, list): return [item for item in value if isinstance(item, str)] return [] def parameter_text( parameters: Any, *, include_names: bool, include_descriptions: bool, include_types: bool, ) -> str: if not isinstance(parameters, dict): return "" parts: list[str] = [] for name in sorted(parameters): value = parameters[name] if not isinstance(value, dict): continue if include_names: parts.append(name) for key in ["property"]: item = value.get(key) if isinstance(item, str): parts.append(item) enum = value.get("enum") if isinstance(enum, list): parts.extend( str(item) for item in enum if isinstance(item, (str, int, float)) ) if include_descriptions: for key in ["description", "interface_description"]: item = value.get(key) if isinstance(item, str): parts.append(item) if include_types: for key in ["type", "interface_type"]: item = value.get(key) if isinstance(item, str): parts.append(item) return " ".join(parts) def load_catalog( path: Path, projection: str = DEFAULT_PROJECTION ) -> tuple[str, list[ToolDocument]]: if projection not in PROJECTION_SPECS: raise ValueError(f"unknown projection {projection}") projection_spec = PROJECTION_SPECS[projection] payload = load_catalog_payload(path) tools: list[ToolDocument] = [] for canonical, raw in sorted(payload["tools"].items()): identity_values = [ str(raw.get(key, "")).strip() for key in [ "canonical_path", "cli_path", "primary_cli_path", "name", "cli_name", "group", "product_id", ] if str(raw.get(key, "")).strip() ] if projection_spec.get("aliases"): identity_values.extend(_strings(raw.get("aliases"))) identities = tuple(dict.fromkeys(identity_values)) identity = " ".join(identities) summary = " ".join( str(raw.get(key, "")) for key in ["agent_summary", "title", "display"] ) fields = { "identity": identity if projection_spec.get("identity") else "", "summary": summary if projection_spec.get("summary") else "", "description": ( str(raw.get("description", "")) if projection_spec.get("description") else "" ), "parameters": parameter_text( raw.get("parameters"), include_names=bool(projection_spec.get("parameter_names")), include_descriptions=bool( projection_spec.get("parameter_descriptions") ), include_types=bool(projection_spec.get("parameter_types")), ), } if projection_spec.get("use_when"): fields["use_when"] = " ".join(_strings(raw.get("use_when"))) tools.append( ToolDocument( canonical=canonical, product=str(raw.get("product_id") or canonical.partition(".")[0]), cli_path=str(raw.get("cli_path", "")).strip(), identities=identities, fields=fields, use_when=tuple(_strings(raw.get("use_when"))), avoid_when=tuple(_strings(raw.get("avoid_when"))), examples=tuple(_strings(raw.get("examples"))), ) ) return str(payload.get("source_hash", "")), tools def load_catalog_payload(path: Path) -> dict[str, Any]: """Load either a build-time Catalog shard directory or a test fixture. Production and CI use ``cmd_schema_catalog`` to materialize ``catalog.json`` plus product shards below an ignored temporary directory. A consolidated JSON file remains accepted only so small unit fixtures and old diagnostic artifacts can be inspected without becoming repository delivery inputs. """ if path.is_file(): payload = json.loads(path.read_text(encoding="utf-8")) if not isinstance(payload.get("tools"), dict): raise ValueError(f"Catalog file has no tools object: {path}") return payload envelope_path = path / "catalog.json" tools_dir = path / "tools" if not envelope_path.is_file() or not tools_dir.is_dir(): raise ValueError( "Catalog must be a generated directory containing catalog.json " f"and tools/: {path}" ) envelope = json.loads(envelope_path.read_text(encoding="utf-8")) merged: dict[str, Any] = {} for shard_path in sorted(tools_dir.glob("*.json")): shard = json.loads(shard_path.read_text(encoding="utf-8")) shard_tools = shard.get("tools") if not isinstance(shard_tools, dict): raise ValueError(f"Catalog shard has no tools object: {shard_path}") duplicates = sorted(set(merged).intersection(shard_tools)) if duplicates: raise ValueError( f"Catalog shards duplicate tools {duplicates}: {shard_path}" ) merged.update(shard_tools) expected = envelope.get("catalog", {}).get("agent_metadata", {}).get( "surface_tools" ) if isinstance(expected, int) and expected != len(merged): raise ValueError( f"Catalog shard count {len(merged)} differs from envelope {expected}" ) return { "source_hash": envelope.get("source_hash", ""), "surface_hash": envelope.get("surface_hash", ""), "tools": merged, } def load_workflows(path: Path, known_tools: set[str]) -> list[WorkflowCase]: payload = json.loads(path.read_text(encoding="utf-8")) cases: list[WorkflowCase] = [] for raw in payload["cases"]: required = tuple(raw["required"]) missing = sorted(set(required) - known_tools) if missing: raise ValueError(f"workflow {raw['id']} references missing tools: {missing}") cases.append( WorkflowCase( case_id=raw["id"], query=raw["query"], required=required, subqueries=tuple(raw["subqueries"]), ) ) return cases def tokenize(text: str) -> list[str]: """Tokenize mixed Chinese, English, and CLI identifiers deterministically.""" normalized = unicodedata.normalize("NFKC", text) tokens = [match.group(0).lower() for match in IDENTIFIER.finditer(normalized)] normalized = CAMEL_BOUNDARY.sub(" ", normalized).lower() normalized = normalized.replace("_", " ").replace("-", " ").replace(".", " ") for chunk in TEXT_CHUNK.findall(normalized): if "\u3400" <= chunk[0] <= "\u9fff": if len(chunk) == 1: tokens.append(chunk) else: tokens.extend(chunk[index : index + 2] for index in range(len(chunk) - 1)) else: tokens.append(chunk) return tokens def stable_rank(scores: dict[str, float], limit: int | None = None) -> list[str]: ranked = [item for item in scores.items() if item[1] > 0.0] ranked.sort(key=lambda item: (-item[1], item[0])) if limit is not None: ranked = ranked[:limit] return [canonical for canonical, _ in ranked] class KeywordOverlapRanker: name = "keyword_overlap" def __init__(self, tools: list[ToolDocument]) -> None: self._tokens = { tool.canonical: set( tokenize(" ".join(tool.fields.values())) ) for tool in tools } document_frequency: Counter[str] = Counter() for terms in self._tokens.values(): document_frequency.update(terms) count = len(tools) self._idf = { term: math.log(1.0 + (count + 1.0) / (frequency + 1.0)) for term, frequency in document_frequency.items() } def rank(self, query: str, limit: int | None = None) -> list[str]: terms = set(tokenize(query)) scores = { canonical: math.fsum( self._idf.get(term, 0.0) for term in sorted(terms & document_terms) ) for canonical, document_terms in self._tokens.items() } return stable_rank(scores, limit) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return {query: self.rank(query, limit) for query in dict.fromkeys(queries)} class BM25Index: def __init__( self, documents: dict[str, list[str]], *, k1: float = 0.9, b: float = 0.4, ) -> None: self._documents = {key: Counter(tokens) for key, tokens in documents.items()} self._lengths = {key: sum(counts.values()) for key, counts in self._documents.items()} self._average_length = ( statistics.fmean(self._lengths.values()) if self._lengths else 1.0 ) self._k1 = k1 self._b = b self._count = len(documents) document_frequency: Counter[str] = Counter() for counts in self._documents.values(): document_frequency.update(counts.keys()) self._idf = { term: math.log(1.0 + (self._count - frequency + 0.5) / (frequency + 0.5)) for term, frequency in document_frequency.items() } def scores(self, query: str) -> dict[str, float]: query_terms = Counter(tokenize(query)) result: dict[str, float] = {} for canonical, counts in self._documents.items(): length = max(self._lengths[canonical], 1) score = 0.0 for term, query_frequency in query_terms.items(): frequency = counts.get(term, 0) if frequency == 0: continue denominator = frequency + self._k1 * ( 1.0 - self._b + self._b * length / max(self._average_length, 1.0) ) score += ( self._idf.get(term, 0.0) * (frequency * (self._k1 + 1.0) / denominator) * math.log1p(query_frequency) ) result[canonical] = score return result class TfidfCosineRanker: name = "tfidf_cosine" def __init__(self, tools: list[ToolDocument]) -> None: documents = { tool.canonical: Counter(tokenize(" ".join(tool.fields.values()))) for tool in tools } document_frequency: Counter[str] = Counter() for counts in documents.values(): document_frequency.update(counts.keys()) count = len(documents) self._idf = { term: math.log((count + 1.0) / (frequency + 1.0)) + 1.0 for term, frequency in document_frequency.items() } self._vectors: dict[str, dict[str, float]] = {} self._norms: dict[str, float] = {} for canonical, counts in documents.items(): vector = { term: (1.0 + math.log(frequency)) * self._idf[term] for term, frequency in counts.items() } self._vectors[canonical] = vector self._norms[canonical] = math.sqrt( math.fsum(value * value for value in vector.values()) ) def rank(self, query: str, limit: int | None = None) -> list[str]: counts = Counter(tokenize(query)) query_vector = { term: (1.0 + math.log(frequency)) * self._idf.get(term, 0.0) for term, frequency in counts.items() if term in self._idf } query_norm = math.sqrt( math.fsum(value * value for value in query_vector.values()) ) scores: dict[str, float] = {} if query_norm == 0: return [] for canonical, vector in self._vectors.items(): denominator = query_norm * self._norms[canonical] if denominator == 0: continue scores[canonical] = math.fsum( query_vector[term] * vector.get(term, 0.0) for term in sorted(query_vector) ) / denominator return stable_rank(scores, limit) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return {query: self.rank(query, limit) for query in dict.fromkeys(queries)} class WeightedJaccardRanker: name = "weighted_jaccard" def __init__(self, tools: list[ToolDocument]) -> None: self._tokens = { tool.canonical: set(tokenize(" ".join(tool.fields.values()))) for tool in tools } document_frequency: Counter[str] = Counter() for terms in self._tokens.values(): document_frequency.update(terms) count = len(tools) self._idf = { term: math.log((count + 1.0) / (frequency + 1.0)) + 1.0 for term, frequency in document_frequency.items() } def rank(self, query: str, limit: int | None = None) -> list[str]: query_terms = set(tokenize(query)) scores: dict[str, float] = {} for canonical, document_terms in self._tokens.items(): union = query_terms | document_terms denominator = math.fsum( self._idf.get(term, 0.0) for term in sorted(union) ) if denominator == 0: continue numerator = math.fsum( self._idf.get(term, 0.0) for term in sorted(query_terms & document_terms) ) scores[canonical] = numerator / denominator return stable_rank(scores, limit) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return {query: self.rank(query, limit) for query in dict.fromkeys(queries)} class BM25VariantIndex(BM25Index): def __init__( self, documents: dict[str, list[str]], *, variant: str, k1: float = 0.9, b: float = 0.4, delta: float = 1.0, ) -> None: super().__init__(documents, k1=k1, b=b) self._variant = variant self._delta = delta def scores(self, query: str) -> dict[str, float]: query_terms = Counter(tokenize(query)) result: dict[str, float] = {} for canonical, counts in self._documents.items(): length = max(self._lengths[canonical], 1) length_ratio = length / max(self._average_length, 1.0) score = 0.0 for term, query_frequency in query_terms.items(): frequency = counts.get(term, 0) if frequency == 0: continue if self._variant == "bm25l": normalized = frequency / (1.0 - self._b + self._b * length_ratio) contribution = ( (self._k1 + 1.0) * (normalized + self._delta) / (self._k1 + normalized + self._delta) ) elif self._variant == "bm25plus": denominator = frequency + self._k1 * ( 1.0 - self._b + self._b * length_ratio ) contribution = ( frequency * (self._k1 + 1.0) / denominator + self._delta ) else: # pragma: no cover - constructor-owned invariant raise ValueError(f"unsupported BM25 variant {self._variant}") score += ( self._idf.get(term, 0.0) * contribution * math.log1p(query_frequency) ) result[canonical] = score return result class BM25LRanker: name = "bm25l" def __init__(self, tools: list[ToolDocument]) -> None: self._index = BM25VariantIndex( { tool.canonical: tokenize(" ".join(tool.fields.values())) for tool in tools }, variant="bm25l", ) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return { query: stable_rank(self._index.scores(query), limit) for query in dict.fromkeys(queries) } class BM25PlusRanker: name = "bm25plus" def __init__(self, tools: list[ToolDocument]) -> None: self._index = BM25VariantIndex( { tool.canonical: tokenize(" ".join(tool.fields.values())) for tool in tools }, variant="bm25plus", ) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return { query: stable_rank(self._index.scores(query), limit) for query in dict.fromkeys(queries) } class BM25Ranker: name = "bm25_unfielded" def __init__(self, tools: list[ToolDocument]) -> None: self._index = BM25Index( { tool.canonical: tokenize(" ".join(tool.fields.values())) for tool in tools } ) def rank(self, query: str, limit: int | None = None) -> list[str]: return stable_rank(self._index.scores(query), limit) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return {query: self.rank(query, limit) for query in dict.fromkeys(queries)} class BM25FRanker: name = "bm25_field_weighted" def __init__(self, tools: list[ToolDocument]) -> None: self._indexes = { field: BM25Index( {tool.canonical: tokenize(tool.fields[field]) for tool in tools} ) for field in FIELD_WEIGHTS } self._identities = { tool.canonical: { unicodedata.normalize("NFKC", value).strip().lower() for value in tool.identities if value.strip() } for tool in tools } self._tools = tools def rank(self, query: str, limit: int | None = None) -> list[str]: scores: defaultdict[str, float] = defaultdict(float) for field, index in self._indexes.items(): weight = FIELD_WEIGHTS[field] for canonical, score in index.scores(query).items(): scores[canonical] += weight * score normalized_query = unicodedata.normalize("NFKC", query).strip().lower() for tool in self._tools: canonical = tool.canonical identity_values = self._identities[canonical] if normalized_query == canonical.lower(): scores[canonical] += 1_000_000.0 elif any(normalized_query == value for value in identity_values): scores[canonical] += 100_000.0 return stable_rank(scores, limit) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: return {query: self.rank(query, limit) for query in dict.fromkeys(queries)} class SentenceTransformerDenseRanker: name = "dense" def __init__(self, tools: list[ToolDocument], model_name: str, batch_size: int) -> None: try: import numpy as np from sentence_transformers import SentenceTransformer except ImportError as exc: # pragma: no cover - environment dependent raise RuntimeError( "--dense-model requires sentence-transformers; use the uv command " "shown in this script's docstring" ) from exc self._np = np self._model_name = model_name self._model = SentenceTransformer(model_name) self._batch_size = batch_size self._canonicals = [tool.canonical for tool in tools] documents = [self._document_input(tool.dense_text) for tool in tools] self._document_vectors = self._model.encode( documents, batch_size=batch_size, normalize_embeddings=True, show_progress_bar=True, convert_to_numpy=True, ) def _query_input(self, text: str) -> str: if "e5" in self._model_name.lower(): return f"query: {text}" return text def _document_input(self, text: str) -> str: if "e5" in self._model_name.lower(): return f"passage: {text}" return text def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: unique_queries = list(dict.fromkeys(queries)) vectors = self._model.encode( [self._query_input(query) for query in unique_queries], batch_size=self._batch_size, normalize_embeddings=True, show_progress_bar=True, convert_to_numpy=True, ) similarities = vectors @ self._document_vectors.T result: dict[str, list[str]] = {} for row, query in zip(similarities, unique_queries): order = self._np.lexsort( ( self._np.asarray(self._canonicals, dtype=object), -row, ) ) if limit is not None: order = order[:limit] result[query] = [self._canonicals[int(index)] for index in order] return result class FastEmbedDenseRanker: name = "dense" def __init__(self, tools: list[ToolDocument], model_name: str, batch_size: int) -> None: try: import numpy as np from fastembed import TextEmbedding except ImportError as exc: # pragma: no cover - environment dependent raise RuntimeError( "--dense-backend fastembed requires fastembed; use the uv command " "shown in this script's docstring" ) from exc self._np = np self._model = TextEmbedding(model_name=model_name) self._batch_size = batch_size self._canonicals = [tool.canonical for tool in tools] document_vectors = list( self._model.passage_embed( [tool.dense_text for tool in tools], batch_size=batch_size ) ) self._document_vectors = self._normalize(np.asarray(document_vectors)) def _normalize(self, vectors: Any) -> Any: norms = self._np.linalg.norm(vectors, axis=1, keepdims=True) return vectors / self._np.maximum(norms, 1e-12) def rank_many( self, queries: Iterable[str], limit: int | None = None ) -> dict[str, list[str]]: unique_queries = list(dict.fromkeys(queries)) query_vectors = list( self._model.query_embed(unique_queries, batch_size=self._batch_size) ) vectors = self._normalize(self._np.asarray(query_vectors)) similarities = vectors @ self._document_vectors.T result: dict[str, list[str]] = {} for row, query in zip(similarities, unique_queries): order = self._np.lexsort( ( self._np.asarray(self._canonicals, dtype=object), -row, ) ) if limit is not None: order = order[:limit] result[query] = [self._canonicals[int(index)] for index in order] return result def rrf_rank( rankings: Iterable[list[str]], limit: int | None = None, k: float = 60.0 ) -> list[str]: scores: defaultdict[str, float] = defaultdict(float) for ranking in rankings: for rank, canonical in enumerate(ranking, start=1): scores[canonical] += 1.0 / (k + rank) return stable_rank(scores, limit) def reciprocal_rank(ranking: list[str], gold: str) -> float: try: return 1.0 / (ranking.index(gold) + 1) except ValueError: return 0.0 def rank_position(ranking: list[str], gold: str) -> int | None: try: return ranking.index(gold) + 1 except ValueError: return None def evaluate_single_cases( cases: list[SingleCase], rankings: dict[str, list[str]] ) -> tuple[dict[str, Any], list[bool]]: reciprocal_ranks: list[float] = [] positions: list[int] = [] top1: list[bool] = [] top5: list[bool] = [] top10: list[bool] = [] empty = 0 per_product: defaultdict[str, list[bool]] = defaultdict(list) for case in cases: ranking = rankings[case.query] if not ranking: empty += 1 position = rank_position(ranking, case.gold) if position is not None: positions.append(position) reciprocal_ranks.append(reciprocal_rank(ranking, case.gold)) hit1 = position is not None and position <= 1 hit5 = position is not None and position <= 5 hit10 = position is not None and position <= 10 top1.append(hit1) top5.append(hit5) top10.append(hit10) per_product[case.product].append(hit5) product_recall = { product: sum(values) / len(values) for product, values in sorted(per_product.items()) } metrics: dict[str, Any] = { "cases": len(cases), "recall_at_1": sum(top1) / len(cases), "recall_at_5": sum(top5) / len(cases), "recall_at_10": sum(top10) / len(cases), "mrr": statistics.fmean(reciprocal_ranks), "median_rank_when_found": statistics.median(positions) if positions else None, "empty_result_rate": empty / len(cases), "worst_product_recall_at_5": dict( sorted(product_recall.items(), key=lambda item: (item[1], item[0]))[:5] ), "product_recall_at_5": product_recall, } return metrics, top5 def evaluate_forbidden( cases: list[SingleCase], rankings: dict[str, list[str]] ) -> dict[str, Any]: ranks = [rank_position(rankings[case.query], case.gold) for case in cases] return { "cases": len(cases), "forbidden_at_1": sum(rank == 1 for rank in ranks) / len(cases), "forbidden_at_5": sum(rank is not None and rank <= 5 for rank in ranks) / len(cases), "forbidden_at_10": sum(rank is not None and rank <= 10 for rank in ranks) / len(cases), } def evaluate_workflows( cases: list[WorkflowCase], rankings: dict[str, list[str]], top_k: int = 5 ) -> dict[str, Any]: complete = 0 recalls: list[float] = [] ndcgs: list[float] = [] details: list[dict[str, Any]] = [] for case in cases: selected = rankings[case.query][:top_k] found = [tool for tool in case.required if tool in selected] missing = [tool for tool in case.required if tool not in selected] complete += not missing recalls.append(len(found) / len(case.required)) dcg = math.fsum( 1.0 / math.log2(rank + 2.0) for rank, canonical in enumerate(selected) if canonical in case.required ) ideal_count = min(len(case.required), top_k) ideal_dcg = math.fsum( 1.0 / math.log2(rank + 2.0) for rank in range(ideal_count) ) ndcgs.append(dcg / ideal_dcg if ideal_dcg else 0.0) details.append( { "id": case.case_id, "required": list(case.required), "top_k": selected, "missing": missing, } ) return { "cases": len(cases), "comprehensiveness_at_5": complete / len(cases), "mean_required_recall_at_5": statistics.fmean(recalls), "mean_ndcg_at_5": statistics.fmean(ndcgs), "details": details, } def decompose_workflow_rankings( cases: list[WorkflowCase], rankings: dict[str, list[str]], *, top_k: int = 5, ) -> dict[str, list[str]]: """Round-robin per-step rankings for a manually reviewed upper bound.""" result: dict[str, list[str]] = {} for case in cases: selected: list[str] = [] depth = 0 while len(selected) < top_k: added = False for subquery in case.subqueries: ranking = rankings[subquery] if depth >= len(ranking): continue canonical = ranking[depth] if canonical not in selected: selected.append(canonical) added = True if len(selected) == top_k: break if not added: break depth += 1 result[case.query] = selected return result def diagnostic_examples( cases: list[SingleCase], rankings: dict[str, list[str]], *, forbidden: bool = False ) -> list[dict[str, Any]]: examples: list[dict[str, Any]] = [] for case in cases: top5 = rankings[case.query][:5] position = rank_position(rankings[case.query], case.gold) selected = position == 1 if forbidden else position is None or position > 5 if not selected: continue examples.append( { "query": case.query, "gold" if not forbidden else "forbidden": case.gold, "top_5": top5, } ) if len(examples) == 8: break return examples def paired_change_examples( cases: list[SingleCase], baseline: dict[str, list[str]], candidate: dict[str, list[str]], ) -> dict[str, list[dict[str, Any]]]: rescued: list[dict[str, Any]] = [] regressed: list[dict[str, Any]] = [] for case in cases: base_hit = case.gold in baseline[case.query][:5] candidate_hit = case.gold in candidate[case.query][:5] if base_hit == candidate_hit: continue record = { "query": case.query, "gold": case.gold, "baseline_top_5": baseline[case.query][:5], "candidate_top_5": candidate[case.query][:5], } target = rescued if candidate_hit else regressed if len(target) < 8: target.append(record) return {"rescued": rescued, "regressed": regressed} def paired_bootstrap_interval( baseline: list[bool], candidate: list[bool], samples: int = 2_000 ) -> dict[str, float]: if len(baseline) != len(candidate): raise ValueError("paired samples differ in length") random_source = random.Random(20260812) differences: list[float] = [] count = len(baseline) observed = statistics.fmean(candidate) - statistics.fmean(baseline) for _ in range(samples): indexes = [random_source.randrange(count) for _ in range(count)] difference = statistics.fmean(candidate[index] for index in indexes) - statistics.fmean( baseline[index] for index in indexes ) differences.append(difference) differences.sort() return { "recall_at_5_delta": observed, "bootstrap_95_low": differences[int(samples * 0.025)], "bootstrap_95_high": differences[int(samples * 0.975)], } def paired_product_cluster_bootstrap_interval( cases: list[SingleCase], baseline: list[bool], candidate: list[bool], samples: int = 2_000, ) -> dict[str, float]: """Bootstrap macro-product paired R@5 differences. Cases written for the same product share authors, vocabulary, and sibling tools, so treating all rows as independent overstates cross-product certainty. Each draw resamples product clusters and averages their paired mean differences with equal product weight. """ if len(cases) != len(baseline) or len(baseline) != len(candidate): raise ValueError("clustered paired samples differ in length") clusters: defaultdict[str, list[int]] = defaultdict(list) for index, case in enumerate(cases): clusters[case.product].append(index) products = sorted(clusters) if not products: raise ValueError("clustered paired samples are empty") product_differences = { product: statistics.fmean( float(candidate[index]) - float(baseline[index]) for index in clusters[product] ) for product in products } observed = statistics.fmean(product_differences.values()) random_source = random.Random(20260812) differences: list[float] = [] for _ in range(samples): selected = [ products[random_source.randrange(len(products))] for _ in range(len(products)) ] differences.append( statistics.fmean(product_differences[product] for product in selected) ) differences.sort() return { "macro_product_recall_at_5_delta": observed, "bootstrap_95_low": differences[int(samples * 0.025)], "bootstrap_95_high": differences[int(samples * 0.975)], "product_clusters": len(products), } def query_language_slice(query: str) -> str: has_cjk = any("\u3400" <= character <= "\u9fff" for character in query) has_ascii_word = any(character.isascii() and character.isalnum() for character in query) if has_cjk and has_ascii_word: return "mixed_chinese_ascii" if has_cjk: return "chinese_only" return "non_chinese" def evaluate_language_slices( cases: list[SingleCase], rankings: dict[str, list[str]] ) -> dict[str, dict[str, Any]]: result: dict[str, dict[str, Any]] = {} for slice_name in ["chinese_only", "mixed_chinese_ascii", "non_chinese"]: selected = [ case for case in cases if query_language_slice(case.query) == slice_name ] if selected: result[slice_name] = evaluate_single_cases(selected, rankings)[0] return result def percentile(values: list[float], fraction: float) -> float: if not values: raise ValueError("percentile input is empty") ordered = sorted(values) index = max(0, min(len(ordered) - 1, math.ceil(fraction * len(ordered)) - 1)) return ordered[index] def measure_single_query_latency( ranker: Ranker, queries: list[str], *, samples: int, ) -> dict[str, float | int]: """Measure warmed one-query Python latency without claiming production SLA.""" if samples <= 0: return {} unique_queries = list(dict.fromkeys(queries)) if not unique_queries: raise ValueError("latency query set is empty") warmup = min(20, samples) for index in range(warmup): ranker.rank_many([unique_queries[index % len(unique_queries)]], limit=5) durations: list[float] = [] for index in range(samples): query = unique_queries[index % len(unique_queries)] started = time.perf_counter_ns() ranker.rank_many([query], limit=5) durations.append((time.perf_counter_ns() - started) / 1_000_000.0) return { "samples": samples, "p50": statistics.median(durations), "p95": percentile(durations, 0.95), "p99": percentile(durations, 0.99), "max": max(durations), } def evaluate_projection_ablation( catalog_path: Path, intent: list[SingleCase], forbidden: list[SingleCase], workflows: list[WorkflowCase], all_queries: list[str], ) -> dict[str, Any]: """Compare cumulative field projections with identical ranker contracts.""" result: dict[str, Any] = { "order": list(PROJECTION_SPECS), "projections": {}, "paired_vs_proxy_v1": {}, "paired_product_cluster_vs_proxy_v1": {}, } success_vectors: dict[tuple[str, str], list[bool]] = {} for projection_name, projection_spec in PROJECTION_SPECS.items(): _, tools = load_catalog(catalog_path, projection_name) projection_result: dict[str, Any] = { "fields": sorted( key for key, enabled in projection_spec.items() if enabled ), "methods": {}, } for ranker_type in [TfidfCosineRanker, BM25Ranker]: started = time.perf_counter() ranker = ranker_type(tools) build_seconds = time.perf_counter() - started started = time.perf_counter() rankings = ranker.rank_many(all_queries) batch_seconds = time.perf_counter() - started intent_metrics, success = evaluate_single_cases(intent, rankings) success_vectors[(projection_name, ranker.name)] = success projection_result["methods"][ranker.name] = { "index_build_seconds": build_seconds, "offline_batch_seconds": batch_seconds, "intent": intent_metrics, "intent_language_slices": evaluate_language_slices( intent, rankings ), "forbidden": evaluate_forbidden(forbidden, rankings), "workflow": evaluate_workflows(workflows, rankings), } result["projections"][projection_name] = projection_result for projection_name in PROJECTION_SPECS: if projection_name == DEFAULT_PROJECTION: continue result["paired_vs_proxy_v1"][projection_name] = {} result["paired_product_cluster_vs_proxy_v1"][projection_name] = {} for method in ["tfidf_cosine", "bm25_unfielded"]: baseline = success_vectors[(DEFAULT_PROJECTION, method)] candidate = success_vectors[(projection_name, method)] result["paired_vs_proxy_v1"][projection_name][method] = ( paired_bootstrap_interval(baseline, candidate) ) result["paired_product_cluster_vs_proxy_v1"][projection_name][method] = ( paired_product_cluster_bootstrap_interval( intent, baseline, candidate ) ) return result def build_cases( tools: list[ToolDocument], workflows_path: Path ) -> tuple[list[SingleCase], list[SingleCase], list[SingleCase], list[WorkflowCase]]: intent: list[SingleCase] = [] forbidden: list[SingleCase] = [] identity: list[SingleCase] = [] for tool in tools: intent.extend( SingleCase(query=query, gold=tool.canonical, product=tool.product) for query in tool.use_when ) forbidden.extend( SingleCase(query=query, gold=tool.canonical, product=tool.product) for query in tool.avoid_when ) identity.append( SingleCase( query=tool.canonical, gold=tool.canonical, product=tool.product, kind="canonical", ) ) if tool.cli_path: identity.append( SingleCase( query=tool.cli_path, gold=tool.canonical, product=tool.product, kind="cli_path", ) ) workflows = load_workflows(workflows_path, {tool.canonical for tool in tools}) return intent, forbidden, identity, workflows def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--catalog", type=Path, default=DEFAULT_CATALOG, help=( "build-time Catalog shard directory from cmd_schema_catalog; " "a consolidated JSON file is accepted for diagnostic fixtures" ), ) parser.add_argument("--workflows", type=Path, default=DEFAULT_WORKFLOWS) parser.add_argument("--dense-model", default="") parser.add_argument( "--dense-backend", choices=["fastembed", "sentence-transformers"], default="fastembed", ) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--rrf-depth", type=int, default=100) parser.add_argument( "--projection", choices=list(PROJECTION_SPECS), default=DEFAULT_PROJECTION, ) parser.add_argument( "--skip-projection-ablation", action="store_true", help="skip cumulative field projection comparisons", ) parser.add_argument( "--latency-samples", type=int, default=0, help="measure warmed Python one-query latency for each base ranker", ) parser.add_argument("--output", type=Path) args = parser.parse_args() source_hash, tools = load_catalog(args.catalog, args.projection) intent, forbidden, identity, workflows = build_cases(tools, args.workflows) all_queries = [ case.query for case in [*intent, *forbidden, *identity] ] + [ query for case in workflows for query in (case.query, *case.subqueries) ] rankers: list[Ranker] = [] build_seconds: dict[str, float] = {} for ranker_type in [ KeywordOverlapRanker, TfidfCosineRanker, WeightedJaccardRanker, BM25Ranker, BM25LRanker, BM25PlusRanker, BM25FRanker, ]: started = time.perf_counter() ranker = ranker_type(tools) build_seconds[ranker.name] = time.perf_counter() - started rankers.append(ranker) if args.dense_model: started = time.perf_counter() if args.dense_backend == "fastembed": dense_ranker: Ranker = FastEmbedDenseRanker( tools, args.dense_model, args.batch_size ) else: dense_ranker = SentenceTransformerDenseRanker( tools, args.dense_model, args.batch_size ) build_seconds[dense_ranker.name] = time.perf_counter() - started rankers.append(dense_ranker) result: dict[str, Any] = { "protocol": { "catalog": str(args.catalog), "catalog_source_hash": source_hash, "tool_count": len(tools), "intent_cases": len(intent), "forbidden_cases": len(forbidden), "identity_cases": len(identity), "workflow_cases": len(workflows), "index_excludes": [ field for field in ["use_when", "avoid_when", "examples"] if field != "use_when" or not PROJECTION_SPECS[args.projection].get("use_when") ], "projection_version": args.projection, "projection_fields": sorted( key for key, enabled in PROJECTION_SPECS[args.projection].items() if enabled ), "intent_qrels_limitation": ( "use_when queries and indexed agent_summary text are authored in " "the same reviewed selection files; this is a proxy benchmark, " "not an independent release test set" ), "field_weights": FIELD_WEIGHTS, "dense_model": args.dense_model or None, "dense_backend": args.dense_backend if args.dense_model else None, "rrf_depth": args.rrf_depth, }, "methods": {}, "paired_vs_bm25_field_weighted": {}, "paired_product_cluster_vs_bm25_field_weighted": {}, "paired_vs_bm25_unfielded": {}, "paired_product_cluster_vs_bm25_unfielded": {}, } all_rankings: dict[str, dict[str, list[str]]] = {} success_vectors: dict[str, list[bool]] = {} for ranker in rankers: started = time.perf_counter() rankings = ranker.rank_many(all_queries) duration = time.perf_counter() - started all_rankings[ranker.name] = rankings intent_metrics, success = evaluate_single_cases(intent, rankings) success_vectors[ranker.name] = success identity_metrics, _ = evaluate_single_cases(identity, rankings) identity_by_kind = { kind: evaluate_single_cases( [case for case in identity if case.kind == kind], rankings )[0] for kind in ["canonical", "cli_path"] } result["methods"][ranker.name] = { "index_build_seconds": build_seconds[ranker.name], "offline_batch_seconds": duration, "single_query_latency_ms": measure_single_query_latency( ranker, [case.query for case in intent], samples=args.latency_samples, ), "intent": intent_metrics, "intent_language_slices": evaluate_language_slices(intent, rankings), "forbidden": evaluate_forbidden(forbidden, rankings), "identity": identity_metrics, "identity_by_kind": identity_by_kind, "workflow": evaluate_workflows(workflows, rankings), "diagnostics": { "intent_misses": diagnostic_examples(intent, rankings), "forbidden_at_1": diagnostic_examples( forbidden, rankings, forbidden=True ), }, } result["lexical_rrf_sweep"] = [] for depth in [20, 50, 100]: for rrf_k in [10.0, 20.0, 60.0]: lexical_rankings = { query: rrf_rank( [ all_rankings["tfidf_cosine"][query][:depth], all_rankings["bm25_unfielded"][query][:depth], ], k=rrf_k, ) for query in dict.fromkeys(all_queries) } metrics, _ = evaluate_single_cases(intent, lexical_rankings) result["lexical_rrf_sweep"].append( { "arms": ["tfidf_cosine", "bm25_unfielded"], "depth": depth, "k": rrf_k, "intent": metrics, "forbidden": evaluate_forbidden(forbidden, lexical_rankings), "workflow": evaluate_workflows(workflows, lexical_rankings), } ) lexical_hybrid_name = "lexical_rrf_tfidf_bm25" lexical_hybrid_rankings = { query: rrf_rank( [ all_rankings["tfidf_cosine"][query][: args.rrf_depth], all_rankings["bm25_unfielded"][query][: args.rrf_depth], ] ) for query in dict.fromkeys(all_queries) } all_rankings[lexical_hybrid_name] = lexical_hybrid_rankings intent_metrics, success = evaluate_single_cases(intent, lexical_hybrid_rankings) success_vectors[lexical_hybrid_name] = success result["methods"][lexical_hybrid_name] = { "index_build_seconds": None, "offline_batch_seconds": None, "single_query_latency_ms": {}, "intent": intent_metrics, "intent_language_slices": evaluate_language_slices( intent, lexical_hybrid_rankings ), "forbidden": evaluate_forbidden(forbidden, lexical_hybrid_rankings), "identity": evaluate_single_cases(identity, lexical_hybrid_rankings)[0], "identity_by_kind": { kind: evaluate_single_cases( [case for case in identity if case.kind == kind], lexical_hybrid_rankings, )[0] for kind in ["canonical", "cli_path"] }, "workflow": evaluate_workflows(workflows, lexical_hybrid_rankings), "diagnostics": { "intent_misses": diagnostic_examples(intent, lexical_hybrid_rankings), "forbidden_at_1": diagnostic_examples( forbidden, lexical_hybrid_rankings, forbidden=True ), }, } if "dense" in all_rankings: dense_ranks = all_rankings["dense"] result["dense_rrf_sweep"] = [] for sparse_name in [ "tfidf_cosine", "bm25_unfielded", "bm25_field_weighted", ]: sparse_ranks = all_rankings[sparse_name] for depth in [20, 50, 100]: for rrf_k in [10.0, 20.0, 60.0]: sweep_rankings = { query: rrf_rank( [ sparse_ranks[query][:depth], dense_ranks[query][:depth], ], k=rrf_k, ) for query in dict.fromkeys(all_queries) } sweep_metrics, _ = evaluate_single_cases( intent, sweep_rankings ) result["dense_rrf_sweep"].append( { "arms": [sparse_name, "dense"], "depth": depth, "k": rrf_k, "intent": sweep_metrics, "forbidden": evaluate_forbidden( forbidden, sweep_rankings ), "workflow": evaluate_workflows( workflows, sweep_rankings ), } ) for hybrid_name, sparse_name in [ ("hybrid_rrf_tfidf", "tfidf_cosine"), ("hybrid_rrf_fielded", "bm25_field_weighted"), ("hybrid_rrf_unfielded", "bm25_unfielded"), ]: sparse_ranks = all_rankings[sparse_name] hybrid_rankings = { query: rrf_rank( [ sparse_ranks[query][: args.rrf_depth], dense_ranks[query][: args.rrf_depth], ] ) for query in dict.fromkeys(all_queries) } all_rankings[hybrid_name] = hybrid_rankings intent_metrics, success = evaluate_single_cases(intent, hybrid_rankings) success_vectors[hybrid_name] = success identity_metrics, _ = evaluate_single_cases(identity, hybrid_rankings) identity_by_kind = { kind: evaluate_single_cases( [case for case in identity if case.kind == kind], hybrid_rankings )[0] for kind in ["canonical", "cli_path"] } result["methods"][hybrid_name] = { "index_build_seconds": None, "offline_batch_seconds": None, "intent": intent_metrics, "intent_language_slices": evaluate_language_slices(intent, hybrid_rankings), "forbidden": evaluate_forbidden(forbidden, hybrid_rankings), "identity": identity_metrics, "identity_by_kind": identity_by_kind, "workflow": evaluate_workflows(workflows, hybrid_rankings), "diagnostics": { "intent_misses": diagnostic_examples(intent, hybrid_rankings), "forbidden_at_1": diagnostic_examples( forbidden, hybrid_rankings, forbidden=True ), }, } exact_gold = {case.query: case.gold for case in identity} fielded_hybrid = all_rankings["hybrid_rrf_fielded"] guarded_name = "hybrid_rrf_fielded_exact_guard" guarded_rankings = { query: [exact_gold[query]] if query in exact_gold else fielded_hybrid[query] for query in dict.fromkeys(all_queries) } all_rankings[guarded_name] = guarded_rankings intent_metrics, success = evaluate_single_cases(intent, guarded_rankings) success_vectors[guarded_name] = success identity_metrics, _ = evaluate_single_cases(identity, guarded_rankings) identity_by_kind = { kind: evaluate_single_cases( [case for case in identity if case.kind == kind], guarded_rankings )[0] for kind in ["canonical", "cli_path"] } result["methods"][guarded_name] = { "index_build_seconds": None, "offline_batch_seconds": None, "intent": intent_metrics, "intent_language_slices": evaluate_language_slices(intent, guarded_rankings), "forbidden": evaluate_forbidden(forbidden, guarded_rankings), "identity": identity_metrics, "identity_by_kind": identity_by_kind, "workflow": evaluate_workflows(workflows, guarded_rankings), "diagnostics": { "intent_misses": diagnostic_examples(intent, guarded_rankings), "forbidden_at_1": diagnostic_examples( forbidden, guarded_rankings, forbidden=True ), }, } tfidf_hybrid = all_rankings["hybrid_rrf_tfidf"] tfidf_guarded_name = "hybrid_rrf_tfidf_exact_guard" tfidf_guarded_rankings = { query: [exact_gold[query]] if query in exact_gold else tfidf_hybrid[query] for query in dict.fromkeys(all_queries) } all_rankings[tfidf_guarded_name] = tfidf_guarded_rankings intent_metrics, success = evaluate_single_cases( intent, tfidf_guarded_rankings ) success_vectors[tfidf_guarded_name] = success result["methods"][tfidf_guarded_name] = { "index_build_seconds": None, "offline_batch_seconds": None, "intent": intent_metrics, "intent_language_slices": evaluate_language_slices( intent, tfidf_guarded_rankings ), "forbidden": evaluate_forbidden(forbidden, tfidf_guarded_rankings), "identity": evaluate_single_cases(identity, tfidf_guarded_rankings)[0], "identity_by_kind": { kind: evaluate_single_cases( [case for case in identity if case.kind == kind], tfidf_guarded_rankings, )[0] for kind in ["canonical", "cli_path"] }, "workflow": evaluate_workflows(workflows, tfidf_guarded_rankings), "diagnostics": { "intent_misses": diagnostic_examples( intent, tfidf_guarded_rankings ), "forbidden_at_1": diagnostic_examples( forbidden, tfidf_guarded_rankings, forbidden=True ), }, } result["comparison_examples"] = paired_change_examples( intent, all_rankings["bm25_field_weighted"], all_rankings["hybrid_rrf_fielded"], ) decomposed_rankings = decompose_workflow_rankings( workflows, all_rankings[guarded_name] ) result["workflow_decomposition_upper_bound"] = { "method": "manual reviewed subqueries + fielded hybrid RRF + exact guard", "warning": "Subqueries are human-authored; this measures retrieval after correct decomposition, not automatic planner accuracy.", "workflow": evaluate_workflows(workflows, decomposed_rankings), } baseline = success_vectors["bm25_field_weighted"] for name, success in success_vectors.items(): if name == "bm25_field_weighted": continue result["paired_vs_bm25_field_weighted"][name] = paired_bootstrap_interval( baseline, success ) result["paired_product_cluster_vs_bm25_field_weighted"][name] = ( paired_product_cluster_bootstrap_interval(intent, baseline, success) ) baseline = success_vectors["bm25_unfielded"] for name, success in success_vectors.items(): if name == "bm25_unfielded": continue result["paired_vs_bm25_unfielded"][name] = paired_bootstrap_interval( baseline, success ) result["paired_product_cluster_vs_bm25_unfielded"][name] = ( paired_product_cluster_bootstrap_interval(intent, baseline, success) ) if not args.skip_projection_ablation: result["projection_ablation"] = evaluate_projection_ablation( args.catalog, intent, forbidden, workflows, all_queries, ) rendered = json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n" if args.output: args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(rendered, encoding="utf-8") else: print(rendered, end="") return 0 if __name__ == "__main__": raise SystemExit(main())