| |
| """Aggregate MCP scaling experiment metrics from Biomni/bioagent-bench run dirs.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| import time |
| from functools import lru_cache |
| from pathlib import Path |
| from typing import Any |
|
|
| DEFAULT_EXP_DIR = Path(__file__).resolve().parents[1] |
| DEFAULT_GOLD = DEFAULT_EXP_DIR / "gold_tools.json" |
| DEFAULT_MANIFEST = DEFAULT_EXP_DIR / "configs" / "mcp_scale_manifest.json" |
| DEFAULT_DATASET_ROOT = Path("/225040511/project/bioagent-bench/dataset") |
| DEFAULT_PRIMARY_K = 10 |
|
|
| TYPE_MISMATCH_PATTERNS = [ |
| r"typeerror", |
| r"valueerror", |
| r"validationerror", |
| r"pydantic", |
| r"missing required positional argument", |
| r"unexpected keyword argument", |
| r"unsupported operand type", |
| r"input should be", |
| r"should be of type", |
| r"must be (?:a|an)?\s*(?:int|integer|float|bool|boolean|string|str|list|dict|path|file)", |
| r"expected .*?(?:int|integer|float|bool|boolean|string|str|list|dict|path|file)", |
| r"invalid literal for int", |
| r"cannot be interpreted as an integer", |
| r"is not iterable", |
| ] |
|
|
| CONSTRAINT_VIOLATION_PATTERNS = [ |
| r"\bpip install\b", |
| r"\bconda install\b", |
| r"\bapt-get\b", |
| r"\bwget\b", |
| r"\bcurl\b", |
| r"\bgit clone\b", |
| r"https?://", |
| ] |
|
|
|
|
| def load_json(path: Path, default: Any = None) -> Any: |
| if not path.exists(): |
| return default |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def load_text(path: Path, default: str = "") -> str: |
| if not path.exists(): |
| return default |
| return path.read_text(encoding="utf-8", errors="ignore") |
|
|
|
|
| def load_yaml_if_available(path: Path, default: Any = None) -> Any: |
| if not path.exists(): |
| return default |
| try: |
| import yaml |
| except ModuleNotFoundError: |
| return default |
| return yaml.safe_load(path.read_text(encoding="utf-8")) or default |
|
|
|
|
| def normalize(text: str) -> str: |
| return re.sub(r"[^a-z0-9]+", "", text.lower()) |
|
|
|
|
| def extract_names(obj: Any) -> list[str]: |
| names: list[str] = [] |
| if isinstance(obj, dict): |
| for key, value in obj.items(): |
| if key in {"name", "tool_name", "server", "module"} and isinstance(value, str): |
| names.append(value) |
| names.extend(extract_names(value)) |
| elif isinstance(obj, list): |
| for item in obj: |
| names.extend(extract_names(item)) |
| elif isinstance(obj, str): |
| for token in re.findall(r"\b[a-zA-Z][a-zA-Z0-9_.-]{2,}\b", obj): |
| if any( |
| hint in token.lower() |
| for hint in ( |
| "mcp", |
| "tool", |
| "kraken", |
| "kallisto", |
| "salmon", |
| "scanpy", |
| "star", |
| "bwa", |
| "bcftools", |
| "samtools", |
| "deseq", |
| "kaiju", |
| "megahit", |
| "gatk", |
| "freebayes", |
| ) |
| ): |
| names.append(token) |
| return names |
|
|
|
|
| def recall_at_k(candidates: list[str], gold: list[str], k: int) -> float: |
| if not gold: |
| return 0.0 |
| top = {normalize(x) for x in candidates[:k]} |
| hits = 0 |
| for item in gold: |
| n = normalize(item) |
| if any(n in candidate or candidate in n for candidate in top): |
| hits += 1 |
| return hits / len(gold) |
|
|
|
|
| def gold_hit(candidates: list[str], gold: list[str]) -> bool: |
| return recall_at_k(candidates, gold, max(len(candidates), 1)) > 0 |
|
|
|
|
| def mean_or_none(values: list[float | int | None]) -> float | None: |
| numeric = [float(value) for value in values if isinstance(value, (int, float))] |
| if not numeric: |
| return None |
| return sum(numeric) / len(numeric) |
|
|
|
|
| def workflow_validity(evidence: list[str], workflow: list[dict[str, Any]]) -> dict[str, Any]: |
| found = [] |
| cursor = 0 |
| normalized_evidence = [normalize(x) for x in evidence] |
| for step in workflow: |
| patterns = step.get("tool_patterns", []) |
| matched_at = None |
| for i in range(cursor, len(normalized_evidence)): |
| row = normalized_evidence[i] |
| if any(normalize(pattern) in row for pattern in patterns): |
| matched_at = i |
| break |
| if matched_at is not None: |
| found.append(step.get("name", "step")) |
| cursor = matched_at + 1 |
| return { |
| "valid": len(found) == len(workflow) and bool(workflow), |
| "steps_matched": len(found), |
| "steps_total": len(workflow), |
| "matched_step_names": found, |
| } |
|
|
|
|
| def execution_success(run_dir: Path) -> bool: |
| validation = load_json(run_dir / "output_validation.json", {}) |
| if isinstance(validation, dict): |
| if validation.get("success") is True or validation.get("all_outputs_exist") is True: |
| return True |
| outputs = validation.get("outputs") or validation.get("output_paths") |
| if isinstance(outputs, list) and outputs: |
| return all(bool(x.get("exists")) for x in outputs if isinstance(x, dict)) |
| summary = load_json(run_dir / "run_summary.json", {}) |
| if isinstance(summary, dict): |
| if summary.get("success") is True: |
| return True |
| if summary.get("status") in {"success", "completed"}: |
| return True |
| metadata = load_json(run_dir / "run_metadata.json", {}) |
| for output in metadata.get("output_paths", []) if isinstance(metadata, dict) else []: |
| if not Path(output).exists(): |
| return False |
| return bool(metadata.get("output_paths")) if isinstance(metadata, dict) else False |
|
|
|
|
| def rough_token_count(text: str) -> int: |
| return max(1, len(re.findall(r"\S+", text))) |
|
|
|
|
| def discover_task_ids(dataset_root: Path) -> list[str]: |
| if not dataset_root.exists(): |
| return [] |
| return sorted(p.name for p in dataset_root.iterdir() if p.is_dir()) |
|
|
|
|
| def collect_task_runs(runs_root: Path, task_id: str) -> list[Path]: |
| if not runs_root.exists(): |
| return [] |
| pattern = re.compile(rf"^{re.escape(task_id)}_(\d{{8}}_\d{{6}})$") |
| matched: list[tuple[str, Path]] = [] |
| for run_dir in runs_root.iterdir(): |
| if not run_dir.is_dir(): |
| continue |
| match = pattern.match(run_dir.name) |
| if match: |
| matched.append((match.group(1), run_dir)) |
| matched.sort(key=lambda x: x[0]) |
| return [path for _, path in matched] |
|
|
|
|
| @lru_cache(maxsize=16) |
| def load_mcp_tool_universe(config_path: str) -> list[str]: |
| cfg = load_yaml_if_available(Path(config_path), default={}) |
| if not isinstance(cfg, dict): |
| return [] |
| names: list[str] = [] |
| for server_meta in (cfg.get("mcp_servers") or {}).values(): |
| if not isinstance(server_meta, dict): |
| continue |
| for tool in server_meta.get("tools", []): |
| if isinstance(tool, dict) and tool.get("name"): |
| names.append(tool["name"]) |
| return names |
|
|
|
|
| def selected_tool_names(retrieval_plan: dict[str, Any]) -> list[str]: |
| names = retrieval_plan.get("selected_resource_names", {}).get("tools", []) |
| if names: |
| return [str(name) for name in names if name] |
| tool_items = retrieval_plan.get("selected_resources", {}).get("tools", []) |
| return [item.get("name") for item in tool_items if isinstance(item, dict) and item.get("name")] |
|
|
|
|
| def registered_tool_names(retrieval_plan: dict[str, Any]) -> list[str]: |
| return [str(name) for name in retrieval_plan.get("registered_tool_names", []) if name] |
|
|
|
|
| def stored_selected_tool_count(retrieval_plan: dict[str, Any]) -> int | None: |
| count = retrieval_plan.get("selected_tool_count") |
| return count if isinstance(count, int) else None |
|
|
|
|
| def retrieval_metadata_available(retrieval_plan: dict[str, Any]) -> bool: |
| if selected_tool_names(retrieval_plan): |
| return True |
| if registered_tool_names(retrieval_plan): |
| return True |
| planning_context = retrieval_plan.get("planning_context_text") |
| if isinstance(planning_context, str) and planning_context.strip(): |
| return True |
| query_context = retrieval_plan.get("query_context") |
| if isinstance(query_context, dict) and query_context: |
| return True |
| return False |
|
|
|
|
| def retrieval_candidates_from_plan(retrieval_plan: dict[str, Any]) -> list[str]: |
| candidates: list[str] = [] |
| candidates.extend(selected_tool_names(retrieval_plan)) |
| candidates.extend(registered_tool_names(retrieval_plan)) |
| planning_context = retrieval_plan.get("planning_context_text") |
| if isinstance(planning_context, str) and planning_context.strip(): |
| candidates.extend(extract_names(planning_context)) |
| query_context = retrieval_plan.get("query_context") |
| if isinstance(query_context, dict) and query_context: |
| candidates.extend(extract_names(query_context)) |
| return sorted(dict.fromkeys(name for name in candidates if name)) |
|
|
|
|
| def planning_context_fallback(retrieval_plan: dict[str, Any], metadata: dict[str, Any], run_dir: Path) -> str: |
| planning_context_text = retrieval_plan.get("planning_context_text") |
| if isinstance(planning_context_text, str) and planning_context_text.strip(): |
| return planning_context_text |
|
|
| query_context = retrieval_plan.get("query_context") |
| if isinstance(query_context, dict) and query_context: |
| return json.dumps(query_context, ensure_ascii=False, default=str) |
|
|
| query = retrieval_plan.get("query") |
| if isinstance(query, str) and query.strip(): |
| return query |
|
|
| metadata_query = metadata.get("query") if isinstance(metadata, dict) else None |
| if isinstance(metadata_query, str) and metadata_query.strip(): |
| return metadata_query |
|
|
| task_query = load_text(run_dir / "task_query.txt").strip() |
| if task_query: |
| return task_query |
|
|
| return "" |
|
|
|
|
| def non_human_log_text(execution_log: dict[str, Any], execution_text_path: Path, final_text: str) -> str: |
| entries = execution_log.get("log_entries", []) if isinstance(execution_log, dict) else [] |
| filtered = [] |
| for entry in entries: |
| text = str(entry) |
| if "================================ Human Message" in text: |
| continue |
| filtered.append(text) |
| if filtered: |
| return "\n\n".join(filtered) + "\n" + final_text |
| if execution_text_path.exists(): |
| return execution_text_path.read_text(encoding="utf-8", errors="ignore") + "\n" + final_text |
| return final_text |
|
|
|
|
| def matched_tool_names(text: str, candidate_names: list[str]) -> list[str]: |
| haystack = normalize(text) |
| matched = [] |
| for name in candidate_names: |
| token = normalize(name) |
| if token and token in haystack: |
| matched.append(name) |
| return sorted(set(matched)) |
|
|
|
|
| def execution_attempt_count(text: str, executed_tools: list[str]) -> int: |
| counts = [ |
| len(re.findall(r"<execute>", text, flags=re.IGNORECASE)), |
| len(re.findall(r"command_executed", text, flags=re.IGNORECASE)), |
| len(executed_tools), |
| ] |
| return max(1, *counts) |
|
|
|
|
| def pattern_rate(text: str, patterns: list[str], denominator: int) -> tuple[float, int]: |
| count = 0 |
| for pattern in patterns: |
| count += len(re.findall(pattern, text, flags=re.IGNORECASE)) |
| return min(1.0, count / max(1, denominator)), count |
|
|
|
|
| def biological_constraint_rate(text: str, metadata: dict[str, Any], denominator: int) -> tuple[float, int]: |
| count = 0 |
| for pattern in CONSTRAINT_VIOLATION_PATTERNS: |
| count += len(re.findall(pattern, text, flags=re.IGNORECASE)) |
|
|
| task_id = metadata.get("task_id", "") if isinstance(metadata, dict) else "" |
| if task_id: |
| sibling_pattern = rf"/225040511/project/bioagent-bench/dataset/(?!{re.escape(task_id)})([a-zA-Z0-9_.-]+)" |
| count += len(re.findall(sibling_pattern, text)) |
| count += len(re.findall(r"/225040511/project/bioagent-bench/dataset/.+?/results", text)) |
| return min(1.0, count / max(1, denominator)), count |
|
|
|
|
| def summarize_run(run_dir: Path, gold: dict[str, Any], ks: list[int], primary_k: int) -> dict[str, Any]: |
| metadata = load_json(run_dir / "run_metadata.json", {}) |
| task_id = metadata.get("task_id") if isinstance(metadata, dict) else None |
| if not task_id: |
| task_id = re.sub(r"_20\d{6}_\d{6}$", "", run_dir.name) |
| gold_entry = gold.get(task_id, {}) |
|
|
| retrieval_plan = load_json(run_dir / "retrieval_plan.json", {}) |
| execution_log = load_json(run_dir / "execution_log.json", {}) |
| execution_text_path = run_dir / "execution_log.txt" |
| final_text = (run_dir / "final_answer.txt").read_text(encoding="utf-8", errors="ignore") if (run_dir / "final_answer.txt").exists() else "" |
|
|
| retrieval_candidates = retrieval_candidates_from_plan(retrieval_plan) |
| retrieval_available = retrieval_metadata_available(retrieval_plan) |
| filtered_execution_text = non_human_log_text(execution_log, execution_text_path, final_text) |
| execution_evidence = extract_names(execution_log) + re.findall( |
| r"\b[a-zA-Z][a-zA-Z0-9_.-]{2,}\b", |
| filtered_execution_text, |
| ) |
|
|
| gold_tools = gold_entry.get("gold_tools", []) |
| gold_servers = gold_entry.get("gold_servers", []) |
| all_gold = [*gold_tools, *gold_servers] |
| recall = { |
| str(k): (recall_at_k(retrieval_candidates, all_gold, k) if retrieval_available else None) |
| for k in ks |
| } |
|
|
| planning_context_text = planning_context_fallback(retrieval_plan, metadata, run_dir) |
|
|
| timing = metadata.get("timing", {}) if isinstance(metadata, dict) else {} |
| metrics = metadata.get("metrics", {}) if isinstance(metadata, dict) else {} |
| planning_latency = ( |
| retrieval_plan.get("planning_latency_seconds") |
| or timing.get("planning_latency_seconds") |
| or metrics.get("planning_latency") |
| or metrics.get("plan_seconds") |
| ) |
|
|
| mcp_config_path = retrieval_plan.get("mcp_config") or metadata.get("mcp_config") |
| mcp_tool_universe = load_mcp_tool_universe(str(mcp_config_path)) if mcp_config_path else [] |
| selected_tools = selected_tool_names(retrieval_plan) |
| registered_tools = registered_tool_names(retrieval_plan) |
| retrieved_tool_candidates = set(selected_tools or registered_tools) |
| executed_mcp_tools = matched_tool_names(filtered_execution_text, mcp_tool_universe) if mcp_tool_universe else [] |
| hallucinated_tools = sorted(set(executed_mcp_tools) - retrieved_tool_candidates) |
| hallucinated_tool_rate = ( |
| len(hallucinated_tools) / len(executed_mcp_tools) |
| if executed_mcp_tools and retrieved_tool_candidates |
| else None |
| ) |
|
|
| attempts = execution_attempt_count(filtered_execution_text, executed_mcp_tools) |
| data_type_mismatch_rate, data_type_mismatch_events = pattern_rate(filtered_execution_text, TYPE_MISMATCH_PATTERNS, attempts) |
| biological_constraint_error_rate, biological_constraint_events = biological_constraint_rate( |
| filtered_execution_text, |
| metadata if isinstance(metadata, dict) else {}, |
| attempts, |
| ) |
|
|
| workflow = workflow_validity(execution_evidence, gold_entry.get("workflow", [])) |
| execution_ok = execution_success(run_dir) |
|
|
| return { |
| "task_id": task_id, |
| "run_dir": str(run_dir), |
| "scale_label": run_dir.parent.name, |
| "retrieval_metadata_available": retrieval_available, |
| "retrieval_recall_at_k": recall, |
| "retrieval_recall_at_primary_k": recall.get(str(primary_k)), |
| "retrieval_recall_k": primary_k, |
| "tool_selection_accuracy": gold_hit(execution_evidence, all_gold), |
| "workflow_validity": workflow, |
| "execution_success": execution_ok, |
| "context_tokens_rough": rough_token_count(planning_context_text), |
| "planning_latency_seconds": planning_latency, |
| "retrieved_candidate_count": len(retrieval_candidates), |
| "execution_evidence_count": len(execution_evidence), |
| "selected_tool_count": len(selected_tools) if selected_tools else stored_selected_tool_count(retrieval_plan) or 0, |
| "registered_tool_count": len(registered_tools), |
| "executed_mcp_tools": executed_mcp_tools, |
| "hallucinated_tools": hallucinated_tools, |
| "hallucinated_tool_rate": hallucinated_tool_rate, |
| "data_type_mismatch_rate": data_type_mismatch_rate, |
| "data_type_mismatch_events": data_type_mismatch_events, |
| "biological_constraint_error_rate": biological_constraint_error_rate, |
| "biological_constraint_error_events": biological_constraint_events, |
| "execution_attempts_estimate": attempts, |
| } |
|
|
|
|
| def aggregate_results(results: list[dict[str, Any]], ks: list[int], primary_k: int) -> dict[str, Any]: |
| aggregate: dict[str, Any] = {"run_count": len(results), "retrieval_recall_k": primary_k} |
| if not results: |
| return aggregate |
|
|
| for k in ks: |
| vals = [r["retrieval_recall_at_k"][str(k)] for r in results] |
| aggregate[f"mean_retrieval_recall@{k}"] = mean_or_none(vals) |
| aggregate["tool_selection_accuracy"] = sum(r["tool_selection_accuracy"] for r in results) / len(results) |
| aggregate["workflow_validity"] = sum(r["workflow_validity"]["valid"] for r in results) / len(results) |
| aggregate["execution_success_rate"] = sum(r["execution_success"] for r in results) / len(results) |
| aggregate["mean_context_tokens_rough"] = sum(r["context_tokens_rough"] for r in results) / len(results) |
| latencies = [r["planning_latency_seconds"] for r in results if isinstance(r["planning_latency_seconds"], (int, float))] |
| aggregate["mean_planning_latency_seconds"] = sum(latencies) / len(latencies) if latencies else None |
| aggregate["retrieval_metadata_coverage"] = ( |
| sum(r["retrieval_metadata_available"] for r in results) / len(results) |
| ) |
| aggregate["mean_hallucinated_tool_rate"] = mean_or_none([r["hallucinated_tool_rate"] for r in results]) |
| aggregate["mean_data_type_mismatch_rate"] = mean_or_none([r["data_type_mismatch_rate"] for r in results]) |
| aggregate["mean_biological_constraint_error_rate"] = mean_or_none( |
| [r["biological_constraint_error_rate"] for r in results] |
| ) |
|
|
| aggregate["Retrieval Recall@k"] = aggregate.get(f"mean_retrieval_recall@{primary_k}") |
| aggregate["Workflow Validity"] = aggregate["workflow_validity"] |
| aggregate["Execution Success Rate"] = aggregate["execution_success_rate"] |
| aggregate["Context Tokens"] = aggregate["mean_context_tokens_rough"] |
| aggregate["Planning Latency"] = aggregate["mean_planning_latency_seconds"] |
| aggregate["Hallucinated Tool Rate"] = aggregate["mean_hallucinated_tool_rate"] |
| aggregate["Data-Type Mismatch Rate"] = aggregate["mean_data_type_mismatch_rate"] |
| aggregate["Biological Constraint Error Rate"] = aggregate["mean_biological_constraint_error_rate"] |
| return aggregate |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--runs-root", type=Path, required=True, help="Directory containing tier/task run dirs.") |
| parser.add_argument( |
| "--dataset-root", |
| type=Path, |
| default=DEFAULT_DATASET_ROOT, |
| help="Benchmark dataset root used to enumerate task IDs for per-task outputs.", |
| ) |
| parser.add_argument("--gold", type=Path, default=DEFAULT_GOLD) |
| parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST) |
| parser.add_argument("--out", type=Path, default=DEFAULT_EXP_DIR / "results" / "metrics_summary.json") |
| parser.add_argument( |
| "--per-task-dir", |
| type=Path, |
| default=DEFAULT_EXP_DIR / "results" / "by_task", |
| help="Directory where one metrics JSON per task is written.", |
| ) |
| parser.add_argument("--k", type=int, nargs="+", default=[1, 3, 5, 10, 20, 50]) |
| parser.add_argument("--primary-k", type=int, default=DEFAULT_PRIMARY_K) |
| args = parser.parse_args() |
|
|
| started = time.time() |
| gold = load_json(args.gold, {}) |
| manifest = load_json(args.manifest, {}) |
| task_ids = discover_task_ids(args.dataset_root) |
| args.per_task_dir.mkdir(parents=True, exist_ok=True) |
|
|
| results: list[dict[str, Any]] = [] |
| task_index = [] |
| for task_id in task_ids: |
| task_run_dirs = collect_task_runs(args.runs_root, task_id) |
| task_results = [summarize_run(run_dir, gold, args.k, args.primary_k) for run_dir in task_run_dirs] |
| results.extend(task_results) |
| task_payload = { |
| "evaluated_at_unix": started, |
| "task_id": task_id, |
| "runs_root": str(args.runs_root), |
| "dataset_root": str(args.dataset_root), |
| "manifest": manifest, |
| "aggregate": aggregate_results(task_results, args.k, args.primary_k), |
| "results": task_results, |
| } |
| task_out = args.per_task_dir / f"{task_id}.json" |
| task_out.write_text(json.dumps(task_payload, indent=2, ensure_ascii=False), encoding="utf-8") |
| task_index.append({"task_id": task_id, "run_count": len(task_results), "output": str(task_out)}) |
|
|
| aggregate = aggregate_results(results, args.k, args.primary_k) |
| payload = { |
| "evaluated_at_unix": started, |
| "runs_root": str(args.runs_root), |
| "dataset_root": str(args.dataset_root), |
| "manifest": manifest, |
| "tasks": task_index, |
| "aggregate": aggregate, |
| "results": results, |
| } |
| args.out.parent.mkdir(parents=True, exist_ok=True) |
| args.out.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8") |
| print(json.dumps(payload["aggregate"], indent=2, ensure_ascii=False)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
