mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-29 19:37:54 +08:00
严格检查全部宿主源码,清零类型错误预算并移除按模块类型豁免。 按原因、所属方法、错误码及语法指纹锁定必要兼容例外和站点类型桩, CI 阻断新增忽略、检查范围退化、本地桩替代真实接口及新的类型错误。 修复空值、异步和供应方类型边界,保持插件与 API 兼容合同, 并补充存储、订阅、通知、Agent 和工作流边界的回归覆盖。
959 lines
40 KiB
Python
959 lines
40 KiB
Python
"""
|
||
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)
|