Agent/agent/trace/delegate.py

"""
Delegate 工具 - 委托任务给子 Agent

将大任务委托给独立的 Sub-Trace 执行，获得完整权限。
"""

from typing import Optional, Dict, Any
from datetime import datetime

from .models import Trace, Message
from .trace_id import generate_sub_trace_id
from .goal_models import Goal


async def delegate_tool(
    current_trace_id: str,
    current_goal_id: str,
    task: str,
    store=None,
    run_agent=None
) -> str:
    """
    将任务委托给独立的 Sub-Agent

    Args:
        current_trace_id: 当前主 Trace ID
        current_goal_id: 当前 Goal ID
        task: 委托的任务描述
        store: TraceStore 实例
        run_agent: 运行 Agent 的函数

    Returns:
        任务执行结果摘要

    Example:
        >>> result = await delegate_tool(
        ...     current_trace_id="abc123",
        ...     current_goal_id="3",
        ...     task="实现用户登录功能",
        ...     store=store,
        ...     run_agent=run_agent_func
        ... )
    """
    if not store:
        raise ValueError("store parameter is required")
    if not run_agent:
        raise ValueError("run_agent parameter is required")

    # 1. 创建 agent_call Goal
    await store.update_goal(current_trace_id, current_goal_id,
                           type="agent_call",
                           agent_call_mode="delegate",
                           status="in_progress")

    # 2. 生成 Sub-Trace ID
    sub_trace_id = generate_sub_trace_id(current_trace_id, "delegate")

    # 3. 创建 Sub-Trace
    sub_trace = Trace(
        trace_id=sub_trace_id,
        mode="agent",
        task=task,
        parent_trace_id=current_trace_id,
        parent_goal_id=current_goal_id,
        agent_type="delegate",
        context={
            # delegate 模式：完整权限
            "allowed_tools": None,  # None = 所有工具
            "max_turns": 50
        },
        status="running",
        created_at=datetime.now()
    )

    # 保存 Sub-Trace
    await store.create_trace(sub_trace)

    # 更新主 Goal 的 sub_trace_ids
    await store.update_goal(current_trace_id, current_goal_id, sub_trace_ids=[sub_trace_id])

    # 推送 sub_trace_started 事件
    await store.append_event(current_trace_id, "sub_trace_started", {
        "trace_id": sub_trace_id,
        "parent_trace_id": current_trace_id,
        "parent_goal_id": current_goal_id,
        "agent_type": "delegate",
        "task": task
    })

    # 4. 执行 Sub-Trace
    try:
        result = await run_agent(sub_trace)

        # 获取 Sub-Trace 的最终状态
        updated_trace = await store.get_trace(sub_trace_id)

        if isinstance(result, dict):
            summary = result.get("summary", "任务完成")
        else:
            summary = "任务完成"

        # 推送 sub_trace_completed 事件
        await store.append_event(current_trace_id, "sub_trace_completed", {
            "trace_id": sub_trace_id,
            "status": "completed",
            "summary": summary,
            "stats": {
                "total_messages": updated_trace.total_messages if updated_trace else 0,
                "total_tokens": updated_trace.total_tokens if updated_trace else 0,
                "total_cost": updated_trace.total_cost if updated_trace else 0
            }
        })

        # 5. 完成主 Goal
        await store.update_goal(current_trace_id, current_goal_id,
                               status="completed",
                               summary=f"已委托完成: {task}")

        # 格式化返回结果
        return f"""## 委托任务完成

**任务**: {task}

**结果**: {summary}

**统计**:
- 消息数: {updated_trace.total_messages if updated_trace else 0}
- Tokens: {updated_trace.total_tokens if updated_trace else 0}
- 成本: ${updated_trace.total_cost if updated_trace else 0:.4f}
"""

    except Exception as e:
        # 推送失败事件
        await store.append_event(current_trace_id, "sub_trace_completed", {
            "trace_id": sub_trace_id,
            "status": "failed",
            "error": str(e)
        })

        # 更新主 Goal 为失败
        await store.update_goal(current_trace_id, current_goal_id,
                               status="failed",
                               summary=f"委托任务失败: {str(e)}")

        return f"""## 委托任务失败

**任务**: {task}

**错误**: {str(e)}
"""


def create_delegate_tool_schema() -> Dict[str, Any]:
    """
    创建 delegate 工具的 JSON Schema

    Returns:
        工具的 JSON Schema
    """
    return {
        "type": "function",
        "function": {
            "name": "delegate",
            "description": "将大任务委托给独立的 Sub-Agent 执行。Sub-Agent 拥有完整权限，适合执行复杂的、需要多步骤的任务。",
            "parameters": {
                "type": "object",
                "properties": {
                    "task": {
                        "type": "string",
                        "description": "要委托的任务描述，应该清晰具体"
                    }
                },
                "required": ["task"]
            }
        }
    }