chendandan_homework/scripts/test_resource_graph.py

from __future__ import annotations

"""测试资源查询阶段子图：验证航班和酒店搜索的完整流程（调用真实MCP服务）。"""

import asyncio

from rich.console import Console

from app.graph.travel_builder import (
    build_resource_graph,
)


console = Console()


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

    result = await graph.ainvoke(
        {
            "user_query": (
                "我和妻子计划2026年8月10日"
                "从上海去成都，8月13日返回，"
                "总预算8000元。"
                "喜欢熊猫、自然景观和川菜，"
                "不想安排太赶。"
                "酒店希望靠近地铁，"
                "每晚不超过600元。"
                "机票不限制是否中转。"
            ),
            "errors": [],
            "missing_fields": [],
        }
    )

    console.rule("[bold]最终状态摘要")

    console.print(result["final_answer"])

    console.print(
        {
            "出发机场": result.get(
                "origin_airports"
            ),
            "目的机场": result.get(
                "destination_airports"
            ),
            "航班候选数": (
                result.get(
                    "flight_search_result",
                    {},
                ).get("total_count")
            ),
            "酒店候选数": (
                result.get(
                    "hotel_search_result",
                    {},
                ).get("total_count")
            ),
            "错误": result.get("errors"),
        }
    )


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