Agent/examples/subagent_example.py

"""
Sub-Agent 使用示例

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

import asyncio
import os
from agent import AgentRunner, AgentConfig
from agent.llm import create_gemini_llm_call
from agent.agent_registry import get_agent_registry
from agent.models.agent import AgentDefinition


async def example_basic_subagent():
    """示例 1: 使用内置 Sub-Agent"""
    print("=== 示例 1: 使用内置 Sub-Agent ===\n")

    # 配置主 Agent
    config = AgentConfig(
        agent_type="primary",
        max_iterations=20,
        skills_dir="./skills",
    )

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

    # 主 Agent 会自动知道如何使用 Task 工具启动 Sub-Agent
    task = """
    分析这个 Python 项目的架构：
    1. 找出所有主要的模块和它们的职责
    2. 识别核心的数据流
    3. 列出使用的外部依赖

    请使用 explore sub-agent 来探索代码库。
    """

    async for event in runner.run(task=task, model="gemini-2.0-flash-exp"):
        if event.type == "tool_call" and event.data.get("tool") == "task":
            print(f"🚀 启动 Sub-Agent: {event.data['args']['subagent_type']}")
            print(f"   任务: {event.data['args']['description']}")

        elif event.type == "tool_result" and event.data.get("tool") == "task":
            metadata = event.data.get("metadata", {})
            print(f"✅ Sub-Agent 完成")
            print(f"   Sub-Trace ID: {metadata.get('sub_trace_id')}")
            print(f"   工具调用: {metadata.get('tool_summary')}")

        elif event.type == "conclusion":
            print(f"\n📋 最终结果:\n{event.data['content']}")


async def example_custom_subagent():
    """示例 2: 定义和使用自定义 Sub-Agent"""
    print("\n=== 示例 2: 自定义 Sub-Agent ===\n")

    # 1. 定义自定义 Sub-Agent
    security_scanner = AgentDefinition(
        name="security-scanner",
        description="安全扫描专家，查找常见漏洞和安全问题",
        mode="subagent",
        allowed_tools=[
            "read_file",
            "search_code",
            "list_files",
        ],
        denied_tools=[
            "write_file",
            "edit_file",
            "execute_bash",
        ],
        system_prompt="""你是一个安全扫描专家。专注于：

1. 查找常见安全漏洞：
   - SQL 注入
   - XSS 攻击
   - CSRF 漏洞
   - 不安全的序列化
   - 路径遍历

2. 检查敏感信息泄露：
   - 硬编码的密钥和密码
   - API token
   - 数据库凭据

3. 依赖项安全：
   - 过时的依赖
   - 已知漏洞的包

输出格式：
- **漏洞类型**: [类型名称]
- **严重程度**: [高/中/低]
- **位置**: [文件路径:行号]
- **描述**: [详细说明]
- **修复建议**: [如何修复]
""",
        max_iterations=25,
        temperature=0.2,  # 更精确的分析
    )

    # 2. 注册到全局注册表
    get_agent_registry().register(security_scanner)
    print(f"✓ 注册自定义 Sub-Agent: {security_scanner.name}")

    # 3. 使用自定义 Sub-Agent
    config = AgentConfig(
        agent_type="primary",
        max_iterations=20,
    )

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

    task = """
    对这个项目进行安全审计，使用 security-scanner sub-agent。
    重点关注：
    1. 认证和授权相关代码
    2. 数据库交互代码
    3. 用户输入处理
    """

    async for event in runner.run(task=task, model="gemini-2.0-flash-exp"):
        if event.type == "tool_call" and event.data.get("tool") == "task":
            print(f"🔍 启动安全扫描...")

        elif event.type == "conclusion":
            print(f"\n🔒 安全审计报告:\n{event.data['content']}")


async def example_multiple_subagents():
    """示例 3: 并行使用多个 Sub-Agent"""
    print("\n=== 示例 3: 并行使用多个 Sub-Agent ===\n")

    config = AgentConfig(
        agent_type="primary",
        max_iterations=30,
    )

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

    task = """
    完整分析这个项目：

    1. 使用 explore sub-agent 探索代码结构
    2. 使用 analyst sub-agent 分析技术栈和依赖
    3. 综合两个分析结果，给出项目概览

    可以并行启动多个 sub-agent 来加快速度。
    """

    subagent_results = {}

    async for event in runner.run(task=task, model="gemini-2.0-flash-exp"):
        if event.type == "tool_call" and event.data.get("tool") == "task":
            agent_type = event.data['args']['subagent_type']
            print(f"🔄 启动 Sub-Agent: {agent_type}")

        elif event.type == "tool_result" and event.data.get("tool") == "task":
            metadata = event.data.get("metadata", {})
            agent_type = metadata.get("subagent_type")
            subagent_results[agent_type] = {
                "sub_trace_id": metadata.get("sub_trace_id"),
                "summary": metadata.get("tool_summary"),
            }
            print(f"✅ {agent_type} 完成: {metadata.get('tool_summary')}")

        elif event.type == "conclusion":
            print(f"\n📊 综合分析报告:\n{event.data['content']}")
            print(f"\n📈 Sub-Agent 使用统计:")
            for agent_type, result in subagent_results.items():
                print(f"  - {agent_type}: {result['summary']}")


async def example_subagent_config_file():
    """示例 4: 从配置文件加载 Sub-Agent"""
    print("\n=== 示例 4: 配置文件方式 ===\n")

    # 1. 从配置文件加载
    # 假设有 sub_agents.json:
    # {
    #   "agents": {
    #     "code-reviewer": {
    #       "description": "代码审查专家",
    #       "mode": "subagent",
    #       "allowed_tools": ["read_file", "search_code"],
    #       "system_prompt": "你是代码审查专家..."
    #     }
    #   }
    # }

    registry = get_agent_registry()

    # 如果配置文件存在
    config_path = "sub_agents.json"
    if os.path.exists(config_path):
        registry.load_from_config(config_path)
        print(f"✓ 从 {config_path} 加载 Agent 配置")

    # 2. 列出所有可用的 Sub-Agent
    print("\n可用的 Sub-Agent:")
    for agent in registry.list_subagents():
        print(f"  - {agent.name}: {agent.description}")


async def main():
    """运行所有示例"""
    # 示例 1: 基本使用
    await example_basic_subagent()

    # 示例 2: 自定义 Sub-Agent
    await example_custom_subagent()

    # 示例 3: 并行使用多个 Sub-Agent
    await example_multiple_subagents()

    # 示例 4: 配置文件方式
    await example_subagent_config_file()


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