chendandan_homework/scripts/test_planner_graph.py

from __future__ import annotations

"""测试行程规划阶段子图：验证ItineraryPlannerAgent生成多日行程的流程。"""

import asyncio

from rich.console import Console
from rich.table import Table

from app.graph.planner_builder import (
    build_planner_graph,
)


console = Console()


async def main() -> None:
    graph = await build_planner_graph()

    result = await graph.ainvoke(
        {
            "user_query": (
                "我和妻子计划2026年8月10日"
                "从上海去成都，8月14日返回，"
                "总预算5000元。"
                "喜欢熊猫、自然景观和川菜，"
                "希望行程轻松一些。"
                "酒店每晚不超过300元，"
                "评分不低于3，"
                "并且希望靠近地铁。"
                "机票价格和便利性综合考虑。"
            ),
            "errors": [],
            "missing_fields": [],
            "planning_attempts": 0,
            "revision_feedback": [],
        }
    )

    console.rule("[bold]Planner工作流结果")
    console.print(result["final_answer"])

    plan = result.get("itinerary_plan")

    if plan is None:
        console.print(
            {
                "错误": result.get("errors"),
            }
        )
        return

    console.rule("[bold]方案摘要")

    console.print(
        {
            "标题": plan.title,
            "概述": plan.overview,
            "规划次数": result.get(
                "planning_attempts"
            ),
            "航班组合": (
                plan.selected_flight
                .combination_id
            ),
            "航班报价": (
                plan.selected_flight
                .quoted_price
            ),
            "酒店": plan.selected_hotel.name,
            "酒店ID": (
                plan.selected_hotel.hotel_id
            ),
            "住宿晚数": (
                plan.selected_hotel.nights
            ),
        }
    )

    for day in plan.days:
        table = Table(
            title=(
                f"第{day.day_index}天 "
                f"{day.date}｜{day.theme}"
            )
        )

        table.add_column("时间")
        table.add_column("类型")
        table.add_column("活动")
        table.add_column("地址")
        table.add_column("交通")
        table.add_column("备注")

        for activity in day.activities:
            table.add_row(
                (
                    f"{activity.start_time}"
                    f"—{activity.end_time}"
                ),
                activity.activity_type,
                activity.name,
                activity.address or "",
                activity.transport_mode or "",
                activity.notes or "",
            )

        console.print(table)

    if plan.assumptions:
        console.rule("[bold yellow]假设")

        for item in plan.assumptions:
            console.print(f"- {item}")

    if plan.warnings:
        console.rule("[bold yellow]警告")

        for item in plan.warnings:
            console.print(f"- {item}")


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