Agent/examples/create/run_inspiration_test.py

"""
测试灵感点生成选题
"""
import asyncio
import sys
from pathlib import Path

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

from dotenv import load_dotenv
load_dotenv()

from agent.core.runner import AgentRunner, RunConfig
from agent.llm import create_openrouter_llm_call
from agent.llm.prompts import SimplePrompt
from agent.trace import FileSystemTraceStore, Message, Trace
from agent.utils import setup_logging
from agent.cli import InteractiveController

# 导入工具
import examples.create.tool  # noqa: F401

# 导入项目配置
from config import DEBUG, LOG_LEVEL, RUN_CONFIG, TRACE_STORE_PATH
import logging

logger = logging.getLogger(__name__)

async def main():
    setup_logging(level=LOG_LEVEL)

    base_dir = Path(__file__).parent
    prompt_path = base_dir / "test_inspiration.prompt"

    # 加载 prompt
    prompt = SimplePrompt(prompt_path)

    # 注入流程文档
    process_path = base_dir / "PRD" / "inspiration_to_topic_process_v2.md"
    if process_path.exists():
        process_content = process_path.read_text(encoding="utf-8")
        if "system" in prompt._messages:
            prompt._messages["system"] = prompt._messages["system"].replace(
                "{inspiration_to_topic_process_v2}", process_content
            )

    messages = prompt.build_messages()

    # 创建 runner
    store = FileSystemTraceStore(base_path=TRACE_STORE_PATH)
    runner = AgentRunner(
        trace_store=store,
        llm_call=create_openrouter_llm_call(model=f"anthropic/{RUN_CONFIG.model}"),
        debug=DEBUG,
    )

    # 创建交互控制器
    interactive = InteractiveController(
        runner=runner,
        store=store,
        enable_stdin_check=True,
    )

    # 运行
    config = RUN_CONFIG

    current_trace_id = None
    current_sequence = 0
    should_exit = False

    logger.info("💡 交互提示：")
    logger.info("   - 执行过程中输入 'p' 或 'pause' 暂停并进入交互模式")
    logger.info("   - 执行过程中输入 'q' 或 'quit' 停止执行")
    logger.info("")

    try:
        async for item in runner.run(messages=messages, config=config):
            # 检查用户输入
            cmd = interactive.check_stdin()
            if cmd == "pause":
                logger.info("\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
                    break
                if menu_result["action"] == "continue":
                    new_messages = menu_result.get("messages", [])
                    if new_messages:
                        # 用户输入了反馈，在当前 trace 上继续执行
                        continue_config = RunConfig(
                            model=config.model,
                            trace_id=current_trace_id,
                            after_sequence=current_sequence
                        )
                        async for item2 in runner.run(messages=new_messages, config=continue_config):
                            if isinstance(item2, Message):
                                current_sequence = item2.sequence
                                print(item2)
                    break

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

            if isinstance(item, Trace):
                current_trace_id = item.trace_id
                # 如果 trace 完成，进入交互模式等待反馈
                if item.status == "completed":
                    logger.info("\n✅ Agent 执行完成，等待用户反馈...")
                    logger.info("💡 按 'p' 键进入交互模式，输入反馈或输入 'q' 退出")

                    # 等待用户输入
                    while True:
                        await asyncio.sleep(0.5)
                        cmd = interactive.check_stdin()
                        if cmd == "pause":
                            menu_result = await interactive.show_menu(current_trace_id, current_sequence)
                            if menu_result["action"] == "stop":
                                should_exit = True
                                break
                            if menu_result["action"] == "continue":
                                new_messages = menu_result.get("messages", [])
                                if new_messages:
                                    # 用户输入了反馈，在当前 trace 上继续执行
                                    continue_config = RunConfig(
                                        model=config.model,
                                        trace_id=current_trace_id,
                                        after_sequence=current_sequence
                                    )
                                    async for item2 in runner.run(messages=new_messages, config=continue_config):
                                        if isinstance(item2, Message):
                                            current_sequence = item2.sequence
                                            print(item2)
                                break
                        elif cmd == "quit":
                            should_exit = True
                            break

                    if should_exit:
                        break

            elif isinstance(item, Message):
                current_sequence = item.sequence

            print(item)

    except KeyboardInterrupt:
        logger.info("\n\n用户中断 (Ctrl+C)")
        if current_trace_id:
            await runner.stop(current_trace_id)

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