Agent/examples/production_restore/run.py

"""
图像还原流程 - 使用框架交互功能

使用 Agent 模式 + Skills + 框架交互控制器
读取 input/ 中的 pipeline.json，逐阶段执行图像还原。

功能：
1. 使用框架提供的 InteractiveController
2. 使用配置文件管理运行参数
3. 支持命令行随时打断（输入 'p' 暂停，'q' 退出）
4. 暂停后可插入干预消息
5. 支持触发经验总结
6. 查看当前 GoalTree
7. 支持通过 --trace <ID> 恢复已有 Trace 继续执行
"""

import argparse
import os
import sys
import asyncio
from pathlib import Path

# Clash Verge TUN 模式兼容：禁止 httpx/urllib 自动检测系统 HTTP 代理
os.environ.setdefault("no_proxy", "*")

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

from dotenv import load_dotenv
load_dotenv()

from cyber_agent.llm.prompts import SimplePrompt
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.trace import (
    FileSystemTraceStore,
    Trace,
    Message,
)
from cyber_agent.llm import create_qwen_llm_call
from cyber_agent.cli import InteractiveController
from cyber_agent.utils import setup_logging
from cyber_agent.tools.builtin.browser.baseClass import init_browser_session, kill_browser_session

# 导入自定义工具（触发 @tool 注册）
from cyber_agent.tools.builtin.toolhub import toolhub_health, toolhub_search, toolhub_call, image_uploader, image_downloader  # noqa: F401
from cyber_agent.tools.builtin.knowledge import requirement_search, requirement_list  # noqa: F401
from evaluate_tool import evaluate_image  # noqa: F401
from examples.production_restore.tools.category_query import search_category_tree, extract_requirements_from_table  # noqa: F401

# 导入项目配置
from config import RUN_CONFIG, SKILLS_DIR, TRACE_STORE_PATH, DEBUG, LOG_LEVEL, LOG_FILE, BROWSER_TYPE, HEADLESS, INPUT_DIR, OUTPUT_DIR, FEATURES_DIR


async def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="图像还原 (Agent 模式 + 交互增强)")
    parser.add_argument(
        "--trace", type=str, default=None,
        help="已有的 Trace ID，用于恢复继续执行（不指定则新建）",
    )
    args = parser.parse_args()

    # 路径配置
    base_dir = Path(__file__).parent
    project_root = base_dir.parent.parent
    prompt_path = base_dir / "requirement.prompt"
    output_dir = project_root / OUTPUT_DIR
    output_dir.mkdir(parents=True, exist_ok=True)

    # 1. 配置日志
    setup_logging(level=LOG_LEVEL, file=LOG_FILE)

    # 2. 加载项目级 presets
    print("2. 加载 presets...")
    presets_path = base_dir / "presets.json"
    if presets_path.exists():
        from cyber_agent.core.presets import load_presets_from_json
        load_presets_from_json(str(presets_path))
        print(f"   - 已加载项目 presets")
    else:
        print(f"   - 未找到 presets.json，跳过")

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

    # 4. 构建任务消息
    print("4. 构建任务消息...")
    print(f"   - 输入目录: {INPUT_DIR}")
    print(f"   - 输出目录: {OUTPUT_DIR}")
    messages = prompt.build_messages(input_dir=INPUT_DIR, output_dir=OUTPUT_DIR, features_dir=FEATURES_DIR)

    # 5. 初始化浏览器
    browser_mode_names = {"cloud": "云浏览器", "local": "本地浏览器", "container": "容器浏览器"}
    browser_mode_name = browser_mode_names.get(BROWSER_TYPE, BROWSER_TYPE)
    print(f"5. 正在初始化{browser_mode_name}...")
    await init_browser_session(
        browser_type=BROWSER_TYPE,
        headless=HEADLESS,
        url="https://www.google.com/",
        profile_name=""
    )
    print(f"   ✅ {browser_mode_name}初始化完成\n")

    # 6. 创建 Agent Runner
    print("6. 创建 Agent Runner...")
    print(f"   - Skills 目录: {SKILLS_DIR}")

    # 从 prompt 的 frontmatter 中提取模型配置（优先于 config.py）
    prompt_model = prompt.config.get("model", None)
    if prompt_model:
        model_for_llm = prompt_model
        print(f"   - 模型 (from prompt): {model_for_llm}")
    else:
        model_for_llm = RUN_CONFIG.model
        print(f"   - 模型 (from config): {model_for_llm}")

    store = FileSystemTraceStore(base_path=TRACE_STORE_PATH)
    runner = AgentRunner(
        trace_store=store,
        llm_call=create_qwen_llm_call(model=model_for_llm),
        skills_dir=SKILLS_DIR,
        debug=DEBUG
    )

    # 7. 创建交互控制器
    interactive = InteractiveController(
        runner=runner,
        store=store,
        enable_stdin_check=True
    )
    # 将 stdin 检查回调注入 runner，供子 agent 执行期间使用
    runner.stdin_check = interactive.check_stdin

    # 8. 任务信息
    task_name = RUN_CONFIG.name or base_dir.name
    print("=" * 60)
    print(f"{task_name}")
    print("=" * 60)
    print("💡 交互提示：")
    print("   - 执行过程中输入 'p' 或 'pause' 暂停并进入交互模式")
    print("   - 执行过程中输入 'q' 或 'quit' 停止执行")
    print("=" * 60)
    print()

    # 9. 判断是新建还是恢复
    resume_trace_id = args.trace
    if resume_trace_id:
        existing_trace = await store.get_trace(resume_trace_id)
        if not existing_trace:
            print(f"\n错误: Trace 不存在: {resume_trace_id}")
            sys.exit(1)
        print(f"恢复已有 Trace: {resume_trace_id[:8]}...")
        print(f"   - 状态: {existing_trace.status}")
        print(f"   - 消息数: {existing_trace.total_messages}")
        print(f"\n💡 提示：恢复 Trace 时会先进入交互菜单，您可以选择从指定消息续跑")
    else:
        print(f"启动新 Agent...")

    print()

    final_response = ""
    current_trace_id = resume_trace_id
    current_sequence = 0
    should_exit = False

    try:
        # 配置
        run_config = RUN_CONFIG
        if resume_trace_id:
            initial_messages = None
            run_config.trace_id = resume_trace_id
        else:
            initial_messages = messages
            run_config.name = f"{task_name}：还原任务"

        while not should_exit:
            if current_trace_id:
                run_config.trace_id = current_trace_id

            final_response = ""

            # 如果是恢复 trace 或 trace 已完成/失败且没有新消息，进入交互菜单
            if current_trace_id and initial_messages is None:
                check_trace = await store.get_trace(current_trace_id)
                if check_trace:
                    # 显示 trace 状态
                    if check_trace.status == "completed":
                        print(f"\n[Trace] ✅ 已完成")
                        print(f"  - Total messages: {check_trace.total_messages}")
                        print(f"  - Total cost: ${check_trace.total_cost:.4f}")
                    elif check_trace.status == "failed":
                        print(f"\n[Trace] ❌ 已失败: {check_trace.error_message}")
                    elif check_trace.status == "stopped":
                        print(f"\n[Trace] ⏸️ 已停止")
                        print(f"  - Total messages: {check_trace.total_messages}")
                    else:
                        print(f"\n[Trace] 📊 状态: {check_trace.status}")
                        print(f"  - Total messages: {check_trace.total_messages}")

                    current_sequence = check_trace.head_sequence

                    menu_result = await interactive.show_menu(current_trace_id, current_sequence)

                    if menu_result["action"] == "stop":
                        break
                    elif menu_result["action"] == "continue":
                        new_messages = menu_result.get("messages", [])
                        if new_messages:
                            initial_messages = new_messages
                            run_config.after_sequence = menu_result.get("after_sequence")
                        else:
                            initial_messages = []
                            run_config.after_sequence = None
                        continue
                    break

            # 如果没有进入菜单（新建 trace），设置初始消息
            if initial_messages is None:
                initial_messages = []

            print(f"{'▶️ 开始执行...' if not current_trace_id else '▶️ 继续执行...'}")

            # 执行 Agent
            paused = False
            try:
                async for item in runner.run(messages=initial_messages, config=run_config):
                    # 检查用户中断
                    cmd = interactive.check_stdin()
                    if cmd == 'pause':
                        print("\n⏸️ 正在暂停执行...")
                        if current_trace_id:
                            await runner.stop(current_trace_id)
                        await asyncio.sleep(0.5)

                        menu_result = await interactive.show_menu(current_trace_id, current_sequence)

                        if menu_result["action"] == "stop":
                            should_exit = True
                            paused = True
                            break
                        elif menu_result["action"] == "continue":
                            new_messages = menu_result.get("messages", [])
                            if new_messages:
                                initial_messages = new_messages
                                after_seq = menu_result.get("after_sequence")
                                if after_seq is not None:
                                    run_config.after_sequence = after_seq
                                paused = True
                                break
                            else:
                                initial_messages = []
                                run_config.after_sequence = None
                                paused = True
                                break

                    elif cmd == 'quit':
                        print("\n🛑 用户请求停止...")
                        if current_trace_id:
                            await runner.stop(current_trace_id)
                        should_exit = True
                        break

                    # 处理 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"\n[Trace] ✅ 完成")
                            print(f"  - Total messages: {item.total_messages}")
                            print(f"  - Total cost: ${item.total_cost:.4f}")
                        elif item.status == "failed":
                            print(f"\n[Trace] ❌ 失败: {item.error_message}")
                        elif item.status == "stopped":
                            print(f"\n[Trace] ⏸️ 已停止")

                    # 处理 Message 对象
                    elif isinstance(item, Message):
                        current_sequence = item.sequence

                        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"\n[Response] Agent 回复：")
                                    print(text)
                                elif text:
                                    preview = text[:150] + "..." if len(text) > 150 else text
                                    print(f"[Assistant] {preview}")

                        elif item.role == "tool":
                            content = item.content
                            tool_name = "unknown"
                            if isinstance(content, dict):
                                tool_name = content.get("tool_name", "unknown")

                            if item.description and item.description != tool_name:
                                desc = item.description[:80] if len(item.description) > 80 else item.description
                                print(f"[Tool Result] ✅ {tool_name}: {desc}...")
                            else:
                                print(f"[Tool Result] ✅ {tool_name}")

            except Exception as e:
                print(f"\n执行出错: {e}")
                import traceback
                traceback.print_exc()

            if paused:
                if should_exit:
                    break
                continue

            if should_exit:
                break

            # Runner 退出后显示交互菜单
            if current_trace_id:
                menu_result = await interactive.show_menu(current_trace_id, current_sequence)

                if menu_result["action"] == "stop":
                    break
                elif menu_result["action"] == "continue":
                    new_messages = menu_result.get("messages", [])
                    if new_messages:
                        initial_messages = new_messages
                        run_config.after_sequence = menu_result.get("after_sequence")
                    else:
                        initial_messages = []
                        run_config.after_sequence = None
                    continue
            break

    except KeyboardInterrupt:
        print("\n\n用户中断 (Ctrl+C)")
        if current_trace_id:
            await runner.stop(current_trace_id)
    finally:
        # 清理浏览器会话
        try:
            await kill_browser_session()
        except Exception:
            pass

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

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

    # 可视化提示
    if current_trace_id:
        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())