Files
MoviePilot/app/agent/middleware/selection.py
jxxghp 6d4ce83837 fix(types): clear host mypy debt and prevent typing regressions
严格检查全部宿主源码,清零类型错误预算并移除按模块类型豁免。
按原因、所属方法、错误码及语法指纹锁定必要兼容例外和站点类型桩,
CI 阻断新增忽略、检查范围退化、本地桩替代真实接口及新的类型错误。

修复空值、异步和供应方类型边界,保持插件与 API 兼容合同,
并补充存储、订阅、通知、Agent 和工作流边界的回归覆盖。
2026-10-21 22:03:18 +08:00

959 lines
40 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
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"(?<![a-z0-9]){re.escape(term)}(?![a-z0-9])", text))
@classmethod
def _discovery_score(cls, tool: BaseTool, search: str, keywords: list[str], capabilities: set[str]) -> 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 (
"<tool_discovery_availability>\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"
"</tool_discovery_availability>"
)
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)