Agent/examples/subagent_unified/run.py

"""
统一 subagent 工具集成测试（mock LLM）。
"""

import asyncio
import os
import sys
from pathlib import Path
from tempfile import TemporaryDirectory

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from agent.core.runner import AgentRunner
from agent.trace.store import FileSystemTraceStore
from agent.trace.models import Trace
from agent.trace.goal_models import GoalTree
from agent.tools.builtin.subagent import subagent


async def mock_llm_call(messages, model="gpt-4o", tools=None, **kwargs):
    last_user = ""
    for msg in reversed(messages):
        if msg.get("role") == "user":
            last_user = str(msg.get("content", ""))
            break

    if "# 评估任务" in last_user:
        content = "## 评估结论\n通过\n\n## 评估理由\n满足需求。"
    elif "# 探索任务" in last_user:
        content = "探索完成：建议优先采用方案 1。"
    else:
        content = "委托任务已完成。"

    return {
        "content": content,
        "tool_calls": None,
        "finish_reason": "stop",
        "prompt_tokens": 10,
        "completion_tokens": 10,
        "cost": 0.0,
    }


async def run_case():
    with TemporaryDirectory(prefix="subagent-unified-") as tmp_dir:
        store = FileSystemTraceStore(base_path=tmp_dir)
        runner = AgentRunner(trace_store=store, llm_call=mock_llm_call)

        # 创建主 Trace 与 GoalTree（供 subagent 作为父上下文）
        main_trace = Trace(
            trace_id="main-trace",
            mode="agent",
            task="主任务",
            agent_type="default",
            status="running",
        )
        await store.create_trace(main_trace)
        goal_tree = GoalTree(mission="主任务")
        new_goals = goal_tree.add_goals(["实现主流程"])
        goal_tree.focus(new_goals[0].id)
        await store.update_goal_tree(main_trace.trace_id, goal_tree)

        context = {
            "store": store,
            "trace_id": main_trace.trace_id,
            "goal_id": new_goals[0].id,
            "runner": runner,
        }

        # 1) delegate
        delegate_result = await subagent(
            mode="delegate",
            task="实现用户登录功能",
            context=context,
        )
        assert delegate_result["status"] == "completed", delegate_result
        assert delegate_result["summary"], delegate_result
        delegate_trace = await store.get_trace(delegate_result["sub_trace_id"])
        assert delegate_trace is not None
        assert delegate_trace.parent_trace_id == main_trace.trace_id
        assert delegate_trace.parent_goal_id == new_goals[0].id

        # 2) explore
        explore_result = await subagent(
            mode="explore",
            branches=["JWT 方案", "Session 方案"],
            background="请比较维护成本和安全性。",
            context=context,
        )
        assert explore_result["status"] == "completed", explore_result
        assert "探索" in explore_result["summary"], explore_result

        # 3) evaluate
        evaluate_result = await subagent(
            mode="evaluate",
            target_goal_id=new_goals[0].id,
            evaluation_input={"actual_result": "已实现登录接口并通过单元测试"},
            requirements="请评估是否满足安全和可维护性要求。",
            context=context,
        )
        assert evaluate_result["status"] == "completed", evaluate_result
        assert "评估结论" in evaluate_result["summary"], evaluate_result

        # 4) continue_from
        continue_result = await subagent(
            mode="delegate",
            task="继续补充边界条件处理",
            continue_from=delegate_result["sub_trace_id"],
            context=context,
        )
        assert continue_result["status"] == "completed", continue_result
        assert continue_result["sub_trace_id"] == delegate_result["sub_trace_id"]
        assert continue_result["continue_from"] is True

        print("✅ unified subagent tests passed")
        print(f"delegate: {delegate_result['sub_trace_id']}")
        print(f"explore : {explore_result['sub_trace_id']}")
        print(f"evaluate: {evaluate_result['sub_trace_id']}")


def main():
    asyncio.run(run_case())


if __name__ == "__main__":
    main()