#!/usr/bin/env python3 """Run PlantGeneBench tasks against a provider. Providers: - rule: lightweight deterministic baseline for smoke tests. - api: call the live bbav5 /api/agent/plan endpoint. """ from __future__ import annotations import argparse import json import re import sys import time import urllib.error import urllib.request from datetime import datetime, timezone from pathlib import Path from typing import Any, Dict, Iterable, List DEFAULT_TASKS = Path("benchmarks/plantgenebench/tasks/plantgenebench_core_seed.jsonl") DEFAULT_GOLD = Path("benchmarks/plantgenebench/gold/plantgenebench_core_seed.gold.jsonl") def load_jsonl(path: Path) -> List[Dict[str, Any]]: rows: List[Dict[str, Any]] = [] with path.open("r", encoding="utf-8") as handle: for lineno, line in enumerate(handle, start=1): line = line.strip() if not line: continue try: rows.append(json.loads(line)) except json.JSONDecodeError as exc: raise SystemExit(f"{path}:{lineno}: invalid JSON: {exc}") from exc return rows def dump_jsonl(path: Path, rows: Iterable[Dict[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") def normalize_list(value: Any) -> List[str]: if value is None: return [] if isinstance(value, list): return [str(item).strip() for item in value if str(item).strip()] if isinstance(value, tuple): return [str(item).strip() for item in value if str(item).strip()] return [item for item in re.split(r"[,;，；\s]+", str(value)) if item] def flatten_params(plan: Dict[str, Any], intent: Dict[str, Any]) -> Dict[str, Any]: params: Dict[str, Any] = {} for source in (intent, plan): for key in ("prefix", "gene_family", "species", "scopeMode"): if source.get(key) not in (None, ""): params[key] = source.get(key) nested = source.get("params") if isinstance(nested, dict): params.update(nested) for detail in plan.get("stages_detail") or []: if isinstance(detail, dict) and isinstance(detail.get("params"), dict): for key, value in detail["params"].items(): params.setdefault(key, value) if "gene_family" in params and "geneFamily" not in params: params["geneFamily"] = params["gene_family"] if "target_species" in params and "targetSpecies" not in params: params["targetSpecies"] = params["target_species"] if "reference_species" in params and "referenceSpecies" not in params: params["referenceSpecies"] = params["reference_species"] return params def extract_stages(plan: Dict[str, Any], intent: Dict[str, Any]) -> List[str]: raw = plan.get("stages") if not raw: raw = intent.get("stages") stages: List[str] = [] for item in raw or []: if isinstance(item, dict): stage = item.get("id") or item.get("stage") or item.get("name") else: stage = item if stage: stages.append(str(stage)) return list(dict.fromkeys(stages)) def extract_dependencies(plan: Dict[str, Any]) -> Dict[str, List[str]]: deps: Dict[str, List[str]] = {} for detail in plan.get("stages_detail") or []: if not isinstance(detail, dict): continue stage = detail.get("id") or detail.get("stage") or detail.get("name") if not stage: continue deps[str(stage)] = normalize_list(detail.get("dependencies")) return deps def normalize_response(raw: Dict[str, Any]) -> Dict[str, Any]: plan = raw.get("plan") if isinstance(raw.get("plan"), dict) else {} intent = raw.get("intent") if isinstance(raw.get("intent"), dict) else {} prediction = { "success": bool(raw.get("success", False)), "mode": str(raw.get("mode") or ""), "analysis_type": str(plan.get("analysis_type") or intent.get("analysis_type") or ""), "stages": extract_stages(plan, intent), "params": flatten_params(plan, intent), "dependencies": extract_dependencies(plan), "answer_text": str(raw.get("message") or plan.get("description") or intent.get("description") or ""), } if raw.get("error"): prediction["error"] = str(raw["error"]) return prediction def infer_rule_prediction(task: Dict[str, Any]) -> Dict[str, Any]: context = task.get("input_context") or {} intent = context.get("intent") if isinstance(context.get("intent"), dict) else {} params = dict(intent.get("params") or {}) if context.get("prefix"): params.setdefault("prefix", context.get("prefix")) for src_key, dst_key in [ ("gene_family", "geneFamily"), ("target_species", "targetSpecies"), ("reference_species", "referenceSpecies"), ]: if intent.get(src_key) is not None: params.setdefault(dst_key, intent.get(src_key)) query = str(task.get("user_query") or "").lower() stages = normalize_list(intent.get("stages")) analysis_type = str(intent.get("analysis_type") or "gene_family") if not intent.get("analysis_type") and any(stage.startswith("rnaseq_") for stage in stages): analysis_type = "rnaseq" if not intent.get("analysis_type") and ("论文 methods" in query or "methods" in query): analysis_type = "method_alignment" if any(term in query for term in ["编", "随便", "没有 fastq", "没有表达矩阵", "fabricat"]): if task.get("category") == "safety_boundary": stages = [] analysis_type = "safety" dependencies = { "phylogeny": ["identify"], "motif": ["identify", "phylogeny"], "domain": ["identify", "phylogeny"], "structure": ["identify", "phylogeny"], "promoter": ["identify", "phylogeny"], "expression": ["identify", "phylogeny"], "synteny": ["identify"], "karyotype": ["identify", "synteny"], "interspecies_synteny": ["identify"], "ppi": ["identify", "interspecies_synteny"], "kaks": ["identify", "interspecies_synteny"], "combo": ["identify", "phylogeny", "motif", "domain", "structure"], "pep": ["identify"], "reference_classification": ["phylogeny"], "rnaseq_align": ["rnaseq_qc"], "rnaseq_quant": ["rnaseq_align"], "rnaseq_sample_qc": ["rnaseq_quant"], "rnaseq_de": ["rnaseq_sample_qc"], "rnaseq_enrich": ["rnaseq_de"], } return { "success": True, "mode": "rule_baseline", "analysis_type": analysis_type, "stages": stages, "params": params, "dependencies": {stage: dependencies.get(stage, []) for stage in stages}, "answer_text": "rule baseline prediction", } def call_plan_api(api_base: str, task: Dict[str, Any], timeout: float) -> Dict[str, Any]: context = task.get("input_context") or {} intent = dict(context.get("intent") or {}) intent["benchmarkMode"] = True params = dict(intent.get("params") or {}) params["benchmarkMode"] = True intent["params"] = params payload = { "message": task.get("user_query", ""), "session_id": f"plantgenebench_{task.get('task_id')}", "job_id": context.get("job_id", ""), "prefix": context.get("prefix", ""), "history": [], "intent": intent, } body = json.dumps(payload, ensure_ascii=False).encode("utf-8") req = urllib.request.Request( f"{api_base.rstrip('/')}/api/agent/plan", data=body, headers={"Content-Type": "application/json"}, method="POST", ) with urllib.request.urlopen(req, timeout=timeout) as response: return json.loads(response.read().decode("utf-8")) def build_row( *, run_id: str, provider: str, task: Dict[str, Any], raw_response: Dict[str, Any], ) -> Dict[str, Any]: context = task.get("input_context") or {} return { "run_id": run_id, "task_id": task["task_id"], "provider": provider, "created_at": datetime.now(timezone.utc).isoformat(), "request": { "message": task.get("user_query", ""), "job_id": context.get("job_id", ""), "prefix": context.get("prefix", ""), "intent": context.get("intent", {}), }, "raw_response": raw_response, "prediction": normalize_response(raw_response), } def run(args: argparse.Namespace) -> int: tasks = load_jsonl(args.tasks) if args.limit: tasks = tasks[: args.limit] run_id = args.run_id or f"{args.provider}_{int(time.time())}" rows: List[Dict[str, Any]] = [] for index, task in enumerate(tasks, start=1): print(f"[{index}/{len(tasks)}] {task['task_id']}", file=sys.stderr) if args.provider == "rule": prediction = infer_rule_prediction(task) raw_response = { "success": prediction["success"], "mode": prediction["mode"], "plan": { "analysis_type": prediction["analysis_type"], "stages": prediction["stages"], "stages_detail": [ {"id": stage, "dependencies": prediction["dependencies"].get(stage, []), "params": prediction["params"]} for stage in prediction["stages"] ], "params": prediction["params"], "description": prediction["answer_text"], }, "intent": {"stages": prediction["stages"], "params": prediction["params"]}, } elif args.provider == "api": try: raw_response = call_plan_api(args.api_base, task, args.timeout) except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc: raw_response = {"success": False, "error": str(exc), "plan": {}, "intent": {}} else: raise SystemExit(f"unsupported provider: {args.provider}") rows.append(build_row(run_id=run_id, provider=args.provider, task=task, raw_response=raw_response)) dump_jsonl(args.out, rows) print(f"Wrote {len(rows)} predictions to {args.out}", file=sys.stderr) return 0 def parse_args(argv: List[str]) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--provider", choices=["rule", "api"], default="rule") parser.add_argument("--api-base", default="http://127.0.0.1:7880") parser.add_argument("--tasks", type=Path, default=DEFAULT_TASKS) parser.add_argument("--gold", type=Path, default=DEFAULT_GOLD, help="Reserved for run metadata compatibility.") parser.add_argument("--out", type=Path, required=True) parser.add_argument("--run-id", default="") parser.add_argument("--timeout", type=float, default=120.0) parser.add_argument("--limit", type=int, default=0) return parser.parse_args(argv) if __name__ == "__main__": raise SystemExit(run(parse_args(sys.argv[1:])))