mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-11-11 16:21:09 +08:00
167 lines
8.1 KiB
Python
167 lines
8.1 KiB
Python
"""为当前会话保留有界工具结果，并按字符位置继续读取。"""
|
|
|
|
import json
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections import OrderedDict
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
from typing import Annotated, Any
|
|
|
|
from langchain.agents.middleware.types import AgentMiddleware, ToolCallRequest
|
|
from langchain.tools import ToolRuntime
|
|
from langchain_core.messages import ToolMessage
|
|
from langchain_core.tools import StructuredTool, create_schema_from_function
|
|
from pydantic import Field
|
|
|
|
from app.agent.code.authority import RPC_RESULT_KEY
|
|
from app.agent.policy.contracts import ToolPolicyContext
|
|
from app.agent.tools.base import TOOL_RESULT_RECORDER, format_tool_result_for_agent
|
|
from app.agent.tools.result import inspect_tool_result
|
|
|
|
READ_TOOL_RESULT_NAME = "read_tool_result"
|
|
MAX_RESULT_BYTES = 1024 * 1024
|
|
MAX_TOTAL_BYTES = 4 * MAX_RESULT_BYTES
|
|
MAX_RESULTS = 8
|
|
RESULT_TTL_SECONDS = 15 * 60
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _StoredResult:
|
|
"""内存中短期保留原始授权结果，不复制到日志、记忆或磁盘。"""
|
|
|
|
thread_id: str
|
|
text: str
|
|
byte_size: int
|
|
tool_name: str
|
|
expires_at: float
|
|
requires_admin: bool
|
|
|
|
|
|
# follow_imports=skip 不分析第三方基类，仅在框架继承边界忽略 misc。
|
|
class ToolOutputMiddleware(AgentMiddleware): # type: ignore[misc]
|
|
"""包装当前图的工具结果；读取按会话与当前管理员身份双重隔离。"""
|
|
|
|
def __init__(self, context: ToolPolicyContext) -> None:
|
|
"""注册内部续读工具，缓存容量和寿命均受固定上限约束。"""
|
|
self.context = context
|
|
self._results: OrderedDict[str, _StoredResult] = OrderedDict()
|
|
self._lock = threading.RLock()
|
|
self.tools = [StructuredTool.from_function(
|
|
name=READ_TOOL_RESULT_NAME,
|
|
description=(
|
|
"Read the next page of an archived tool result using its result_id and next_offset. "
|
|
"Offsets count Unicode characters. Results expire after 15 minutes, on eviction, "
|
|
"or when this conversation graph is rebuilt. Never guess result IDs."
|
|
),
|
|
coroutine=self._read_result,
|
|
args_schema=create_schema_from_function(
|
|
"ReadToolResultInput", self._read_result, filter_args=["runtime"],
|
|
),
|
|
tags=["agent_tool", "read"],
|
|
)]
|
|
object.__setattr__(self.tools[0], "_agent_tool_source", "middleware:output")
|
|
|
|
@staticmethod
|
|
def _thread_id(runtime: ToolRuntime) -> str:
|
|
"""使用图的真实线程身份，避免同一中间件被不同图线程误复用。"""
|
|
return str(runtime.config.get("configurable", {}).get("thread_id") or "")
|
|
|
|
def _prune(self) -> None:
|
|
"""调用方持锁时删除到期结果，并按写入顺序执行 UTF-8 字节容量淘汰。"""
|
|
now = time.monotonic()
|
|
for result_id, result in list(self._results.items()):
|
|
if result.expires_at <= now:
|
|
del self._results[result_id]
|
|
size = sum(result.byte_size for result in self._results.values())
|
|
while self._results and (len(self._results) > MAX_RESULTS or size > MAX_TOTAL_BYTES):
|
|
_, result = self._results.popitem(last=False)
|
|
size -= result.byte_size
|
|
|
|
def _store(self, thread_id: str, tool_name: str, text: str, requires_admin: bool) -> dict[str, Any]:
|
|
"""只保存完整的有界结果，超过容量时明确说明无法续读。"""
|
|
if not thread_id:
|
|
return {"result_unavailable": "thread_unavailable"}
|
|
# 先用字符数下界拒绝超大文本，避免为明显无法归档的输出再分配编码副本。
|
|
if len(text) > MAX_RESULT_BYTES or (byte_size := len(text.encode("utf-8"))) > MAX_RESULT_BYTES:
|
|
return {"result_unavailable": "result_too_large", "result_limit_bytes": MAX_RESULT_BYTES}
|
|
result_id = uuid.uuid4().hex
|
|
with self._lock:
|
|
self._results[result_id] = _StoredResult(
|
|
thread_id=thread_id, text=text, byte_size=byte_size, tool_name=tool_name,
|
|
expires_at=time.monotonic() + RESULT_TTL_SECONDS,
|
|
requires_admin=requires_admin,
|
|
)
|
|
self._prune()
|
|
return {
|
|
"result_id": result_id,
|
|
"read_tool": READ_TOOL_RESULT_NAME,
|
|
"expires_in_seconds": RESULT_TTL_SECONDS,
|
|
"offset_unit": "unicode_characters",
|
|
}
|
|
|
|
async def _read_result(
|
|
self,
|
|
result_id: Annotated[str, Field(min_length=32, max_length=32)],
|
|
runtime: ToolRuntime,
|
|
offset: Annotated[int, Field(ge=0)] = 0,
|
|
limit: Annotated[int, Field(ge=1, le=16000)] = 4000,
|
|
) -> str:
|
|
"""按显式游标续读；过期、越权和未知 ID 使用同一不可用响应。"""
|
|
with self._lock:
|
|
self._prune()
|
|
result = self._results.get(result_id)
|
|
if (
|
|
not self._thread_id(runtime) or result is None or result.thread_id != self._thread_id(runtime)
|
|
or (result.requires_admin and not self.context.agent_context.get("is_admin"))
|
|
):
|
|
return json.dumps({"success": False, "error": "result_unavailable"})
|
|
if offset > len(result.text):
|
|
return json.dumps({"success": False, "error": "offset_out_of_range", "total_chars": len(result.text)})
|
|
end = min(offset + limit, len(result.text))
|
|
return json.dumps({
|
|
"success": True, "result_id": result_id, "tool_name": result.tool_name,
|
|
"offset": offset, "next_offset": end if end < len(result.text) else None,
|
|
"total_chars": len(result.text), "content": result.text[offset:end],
|
|
}, ensure_ascii=False)
|
|
|
|
async def awrap_tool_call(
|
|
self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], Awaitable[Any]],
|
|
) -> Any:
|
|
"""让内置工具在截断前归档，并覆盖返回大文本的外部 MCP 工具。"""
|
|
if request.tool_call.get("name") == READ_TOOL_RESULT_NAME:
|
|
return await handler(request)
|
|
if request.tool_call.get('name') == 'execute_code' and getattr(request.tool, '_agent_tool_source', None) == 'builtin':
|
|
return await handler(request)
|
|
# 历史工具已经按消息/窗口预算裁剪；再次套通用预览会丢失锚点与书挡结构。
|
|
if request.tool_call.get("name") == "session_search" and getattr(request.tool, "_agent_tool_source", None) == "middleware:recall":
|
|
return await handler(request)
|
|
# 学习工具已限制单文件、条目及批次；预览截断会破坏 JSON 和写前读取证据。
|
|
if (request.tool_call.get('name') in {'memory', 'skill_manage', 'skill_view', 'skills_list'}
|
|
and getattr(getattr(request, 'tool', None), '_agent_tool_source', None) == 'middleware:learning'):
|
|
return await handler(request)
|
|
thread_id = self._thread_id(request.runtime)
|
|
requires_admin = bool(self.context.agent_context.get("is_admin"))
|
|
def record(name: str, text: str) -> dict[str, Any]:
|
|
"""Python 聚合读取完整结果，同时保持普通工具的有界展示和归档引用。"""
|
|
state = getattr(request, 'state', None)
|
|
raw = state.get(RPC_RESULT_KEY) if isinstance(state, dict) else None
|
|
if isinstance(raw, dict):
|
|
raw[name] = text
|
|
return self._store(thread_id, name, text, requires_admin)
|
|
|
|
token = TOOL_RESULT_RECORDER.set(record)
|
|
try:
|
|
result = await handler(request)
|
|
if isinstance(result, ToolMessage) and isinstance(result.content, str):
|
|
content = format_tool_result_for_agent(result.content, tool_name=result.name)
|
|
if content != result.content:
|
|
payload = json.loads(content)
|
|
payload["execution_outcome"] = inspect_tool_result(result).value
|
|
content = json.dumps(payload, ensure_ascii=False, indent=2)
|
|
return result.model_copy(update={"content": content})
|
|
return result
|
|
finally:
|
|
TOOL_RESULT_RECORDER.reset(token)
|