""" MoviePilot 自定义工具筛选中间件。 """ import json import re from collections.abc import Awaitable, Callable from dataclasses import dataclass, replace from typing import Annotated, Any, NotRequired, cast from langchain.agents.middleware.tool_selection import ( DEFAULT_SYSTEM_PROMPT, LLMToolSelectorMiddleware, ) from langchain.agents.middleware.types import ( AgentState, ModelRequest, ModelResponse, PrivateStateAttr, # noqa ResponseT, ) from langchain.tools import ToolRuntime from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.messages.utils import count_tokens_approximately from langchain_core.runnables import RunnableConfig from langchain_core.tools import BaseTool, StructuredTool, create_schema_from_function from langgraph.runtime import Runtime from langgraph.types import Command, Overwrite from pydantic import Field from typing_extensions import TypedDict # noqa from app.agent.llm.helper import LLMHelper from app.agent.middleware.usage import UsageMiddleware from app.agent.middleware.utils import append_to_system_message from app.agent.tools.tags import ToolTag from app.runtime.log import logger MIN_SELECTED_TOOL_COUNT = 4 RECENT_SELECTION_CONTEXT_MESSAGE_LIMIT = 6 RECENT_SELECTION_CONTEXT_MAX_CHARS = 6000 RECENT_SELECTION_CONTEXT_TRUNCATION_PREFIX = "..." TOOL_DISCOVERY_NAME = "search_tools" TOOL_DISCOVERY_MAX_RESULTS = 8 TOOL_DISCOVERY_DESCRIPTION_CHARS = 400 TOOL_DISCOVERY_MAX_TAGS = 6 TOOL_DISCOVERY_TAG_CHARS = 48 TOOL_DISCOVERY_WINDOW_SIZE = 16 TOOL_DISCOVERY_SCHEMA_TOKENS = 4096 TOOL_DISCOVERY_SCHEMA_WINDOW_FRACTION = 0.10 TOOL_DISCOVERY_INPUT_FRACTION = 0.85 TOOL_DISCOVERY_STOP_WORDS = frozenset({ "the", "a", "an", "and", "or", "for", "to", "of", "in", "on", "with", "please", "find", "show", "help", "me", "my", "tool", "tools", }) TOOL_DISCOVERY_ALIASES = { ToolTag.Media.value: ("media", "movie", "movies", "tv", "电影", "电视剧", "剧集"), ToolTag.Resource.value: ("resource", "torrent", "种子", "资源"), ToolTag.Site.value: ("site", "sites", "tracker", "站点"), ToolTag.Subscription.value: ("subscription", "subscriptions", "subscribe", "订阅", "追剧"), ToolTag.Download.value: ("download", "downloads", "downloader", "下载"), ToolTag.Library.value: ("library", "mediaserver", "媒体库", "媒体服务器"), ToolTag.Transfer.value: ("transfer", "organize", "整理", "转移"), ToolTag.System.value: ("system", "系统"), ToolTag.Settings.value: ("settings", "configuration", "配置", "设置"), ToolTag.Plugin.value: ("plugin", "plugins", "插件"), ToolTag.Workflow.value: ("workflow", "workflows", "工作流"), ToolTag.Scheduler.value: ("scheduler", "schedule", "定时", "调度"), ToolTag.AgentTask.value: ("agent_task", "智能体任务", "后台任务"), ToolTag.File.value: ("file", "files", "文件"), ToolTag.Directory.value: ("directory", "directories", "folder", "folders", "目录", "文件夹"), ToolTag.Web.value: ("web", "internet", "网页", "互联网"), ToolTag.Command.value: ("command", "shell", "命令", "终端"), ToolTag.FilterRule.value: ("filter", "过滤", "筛选规则"), ToolTag.Persona.value: ("persona", "人格"), ToolTag.Recommendation.value: ("recommendation", "recommend", "推荐"), ToolTag.Metadata.value: ("metadata", "scrape", "元数据", "刮削"), ToolTag.Skill.value: ("skill", "skills", "技能"), } TOOL_GROUP_EXCLUDED_TAGS = frozenset( { ToolTag.AgentTool.value, ToolTag.Read.value, ToolTag.Write.value, ToolTag.Admin.value, ToolTag.Message.value, ToolTag.UserInteraction.value, ToolTag.TerminalResponse.value, } ) MOVIEPILOT_TOOL_SELECTION_HINT = """ MoviePilot tool-chain hints: - Tools with the same capability tag belong to the same functional group. - For multi-step MoviePilot tasks, keep same-tag tools together when relevant. - Prefer selecting likely next-step tools in the same capability group instead of selecting only the first tool. """ def _merge_discovered_tool_names(current: list[str], added: list[str]) -> list[str]: """ 按图更新顺序合并最近发现窗口，重复发现刷新优先级且窗口外能力自动淘汰。 """ latest = list(dict.fromkeys(added)) return [*[name for name in current if name not in latest], *latest][-TOOL_DISCOVERY_WINDOW_SIZE:] class ToolSelectionState(AgentState): """ 工具筛选中间件私有状态。 :ivar selected_tool_names: 本条用户请求首轮筛选得到的工具列表；缺失或 None 表示尚未完成筛选 :ivar discovered_tool_names: 图归并器维护的已发现工具增量，新用户请求开始时清空 """ selected_tool_names: NotRequired[Annotated[list[str] | None, PrivateStateAttr]] discovered_tool_names: Annotated[list[str], PrivateStateAttr, _merge_discovered_tool_names] class ToolSelectionStateUpdate(TypedDict): """ 工具筛选中间件状态更新项。 :ivar selected_tool_names: 要更新的筛选结果；None 表示重置为尚未选择 :ivar discovered_tool_names: 需要整体覆盖工具增量时提供的 Overwrite；缺失时保留原值 """ selected_tool_names: list[str] | None discovered_tool_names: NotRequired[Overwrite] @dataclass(frozen=True) class _ToolSelectionAttempt: """ 工具筛选尝试结果，用于统一记录最终日志。 :param request: 应用本次工具选择后的模型请求 :param selected_tool_names: 本次选择的工具名称列表 :param status: 选择过程的成功、回退或失败状态标识 :param detail: 选择原因或回退说明；无需补充时为空字符串 """ request: ModelRequest[Any] selected_tool_names: list[str] status: str detail: str = "" class ToolSelectorMiddleware(LLMToolSelectorMiddleware): """ 使用 provider-neutral JSON 提示执行工具筛选。 LangChain 默认会通过 `with_structured_output()` 走 provider-specific 的 结构化输出能力，不同 OpenAI/Anthropic 兼容端点对 `response_format`、 JSON schema 和工具绑定的支持并不一致。工具筛选只是 Agent 执行前的 辅助优化，失败时也会恢复使用全部工具，因此这里统一使用文本提示约束 模型返回 `{"tools": [...]}` 并手动解析，避免在筛选阶段引入额外兼容分支。 另外，LangChain 原生工具筛选挂在 `wrap_model_call` 上，会在同一条用户请求 的每次“模型回合”前都重新筛选一次工具。对于会多轮调用工具的复杂任务， 这会重复消耗一次额外的 LLM 调用。这里改成： - `abefore_agent()`：在本轮 Agent 执行开始时筛选一次； - `awrap_model_call()`：复用首轮结果，合入本轮按需发现的额外工具。 """ state_schema = ToolSelectionState def __init__( self, model: BaseChatModel | str | None = None, system_prompt: str = DEFAULT_SYSTEM_PROMPT, selection_tools: list[Any] | None = None, max_tools: int | None = None, always_include: list[str] | None = None, enable_discovery: bool = False, ) -> None: """ 配置首轮筛选与按需发现，新增能力受最近窗口和额外参数预算约束。 """ super().__init__( model=model, system_prompt=self._append_tool_selection_hint(system_prompt), max_tools=max_tools, always_include=always_include, ) self.selection_tools = list(selection_tools or []) self.always_include: list[str] = list(always_include or []) self.tools = [self._build_discovery_tool()] if enable_discovery else [] if self.tools: self.selection_tools.extend(self.tools) self.always_include = [*self.always_include, TOOL_DISCOVERY_NAME] def _build_discovery_tool(self) -> StructuredTool: """ 注册仅供当前 Agent 使用的发现工具，保持目录身份与模型参数可序列化。 """ tool = StructuredTool.from_function( coroutine=self._asearch_tools, name=TOOL_DISCOVERY_NAME, description=( "Find and enable tools missing from the current tool list. Search by exact tool name " "or Chinese/English capability keywords. Matches are candidates for the next model call, " "subject to the latest 16 discoveries and an additional schema/context budget. " "The next model call includes an authoritative availability notice with reasons. " "Use when the task changes or a skill references an unavailable tool. " "This searches only the current authorized catalog; no match enables nothing." ), args_schema=create_schema_from_function( TOOL_DISCOVERY_NAME, self._asearch_tools, filter_args=["runtime"], ), tags=[ToolTag.AgentTool.value, ToolTag.Read.value], ) object.__setattr__(tool, "_agent_tool_source", "middleware:selection") return tool @staticmethod def _matches_discovery_term(term: str, text: str) -> bool: """ 中文短语允许句内匹配，英文使用完整单词避免把子串误认成能力。 """ if re.search(r"[\u4e00-\u9fff]", term): return term in text return bool(re.search(rf"(? int: """ 按工具名、能力标签和说明排序，优先保留精确名称命中的能力。 """ name = tool.name.casefold() description = str(tool.description or "").casefold() tool_tags = cls._normalize_tool_tags(tool) tags = " ".join(tool_tags).casefold() if search == name: return 10000 return 100 * len(capabilities.intersection(tool_tags)) + sum( 10 * cls._matches_discovery_term(keyword, name) + 5 * cls._matches_discovery_term(keyword, tags) + cls._matches_discovery_term(keyword, description) for keyword in keywords ) def _find_discovery_tools(self, search: str, limit: int) -> list[BaseTool]: """ 只在当前授权目录内离线匹配关键词，空白或无关查询不会展开全量工具。 """ query = search.strip().casefold() keywords = [word for word in dict.fromkeys(re.findall(r"[^\W_]+", query)) if word not in TOOL_DISCOVERY_STOP_WORDS] if not keywords: return [] capabilities = { tag for tag, aliases in TOOL_DISCOVERY_ALIASES.items() if any(self._matches_discovery_term(alias, query) for alias in aliases) } ranked_tools = [ (self._discovery_score(tool, query, keywords, capabilities), tool) for tool in self.selection_tools if not isinstance(tool, dict) and tool.name != TOOL_DISCOVERY_NAME ] ranked_tools.sort(key=lambda item: (-item[0], item[1].name)) return [tool for score, tool in ranked_tools if score > 0][:limit] async def _asearch_tools( self, search: Annotated[str, Field( min_length=1, max_length=160, description="Tool name or concise keywords matching descriptions or capability tags.", )], runtime: ToolRuntime, limit: Annotated[int, Field( ge=1, le=TOOL_DISCOVERY_MAX_RESULTS, description="Maximum number of matching tools to enable for subsequent model calls.", )] = 5, ) -> Command[Any]: """ 返回有界工具目录并以图状态增量启用结果，不修改共享中间件实例。 """ matched_tools = self._find_discovery_tools(search, limit) baseline_names = set(self.always_include) | set(getattr(runtime, "state", {}).get("selected_tool_names") or []) catalog = [ { "name": tool.name, "description": str(tool.description or "")[:TOOL_DISCOVERY_DESCRIPTION_CHARS], "tags": [ tag[:TOOL_DISCOVERY_TAG_CHARS] for tag in self._normalize_tool_tags(tool)[:TOOL_DISCOVERY_MAX_TAGS] ], "status": "already_selected" if tool.name in baseline_names else "pending_budget_check", } for tool in matched_tools ] payload = { "tools": catalog, "message": ( "Candidates registered. The next model call reports which tools fit the recent discovery window " "and schema/context budget. Earlier discoveries may be evicted; search an exact name to refresh it." if matched_tools else "No matching tools. Try a precise tool name or different keywords." ), } return Command(update={ "discovered_tool_names": [tool.name for tool in reversed(matched_tools) if tool.name not in baseline_names], "messages": [ToolMessage( content=json.dumps(payload, ensure_ascii=False), tool_call_id=runtime.tool_call_id, name=TOOL_DISCOVERY_NAME, )], }) def _selection_state_update(self, selected_tool_names: list[str] | None) -> ToolSelectionStateUpdate: """ 写入首轮筛选结果并清空前一用户请求的工具发现状态。 """ update = ToolSelectionStateUpdate(selected_tool_names=selected_tool_names) if self.tools: update["discovered_tool_names"] = Overwrite([]) return update @staticmethod def _discovery_availability_notice(statuses: dict[str, str]) -> str: """ 给模型提供当前绑定结果，覆盖发现工具此前返回的候选状态。 """ return ( "\n" "Authoritative availability for this model call (supersedes pending search_tools results):\n" f"{json.dumps(statuses, ensure_ascii=False)}\n" f"Only the latest {TOOL_DISCOVERY_WINDOW_SIZE} discovered tools are candidates. " "Older discoveries are evicted unless initially selected or mandatory. Search an exact name to refresh it. " "Only enabled/already_selected tools are available; budget-disabled tools may fit after context shrinks.\n" "" ) def _discovery_budget(self, request: ModelRequest[Any], statuses: dict[str, str]) -> tuple[int, int]: """ 复用最终请求估算器，为额外工具预留上下文余量并限制结构化参数总成本。 """ budget_request = request.override(system_message=append_to_system_message( request.system_message, self._discovery_availability_notice(statuses), )) budget = UsageMiddleware.estimate_request(budget_request) schema_budget = TOOL_DISCOVERY_SCHEMA_TOKENS context_budget = schema_budget context_window = budget.get("context_window_tokens") if isinstance(context_window, int) and context_window > 0: schema_budget = min(schema_budget, int(context_window * TOOL_DISCOVERY_SCHEMA_WINDOW_FRACTION)) context_budget = max(0, int(context_window * TOOL_DISCOVERY_INPUT_FRACTION) - budget["estimated_input_tokens"]) return schema_budget, context_budget def _budget_discovery_tools(self, request: ModelRequest[Any], candidates: list[BaseTool], statuses: dict[str, str]) -> list[str]: """ 按最近发现优先接纳可容纳的工具，估算失败时只禁用额外发现能力。 """ try: schema_budget, context_budget = self._discovery_budget(request, statuses) except Exception as error: logger.warning(f"工具发现预算估算失败，保留首轮工具: {type(error).__name__}") statuses.update({tool.name: "budget_estimate_unavailable" for tool in candidates}) return [] enabled = [] for tool in candidates: try: cost = count_tokens_approximately([], tools=[tool], use_usage_metadata_scaling=False) except Exception as error: logger.warning(f"工具参数预算无法估算: {tool.name}, {type(error).__name__}") statuses[tool.name] = "budget_estimate_unavailable" continue if cost > schema_budget: statuses[tool.name] = "schema_budget_exceeded" continue if cost > context_budget: statuses[tool.name] = "context_budget_exceeded" continue statuses[tool.name] = "enabled" enabled.append(tool.name) schema_budget -= cost context_budget -= cost return enabled def _bind_discovery_tools(self, request: ModelRequest[Any], selected_tool_names: list[str]) -> ModelRequest[Any]: """ 从当前请求精确工具实例中绑定预算内发现项，并保留初选、强制和供应商工具。 """ baseline_names = list(dict.fromkeys([*selected_tool_names, *self.always_include])) discovered_names = list(reversed(cast(ToolSelectionState, request.state).get("discovered_tool_names", [])[-TOOL_DISCOVERY_WINDOW_SIZE:])) baseline = self._apply_selected_tools(request, baseline_names) if not discovered_names: return baseline current_tools = {tool.name: tool for tool in request.tools if not isinstance(tool, dict)} statuses = { name: "already_selected" if name in baseline_names and name in current_tools else "unavailable_in_current_catalog" for name in discovered_names } candidates = [current_tools[name] for name in discovered_names if name not in baseline_names and name in current_tools] # 先按最长预算拒绝文案估算提示成本，实际启用文案只会缩短。 statuses.update({tool.name: "budget_estimate_unavailable" for tool in candidates}) enabled = self._budget_discovery_tools(baseline, candidates, statuses) bound = self._apply_selected_tools(request, list(dict.fromkeys([ *selected_tool_names, *enabled, *self.always_include, ]))) return bound.override(system_message=append_to_system_message( request.system_message, self._discovery_availability_notice(statuses), )) @classmethod def _render_recent_conversation_context( cls, messages: list[Any], ) -> tuple[str, int]: """ 渲染最近对话上下文，供工具筛选模型理解多轮追问。 """ rendered_messages = [] for message in messages: if isinstance(message, HumanMessage): role = "User" elif isinstance(message, AIMessage): role = "Assistant" else: continue content = LLMHelper.extract_text_content(message.content).strip() if not content: continue rendered_messages.append(f"{role}: {content}") recent_messages = rendered_messages[-RECENT_SELECTION_CONTEXT_MESSAGE_LIMIT:] context = "\n\n".join(recent_messages) if len(context) > RECENT_SELECTION_CONTEXT_MAX_CHARS: context = ( f"{RECENT_SELECTION_CONTEXT_TRUNCATION_PREFIX}" f"{context[-RECENT_SELECTION_CONTEXT_MAX_CHARS:]}" ) return context, len(recent_messages) @classmethod def _build_contextual_user_message( cls, messages: list[Any], last_user_message: HumanMessage, ) -> HumanMessage: """ 根据最近对话构造工具筛选专用用户消息。 """ context, message_count = cls._render_recent_conversation_context(messages) if message_count <= 1: return last_user_message return HumanMessage( content=( "Recent conversation context for tool selection:\n" f"{context}\n\n" "Select tools for the latest user instruction. Use prior assistant " "messages and earlier user requests when the latest user message " "depends on previous context." ) ) def _prepare_selection_request( self, request: ModelRequest[Any], ) -> Any | None: """ 准备带最近对话上下文的工具筛选请求。 """ selection_request = super()._prepare_selection_request(request) if selection_request is None: return None contextual_user_message = self._build_contextual_user_message( messages=request.messages, last_user_message=selection_request.last_user_message, ) if contextual_user_message is selection_request.last_user_message: return selection_request return replace(selection_request, last_user_message=contextual_user_message) @staticmethod def _append_tool_selection_hint(system_prompt: str) -> str: """ 追加 MoviePilot 工具组选择提示，避免复杂链路只选中首个工具。 """ if "MoviePilot tool-chain hints:" in system_prompt: return system_prompt return f"{system_prompt.rstrip()}{MOVIEPILOT_TOOL_SELECTION_HINT}" def _get_tool_selection_limit(self, valid_tool_names: list[str]) -> int: """ 计算补齐筛选结果时允许使用的工具数量上限。 """ if self.max_tools: return min(self.max_tools, len(valid_tool_names)) return len(valid_tool_names) @staticmethod def _normalize_tool_tags(tool: BaseTool) -> list[str]: """ 读取工具的业务标签，过滤掉无法表达工具组的通用标签。 """ tags = getattr(tool, "tags", None) or [] if isinstance(tags, str): tags = [tags] normalized_tags = [] for tag in tags: tag_value = getattr(tag, "value", tag) if not tag_value: continue tag_name = str(tag_value) if tag_name in TOOL_GROUP_EXCLUDED_TAGS or tag_name in normalized_tags: continue normalized_tags.append(tag_name) return normalized_tags @classmethod def _build_tool_groups( cls, available_tools: list[BaseTool], valid_tool_names: list[str], ) -> list[tuple[str, list[str]]]: """ 根据工具标签构造能力组，保留当前工具列表中的稳定顺序。 """ valid_tool_set = set(valid_tool_names) tool_groups: dict[str, list[str]] = {} for tool in available_tools: tool_name = getattr(tool, "name", None) if not tool_name or tool_name not in valid_tool_set: continue for tag in cls._normalize_tool_tags(tool): group_tool_names = tool_groups.setdefault(tag, []) if tool_name not in group_tool_names: group_tool_names.append(tool_name) return [ (tag, tool_names) for tag, tool_names in tool_groups.items() if len(tool_names) > 1 ] @classmethod def _get_matched_tool_groups( cls, selected_names: list[str], available_tools: list[BaseTool], valid_tool_names: list[str], ) -> list[tuple[str, list[str]]]: """ 返回已选工具命中的标签能力组。 """ groups_by_tag = { tag: tool_names for tag, tool_names in cls._build_tool_groups( available_tools=available_tools, valid_tool_names=valid_tool_names, ) } tools_by_name = { tool.name: tool for tool in available_tools if getattr(tool, "name", None) } matched_groups: list[tuple[str, list[str]]] = [] seen_tags = set() for tool_name in selected_names: tool = tools_by_name.get(tool_name) if not tool: continue for tag in cls._normalize_tool_tags(tool): if tag in seen_tags or tag not in groups_by_tag: continue matched_groups.append((tag, groups_by_tag[tag])) seen_tags.add(tag) return matched_groups def _complete_low_count_selection( self, selected_tool_names: list[str], valid_tool_names: list[str], available_tools: list[BaseTool], ) -> list[str]: """ 当模型只选出极少工具时，按工具标签补齐同组工具。 工具标签是工具自身声明的能力归属。这里只补齐已经命中的标签组， 不会把所有工具组都展开。 """ limit = self._get_tool_selection_limit(valid_tool_names) selected_names = [ tool_name for tool_name in selected_tool_names if tool_name in valid_tool_names ] selected_set = set(selected_names) valid_tool_set = set(valid_tool_names) completed_names = list(selected_names) matched_groups = self._get_matched_tool_groups( selected_names=selected_names, available_tools=available_tools, valid_tool_names=valid_tool_names, ) if not matched_groups: return completed_names[:limit] matched_group_tool_names = { tool_name for _, group_tool_names in matched_groups for tool_name in group_tool_names } target_count = min( max(MIN_SELECTED_TOOL_COUNT, len(matched_group_tool_names)), limit, ) if len(selected_names) >= target_count: return selected_names[:limit] for _, group_tool_names in matched_groups: for tool_name in group_tool_names: if tool_name in selected_set or tool_name not in valid_tool_set: continue completed_names.append(tool_name) selected_set.add(tool_name) if len(completed_names) >= target_count: return completed_names[:limit] return completed_names[:limit] def _process_selection_response( self, response: dict[str, Any], available_tools: list[BaseTool], valid_tool_names: list[str], request: ModelRequest[Any], *, apply_max_tools: bool = True, ) -> ModelRequest[Any]: """ 处理工具筛选响应，并在正常空结果时禁用可筛选工具。 """ if response.get("tools") == []: always_included_tools: list[BaseTool] = [ tool for tool in request.tools if not isinstance(tool, dict) and tool.name in self.always_include ] provider_tools = [tool for tool in request.tools if isinstance(tool, dict)] return request.override(tools=[*always_included_tools, *provider_tools]) response["tools"] = self._complete_low_count_selection( selected_tool_names=[ tool_name for tool_name in response.get("tools", []) if isinstance(tool_name, str) ], valid_tool_names=valid_tool_names, available_tools=available_tools, ) modified_request = super()._process_selection_response( response, available_tools, valid_tool_names, request, apply_max_tools=apply_max_tools, ) return modified_request @staticmethod def _parse_json_object(text: str) -> dict[str, Any]: """ 解析模型返回的 JSON。 不同模型可能偶发输出 Markdown 围栏或前后说明文本，因此这里从 响应中提取第一个 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 @classmethod def _render_tool_list(cls, available_tools: list[Any]) -> str: """ 把工具名和描述渲染成稳定的文本列表。 """ lines = [] for tool in available_tools: tags = cls._normalize_tool_tags(tool) tag_text = f" [group tags: {', '.join(tags)}]" if tags else "" lines.append(f"- {tool.name}{tag_text}: {tool.description}") return "\n".join(lines) @classmethod def _render_tool_groups(cls, available_tools: list[BaseTool]) -> str: """ 把当前可用工具按标签渲染成能力组提示。 """ valid_tool_names = [ tool.name for tool in available_tools if getattr(tool, "name", None) ] groups = cls._build_tool_groups( available_tools=available_tools, valid_tool_names=valid_tool_names, ) if not groups: return "" rendered_groups = "\n".join( f"- {tag}: {', '.join(tool_names)}" for tag, tool_names in groups ) return f"Capability groups from tool tags:\n{rendered_groups}\n\n" def _build_json_selection_prompt(self, selection_request: Any) -> str: """ 生成显式 JSON 输出提示。 使用纯提示约束可覆盖更多兼容端点，避免在工具筛选阶段依赖某个 provider 专属的 `response_format` 或 schema 能力。 """ limit_instruction = "" if self.max_tools: limit_instruction = f"- Select up to {self.max_tools} tools. Return an empty array if no tools are relevant." 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" "- Tools sharing the same capability tag are in the same group; include same-group tools together when relevant.\n" f"{limit_instruction}\n" "- Do not add explanations, markdown, or extra keys.\n\n" f"{self._render_tool_groups(selection_request.available_tools)}" "Allowed tools:\n" f"{self._render_tool_list(selection_request.available_tools)}" ) def _normalize_selection_response(self, response: Any) -> dict[str, list[str]]: """ 解析并标准化显式 JSON 模式的工具筛选结果。 """ content = getattr(response, "content", response) text = LLMHelper.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_json_prompt( self, selection_request: Any ) -> dict[str, list[str]]: """ 使用 JSON 提示执行异步工具筛选。 :param selection_request: LangChain 工具筛选请求 :return: 标准化后的工具名列表 """ logger.debug("工具筛选走 JSON 提示分支") response = await selection_request.model.ainvoke( [ SystemMessage( content=self._build_json_selection_prompt(selection_request) ), selection_request.last_user_message, ] ) return self._normalize_selection_response(response) @staticmethod def _extract_selected_tool_names(request: ModelRequest[Any]) -> list[str]: """ 从已筛选后的请求中提取最终工具名，保留原有顺序。 """ return [tool.name for tool in request.tools if not isinstance(tool, dict)] @staticmethod def _count_request_tools(request: ModelRequest[Any]) -> int: """ 统计当前请求中的 LangChain 工具数量，不包含 provider 原生工具字典。 """ return len([tool for tool in request.tools if not isinstance(tool, dict)]) @classmethod def _log_selection_attempt(cls, attempt: _ToolSelectionAttempt) -> None: """ 按工具筛选最终状态记录稳定日志。 """ tool_count = cls._count_request_tools(attempt.request) if attempt.status == "selected": selected_text = ", ".join(attempt.selected_tool_names) or "无有效工具" logger.info(f"工具筛选结果: {selected_text}") return if attempt.status == "failed_fallback": logger.warning( f"工具筛选失败，将恢复使用所有工具（共 {tool_count} 个）: {attempt.detail}" ) return if attempt.status == "skipped": logger.info(f"工具筛选跳过: {attempt.detail}。") return if attempt.status == "reused": selected_text = ", ".join(attempt.selected_tool_names) or "无有效工具" logger.info(f"工具筛选复用已有结果: {selected_text}") @staticmethod def _apply_selected_tools( request: ModelRequest[Any], selected_tool_names: list[str], ) -> ModelRequest[Any]: """ 将已筛选出的工具集应用到当前模型请求。 这里复用首轮筛选及本轮发现的客户端工具名；provider-specific 的 dict 工具仍然 原样保留，避免破坏 LangChain/provider 自身的工具绑定约定。 """ 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[Any] ) -> ModelRequest[Any]: """ 执行一次真实工具筛选，并返回筛选后的请求对象。 这里单独抽成 helper，便于首次筛选后缓存结果，也便于测试覆盖 “首轮筛选，后续复用”的行为。 """ return (await self._aselect_request_once_with_status(request)).request async def _aselect_request_once_with_status( self, request: ModelRequest[Any] ) -> _ToolSelectionAttempt: """ 执行一次真实工具筛选，并携带最终状态供调用方统一记录日志。 """ selection_request = self._prepare_selection_request(request) if selection_request is None: return _ToolSelectionAttempt( request=request, selected_tool_names=self._extract_selected_tool_names(request), status="skipped", detail="没有需要筛选的工具", ) try: response = await self._aselect_tools_with_json_prompt(selection_request) modified_request = self._process_selection_response( response, selection_request.available_tools, selection_request.valid_tool_names, request, ) return _ToolSelectionAttempt( request=modified_request, selected_tool_names=self._extract_selected_tool_names(modified_request), status="selected", ) except Exception as err: return _ToolSelectionAttempt( request=request, selected_tool_names=self._extract_selected_tool_names(request), status="failed_fallback", detail=str(err), ) async def abefore_agent( # noqa self, state: AgentState, runtime: Runtime[Any], # noqa config: RunnableConfig | None = None, ) -> dict[str, Any] | None: """ 在本轮 Agent 执行开始前完成一次真实工具筛选。 这样后续多轮 `model -> tools -> model` 循环都只复用这一次结果， 不会为每次模型回合重复追加一笔 selector LLM 开销。 """ if not self.selection_tools or self.model is None: detail = "没有可筛选工具" if not self.selection_tools else "未配置筛选模型" logger.info(f"工具筛选跳过: {detail}。") return dict(self._selection_state_update(selected_tool_names=None)) selection_request = ModelRequest( model=self.model, tools=list(self.selection_tools), messages=state["messages"], state=state, runtime=runtime, ) attempt = await self._aselect_request_once_with_status(selection_request) self._log_selection_attempt(attempt) selected_tool_names = attempt.selected_tool_names return dict(self._selection_state_update(selected_tool_names=selected_tool_names)) async def awrap_model_call( self, request: ModelRequest[Any], handler: Callable[ [ModelRequest[Any]], Awaitable[ModelResponse[ResponseT]] ], ) -> ModelResponse[ResponseT]: """ 复用首轮筛选并合入本轮按需发现结果，无需追加工具筛选模型调用。 """ selected_tool_names = cast(ToolSelectionState, 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 ): attempt = await self._aselect_request_once_with_status(request) self._log_selection_attempt(attempt) request = attempt.request selected_tool_names = attempt.selected_tool_names cast(ToolSelectionState, request.state)["selected_tool_names"] = selected_tool_names # noqa if selected_tool_names is not None: request = self._bind_discovery_tools(request, selected_tool_names) return await handler(request)