Agent/examples/content_finder/run.py

"""
内容寻找 Agent

使用示例：
    python run.py
"""

import asyncio
import logging
import sys
import os
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from dotenv import load_dotenv
load_dotenv()

from agent import (
    AgentRunner,
    RunConfig,
    FileSystemTraceStore,
    Trace,
    Message,
)
from agent.llm import create_openrouter_llm_call
from agent.llm.prompts import SimplePrompt

# 导入工具（确保工具被注册）
from tools import (
    douyin_search,
    douyin_user_videos,
    get_content_fans_portrait,
    get_account_fans_portrait,
)

# 配置日志
log_dir = Path(__file__).parent / '.cache'
log_dir.mkdir(exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(log_dir / 'agent.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


async def main():
    print("\n" + "=" * 60)
    print("内容寻找 Agent")
    print("=" * 60)
    print("开始执行...\n")

    # 加载 prompt
    prompt_path = Path(__file__).parent / "content_finder.prompt"
    prompt = SimplePrompt(prompt_path)

    # 构建消息
    messages = prompt.build_messages()

    # 初始化
    api_key = os.getenv("OPEN_ROUTER_API_KEY")
    if not api_key:
        raise ValueError("OPEN_ROUTER_API_KEY 未设置，请在 .env 文件中配置")

    model = os.getenv("MODEL", f"anthropic/claude-{prompt.config.get('model', 'sonnet-4.6')}")
    temperature = float(prompt.config.get("temperature", 0.3))
    max_iterations = int(os.getenv("MAX_ITERATIONS", "30"))
    trace_dir = os.getenv("TRACE_DIR", ".cache/traces")
    skills_dir = str(Path(__file__).parent / "skills")

    Path(trace_dir).mkdir(parents=True, exist_ok=True)

    store = FileSystemTraceStore(base_path=trace_dir)

    # 限制工具范围：只使用抖音相关的4个工具
    allowed_tools = [
        "douyin_search",
        "douyin_user_videos",
        "get_content_fans_portrait",
        "get_account_fans_portrait",
    ]

    runner = AgentRunner(
        llm_call=create_openrouter_llm_call(model=model),
        trace_store=store,
        skills_dir=skills_dir,
    )

    config = RunConfig(
        model=model,
        temperature=temperature,
        max_iterations=max_iterations,
        tools=allowed_tools,  # 限制工具范围
    )

    # 执行
    try:
        async for item in runner.run(messages=messages, config=config):
            if isinstance(item, Trace):
                if item.status == "completed":
                    print(f"\n[完成] trace_id={item.trace_id}")
                elif item.status == "failed":
                    print(f"\n[失败] {item.error_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:
                            print(f"\n{text}")
                        elif text:
                            print(f"[思考] {text[:100]}..." if len(text) > 100 else f"[思考] {text}")
                        if tool_calls:
                            for tc in tool_calls:
                                tool_name = tc.get("function", {}).get("name", "unknown")
                                print(f"[工具] {tool_name}")
                    elif isinstance(content, str) and content:
                        print(f"\n{content}")

                elif item.role == "tool":
                    content = item.content
                    if isinstance(content, dict):
                        tool_name = content.get("tool_name", "unknown")
                        print(f"[结果] {tool_name} ✓")

    except KeyboardInterrupt:
        print("\n用户中断")
    except Exception as e:
        logger.error(f"执行失败: {e}", exc_info=True)
        print(f"\n执行失败: {e}")
        sys.exit(1)


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