SupplyAgent/find_agent_v2/test_entry.py

"""Local CLI for preparing or executing one database-backed v2 test run."""

from __future__ import annotations

import argparse
import json

from find_agent_v2.demand_context import prepare_v2_demand_run
from find_agent_v2.runner import run_find_agent_v2, run_prepared_find_agent_v2
from find_agent_v2.service import get_find_agent_v2_service


def main() -> None:
    parser = argparse.ArgumentParser(description="准备或执行一条寻找 Agent v2 本地测试任务")
    parser.add_argument("--demand-grade-id", type=int, default=None)
    parser.add_argument("--index", type=int, default=0, help="未指定需求 ID 时选择最新 S 级上下文序号")
    parser.add_argument("--run-id", default=None)
    parser.add_argument("--existing-run-id", default=None, help="读取并执行已准备的 v2 run")
    parser.add_argument("--execute", action="store_true", help="准备后立即调用模型和外部搜索")
    parser.add_argument("--model", default=None)
    parser.add_argument("--resume", action="store_true", help="从 terminal run 的下一轮恢复")
    args = parser.parse_args()

    if args.existing_run_id:
        run = get_find_agent_v2_service().require_run(args.existing_run_id)
        prepared = None
        output: dict = {"prepared": {
            "run_id": run["run_id"],
            "demand_grade_id": run["demand_grade_id"],
            "demand_name": run["demand_word"],
            "status": run["status"],
        }, "executed": False}
    else:
        prepared = prepare_v2_demand_run(
            demand_grade_id=args.demand_grade_id,
            index=args.index,
            run_id=args.run_id,
        )
        output = {"prepared": prepared.summary(), "executed": False}
    if args.execute:
        result = (
            run_prepared_find_agent_v2(
                args.existing_run_id, model=args.model, resume=args.resume,
            )
            if args.existing_run_id
            else run_find_agent_v2(
                prepared.user_input,
                run_id=prepared.run_id,
                model=args.model,
            )
        )
        output.update({"executed": True, "result": {
            "status": result.status,
            "succeeded": result.succeeded,
            "valid_primary_count": result.valid_primary_count,
            "stop_reason": result.stop_reason,
            "final_output": result.final_output,
        }})
    print(json.dumps(output, ensure_ascii=False, default=str, indent=2))


if __name__ == "__main__":
    main()