chendandan_homework/scripts/test_requirement_graph.py

from __future__ import annotations

"""测试需求解析阶段子图：验证自然语言解析为TravelRequest的完整流程与澄清分支。"""

import asyncio

from rich.console import Console

from app.graph.builder import (
    build_requirement_graph,
)


console = Console()


async def run_case(
    graph,
    user_query: str,
) -> None:
    console.rule("[bold]用户输入")
    console.print(user_query)

    result = await graph.ainvoke(
        {
            "user_query": user_query,
            "errors": [],
            "missing_fields": [],
        }
    )

    console.rule("[bold]Graph输出")
    console.print(result["final_answer"])

    request = result.get("travel_request")

    if request is not None:
        console.rule("[bold]结构化需求")
        console.print_json(
            request.model_dump_json(
                indent=2,
            )
        )


async def main() -> None:
    graph = build_requirement_graph()

    await run_case(
        graph,
        (
            "我和妻子计划2026年8月10日"
            "从上海去成都，5天4晚，"
            "总预算5000元。喜欢熊猫、"
            "自然景观和川菜，不想安排太赶。"
            "酒店希望靠近地铁，最好住在市中心，住宿每晚不超过"
            "300元。机票价格相差不大的情况下，优先考虑落地双流机场"
            "只优先考虑直达。"
        ),
    )

    # 可选测试：信息不完整时是否追问。
    await run_case(
        graph,
        "我想从上海去成都旅游，喜欢美食。",
    )


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