BioScientist/agent_system/main.py

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
from pathlib import Path

from .orchestrator import DualModeAgentSystem


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser("Bioinfo Dual-Mode Agent System")
    parser.add_argument("--project_root", type=str, default=str(Path(__file__).resolve().parents[1]))
    parser.add_argument(
        "--execution_backend",
        type=str,
        choices=["docker", "python"],
        default="docker",
        help="Tool execution backend. docker=run tool calls inside container images; python=run in host process.",
    )
    sub = parser.add_subparsers(dest="mode", required=True)

    p_exec = sub.add_parser("execute")
    p_exec.add_argument("--task", type=str, required=True)
    p_exec.add_argument("--task_scope", type=str, required=True)
    p_exec.add_argument("--input_manifest", type=str, required=True, help="JSON string")

    sub.add_parser("reflect")

    p_consult = sub.add_parser("consult")
    p_consult.add_argument("--user_goal", type=str, required=True)
    p_consult.add_argument("--task_scope", type=str, required=True)

    p_auto = sub.add_parser("autopilot")
    p_auto.add_argument("--user_goal", type=str, required=True)
    p_auto.add_argument("--data_dir", type=str, required=True)
    p_auto.add_argument("--task_scope", type=str, default="first_pipeline")

    p_cfg = sub.add_parser("propose-config")
    p_cfg.add_argument("--task_scope", type=str, required=True)
    p_cfg.add_argument("--strategy_name", type=str, required=True)
    p_cfg.add_argument("--tools", type=str, required=True, help="JSON list string")
    p_cfg.add_argument("--parameters", type=str, required=True, help="JSON object string")
    p_cfg.add_argument("--rationale", type=str, required=True)

    p_reg = sub.add_parser("register-mcp")
    p_reg.add_argument("--dry_run", action="store_true")

    p_hyp = sub.add_parser("hypothesis-generate")
    p_hyp.add_argument("--user_query", type=str, required=True)
    p_hyp.add_argument("--task_scope", type=str, required=True)
    p_hyp.add_argument("--n", type=int, default=10)
    p_hyp.add_argument("--top_k", type=int, default=5)

    p_loop = sub.add_parser("hypothesis-loop")
    p_loop.add_argument("--user_query", type=str, required=True)
    p_loop.add_argument("--task_scope", type=str, required=True)
    p_loop.add_argument("--n", type=int, default=10)
    p_loop.add_argument("--top_k", type=int, default=5)
    p_loop.add_argument("--validate_top_m", type=int, default=3)
    p_loop.add_argument("--validation_level", type=str, default="L1")
    p_loop.add_argument("--register_mcp", type=str, default="true")
    return parser


def _default_results_dir(project_root: Path) -> Path:
    # Prefer <project_root>/agent_system/results; fall back to package-local results directory.
    candidate = project_root / "agent_system" / "results"
    if candidate.parent.exists():
        return candidate
    return Path(__file__).resolve().parent / "results"


def _persist_result(
    result: dict,
    *,
    mode: str,
    project_root: Path,
) -> Path:
    results_dir = _default_results_dir(project_root)
    results_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    output_path = results_dir / f"{mode}_{ts}.json"
    output_path.write_text(json.dumps(result, indent=2, ensure_ascii=True), encoding="utf-8")
    return output_path


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()
    project_root = Path(args.project_root).resolve()
    system = DualModeAgentSystem(project_root=project_root, execution_backend=args.execution_backend)

    if args.mode == "execute":
        result = system.execute(
            task=args.task,
            input_manifest=json.loads(args.input_manifest),
            task_scope=args.task_scope,
        )
    elif args.mode == "reflect":
        result = system.reflect()
    elif args.mode == "consult":
        result = system.consult(user_goal=args.user_goal, task_scope=args.task_scope)
    elif args.mode == "autopilot":
        result = system.autopilot(
            user_goal=args.user_goal,
            data_dir=args.data_dir,
            task_scope=args.task_scope,
        )
    elif args.mode == "propose-config":
        result = system.propose_config(
            task_scope=args.task_scope,
            strategy_name=args.strategy_name,
            tools=json.loads(args.tools),
            parameters=json.loads(args.parameters),
            rationale=args.rationale,
        )
    elif args.mode == "register-mcp":
        result = system.register_mcp_servers(dry_run=bool(args.dry_run))
    elif args.mode == "hypothesis-generate":
        result = system.propose_hypotheses(
            user_query=args.user_query,
            task_scope=args.task_scope,
            n=args.n,
            top_k=args.top_k,
        )
    elif args.mode == "hypothesis-loop":
        register_mcp = str(args.register_mcp).strip().lower() in ("1", "true", "yes", "y", "on")
        result = system.hypothesis_loop(
            user_query=args.user_query,
            task_scope=args.task_scope,
            n=args.n,
            top_k=args.top_k,
            validate_top_m=args.validate_top_m,
            validation_level=args.validation_level,
            register_mcp=register_mcp,
        )
    else:  # pragma: no cover - argparse already guards
        raise ValueError(f"Unsupported mode: {args.mode}")

    saved_path = _persist_result(
        result,
        mode=args.mode,
        project_root=project_root,
    )
    result = dict(result)
    result["result_file"] = str(saved_path)
    print(json.dumps(result, indent=2, ensure_ascii=True))


if __name__ == "__main__":
    main()