"""MoviePilot 自定义工具筛选中间件。""" import json from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import Annotated, Any, Literal, Union, NotRequired from langchain.agents.middleware.types import ( AgentMiddleware, AgentState, ContextT, ModelRequest, ModelResponse, ResponseT, ) from langchain.agents.middleware.types import ( PrivateStateAttr, # noqa ) from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import HumanMessage from langchain_core.runnables import RunnableConfig from langchain_core.tools import BaseTool from langgraph.runtime import Runtime from pydantic import Field, TypeAdapter from typing_extensions import TypedDict # noqa from app.log import logger DEFAULT_SYSTEM_PROMPT = ( "Your goal is to select the most relevant tools for answering the user's query." ) @dataclass class _SelectionRequest: """Prepared inputs for tool selection.""" available_tools: list[BaseTool] system_message: str last_user_message: HumanMessage model: BaseChatModel valid_tool_names: list[str] def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter[Any]: """Create a structured output schema for tool selection. Args: tools: Available tools to include in the schema. Returns: `TypeAdapter` for a schema where each tool name is a `Literal` with its description. Raises: AssertionError: If `tools` is empty. """ if not tools: msg = "Invalid usage: tools must be non-empty" raise AssertionError(msg) # Create a Union of Annotated Literal types for each tool name with description # For instance: Union[Annotated[Literal["tool1"], Field(description="...")], ...] literals = [ Annotated[Literal[tool.name], Field(description=tool.description)] for tool in tools # noqa ] selected_tool_type = Union[tuple(literals)] # type: ignore[valid-type] # noqa: UP007 description = "Tools to use. Place the most relevant tools first." class ToolSelectionResponse(TypedDict): """Use to select relevant tools.""" tools: Annotated[list[selected_tool_type], Field(description=description)] # type: ignore[valid-type] return TypeAdapter(ToolSelectionResponse) def _render_tool_list(tools: list[BaseTool]) -> str: """Format tools as markdown list. Args: tools: Tools to format. Returns: Markdown string with each tool on a new line. """ return "\n".join(f"- {tool.name}: {tool.description}" for tool in tools) class ToolSelectionState(AgentState): """工具筛选中间件私有状态。""" selected_tool_names: NotRequired[Annotated[list[str] | None, PrivateStateAttr]] """当前这条用户请求首轮筛选得到的工具名列表。""" class ToolSelectionStateUpdate(TypedDict): """工具筛选中间件状态更新项。""" selected_tool_names: list[str] | None class ToolSelectorMiddleware( AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT] ): """ 为 DeepSeek 兼容端点提供更稳妥的工具筛选实现。 LangChain 默认会通过 `with_structured_output()` 走 OpenAI 的 `response_format=json_schema` 路径，但 DeepSeek 官方 OpenAI 兼容端点公开文档 仅保证 `json_object` 模式可用。对于 `deepseek-reasoner`，这会在工具筛选阶段 提前触发 400，导致 Agent 还没真正开始执行工具就失败。 因此这里仅在识别到 DeepSeek 模型/端点时，退回到显式 JSON 输出模式： 1. 使用 `response_format={"type": "json_object"}`； 2. 在提示词中明确约束返回 JSON 结构； 3. 手动解析 `{"tools": [...]}`，其余模型继续沿用 LangChain 默认实现。 另外，LangChain 原生工具筛选挂在 `wrap_model_call` 上，会在同一条用户请求 的每次“模型回合”前都重新筛选一次工具。对于会多轮调用工具的复杂任务， 这会重复消耗一次额外的 LLM 调用。这里改成： - `abefore_agent()`：在本轮 Agent 执行开始时筛选一次； - `awrap_model_call()`：从 `request.state` 读取首轮筛选结果并复用。 """ state_schema = ToolSelectionState def __init__( self, model: BaseChatModel, system_prompt: str = DEFAULT_SYSTEM_PROMPT, selection_tools: list[Any] | None = None, max_tools: int | None = None, always_include: list[str] | None = None, ) -> None: super().__init__() self.model = model self.system_prompt = system_prompt self.max_tools = max_tools self.always_include = always_include or [] self.selection_tools = selection_tools or [] def _prepare_selection_request( self, request: ModelRequest[ContextT] ) -> _SelectionRequest | None: """Prepare inputs for tool selection. Args: request: the model request. Returns: `SelectionRequest` with prepared inputs, or `None` if no selection is needed. Raises: ValueError: If tools in `always_include` are not found in the request. AssertionError: If no user message is found in the request messages. """ # If no tools available, return None if not request.tools or len(request.tools) == 0: return None # Filter to only BaseTool instances (exclude provider-specific tool dicts) base_tools = [tool for tool in request.tools if not isinstance(tool, dict)] # Validate that always_include tools exist if self.always_include: available_tool_names = {tool.name for tool in base_tools} missing_tools = [ name for name in self.always_include if name not in available_tool_names ] if missing_tools: msg = ( f"Tools in always_include not found in request: {missing_tools}. " f"Available tools: {sorted(available_tool_names)}" ) raise ValueError(msg) # Separate tools that are always included from those available for selection available_tools = [ tool for tool in base_tools if tool.name not in self.always_include ] # If no tools available for selection, return None if not available_tools: return None system_message = self.system_prompt # If there's a max_tools limit, append instructions to the system prompt if self.max_tools is not None: system_message += ( f"\nIMPORTANT: List the tool names in order of relevance, " f"with the most relevant first. " f"If you exceed the maximum number of tools, " f"only the first {self.max_tools} will be used." ) # Get the last user message from the conversation history last_user_message: HumanMessage for message in reversed(request.messages): if isinstance(message, HumanMessage): last_user_message = message break else: msg = "No user message found in request messages" raise AssertionError(msg) model = self.model or request.model valid_tool_names = [tool.name for tool in available_tools] return _SelectionRequest( available_tools=available_tools, system_message=system_message, last_user_message=last_user_message, model=model, valid_tool_names=valid_tool_names, ) def _process_selection_response( self, response: dict[str, Any], available_tools: list[BaseTool], valid_tool_names: list[str], request: ModelRequest[ContextT], ) -> ModelRequest[ContextT]: """Process the selection response and return filtered `ModelRequest`.""" selected_tool_names: list[str] = [] invalid_tool_selections = [] for tool_name in response["tools"]: if tool_name not in valid_tool_names: invalid_tool_selections.append(tool_name) continue # Only add if not already selected and within max_tools limit if tool_name not in selected_tool_names and ( self.max_tools is None or len(selected_tool_names) < self.max_tools ): selected_tool_names.append(tool_name) if invalid_tool_selections: msg = f"Model selected invalid tools: {invalid_tool_selections}" raise ValueError(msg) # Filter tools based on selection and append always-included tools if selected_tool_names: selected_tools: list[BaseTool] = [ tool for tool in available_tools if tool.name in selected_tool_names ] else: # 如果模型筛选结果为空，则不对工具进行裁剪，使用所有可用工具 logger.warning("工具筛选结果为空，将恢复使用所有工具。") selected_tools = available_tools always_included_tools: list[BaseTool] = [ tool for tool in request.tools if not isinstance(tool, dict) and tool.name in self.always_include ] selected_tools.extend(always_included_tools) # Also preserve any provider-specific tool dicts from the original request provider_tools = [tool for tool in request.tools if isinstance(tool, dict)] return request.override(tools=[*selected_tools, *provider_tools]) @staticmethod def _is_deepseek_compatible_model(model: BaseChatModel) -> bool: """ 判断当前模型是否应当走 DeepSeek JSON 兼容分支。 除了官方 `langchain_deepseek`，用户也可能通过 OpenAI-compatible 配置把 DeepSeek 端点接到 `ChatOpenAI`。因此这里同时检查模块名、模型名 和 Base URL，避免只靠单一条件漏判。 """ module_name = type(model).__module__.lower() model_name = ( str(getattr(model, "model_name", "") or getattr(model, "model", "")) .strip() .lower() ) base_url = ( str(getattr(model, "openai_api_base", "") or getattr(model, "api_base", "")) .strip() .lower() ) return ( "deepseek" in module_name or model_name.startswith("deepseek-") or "api.deepseek.com" in base_url ) @staticmethod def _extract_text_content(content: Any) -> str: """ 从模型响应中提取纯文本。 这里不依赖上层 LLMHelper，避免中间件与 LLM 构造逻辑互相耦合。 """ if content is None: return "" if isinstance(content, str): return content if isinstance(content, list): text_parts: list[str] = [] for block in content: if isinstance(block, str): text_parts.append(block) continue if isinstance(block, dict): if block.get("type") == "text" and isinstance( block.get("text"), str ): text_parts.append(block["text"]) continue if not block.get("type") and isinstance(block.get("text"), str): text_parts.append(block["text"]) return "".join(text_parts) if isinstance(content, dict): if content.get("type") == "text" and isinstance(content.get("text"), str): return content["text"] if not content.get("type") and isinstance(content.get("text"), str): return content["text"] return "" @staticmethod def _parse_json_object(text: str) -> dict[str, Any]: """ 解析模型返回的 JSON。 DeepSeek 在 JSON 模式下通常会返回纯 JSON，但这里仍做一层兜底， 兼容模型偶发输出围栏或前后说明文本的情况。 """ stripped_text = text.strip() if not stripped_text: raise ValueError("工具筛选返回了空响应") try: payload = json.loads(stripped_text) if isinstance(payload, dict): return payload except json.JSONDecodeError: pass start = stripped_text.find("{") end = stripped_text.rfind("}") if start == -1 or end == -1 or end <= start: raise ValueError(f"工具筛选返回的内容不是合法 JSON: {stripped_text}") payload = json.loads(stripped_text[start: end + 1]) if not isinstance(payload, dict): raise ValueError("工具筛选 JSON 顶层必须是对象") return payload @staticmethod def _render_tool_list(available_tools: list[Any]) -> str: """把工具名和描述渲染成稳定的文本列表。""" return "\n".join( f"- {tool.name}: {tool.description}" for tool in available_tools ) def _build_deepseek_selection_prompt(self, selection_request: Any) -> str: """ 为 DeepSeek 生成显式 JSON 输出提示。 DeepSeek 官方文档要求在 JSON 输出模式下，提示词中必须明确包含 JSON 约束，否则兼容端点可能返回空内容或无意义输出。 """ limit_instruction = "" if self.max_tools: limit_instruction = f"- Select up to {self.max_tools} tools. IF NO TOOLS ARE RELEVANT, DO NOT RETURN AN EMPTY ARRAY. SELECT THE MOST APPLICABLE ONES TO ENSURE THE REQUEST IS HANDLED." return ( f"{selection_request.system_message}\n\n" "Return the answer in JSON only.\n" 'Use exactly this shape: {"tools": ["tool_name_1", "tool_name_2"]}\n' "Rules:\n" "- The `tools` field must be a JSON array of strings.\n" "- Only use tool names from the allowed list below.\n" "- Order tools by relevance, with the most relevant first.\n" f"{limit_instruction}\n" "- Do not add explanations, markdown, or extra keys.\n\n" "Allowed tools:\n" f"{self._render_tool_list(selection_request.available_tools)}" ) def _normalize_selection_response(self, response: Any) -> dict[str, list[str]]: """ 解析并标准化 DeepSeek JSON 模式的工具筛选结果。 """ content = getattr(response, "content", response) text = self._extract_text_content(content) logger.debug(f"工具筛选原始响应: {text}") payload = self._parse_json_object(text) tools = payload.get("tools") if not isinstance(tools, list): raise ValueError(f"工具筛选 JSON 缺少 `tools` 数组: {payload}") normalized_tools = [ tool_name for tool_name in tools if isinstance(tool_name, str) ] logger.debug(f"工具筛选标准化结果: {normalized_tools}") return {"tools": normalized_tools} async def _aselect_tools_with_deepseek( self, selection_request: Any ) -> dict[str, list[str]]: """ 使用 DeepSeek 兼容的 JSON 输出模式执行异步工具筛选。 """ logger.debug("工具筛选走 DeepSeek JSON 兼容分支") structured_model = selection_request.model.bind( response_format={"type": "json_object"} ) response = await structured_model.ainvoke( [ { "role": "system", "content": self._build_deepseek_selection_prompt(selection_request), }, selection_request.last_user_message, ] ) return self._normalize_selection_response(response) @staticmethod def _extract_selected_tool_names(request: ModelRequest) -> list[str]: """从已筛选后的请求中提取最终工具名，保留原有顺序。""" return [tool.name for tool in request.tools if not isinstance(tool, dict)] @staticmethod def _apply_selected_tools( request: ModelRequest[ContextT], selected_tool_names: list[str], ) -> ModelRequest[ContextT]: """ 将已筛选出的工具集应用到当前模型请求。 这里只复用首次筛选出的客户端工具名；provider-specific 的 dict 工具仍然 原样保留，避免破坏 LangChain/provider 自身的工具绑定约定。 """ if not selected_tool_names: return request current_tools_by_name = { tool.name: tool for tool in request.tools if not isinstance(tool, dict) } selected_tools = [ current_tools_by_name[tool_name] for tool_name in selected_tool_names if tool_name in current_tools_by_name ] provider_tools = [tool for tool in request.tools if isinstance(tool, dict)] return request.override(tools=[*selected_tools, *provider_tools]) async def _aselect_request_once( self, request: ModelRequest[ContextT] ) -> ModelRequest[ContextT]: """ 执行一次真实工具筛选，并返回筛选后的请求对象。 这里单独抽成 helper，便于首次筛选后缓存结果，也便于测试覆盖 “首轮筛选，后续复用”的行为。 """ selection_request = self._prepare_selection_request(request) if selection_request is None: return request if not self._is_deepseek_compatible_model(selection_request.model): captured_request: ModelRequest[ContextT] = request async def _capture_handler( updated_request: ModelRequest[ContextT], ) -> ModelRequest[ContextT]: nonlocal captured_request captured_request = updated_request return updated_request await super().awrap_model_call(request, _capture_handler) return captured_request response = await self._aselect_tools_with_deepseek(selection_request) return self._process_selection_response( response, selection_request.available_tools, selection_request.valid_tool_names, request, ) async def abefore_agent( # noqa self, state: ToolSelectionState, runtime: Runtime, # noqa config: RunnableConfig, ) -> ToolSelectionStateUpdate | None: # ty: ignore[invalid-method-override] """ 在本轮 Agent 执行开始前完成一次真实工具筛选。 这样后续多轮 `model -> tools -> model` 循环都只复用这一次结果， 不会为每次模型回合重复追加一笔 selector LLM 开销。 """ if "selected_tool_names" in state: return None if not self.selection_tools or self.model is None: return ToolSelectionStateUpdate(selected_tool_names=None) selection_request = ModelRequest( model=self.model, tools=list(self.selection_tools), messages=state["messages"], state=state, runtime=runtime, ) modified_request = await self._aselect_request_once(selection_request) selected_tool_names = self._extract_selected_tool_names(modified_request) return ToolSelectionStateUpdate(selected_tool_names=selected_tool_names or None) async def awrap_model_call( self, request: ModelRequest[ContextT], handler: Callable[ [ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]] ], ) -> ModelResponse[ResponseT]: """ 从 state 中读取首次筛选结果，并应用到每次模型回合。 """ selected_tool_names = request.state.get("selected_tool_names") # noqa # 正常路径下，`abefore_agent()` 已经提前写入状态；这里只保留一层兜底， # 兼容直接单测或未来某些绕过 before_agent 的调用场景。 if ( selected_tool_names is None and self.selection_tools and self.model is not None ): request = await self._aselect_request_once(request) selected_tool_names = self._extract_selected_tool_names(request) or None request.state["selected_tool_names"] = selected_tool_names # noqa if selected_tool_names: request = self._apply_selected_tools(request, selected_tool_names) return await handler(request)