chendandan_homework/04_work/main.py

from __future__ import annotations

"""TravelMind CLI入口：读取用户旅行需求（支持多行输入），构建多角色规划图并执行完整工作流。"""

import asyncio
import logging
import sys
from time import perf_counter

from rich.console import Console

from app.graph.planner_builder import (
    build_planner_graph,
)
from app.logging_config import (
    configure_logging,
    new_trace_id,
)


console = Console()

logger = logging.getLogger(__name__)


def get_user_query() -> str:
    """从命令行参数或终端读取旅行需求（支持多行输入，空行结束）。"""

    command_line_query = " ".join(
        sys.argv[1:]
    ).strip()

    if command_line_query:
        return command_line_query

    console.print(
        "[bold]请输入旅行需求"
        "（支持多行，输入空行结束）：[/bold]"
    )

    lines: list[str] = []
    line_count = 0

    while True:
        prefix = "> " if line_count == 0 else "… "
        line = input(prefix)

        stripped = line.strip()

        if not stripped:
            if lines:
                break
            # 第一行就为空，继续等待。
            continue

        lines.append(stripped)
        line_count += 1

    return " ".join(lines)


async def run_workflow(
    user_query: str,
) -> dict:
    """执行完整TravelMind工作流。"""

    graph = await build_planner_graph()

    return await graph.ainvoke(
        {
            "user_query": user_query,
            "errors": [],
            "missing_fields": [],
            "planning_attempts": 0,
            "review_attempts": 0,
            "revision_feedback": [],
            "workflow_status": "running",
        },
        config={
            # 图中存在重新规划循环。
            "recursion_limit": 80,
        },
    )


async def main() -> None:
    """CLI主入口：配置日志、生成trace_id、读取用户需求并执行完整工作流。"""

    configure_logging()

    trace_id = new_trace_id()

    user_query = get_user_query()

    if not user_query:
        console.print(
            "[bold red]旅行需求不能为空。[/bold red]"
        )
        raise SystemExit(1)

    console.rule(
        "[bold blue]TravelMind 多角色旅行规划"
    )

    console.print(
        f"[dim]trace_id：{trace_id}[/dim]"
    )

    started_at = perf_counter()

    logger.info(
        "TravelMind工作流开始，用户输入长度=%s",
        len(user_query),
    )

    try:
        result = await run_workflow(
            user_query
        )

    except Exception as exc:
        elapsed = perf_counter() - started_at

        logger.exception(
            "TravelMind工作流异常，"
            "elapsed_seconds=%.2f",
            elapsed,
        )

        console.print(
            "[bold red]工作流执行失败：[/bold red]"
            f"{type(exc).__name__}: {exc}"
        )

        raise SystemExit(1) from exc

    elapsed = perf_counter() - started_at

    workflow_status = result.get(
        "workflow_status"
    )

    logger.info(
        "TravelMind工作流结束，"
        "status=%s，"
        "planning_attempts=%s，"
        "review_attempts=%s，"
        "elapsed_seconds=%.2f",
        workflow_status,
        result.get("planning_attempts"),
        result.get("review_attempts"),
        elapsed,
    )

    final_answer = result.get(
        "final_answer"
    )

    console.print()

    if final_answer:
        console.print(final_answer)
    else:
        console.print(
            "[bold red]"
            "系统没有生成最终回答。"
            "[/bold red]"
        )

    console.rule("[bold]执行状态")

    console.print(
        {
            "trace_id": trace_id,
            "工作流状态": workflow_status,
            "Planner执行次数": result.get(
                "planning_attempts"
            ),
            "Reviewer执行次数": result.get(
                "review_attempts"
            ),
            "Supervisor决策": result.get(
                "supervisor_decision"
            ),
            "总耗时（秒）": round(
                elapsed,
                2,
            ),
            "系统错误": result.get(
                "errors",
                [],
            ),
        }
    )


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