chendandan_homework/app/graph/map_nodes.py

from __future__ import annotations
"""地图研究节点：调用MapResearchAgent通过高德MCP搜索目的地周边景点、餐饮与天气。"""

import json
import re
from typing import Any, Literal

from langchain_core.messages import AIMessage

from app.agents.map_agent import MapResearchAgent
from app.graph.state import TravelState
from app.schemas.place import MapResearchResult


def _extract_json_from_messages(
    messages: list[Any],
) -> dict[str, Any] | None:
    """从 Agent 最后一条 AI 消息中提取 JSON 代码块。"""

    for message in reversed(messages):
        if not isinstance(message, AIMessage):
            continue

        content = message.content

        if not isinstance(content, str):
            continue

        # 优先匹配 ```json ... ``` 代码块
        json_match = re.search(
            r"```json\s*([\s\S]*?)\s*```",
            content,
        )

        if json_match:
            try:
                return json.loads(
                    json_match.group(1)
                )
            except json.JSONDecodeError:
                continue

        # 尝试匹配 ``` ... ``` 无标注代码块
        code_match = re.search(
            r"```\s*(\{[\s\S]*?\})\s*```",
            content,
        )

        if code_match:
            try:
                return json.loads(
                    code_match.group(1)
                )
            except json.JSONDecodeError:
                continue

        # 最后尝试在整段内容中查找 JSON 对象
        json_match = re.search(
            r"\{[\s\S]*\"city\"[\s\S]*\}",
            content,
        )

        if json_match:
            try:
                return json.loads(
                    json_match.group(0)
                )
            except json.JSONDecodeError:
                continue

    return None


def route_after_resources(
    state: TravelState,
) -> Literal[
    "research_destination",
    "end",
]:
    """资源查询成功后进入地图研究。"""

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

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

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

    if not flight_result or not hotel_result:
        return "end"

    return "research_destination"


def make_map_research_node(
    agent: MapResearchAgent,
):
    """创建地图研究节点。"""

    async def map_research_node(
        state: TravelState,
    ) -> dict[str, Any]:
        """地图研究节点：调用Agent获取高德地图POI数据，从消息中提取JSON结构化结果写入state。"""

        request = state["travel_request"]

        try:
            agent_result = await agent.research(
                request
            )

            messages = agent_result.get(
                "messages",
                [],
            )

            parsed = _extract_json_from_messages(
                messages
            )

            if parsed is None:
                # 调试：打印最后一条消息的内容
                last_content = ""
                for msg in reversed(messages):
                    if isinstance(msg, AIMessage):
                        last_content = str(
                            msg.content
                        )[:500]
                        break

                return {
                    "errors": [
                        "地图Agent未返回有效的JSON结构化结果。"
                    ],
                    "map_messages": messages,
                    "final_answer": (
                        "地图Agent未返回有效结果。"
                        f"最后一条消息预览：{last_content}"
                    ),
                }

            structured_response = (
                MapResearchResult.model_validate(
                    parsed
                )
            )

            attraction_count = len(
                structured_response.attractions
            )

            restaurant_count = len(
                structured_response.restaurants
            )

            weather_count = len(
                structured_response.weather
            )

            return {
                "map_research_result": (
                    structured_response
                ),
                "map_messages": messages,
                "final_answer": (
                    "真实地图资源查询完成："
                    f"获得{attraction_count}个景点候选，"
                    f"{restaurant_count}个餐饮候选，"
                    f"{weather_count}天天气信息。"
                ),
            }

        except Exception as exc:
            message = (
                "地图Agent执行失败："
                f"{type(exc).__name__}: {exc}"
            )

            return {
                "errors": [message],
                "final_answer": message,
            }

    return map_research_node