Agent/examples/feature_extract/run.py

"""
特征提取示例

使用 Agent 模式 + Skills + 多模态支持
"""

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.prompts import SimplePrompt
from agent.runner import AgentRunner
from agent.storage import MemoryTraceStore
from agent.llm.providers.openrouter 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"
    feature_md_path = base_dir / "input_1" / "feature.md"
    image_path = base_dir / "input_1" / "image.png"
    output_dir = base_dir / "output_1"
    output_dir.mkdir(exist_ok=True)

    # Skills 目录
    skills_dir = project_root / "agent" / "skills"

    print("=" * 60)
    print("特征提取任务 (Agent 模式)")
    print("=" * 60)
    print()

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

    # 提取 system prompt 和 user template
    system_prompt = prompt._messages.get("system", "")
    user_template = prompt._messages.get("user", "")

    # 2. 读取特征描述
    print("2. 读取特征描述...")
    with open(feature_md_path, 'r', encoding='utf-8') as f:
        feature_text = f.read()

    # 3. 构建任务文本（包含图片）
    print("3. 构建任务（文本 + 图片）...")

    # 使用 prompt 构建多模态消息
    temp_messages = prompt.build_messages(
        text=feature_text,
        images=image_path
    )

    # 提取用户消息（包含文本和图片）
    user_message_with_image = None
    for msg in temp_messages:
        if msg["role"] == "user":
            user_message_with_image = msg
            break

    if not user_message_with_image:
        raise ValueError("No user message found in prompt")

    print(f"   - 任务已构建（包含图片: {image_path.name}）")

    # 4. 创建 Agent Runner（配置 skills）
    print("4. 创建 Agent Runner...")
    print(f"   - Skills 目录: {skills_dir}")
    print(f"   - 模型: Claude Sonnet 4.5 (via OpenRouter)")

    runner = AgentRunner(
        trace_store=MemoryTraceStore(),
        llm_call=create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5"),
        skills_dir=str(skills_dir),  # 恢复加载 skills，测试 Claude 是否能处理
        debug=True  # 启用 debug，输出到 .trace/
    )

    # 5. Agent 模式执行
    # 注意：使用 OpenRouter 时，模型在创建 llm_call 时已指定
    # 这里传入的 model 参数会被忽略（由 llm_call 内部控制）
    print(f"5. 启动 Agent 模式...")
    print()

    final_response = ""

    async for event in runner.run(
        task="[图片和特征描述已包含在 messages 中]",  # 占位符
        messages=[user_message_with_image],  # 传入包含图片的用户消息
        system_prompt=system_prompt,
        model="anthropic/claude-sonnet-4.5",  # OpenRouter 模型名称
        temperature=float(prompt.config.get('temperature', 0.3)),
        max_iterations=10,
        # tools 参数不传入，测试自动加载内置工具
    ):
        event_type = event.type
        event_data = event.data

        if event_type == "trace_started":
            print(f"[Trace] 开始: {event_data.get('trace_id', '')[:8]}")

        elif event_type == "memory_loaded":
            exp_count = event_data.get('experiences_count', 0)
            if exp_count > 0:
                print(f"[Memory] 加载 {exp_count} 条经验")

        elif event_type == "step_started":
            step_type = event_data.get('step_type', '')
            print(f"[Step] {step_type}...")

        elif event_type == "thought":
            content = event_data.get('content', '')
            if content:
                print(f"[Thought] {content[:100]}...")

        elif event_type == "tool_execution":
            tool_name = event_data.get('tool_name', '')
            print(f"[Tool] 执行 {tool_name}")

        elif event_type == "conclusion":
            final_response = event_data.get('content', '')  # 修正：字段名是 content 不是 response
            print(f"[Conclusion] Agent 完成")

        elif event_type == "trace_completed":
            print(f"[Trace] 完成")
            print(f"  - Total tokens: {event_data.get('total_tokens', 0)}")
            print(f"  - Total cost: ${event_data.get('total_cost', 0.0):.4f}")

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

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

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

    # 8. 提示查看 debug 文件
    print("Debug 文件:")
    print(f"  - 完整可折叠: {Path.cwd() / '.trace' / 'tree.md'}")
    print(f"  - 简洁文本: {Path.cwd() / '.trace' / 'tree.txt'}")


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