Agent/examples/subagent_example.py

"""
Sub-Agent 使用示例

演示如何使用 Sub-Agent 机制处理复杂任务。

注意：本示例中的 AgentDefinition 和 get_agent_registry 尚未实现，
此处仅用于演示未来的设计方向。当前可用的 subagent 功能通过
runner.run(messages, config=RunConfig(...)) 和工具层的 subagent 工具实现。
"""

import asyncio
import os
from agent import AgentRunner
from agent.core.runner import RunConfig
from agent.trace import Trace, Message
from agent.llm import create_gemini_llm_call


async def example_basic_subagent():
    """示例 1: 使用 Agent 执行任务（通过 subagent 工具自动委托子任务）"""
    print("=== 示例 1: 基本 Agent 执行 ===\n")

    runner = AgentRunner(
        llm_call=create_gemini_llm_call(os.getenv("GEMINI_API_KEY")),
    )

    task = """
    分析这个 Python 项目的架构：
    1. 找出所有主要的模块和它们的职责
    2. 识别核心的数据流
    3. 列出使用的外部依赖

    请使用 subagent explore 模式来探索代码库。
    """

    async for item in runner.run(
        messages=[{"role": "user", "content": task}],
        config=RunConfig(
            model="gemini-2.0-flash-exp",
            max_iterations=20,
            name="项目架构分析",
        ),
    ):
        if isinstance(item, Trace):
            if item.status == "running":
                print(f"[Trace] 开始: {item.trace_id[:8]}")
            elif item.status == "completed":
                print(f"[Trace] 完成 (tokens: {item.total_tokens})")
        elif isinstance(item, Message):
            if item.role == "assistant":
                content = item.content
                if isinstance(content, dict):
                    text = content.get("text", "")
                    tool_calls = content.get("tool_calls")
                    if tool_calls:
                        for tc in tool_calls:
                            tool_name = tc.get("function", {}).get("name", "")
                            if tool_name == "subagent":
                                print(f"  启动 Sub-Agent...")
                    elif text:
                        print(f"\n最终结果:\n{text[:500]}")


async def main():
    """运行示例"""
    await example_basic_subagent()


if __name__ == "__main__":
    asyncio.run(main())