Agent/examples/research/run.py

"""
浏览器调研示例

使用 Agent 模式 + 浏览器工具进行网络调研
"""

import os
import sys
import asyncio
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from dotenv import load_dotenv
load_dotenv()

from agent.llm.prompts import SimplePrompt
from agent.core.runner import AgentRunner
from agent.execution import (
    FileSystemTraceStore,
    Trace,
    Message,
)
from agent.llm import create_openrouter_llm_call


async def main():
    # 路径配置
    base_dir = Path(__file__).parent
    project_root = base_dir.parent.parent
    prompt_path = base_dir / "test.prompt"
    output_dir = base_dir / "output"
    output_dir.mkdir(exist_ok=True)

    # Skills 目录（可选：用户自定义 skills）
    # 注意：内置 skills（agent/skills/core.md）会自动加载
    skills_dir = None  # 或者指定自定义 skills 目录，如: project_root / "skills"

    print("=" * 60)
    print("浏览器调研任务 (Agent 模式)")
    print("=" * 60)
    print()

    # 1. 加载 prompt
    print("1. 加载 prompt...")
    prompt = SimplePrompt(prompt_path)

    # 提取配置
    system_prompt = prompt._messages.get("system", "")
    user_task = prompt._messages.get("user", "")
    model_name = prompt.config.get('model', 'gemini-2.5-flash')
    temperature = float(prompt.config.get('temperature', 0.3))

    print(f"   - 任务: {user_task[:80]}...")
    print(f"   - 模型: {model_name}")

    # 2. 构建消息
    print("2. 构建任务消息...")
    messages = prompt.build_messages()

    # 3. 创建 Agent Runner（配置 skills 和浏览器工具）
    print("3. 创建 Agent Runner...")
    print(f"   - Skills 目录: {skills_dir}")
    print(f"   - 模型: {model_name} (via OpenRouter)")

    # 使用 OpenRouter 的 Gemini 模型
    runner = AgentRunner(
        trace_store=FileSystemTraceStore(base_path=".trace"),
        llm_call=create_openrouter_llm_call(model=f"google/{model_name}"),
        skills_dir=skills_dir,
        debug=True  # 启用 debug，输出到 .trace/
    )

    # 4. Agent 模式执行
    print(f"4. 启动 Agent 模式...")
    print()

    final_response = ""
    current_trace_id = None

    async for item in runner.run(
        task=user_task,
        messages=messages,
        system_prompt=system_prompt,
        model=f"google/{model_name}",
        temperature=temperature,
        max_iterations=20,  # 调研任务可能需要更多迭代
    ):
        # 处理 Trace 对象（整体状态变化）
        if isinstance(item, Trace):
            current_trace_id = item.trace_id
            if item.status == "running":
                print(f"[Trace] 开始: {item.trace_id[:8]}")
            elif item.status == "completed":
                print(f"[Trace] 完成")
                print(f"  - Total messages: {item.total_messages}")
                print(f"  - Total tokens: {item.total_tokens}")
                print(f"  - Total cost: ${item.total_cost:.4f}")
            elif item.status == "failed":
                print(f"[Trace] 失败: {item.error_message}")

        # 处理 Message 对象（执行过程）
        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 text and not tool_calls:
                        # 纯文本回复（最终响应）
                        final_response = text
                        print(f"[Response] Agent 完成")
                    elif text:
                        print(f"[Assistant] {text[:100]}...")

                    if tool_calls:
                        for tc in tool_calls:
                            tool_name = tc.get("function", {}).get("name", "unknown")
                            print(f"[Tool Call] {tool_name}")

            elif item.role == "tool":
                content = item.content
                if isinstance(content, dict):
                    tool_name = content.get("tool_name", "unknown")
                    print(f"[Tool Result] {tool_name}")
                if item.description:
                    desc = item.description[:80] if len(item.description) > 80 else item.description
                    print(f"  {desc}...")

    # 5. 输出结果
    print()
    print("=" * 60)
    print("Agent 响应:")
    print("=" * 60)
    print(final_response)
    print("=" * 60)
    print()

    # 6. 保存结果
    output_file = output_dir / "research_result.txt"
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(final_response)

    print(f"✓ 结果已保存到: {output_file}")
    print()

    # 提示使用 API 可视化
    print("=" * 60)
    print("可视化 Step Tree:")
    print("=" * 60)
    print("1. 启动 API Server:")
    print("   python3 api_server.py")
    print()
    print("2. 浏览器访问:")
    print("   http://localhost:8000/api/traces")
    print()
    print(f"3. Trace ID: {current_trace_id}")
    print("=" * 60)


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