289 lines
11 KiB
Python
289 lines
11 KiB
Python
|
#!/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:])))
|
||
|
|
|