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.llm.prompts import SimplePrompt
from agent.core.runner import AgentRunner
from agent.execution import MemoryTraceStore, Trace, Step
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"
    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）
    # 注意：内置 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 和 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=skills_dir,  # 可选：加载额外的用户自定义 skills（内置 skills 会自动加载）
        debug=True  # 启用 debug，输出到 .trace/
    )

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

    final_response = ""

    async for item 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 参数不传入，测试自动加载内置工具
    ):
        # 处理 Trace 对象（整体状态变化）
        if isinstance(item, Trace):
            if item.status == "in_progress":
                print(f"[Trace] 开始: {item.trace_id[:8]}")
            elif item.status == "completed":
                print(f"[Trace] 完成")
                print(f"  - Total tokens: {item.total_tokens}")
                print(f"  - Total cost: ${item.total_cost:.4f}")
            elif item.status == "failed":
                print(f"[Trace] 失败")

        # 处理 Step 对象（执行过程）
        elif isinstance(item, Step):
            if item.step_type == "memory_read":
                exp_count = item.data.get('experiences_count', 0)
                if exp_count > 0:
                    print(f"[Memory] 加载 {exp_count} 条经验")

            elif item.step_type == "thought":
                if item.status == "completed":
                    content = item.data.get('content', '')
                    if content:
                        print(f"[Thought] {content[:100]}...")

            elif item.step_type == "action":
                tool_name = item.data.get('tool_name', '')
                print(f"[Tool] 执行 {tool_name}")

            elif item.step_type == "result":
                tool_name = item.data.get('tool_name', '')
                print(f"[Tool] {tool_name} 完成")

            elif item.step_type == "response":
                final_response = item.data.get('content', '')
                print(f"[Response] Agent 完成")

    # 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())