chendandan_homework/app/graph/selection_builder.py

from __future__ import annotations
"""候选筛选阶段子图：对航班/酒店候选做确定性筛选与排名，为路线评估提供输入。"""

from langgraph.graph import (
    END,
    START,
    StateGraph,
)

from app.agents.requirement_agent import (
    RequirementAgent,
)
from app.agents.resource_agent import (
    ResourceSearchAgent,
)
from app.graph.nodes import (
    clarification_node,
    error_node,
    make_parse_request_node,
    route_after_parse,
)
from app.graph.resource_nodes import (
    make_resource_search_node,
)
from app.graph.selection_nodes import (
    make_candidate_selection_node,
    route_after_resources_to_selection,
)
from app.graph.state import TravelState
from app.mcp_client import load_mcp_tools


async def build_selection_graph():
    """构建资源查询与候选筛选工作流。"""

    tool_bundle = await load_mcp_tools()

    return_flight_tool = next(
        (
            tool
            for tool in tool_bundle.travel_tools
            if tool.name
            == "search_return_flights"
        ),
        None,
    )

    if return_flight_tool is None:
        raise RuntimeError(
            "没有加载到"
            "search_return_flights工具。"
        )

    requirement_agent = RequirementAgent()

    resource_agent = ResourceSearchAgent(
        tools=tool_bundle.travel_tools,
    )

    builder = StateGraph(TravelState)

    builder.add_node(
        "parse_request",
        make_parse_request_node(
            requirement_agent
        ),
    )

    builder.add_node(
        "clarify",
        clarification_node,
    )

    builder.add_node(
        "error",
        error_node,
    )

    builder.add_node(
        "search_resources",
        make_resource_search_node(
            resource_agent
        ),
    )

    builder.add_node(
        "select_candidates",
        make_candidate_selection_node(
            return_flight_tool,
            max_outbound_queries=2,
            max_returns_per_outbound=5,
        ),
    )

    builder.add_edge(
        START,
        "parse_request",
    )

    builder.add_conditional_edges(
        "parse_request",
        route_after_parse,
        {
            "clarify": "clarify",
            "error": "error",
            "ready": "search_resources",
        },
    )

    builder.add_edge("clarify", END)
    builder.add_edge("error", END)

    builder.add_conditional_edges(
        "search_resources",
        route_after_resources_to_selection,
        {
            "select_candidates": (
                "select_candidates"
            ),
            "end": END,
        },
    )

    builder.add_edge(
        "select_candidates",
        END,
    )

    return builder.compile()