chendandan_homework/scripts/test_selection_graph.py

from __future__ import annotations

"""测试候选筛选阶段子图：验证航班筛选、酒店过滤与往返组合排序。"""

import asyncio

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

from app.graph.selection_builder import (
    build_selection_graph,
)


console = Console()


async def main() -> None:
    graph = await build_selection_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"])

    selection = result.get(
        "candidate_selection_result"
    )

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

    console.print(
        {
            "旅行类型": selection.flight_type,
            "去程查询候选数": len(
                selection.outbound_candidates
            ),
            "往返组合数": len(
                selection.round_trip_options
            ),
            "酒店候选数": len(
                selection.hotels
            ),
            "警告": selection.warnings,
        }
    )

    flight_table = Table(
        title="排名靠前的往返组合"
    )

    flight_table.add_column("排名")
    flight_table.add_column("去程航班")
    flight_table.add_column("返程航班")
    flight_table.add_column("接口报价")
    flight_table.add_column("分数")
    flight_table.add_column("排序原因")

    for index, option in enumerate(
        selection.round_trip_options[:5],
        start=1,
    ):
        outbound_numbers = "、".join(
            option.outbound.get(
                "flight_numbers",
                [],
            )
        )

        return_numbers = "、".join(
            option.return_flight.get(
                "flight_numbers",
                [],
            )
        )

        flight_table.add_row(
            str(index),
            outbound_numbers or "未知",
            return_numbers or "未知",
            (
                f"{option.quoted_price:.0f}"
                f" {option.currency}"
                if option.quoted_price is not None
                else "未知"
            ),
            f"{option.score:.2f}",
            "；".join(option.reasons),
        )

    console.print(flight_table)

    hotel_table = Table(
        title="排名靠前的酒店"
    )

    hotel_table.add_column("排名")
    hotel_table.add_column("酒店")
    hotel_table.add_column("每晚价格")
    hotel_table.add_column("评分")
    hotel_table.add_column("分数")
    hotel_table.add_column("排序原因")

    for index, candidate in enumerate(
        selection.hotels[:5],
        start=1,
    ):
        hotel = candidate.hotel

        price = hotel.get("price_per_night")

        hotel_table.add_row(
            str(index),
            str(hotel.get("name", "未知")),
            (
                f"{price} {hotel.get('currency', '')}"
                if price is not None
                else "未知"
            ),
            str(
                hotel.get(
                    "overall_rating",
                    "未知",
                )
            ),
            f"{candidate.score:.2f}",
            "；".join(candidate.reasons),
        )

    console.print(hotel_table)


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