chendandan_homework/app/graph/selection_nodes.py

from __future__ import annotations
"""候选筛选节点：调用CandidateSelectionService对航班/酒店/往返组合做确定性排序。"""

import json
from typing import Any, Literal

from langchain_core.messages import ToolMessage
from langchain_core.tools import BaseTool

from app.graph.state import TravelState
from app.schemas.selection import (
    CandidateSelectionResult,
)
from app.services.candidate_selection_service import (
    CandidateSelectionService,
)


def extract_direct_tool_payload(
    result: Any,
) -> dict[str, Any]:
    """提取直接调用LangChain Tool后的结构化结果。"""

    if isinstance(result, ToolMessage):
        artifact = getattr(
            result,
            "artifact",
            None,
        )

        if isinstance(artifact, dict):
            structured = artifact.get(
                "structured_content"
            )

            if isinstance(structured, dict):
                return structured

            if "flights" in artifact:
                return artifact

        result = result.content

    if isinstance(result, dict):
        structured = result.get(
            "structured_content"
        )

        if isinstance(structured, dict):
            return structured

        if (
            "flights" in result
            or "hotels" in result
            or "suggestions" in result
        ):
            return result

    if isinstance(result, str):
        try:
            parsed = json.loads(result)
        except json.JSONDecodeError as exc:
            raise RuntimeError(
                "MCP工具返回了无法解析的文本。"
            ) from exc

        if isinstance(parsed, dict):
            return parsed

    if isinstance(result, list):
        for block in result:
            if not isinstance(block, dict):
                continue

            text = block.get("text")

            if not isinstance(text, str):
                continue

            try:
                parsed = json.loads(text)
            except json.JSONDecodeError:
                continue

            if isinstance(parsed, dict):
                return parsed

    raise RuntimeError(
        "无法从MCP工具结果中提取结构化数据。"
    )


def get_return_stops_filter(
    max_stops: int | None,
) -> str | None:
    """将内部中转次数转换为MCP查询参数。"""

    if max_stops is None:
        return None

    if max_stops == 0:
        return "nonstop"

    if max_stops == 1:
        return "one_or_fewer"

    if max_stops == 2:
        return "two_or_fewer"

    # 接口没有“三次以内”的精确枚举，
    # 留给本地确定性代码过滤。
    return None


def route_after_resources_to_selection(
    state: TravelState,
) -> Literal[
    "select_candidates",
    "end",
]:
    """资源查询后路由：有结果进入候选筛选，无结果结束工作流。"""

    if state.get("errors"):
        return "end"

    if not state.get("flight_search_result"):
        return "end"

    if not state.get("hotel_search_result"):
        return "end"

    return "select_candidates"


def make_candidate_selection_node(
    return_flight_tool: BaseTool,
    *,
    selector: CandidateSelectionService
    | None = None,
    max_outbound_queries: int = 2,
    max_returns_per_outbound: int = 5,
):
    """创建航班和酒店候选筛选节点。

    为控制接口调用成本，默认只对排名最高的
    两个去程候选查询返程。
    """

    selection_service = (
        selector
        or CandidateSelectionService()
    )

    async def candidate_selection_node(
        state: TravelState,
    ) -> dict[str, Any]:
        """候选筛选节点：对航班和酒店候选排序，生成排名理由与警告。"""

        request = state["travel_request"]

        flight_result = state.get(
            "flight_search_result",
            {},
        )

        hotel_result = state.get(
            "hotel_search_result",
            {},
        )

        raw_flights = flight_result.get(
            "flights",
            [],
        )

        raw_hotels = hotel_result.get(
            "hotels",
            [],
        )

        if not isinstance(raw_flights, list):
            raw_flights = []

        if not isinstance(raw_hotels, list):
            raw_hotels = []

        (
            ranked_hotels,
            hotel_warnings,
        ) = selection_service.rank_hotels(
            raw_hotels,
            request,
            limit=10,
        )

        warnings = list(hotel_warnings)

        # 单程场景不需要调用返程工具。
        if request.return_date is None:
            (
                ranked_flights,
                flight_warnings,
            ) = selection_service.rank_flights(
                raw_flights,
                request,
                limit=10,
            )

            warnings.extend(flight_warnings)

            result = CandidateSelectionResult(
                flight_type="one_way",
                outbound_candidates=(
                    ranked_flights
                ),
                one_way_options=ranked_flights,
                hotels=ranked_hotels,
                warnings=warnings,
            )

            return {
                "candidate_selection_result": result,
                "selection_warnings": warnings,
                "final_answer": (
                    "候选筛选完成："
                    f"保留{len(ranked_flights)}条"
                    "单程航班和"
                    f"{len(ranked_hotels)}家酒店。"
                ),
            }

        (
            outbound_candidates,
            outbound_warnings,
        ) = selection_service.rank_flights(
            raw_flights,
            request,
            limit=max_outbound_queries,
        )

        warnings.extend(outbound_warnings)

        origin_airports = state.get(
            "origin_airports",
            [],
        )

        destination_airports = state.get(
            "destination_airports",
            [],
        )

        combinations: list[
            dict[str, dict[str, Any]]
        ] = []

        stops_filter = get_return_stops_filter(
            request.flight_preferences.max_stops
        )

        for outbound_candidate in (
            outbound_candidates
        ):
            outbound = (
                outbound_candidate.flight
            )

            departure_token = outbound.get(
                "departure_token"
            )

            if not departure_token:
                warnings.append(
                    "去程候选"
                    f"{outbound_candidate.option_id}"
                    "缺少departure_token，"
                    "已跳过返程查询。"
                )

                continue

            tool_args: dict[str, Any] = {
                "departure_token": (
                    departure_token
                ),
                "departure_airports": (
                    origin_airports
                ),
                "arrival_airports": (
                    destination_airports
                ),
                "outbound_date": (
                    request.departure_date
                    .isoformat()
                ),
                "return_date": (
                    request.return_date
                    .isoformat()
                ),
                "adults": request.adults,
                "children": request.children,
                "travel_class": "economy",
                "currency": request.currency,
                "show_hidden": False,
                "deep_search": False,
                "no_cache": False,
            }

            if stops_filter is not None:
                tool_args["stops"] = (
                    stops_filter
                )

            try:
                tool_result = (
                    await return_flight_tool.ainvoke(
                        tool_args
                    )
                )

                payload = (
                    extract_direct_tool_payload(
                        tool_result
                    )
                )

                return_flights = payload.get(
                    "flights",
                    [],
                )

                if not isinstance(
                    return_flights,
                    list,
                ):
                    return_flights = []

                (
                    ranked_returns,
                    return_warnings,
                ) = (
                    selection_service.rank_flights(
                        return_flights,
                        request,
                        limit=(
                            max_returns_per_outbound
                        ),
                        reverse_direction=True,
                        apply_time_preferences=False,
                    )
                )

                warnings.extend(return_warnings)

                for return_candidate in (
                    ranked_returns
                ):
                    combinations.append(
                        {
                            "outbound": outbound,
                            "return_flight": (
                                return_candidate.flight
                            ),
                        }
                    )

            except Exception as exc:
                warnings.append(
                    "返程查询失败："
                    f"{type(exc).__name__}: {exc}"
                )

        round_trip_options = (
            selection_service
            .rank_round_trip_combinations(
                combinations,
                request,
                limit=10,
            )
        )

        result = CandidateSelectionResult(
            flight_type="round_trip",
            outbound_candidates=(
                outbound_candidates
            ),
            round_trip_options=(
                round_trip_options
            ),
            hotels=ranked_hotels,
            warnings=warnings,
        )

        if not round_trip_options:
            message = (
                "没有获得有效的往返航班组合。"
            )

            existing_errors = list(
                state.get("errors", [])
            )

            existing_errors.append(message)

            return {
                "candidate_selection_result": result,
                "selection_warnings": warnings,
                "errors": existing_errors,
                "final_answer": message,
            }

        return {
            "candidate_selection_result": result,
            "selection_warnings": warnings,
            "final_answer": (
                "候选筛选完成："
                f"查询了"
                f"{len(outbound_candidates)}个"
                "去程候选，获得"
                f"{len(round_trip_options)}个"
                "往返组合，并保留"
                f"{len(ranked_hotels)}家酒店。"
            ),
        }

    return candidate_selection_node