chendandan_homework/04_work/scripts/test_route_graph.py

from __future__ import annotations

"""测试路线评估阶段子图：验证酒店位置评估与每日景点路线衔接。"""

import asyncio

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

from app.graph.route_builder import (
    build_route_graph,
)


console = Console()


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

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

    console.rule("[bold]路线评估结果")

    console.print(result["final_answer"])

    route_result = result.get(
        "route_evaluation_result"
    )

    if route_result is None:
        console.print(
            {
                "错误": result.get("errors"),
                "警告": result.get(
                    "route_warnings"
                ),
            }
        )
        return

    console.print(
        {
            "参考景点": (
                route_result
                .reference_attractions
            ),
            "距离调用次数": (
                route_result
                .distance_call_count
            ),
            "周边搜索次数": (
                route_result
                .around_search_call_count
            ),
            "最终酒店ID": (
                route_result
                .selected_hotel_id
            ),
            "警告数": len(
                route_result.warnings
            ),
        }
    )

    table = Table(
        title="酒店综合路线排名"
    )

    table.add_column("排名")
    table.add_column("酒店")
    table.add_column("综合分")
    table.add_column("原酒店分")
    table.add_column("景点平均距离")
    table.add_column("平均驾车时间")
    table.add_column("最近地铁")
    table.add_column("原因")

    for index, hotel in enumerate(
        route_result.evaluated_hotels,
        start=1,
    ):
        average_distance = (
            f"{hotel.average_distance_meters / 1000:.1f}km"
            if hotel.average_distance_meters
            is not None
            else "未知"
        )

        average_duration = (
            f"{hotel.average_duration_seconds / 60:.0f}分钟"
            if hotel.average_duration_seconds
            is not None
            else "未知"
        )

        subway_distance = (
            f"{hotel.nearest_subway_distance_meters:.0f}m"
            if hotel
            .nearest_subway_distance_meters
            is not None
            else "未知"
        )

        table.add_row(
            str(index),
            hotel.hotel_name,
            f"{hotel.final_score:.2f}",
            f"{hotel.base_hotel_score:.2f}",
            average_distance,
            average_duration,
            subway_distance,
            "；".join(hotel.reasons),
        )

    console.print(table)

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

        for warning in route_result.warnings:
            console.print(f"- {warning}")


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