chendandan_homework/scripts/test_research_graph.py

from __future__ import annotations

"""测试地图研究阶段子图：验证高德地图POI搜索与天气查询的完整流程。"""

import asyncio

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

from app.graph.research_builder import (
    build_research_graph,
)


console = Console()


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

    map_result = result.get(
        "map_research_result"
    )

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

    console.print(
        {
            "城市": map_result.city,
            "景点数量": len(
                map_result.attractions
            ),
            "餐饮数量": len(
                map_result.restaurants
            ),
            "天气数量": len(
                map_result.weather
            ),
            "备注": map_result.notes,
        }
    )

    attraction_table = Table(
        title="真实景点候选"
    )

    attraction_table.add_column("名称")
    attraction_table.add_column("地址")
    attraction_table.add_column("坐标")
    attraction_table.add_column("匹配原因")

    for place in map_result.attractions:
        attraction_table.add_row(
            place.name,
            place.address or "未知",
            place.location or "未知",
            place.match_reason or "未说明",
        )

    console.print(attraction_table)

    restaurant_table = Table(
        title="真实餐饮候选"
    )

    restaurant_table.add_column("名称")
    restaurant_table.add_column("地址")
    restaurant_table.add_column("坐标")

    for place in map_result.restaurants:
        restaurant_table.add_row(
            place.name,
            place.address or "未知",
            place.location or "未知",
        )

    console.print(restaurant_table)


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