chendandan_homework/app/agents/resource_agent.py

from __future__ import annotations

"""资源查询Agent：调用航班/酒店MCP工具搜索外部资源，指导工具调用顺序并汇总结果为结构化数据。"""

from typing import Any

from langchain.agents import create_agent
from langchain_core.language_models.chat_models import (
    BaseChatModel,
)
from langchain_core.tools import BaseTool

from app.llm import get_chat_model
from app.schemas.travel_request import TravelRequest


# RESOURCE_AGENT_PROMPT：资源查询的系统提示词，定义航班与酒店搜索的工具调用策略。
RESOURCE_AGENT_PROMPT = """
你是旅行系统中的资源查询Agent。

你只负责查询机场、航班和酒店，不负责生成最终行程。

必须按照以下顺序执行：

1. 使用search_airports查询出发城市的机场。
2. 使用search_airports查询目的城市的机场。
3. 优先选择type为city且城市名称最匹配的结果。
4. 提取该城市包含的全部机场代码。
5. 使用search_flights查询真实航班。
6. 使用search_hotels查询目的城市的真实酒店。
7. 当前阶段不要调用search_return_flights。
8. 用户没有提出限制时，不得自行限制：
   - 是否直飞；
   - 起飞或落地时间；
   - 机场；
   - 机票最高价格；
   - 酒店星级和评分。
9. 不得编造机场代码、价格、航班或酒店。
10. 工具查询结束后，只输出简短的查询完成摘要。

航班查询规则：

- 用户有返程日期时，flight_type使用round_trip。
- departure_airports使用出发城市全部有效机场代码。
- arrival_airports使用目的城市全部有效机场代码。
- 日期必须使用YYYY-MM-DD。
- 成人和儿童数量严格使用用户需求。
- 默认币种使用CNY。

酒店查询规则：

- query使用“目的城市 + 酒店”。
- 入住日期使用出发日期。
- 退房日期使用返程日期。
- 不要自行增加价格、评分或星级限制。
"""


class ResourceSearchAgent:
    """查询机场、航班和酒店的专业Agent。"""

    def __init__(
        self,
        tools: list[BaseTool],
        model: BaseChatModel | None = None,
    ) -> None:
        """初始化资源查询Agent，支持注入自定义模型和MCP工具集。"""
        allowed_names = {
            "search_airports",
            "search_flights",
            "search_hotels",
        }

        selected_tools = [
            tool
            for tool in tools
            if tool.name in allowed_names
        ]

        loaded_names = {
            tool.name
            for tool in selected_tools
        }

        missing_names = (
            allowed_names
            - loaded_names
        )

        if missing_names:
            raise RuntimeError(
                "资源查询Agent缺少工具："
                + "、".join(
                    sorted(missing_names)
                )
            )

        self._agent = create_agent(
            model=model or get_chat_model(),
            tools=selected_tools,
            system_prompt=RESOURCE_AGENT_PROMPT,
        )

    async def search(
        self,
        request: TravelRequest,
    ) -> dict[str, Any]:
        """根据结构化旅行需求调用真实工具。"""

        request_json = request.model_dump_json(
            indent=2,
        )

        user_message = f"""
请根据下面的结构化旅行需求执行真实资源查询：

{request_json}

必须实际调用机场、航班和酒店工具。
不要只根据已有知识直接回答。
"""

        return await self._agent.ainvoke(
            {
                "messages": [
                    {
                        "role": "user",
                        "content": user_message,
                    }
                ]
            }
        )