#!/usr/bin/env python3 import argparse import csv import gzip import json import os import re import shutil import sys from datetime import datetime from pathlib import Path from time import perf_counter PROJECT_ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(PROJECT_ROOT)) DATASET_ROOT = Path(os.getenv("BIOAGENT_BENCH_DATASET_ROOT", PROJECT_ROOT.parent / "bioagent-bench" / "dataset")) METADATA_PATH = Path(os.getenv("BIOAGENT_BENCH_METADATA", PROJECT_ROOT.parent / "bioagent-bench" / "src" / "task_metadata.json")) DEFAULT_OUTPUT_ROOT = PROJECT_ROOT / "bioagent-bench-runs" DEFAULT_MCP_CONFIG = PROJECT_ROOT / "mcp_config_shim.yaml" DEFAULT_EXECUTION_ENV_PREFIX = Path(os.getenv("BIOMNI_EXECUTION_ENV_PREFIX", sys.prefix)) FORBIDDEN_VISIBLE_DIRS = {"biomni_data", "__pycache__"} TASK_OUTPUTS = { "alzheimer-mouse": ["pathway_comparison.csv"], "comparative-genomics": ["cluster_annotation_mapping.csv"], "cystic-fibrosis": ["cf_variants.csv"], "deseq": ["up_regulated_genes.csv"], "evolution": ["variants_shared.csv", "gene_annotations.csv"], "giab": ["predicted.vcf.gz"], "metagenomics": ["phylum_relative_abundances.csv"], "single-cell": ["all_clusters_de_genes.csv"], "transcript-quant": ["truth.tsv"], "viral-metagenomics": ["taxonomy.csv"], } TASK_SCHEMA_RULES = { "alzheimer-mouse": { "pathway_comparison.csv": { "format": "csv", "required_columns": ["pathway", "5xFAD_pvalue", "3xTG_AD_pvalue", "PS3O1S_pvalue"], } }, "comparative-genomics": { "cluster_annotation_mapping.csv": { "format": "csv", "required_columns": ["cluster_number", "consensus_annotation"], } }, "cystic-fibrosis": { "cf_variants.csv": { "format": "csv", "required_columns": [ "chromosome", "position", "variant_id", "reference", "alternate", "gene_name", "gene_id", "annotation", "impact", "transcript_id", "hgvs_c", "hgvs_p", "clinical_significance", "diseases", "review_status", "rs_id", ], } }, "deseq": { "up_regulated_genes.csv": { "format": "csv", "required_columns": ["gene_id", "log2FoldChange", "pvalue", "padj"], } }, "evolution": { "variants_shared.csv": { "format": "csv", "required_columns": ["CHROM", "POS", "REF", "ALT", "GENE", "IMPACT", "EFFECT", "STATUS"], }, "gene_annotations.csv": { "format": "csv", "required_columns": ["CHROM", "POS", "REF", "ALT", "GENE", "IMPACT", "EFFECT", "STATUS"], }, }, "giab": { "predicted.vcf.gz": { "format": "vcf.gz", } }, "metagenomics": { "phylum_relative_abundances.csv": { "format": "csv", "required_columns": ["OTU", "Kingdom", "Phylum", "JP4D", "JC1A"], } }, "single-cell": { "all_clusters_de_genes.csv": { "format": "csv", "required_columns": [ "cluster_id", "predicted_cell_type", "gene_name", "logfoldchanges", "pvals", "pvals_adj", "direction", "abs_logfc", ], } }, "transcript-quant": { "truth.tsv": { "format": "tsv_no_header", } }, "viral-metagenomics": { "taxonomy.csv": { "format": "csv", "required_columns": ["contig_count", "domain", "species"], } }, } TASK_EXTRA_INSTRUCTIONS = { "cystic-fibrosis": ( "For the final row, " "variant_id should be the ClinVar VCF ID for the matching record, and rs_id should be the " "numeric dbSNP identifier from ClinVar INFO when available, without adding an extra 'rs' prefix." ), "giab": ( "The final deliverable must be a bgzip-compatible .vcf.gz file. " "If you also generate index or benchmark helper files, keep them in the same run directory." ), "transcript-quant": ( "The final deliverable must be a two-column tab-separated file with no header line and no extra " "commentary around the table. Each line should be: transcript_idcount." ), } def load_task_metadata(metadata_path: Path) -> list[dict]: return json.loads(metadata_path.read_text(encoding="utf-8")) def build_agent_kwargs(args: argparse.Namespace) -> dict: provider = os.getenv("BIOMNI_LLM_PROVIDER", "").strip().lower() kwargs = { "expected_data_lake_files": [], "rewrite_user_query": args.rewrite_user_query, "dynamic_mcp_registration": args.dynamic_mcp_registration, "use_graph_retriever": args.use_graph_retriever, "use_tool_retriever": args.use_tool_retriever, "timeout_seconds": args.timeout_seconds, "mcp_server_top_k": args.mcp_server_top_k, "mcp_tool_top_k": args.mcp_tool_top_k, } if provider == "deepseek": kwargs.update( { "llm": os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat"), "source": "Custom", "base_url": os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"), "api_key": os.getenv("DEEPSEEK_API_KEY"), } ) elif args.llm: kwargs["llm"] = args.llm if args.source: kwargs["source"] = args.source if args.base_url: kwargs["base_url"] = args.base_url if args.api_key: kwargs["api_key"] = args.api_key return kwargs def list_visible_files(directory: Path, limit: int = 24) -> list[str]: if not directory.exists(): return [] files = [] for path in sorted(directory.rglob("*")): if path.is_file(): rel_path = path.relative_to(directory) rel_parts = rel_path.parts if rel_parts and rel_parts[0] in FORBIDDEN_VISIBLE_DIRS: continue files.append(str(rel_path)) if len(files) >= limit: break return files def build_benchmark_policy(task_meta: dict, task_dir: Path, run_dir: Path) -> str: data_dir = task_dir / "data" ref_dir = task_dir / "reference" results_dir = task_dir / "results" lines = [ "Benchmark data policy:", f"- Allowed input data directory: {data_dir}", f"- Allowed reference directory: {ref_dir if ref_dir.exists() else ''}", f"- Allowed scratch/output directory: {run_dir}", f"- Forbidden truth/results directory: {results_dir}", f"- Forbidden sibling benchmark task directories: {DATASET_ROOT}/", f"- Forbidden generated Biomni cache/runtime directories inside benchmark inputs: {data_dir}/biomni_data and {ref_dir}/biomni_data", "- Do not inspect previous bioagent-bench-runs as data sources.", "- Do not download external databases or install new packages during the benchmark run.", "- You may use installed command-line tools, Python/R packages, and MCP servers as executors, but their inputs must come from the allowed paths above.", ] return "\n".join(lines) def build_execution_guard(task_meta: dict, task_dir: Path, run_dir: Path) -> dict: task_id = task_meta["task_id"] dataset_root = re.escape(str(DATASET_ROOT)) task_id_re = re.escape(task_id) task_dir_re = re.escape(str(task_dir)) run_root_re = re.escape(str(run_dir.parent)) run_name_re = re.escape(run_dir.name) return { "enabled": True, "allowed_roots": [ str(task_dir / "data"), str(task_dir / "reference"), str(run_dir), ], "forbidden_patterns": [ rf"{dataset_root}/(?!{task_id_re}(?:/|$|[\s'\"<>]))[^\s'\"<>]+", rf"{task_dir_re}/results(?:/|$|[^\s'\"<>]*)", rf"{task_dir_re}/(?:data|reference)/biomni_data(?:/|$|[^\s'\"<>]*)", rf"{run_root_re}/(?!{run_name_re}(?:/|$|[\s'\"<>]))[^\s'\"<>]+", rf"os\\.walk\\(['\"]{dataset_root}['\"]\\)", rf"Path\\(['\"]{dataset_root}['\"]\\)\\.rglob", ], "forbidden_substrings": [ "pip install", "conda install", "mamba install", "install.packages(", "BiocManager::install", "http://", "https://", ], "forbidden_commands": [ "wget ", "curl ", "aws s3 cp", "gsutil cp", ], } def build_benchmark_task_context(task_meta: dict, output_paths: list[Path]) -> dict: task_id = task_meta["task_id"] return { "task_id": task_id, "task_name": task_meta.get("name", task_id), "description": task_meta.get("description", ""), "task_prompt": task_meta.get("task_prompt", ""), "extra_instruction": TASK_EXTRA_INSTRUCTIONS.get(task_id, ""), "required_outputs": [path.name for path in output_paths], } def build_delivery_guardrails(task_id: str, output_paths: list[Path]) -> str: rules = TASK_SCHEMA_RULES.get(task_id, {}) lines = [ "Before writing the final , validate the deliverable yourself against these fairness-preserving checks:", "1. The final file names must match the required output paths exactly.", "2. The final file schema must match the requested columns/format exactly.", "3. Do not export a background universe or broad intermediate table when the prompt asks for a filtered/shared/significant final result set.", "4. Do not switch to a different reference coordinate system, taxonomy database, or condition contrast without explicitly proving it is still the task's provided one.", "5. If a tool path fails, do not silently change the biological question, reference space, or output definition just to produce a file.", ] for output_path in output_paths: spec = rules.get(output_path.name, {}) required_columns = spec.get("required_columns", []) if required_columns: lines.append(f"- {output_path.name} required columns: {', '.join(required_columns)}") for warning in spec.get("warnings", []): lines.append(f"- {output_path.name}: {warning}") return "\n".join(lines) def _safe_float(value): try: return float(value) except (TypeError, ValueError): return None def _safe_int(value): try: return int(str(value).strip()) except (TypeError, ValueError): return None def _read_text_preview(path: Path, limit: int = 2000) -> str: if path.suffix == ".gz": with gzip.open(path, "rt", encoding="utf-8", errors="ignore") as handle: return handle.read(limit) return path.read_text(encoding="utf-8", errors="ignore")[:limit] def _find_reference_contigs(task_dir: Path) -> set[str]: contigs: set[str] = set() for path in sorted((task_dir / "reference").glob("*")): if not path.is_file(): continue suffixes = "".join(path.suffixes).lower() if not any(token in suffixes for token in (".fa", ".fasta", ".fna", ".fa.gz", ".fasta.gz", ".fna.gz")): continue try: if path.suffix == ".gz": handle = gzip.open(path, "rt", encoding="utf-8", errors="ignore") else: handle = path.open("r", encoding="utf-8", errors="ignore") with handle: for line in handle: if line.startswith(">"): contigs.add(line[1:].strip().split()[0]) if len(contigs) >= 5000: return contigs except OSError: continue return contigs def validate_output_file(task_id: str, task_dir: Path, path: Path) -> dict: spec = TASK_SCHEMA_RULES.get(task_id, {}).get(path.name, {}) result = { "file": str(path), "exists": path.exists(), "errors": [], "warnings": [], "summary": {}, } if not path.exists(): result["errors"].append("missing_output_file") return result fmt = spec.get("format") if fmt in {"csv", "tsv_no_header"}: delimiter = "\t" if fmt == "tsv_no_header" else "," with path.open("r", encoding="utf-8", errors="ignore", newline="") as handle: rows = list(csv.reader(handle, delimiter=delimiter)) result["summary"]["row_count"] = max(0, len(rows) - (0 if fmt == "tsv_no_header" else 1)) if fmt == "tsv_no_header": if rows and len(rows[0]) != 2: result["errors"].append("expected_two_columns_without_header") else: header = rows[0] if rows else [] result["summary"]["header"] = header required_columns = spec.get("required_columns", []) missing_columns = [col for col in required_columns if col not in header] if missing_columns: result["errors"].append(f"missing_required_columns:{','.join(missing_columns)}") if task_id == "alzheimer-mouse" and result["summary"]["row_count"] > 150: result["warnings"].append("appears_to_export_large_pathway_universe") if task_id == "comparative-genomics" and result["summary"]["row_count"] > 500: result["warnings"].append("appears_to_export_unfiltered_cluster_universe") if task_id == "deseq": try: dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore"))) non_positive = sum( 1 for row in dict_rows if (_safe_float(row.get("log2FoldChange")) is not None and _safe_float(row.get("log2FoldChange")) <= 0) ) if non_positive: result["warnings"].append(f"contains_{non_positive}_non_upregulated_rows") except OSError: pass if task_id == "metagenomics": try: dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore"))) kingdoms = sorted({(row.get("Kingdom") or "").strip() for row in dict_rows if row.get("Kingdom")}) result["summary"]["kingdoms"] = kingdoms if any(k and k != "Bacteria" for k in kingdoms): result["warnings"].append("contains_non_bacterial_rows") sums = {} for sample in ("JP4D", "JC1A"): vals = [_safe_float(row.get(sample)) for row in dict_rows] vals = [v for v in vals if v is not None] if vals: sums[sample] = round(sum(vals), 4) if not (99.0 <= sums[sample] <= 101.0): result["warnings"].append(f"{sample}_relative_abundance_sum_not_near_100") result["summary"]["sample_sums"] = sums except OSError: pass if task_id == "single-cell": try: dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore"))) bad_direction = 0 for row in dict_rows[:5000]: direction = (row.get("direction") or "").strip().lower() logfc = _safe_float(row.get("logfoldchanges")) if logfc is None or direction not in {"up", "down"}: continue if (direction == "up" and logfc < 0) or (direction == "down" and logfc > 0): bad_direction += 1 if bad_direction: result["warnings"].append(f"direction_logfc_mismatch_rows:{bad_direction}") except OSError: pass if task_id == "viral-metagenomics": try: dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore"))) negative_counts = sum( 1 for row in dict_rows if (_safe_int(row.get("contig_count")) is not None and _safe_int(row.get("contig_count")) < 0) ) if negative_counts: result["errors"].append("negative_contig_count") if len(dict_rows) > 25: result["warnings"].append("appears_to_export_overly_broad_taxonomic_summary") except OSError: pass elif fmt == "vcf.gz": preview = _read_text_preview(path) result["summary"]["preview"] = preview[:400] if not preview.startswith("##") and "#CHROM" not in preview: result["errors"].append("vcf_header_not_detected") if task_id in {"evolution", "giab"} and path.exists(): reference_contigs = _find_reference_contigs(task_dir) if reference_contigs: observed_contigs = set() try: if path.suffix == ".gz": handle = gzip.open(path, "rt", encoding="utf-8", errors="ignore") is_vcf = True else: handle = path.open("r", encoding="utf-8", errors="ignore") is_vcf = False with handle: for line in handle: if not line.strip(): continue if is_vcf and line.startswith("#"): continue if path.suffix != ".gz" and line.lower().startswith("chrom,"): continue observed_contigs.add(line.split("\t", 1)[0] if is_vcf else line.split(",", 1)[0]) if len(observed_contigs) >= 100: break except OSError: observed_contigs = set() if observed_contigs and observed_contigs.isdisjoint(reference_contigs): result["warnings"].append("observed_coordinate_system_not_in_reference_headers") result["summary"]["observed_contig_examples"] = sorted(list(observed_contigs))[:10] return result def validate_outputs(task_id: str, task_dir: Path, output_paths: list[Path]) -> dict: file_reports = [validate_output_file(task_id, task_dir, path) for path in output_paths] fatal = [err for report in file_reports for err in report["errors"]] warnings = [warning for report in file_reports for warning in report["warnings"]] return { "passed": not fatal, "file_reports": file_reports, "fatal_errors": fatal, "warnings": warnings, } def build_query( task_meta: dict, task_dir: Path, run_dir: Path, output_paths: list[Path], ) -> str: data_dir = task_dir / "data" ref_dir = task_dir / "reference" output_lines = [f"- {path.name}: {path}" for path in output_paths] data_lines = [f"- {name}" for name in list_visible_files(data_dir)] ref_lines = [f"- {name}" for name in list_visible_files(ref_dir)] if ref_dir.exists() else [] extra = TASK_EXTRA_INSTRUCTIONS.get(task_meta["task_id"], "") policy = build_benchmark_policy(task_meta, task_dir, run_dir) return f""" You are running a bioagent-bench task with local files already prepared. Task ID: {task_meta["task_id"]} Task name: {task_meta["name"]} Benchmark prompt: {task_meta["task_prompt"]} Data background: {task_meta["description"]} Constraints: 1. Use only the benchmark inputs and references explicitly listed below. 2. Do not inspect or use any files under benchmark truth/results directories, sibling task directories, generated biomni_data caches, or previous run outputs. 3. Save the required final deliverables exactly to the paths listed below. 4. Save any intermediate scripts, logs, and scratch outputs inside this run directory: {run_dir} 5. Keep final deliverables in the same schema/format requested by the benchmark prompt. 6. Return a concise final summary after writing the required files. 7. The runner, Python REPL, MCP servers, Rscript, and CLI subprocesses are bound to this conda environment: {os.environ.get("CONDA_PREFIX", DEFAULT_EXECUTION_ENV_PREFIX)}. Do not switch to another conda environment. {policy} Input data directory: {data_dir} Visible input files: {chr(10).join(data_lines) if data_lines else "- "} Reference data directory: {ref_dir if ref_dir.exists() else ""} Visible reference files: {chr(10).join(ref_lines) if ref_lines else "- "} Required final output paths: {chr(10).join(output_lines)} """.strip() def save_json(path: Path, payload: dict) -> None: path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8") def count_tokens(text: str | None) -> int: if not isinstance(text, str) or not text.strip(): return 0 try: import tiktoken return len(tiktoken.get_encoding("cl100k_base").encode(text)) except Exception: return max(1, len(re.findall(r"\S+", text))) def build_token_usage_payload(agent) -> dict[str, int]: usage = dict(getattr(agent, "last_token_usage", {}) or {}) payload: dict[str, int] = {} for key in ("prompt_tokens", "completion_tokens", "total_tokens", "llm_call_count"): value = usage.get(key, 0) try: payload[key] = int(value) except (TypeError, ValueError): payload[key] = 0 return payload def select_tasks_for_shard(task_ids: list[str], shard_index: int | None, shard_count: int | None) -> list[str]: if shard_index is None and shard_count is None: return task_ids if shard_index is None or shard_count is None: raise SystemExit("Provide both --shard-index and --shard-count together.") if shard_count <= 0: raise SystemExit("--shard-count must be > 0.") if shard_index < 0 or shard_index >= shard_count: raise SystemExit("--shard-index must satisfy 0 <= shard_index < shard_count.") return [task_id for idx, task_id in enumerate(task_ids) if idx % shard_count == shard_index] def _build_bound_env(env_prefix: Path) -> dict[str, str]: """Return an environment that resolves Python/CLI tools from one conda env.""" env = os.environ.copy() env_bin = env_prefix / "bin" path_parts = [] for part in env.get("PATH", "").split(os.pathsep): if not part: continue # Avoid silently falling back to another conda env such as bioenv_cli. if "/miniconda3/envs/" in part and Path(part).resolve() != env_bin.resolve(): continue if part not in path_parts: path_parts.append(part) env["PATH"] = os.pathsep.join([str(env_bin), *path_parts]) env["CONDA_PREFIX"] = str(env_prefix) env["CONDA_DEFAULT_ENV"] = env_prefix.name env["CONDA_SHLVL"] = "1" env["PYTHONNOUSERSITE"] = "1" env["BIOMNI_EXECUTION_ENV_PREFIX"] = str(env_prefix) env["BIOMNI_EXECUTION_PYTHON"] = str(env_bin / "python") env.pop("VIRTUAL_ENV", None) return env def bind_process_to_execution_env(env_prefix: Path, *, reexec: bool = True) -> dict[str, str]: """Bind this benchmark runner to biomni_e1 and optionally re-exec into its Python.""" env_prefix = env_prefix.expanduser().resolve() env_python = env_prefix / "bin" / "python" if not env_python.exists(): raise SystemExit(f"Execution environment Python not found: {env_python}") bound_env = _build_bound_env(env_prefix) current_python = Path(sys.executable).resolve() if reexec and current_python != env_python.resolve(): os.execve(str(env_python), [str(env_python), *sys.argv], bound_env) os.environ.clear() os.environ.update(bound_env) return bound_env def _rewrite_csv_header(path: Path, header_map: dict[str, str]) -> bool: if not path.exists() or path.stat().st_size == 0: return False with path.open("r", encoding="utf-8", newline="") as handle: rows = list(csv.reader(handle)) if not rows: return False original = rows[0] rewritten = [header_map.get(col.strip(), header_map.get(col.strip().lower(), col.strip())) for col in original] if rewritten == original: return False with path.open("w", encoding="utf-8", newline="") as handle: writer = csv.writer(handle) writer.writerow(rewritten) writer.writerows(rows[1:]) return True def _normalize_cystic_fibrosis_csv(path: Path) -> dict: report = {"file": str(path), "actions": []} if not path.exists() or path.stat().st_size == 0: report["actions"].append("missing_or_empty") return report with path.open("r", encoding="utf-8", newline="") as handle: reader = csv.DictReader(handle) rows = list(reader) fieldnames = reader.fieldnames or [] if "rs_id" in fieldnames: changed = False for row in rows: value = (row.get("rs_id") or "").strip() if value.lower().startswith("rs") and value[2:].isdigit(): row["rs_id"] = value[2:] changed = True if changed: with path.open("w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) report["actions"].append("stripped_rs_prefix") return report def _normalize_transcript_quant_tsv(path: Path) -> dict: report = {"file": str(path), "actions": []} if not path.exists() or path.stat().st_size == 0: report["actions"].append("missing_or_empty") return report original_lines = path.read_text(encoding="utf-8").splitlines() cleaned = [line.strip() for line in original_lines if line.strip()] if cleaned and cleaned[0].lower().replace(" ", "") in {"transcript_id\tcount", "transcript_id,count"}: cleaned = cleaned[1:] report["actions"].append("removed_header") normalized = [] for line in cleaned: parts = [part.strip() for part in line.replace(",", "\t").split("\t") if part.strip()] if len(parts) >= 2: normalized.append(f"{parts[0]}\t{parts[1]}") if normalized != original_lines: path.write_text("\n".join(normalized) + ("\n" if normalized else ""), encoding="utf-8") report["actions"].append("normalized_two_column_tsv") return report def postprocess_outputs(task_id: str, output_paths: list[Path]) -> list[dict]: """Apply schema-only cleanup that does not read benchmark truth files.""" reports = [] for path in output_paths: report = {"file": str(path), "actions": []} if not path.exists(): report["actions"].append("missing") reports.append(report) continue if task_id == "alzheimer-mouse" and path.name == "pathway_comparison.csv": if _rewrite_csv_header(path, {"pathway": "Pathway"}): report["actions"].append("canonicalized_pathway_header") elif task_id == "cystic-fibrosis" and path.name == "cf_variants.csv": report = _normalize_cystic_fibrosis_csv(path) elif task_id == "transcript-quant" and path.name == "truth.tsv": report = _normalize_transcript_quant_tsv(path) reports.append(report) return reports def run_task(task_meta: dict, args: argparse.Namespace, output_root: Path) -> dict: from biomni.agent import A1 task_id = task_meta["task_id"] task_dir = DATASET_ROOT / task_id timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S") run_dir = output_root / f"{task_id}_{timestamp}" run_dir.mkdir(parents=True, exist_ok=True) output_filenames = TASK_OUTPUTS[task_id] output_paths = [run_dir / name for name in output_filenames] query = build_query(task_meta, task_dir, run_dir, output_paths) agent_kwargs = build_agent_kwargs(args) agent_kwargs["path"] = str(run_dir / "agent_runtime") agent_kwargs["execution_env_prefix"] = str(Path(args.execution_env_prefix).expanduser().resolve()) agent_kwargs["benchmark_guard"] = build_execution_guard(task_meta, task_dir, run_dir) agent_kwargs["benchmark_task_context"] = build_benchmark_task_context(task_meta, output_paths) metadata = { "task_id": task_id, "task_name": task_meta["name"], "run_dir": str(run_dir), "dataset_dir": str(task_dir), "data_dir": str(task_dir / "data"), "reference_dir": str(task_dir / "reference"), "agent_runtime_dir": agent_kwargs["path"], "output_paths": [str(path) for path in output_paths], "agent_kwargs": agent_kwargs, "query": query, "benchmark_policy": build_benchmark_policy(task_meta, task_dir, run_dir), "benchmark_execution_guard": agent_kwargs["benchmark_guard"], "benchmark_task_context": agent_kwargs["benchmark_task_context"], "timestamp_utc": timestamp, "runtime_environment": { "execution_env_prefix": str(Path(args.execution_env_prefix).expanduser().resolve()), "execution_python": os.environ.get("BIOMNI_EXECUTION_PYTHON", sys.executable), "conda_default_env": os.environ.get("CONDA_DEFAULT_ENV"), "conda_prefix": os.environ.get("CONDA_PREFIX"), "path_head": os.environ.get("PATH", "").split(os.pathsep)[:5], }, } save_json(run_dir / "run_metadata.json", metadata) (run_dir / "task_query.txt").write_text(query, encoding="utf-8") agent = A1(**agent_kwargs) if args.use_mcp and args.mcp_graph and Path(args.mcp_graph).exists(): agent.attach_prebuilt_mcp_graph(str(args.mcp_graph), executable_only=args.executable_mcp_only) elif args.use_mcp and args.mcp_config and Path(args.mcp_config).exists(): agent.attach_mcp_catalog(str(args.mcp_config)) run_started = perf_counter() log_entries, answer = agent.go(query) total_runtime_seconds = perf_counter() - run_started token_usage = build_token_usage_payload(agent) planning_context_text = getattr(agent, "last_planning_context_text", None) planning_context_tokens = count_tokens(planning_context_text) (run_dir / "final_answer.txt").write_text(str(answer), encoding="utf-8") save_json(run_dir / "execution_log.json", {"log_entries": log_entries}) (run_dir / "execution_log.txt").write_text("\n\n".join(str(entry) for entry in log_entries), encoding="utf-8") save_json( run_dir / "retrieval_plan.json", { "query_context": getattr(agent, "query_context", {}), "mcp_enabled": args.use_mcp, "graph_enabled": args.use_graph_retriever, "mcp_graph_route": getattr(agent, "last_graph_route", {}), "internal_tool_graph_route": getattr(agent, "last_internal_tool_route", {}), "planning_context_text": planning_context_text, "planning_context_tokens": planning_context_tokens, "planning_context_chars": len(planning_context_text) if isinstance(planning_context_text, str) else 0, "planning_latency_seconds": getattr(agent, "last_retrieval_latency_seconds", None), "total_runtime_seconds": total_runtime_seconds, "token_usage": token_usage, "selected_resources": getattr(agent, "last_selected_resources", {}), "selected_resource_names": getattr(agent, "last_selected_resources_names", {}), "registered_tool_count": len(agent.tool_registry.tools) if hasattr(getattr(agent, "tool_registry", None), "tools") else None, "registered_tool_names": [ tool.get("name") for tool in getattr(getattr(agent, "tool_registry", None), "tools", []) if isinstance(tool, dict) and tool.get("name") ], }, ) postprocess_report = postprocess_outputs(task_id, output_paths) validation_report = validate_outputs(task_id, task_dir, output_paths) save_json( run_dir / "output_validation.json", { "postprocess": postprocess_report, "validator": validation_report, }, ) output_status = [] for path in output_paths: output_status.append( { "path": str(path), "exists": path.exists(), "size_bytes": path.stat().st_size if path.exists() else 0, } ) result = { "task_id": task_id, "run_dir": str(run_dir), "final_answer_path": str(run_dir / "final_answer.txt"), "metadata_path": str(run_dir / "run_metadata.json"), "query_path": str(run_dir / "task_query.txt"), "retrieval_plan_path": str(run_dir / "retrieval_plan.json"), "output_validation_path": str(run_dir / "output_validation.json"), "outputs": output_status, "planning_latency_seconds": getattr(agent, "last_retrieval_latency_seconds", None), "total_runtime_seconds": total_runtime_seconds, "planning_context_tokens": planning_context_tokens, "planning_context_chars": len(planning_context_text) if isinstance(planning_context_text, str) else 0, "validation_passed": validation_report["passed"], "validation_warning_count": len(validation_report["warnings"]), **token_usage, } save_json(run_dir / "run_summary.json", result) metadata["post_run_metrics"] = { "planning_latency_seconds": getattr(agent, "last_retrieval_latency_seconds", None), "total_runtime_seconds": total_runtime_seconds, "planning_context_tokens": planning_context_tokens, "planning_context_chars": len(planning_context_text) if isinstance(planning_context_text, str) else 0, "validation_passed": validation_report["passed"], "validation_warning_count": len(validation_report["warnings"]), "validation_fatal_errors": validation_report["fatal_errors"], **token_usage, } save_json(run_dir / "run_metadata.json", metadata) return result def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run bioagent-bench tasks with Biomanus.") parser.add_argument("--task", action="append", help="Task ID to run. Can be provided multiple times.") parser.add_argument("--all", action="store_true", help="Run all supported tasks.") parser.add_argument("--metadata", default=str(METADATA_PATH)) parser.add_argument("--dataset-root", default=str(DATASET_ROOT)) parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT)) parser.add_argument("--mcp-config", default=str(DEFAULT_MCP_CONFIG)) parser.add_argument( "--mcp-graph", default=None, help="Prebuilt MCP graph directory containing server_catalog.json. Takes precedence over --mcp-config.", ) parser.add_argument("--executable-mcp-only", action="store_true", help="When using --mcp-graph, skip entries without commands.") parser.add_argument("--llm", default=None) parser.add_argument("--source", default=None) parser.add_argument("--base-url", default=None) parser.add_argument("--api-key", default=None) parser.add_argument("--timeout-seconds", type=int, default=1200) parser.add_argument("--mcp-server-top-k", type=int, default=20) parser.add_argument("--mcp-tool-top-k", type=int, default=12) parser.add_argument( "--execution-env-prefix", default=str(DEFAULT_EXECUTION_ENV_PREFIX), help="Conda environment prefix used for the runner, Python REPL, MCP servers, and CLI tools.", ) parser.add_argument( "--no-env-reexec", dest="env_reexec", action="store_false", help="Do not re-exec the benchmark runner with --execution-env-prefix/bin/python.", ) parser.set_defaults(env_reexec=True) parser.add_argument("--rewrite-user-query", action="store_true", default=True) parser.add_argument("--no-rewrite-user-query", dest="rewrite_user_query", action="store_false") parser.add_argument("--dynamic-mcp-registration", action="store_true", default=True) parser.add_argument("--no-dynamic-mcp-registration", dest="dynamic_mcp_registration", action="store_false") parser.add_argument("--use-graph-retriever", action="store_true", default=True) parser.add_argument("--no-graph-retriever", dest="use_graph_retriever", action="store_false") parser.add_argument("--use-mcp", action="store_true", default=True) parser.add_argument("--no-mcp", dest="use_mcp", action="store_false") parser.add_argument("--use-tool-retriever", action="store_true", default=True) parser.add_argument("--no-use-tool-retriever", dest="use_tool_retriever", action="store_false") parser.add_argument("--shard-index", type=int, default=None, help="0-based shard index over the selected task list.") parser.add_argument("--shard-count", type=int, default=None, help="Total number of shards over the selected task list.") return parser.parse_args() def main() -> int: args = parse_args() execution_env = bind_process_to_execution_env(Path(args.execution_env_prefix), reexec=args.env_reexec) global DATASET_ROOT DATASET_ROOT = Path(args.dataset_root) metadata = load_task_metadata(Path(args.metadata)) meta_by_task = {item["task_id"]: item for item in metadata if item["task_id"] in TASK_OUTPUTS} if args.all: task_ids = list(meta_by_task) else: task_ids = args.task or [] task_ids = select_tasks_for_shard(task_ids, args.shard_index, args.shard_count) if not task_ids: if args.shard_index is not None: print( f"No tasks assigned to shard {args.shard_index}/{args.shard_count} under the current selection; exiting." ) return 0 raise SystemExit("Provide --task or use --all.") unsupported = [task_id for task_id in task_ids if task_id not in meta_by_task] if unsupported: raise SystemExit(f"Unsupported task IDs: {unsupported}") output_root = Path(args.output_root) output_root.mkdir(parents=True, exist_ok=True) print( "Execution environment: " f"{execution_env['CONDA_DEFAULT_ENV']} ({execution_env['CONDA_PREFIX']}); " f"python={sys.executable}" ) summaries = [] for task_id in task_ids: print(f"Running task: {task_id}") summary = run_task(meta_by_task[task_id], args, output_root) summaries.append(summary) print(json.dumps(summary, ensure_ascii=False, indent=2)) batch_summary = { "timestamp_utc": datetime.utcnow().strftime("%Y%m%d_%H%M%S"), "shard_index": args.shard_index, "shard_count": args.shard_count, "tasks": summaries, } save_json(output_root / "latest_batch_summary.json", batch_summary) print(f"Batch summary: {output_root / 'latest_batch_summary.json'}") return 0 if __name__ == "__main__": raise SystemExit(main())