149 lines
5.1 KiB
Python
149 lines
5.1 KiB
Python
"""Agent 工具注册与执行边界。"""
|
|
|
|
import inspect
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from time import perf_counter
|
|
from typing import Any, Awaitable, Callable
|
|
|
|
from pydantic import BaseModel, ValidationError
|
|
from jsonschema import Draft202012Validator
|
|
from jsonschema.exceptions import ValidationError as JsonSchemaValidationError
|
|
|
|
from app.contracts import ToolCall, ToolDefinition, ToolResult
|
|
from app.schema_security import reject_external_schema_references
|
|
|
|
ToolExecutor = Callable[[BaseModel, "ToolExecutionContext"], Any | Awaitable[Any]]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ToolExecutionContext:
|
|
run_id: str
|
|
tool_call_id: str | None = None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RegisteredTool:
|
|
definition: ToolDefinition
|
|
arguments_model: type[BaseModel]
|
|
executor: ToolExecutor
|
|
|
|
|
|
class ToolNotFoundError(LookupError):
|
|
pass
|
|
|
|
|
|
class ToolExecutionError(RuntimeError):
|
|
"""Executor 可预期失败，保留领域错误码而不是折叠成通用异常。"""
|
|
|
|
def __init__(self, code: str, message: str) -> None:
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.message = message
|
|
|
|
|
|
class ToolRegistry:
|
|
"""统一校验工具入参并隔离执行异常，避免单个工具击穿 Agent 主循环。"""
|
|
|
|
def __init__(self) -> None:
|
|
self._tools: dict[str, RegisteredTool] = {}
|
|
self._lock = threading.RLock()
|
|
|
|
def register(
|
|
self,
|
|
definition: ToolDefinition,
|
|
arguments_model: type[BaseModel],
|
|
executor: ToolExecutor,
|
|
) -> None:
|
|
Draft202012Validator.check_schema(definition.parameters)
|
|
reject_external_schema_references(definition.parameters)
|
|
with self._lock:
|
|
if definition.name in self._tools:
|
|
raise ValueError(f"Tool already registered: {definition.name}")
|
|
self._tools[definition.name] = RegisteredTool(
|
|
definition=definition,
|
|
arguments_model=arguments_model,
|
|
executor=executor,
|
|
)
|
|
|
|
def unregister(self, name: str) -> None:
|
|
with self._lock:
|
|
self._tools.pop(name, None)
|
|
|
|
def contains(self, name: str) -> bool:
|
|
with self._lock:
|
|
return name in self._tools
|
|
|
|
def get(self, name: str) -> RegisteredTool:
|
|
with self._lock:
|
|
try:
|
|
return self._tools[name]
|
|
except KeyError as exc:
|
|
raise ToolNotFoundError(name) from exc
|
|
|
|
def definitions(self, allowed: list[str] | None = None) -> list[ToolDefinition]:
|
|
names = set(allowed) if allowed is not None else None
|
|
with self._lock:
|
|
return [
|
|
item.definition.model_copy(deep=True)
|
|
for name, item in self._tools.items()
|
|
if names is None or name in names
|
|
]
|
|
|
|
async def execute(self, call: ToolCall, context: ToolExecutionContext) -> ToolResult:
|
|
started = perf_counter()
|
|
try:
|
|
registered = self.get(call.name)
|
|
except ToolNotFoundError:
|
|
return ToolResult(
|
|
tool_call_id=call.tool_call_id,
|
|
name=call.name,
|
|
success=False,
|
|
error_code="TOOL_NOT_FOUND",
|
|
error_message=f"Tool is not registered: {call.name}",
|
|
)
|
|
|
|
try:
|
|
# JSON Schema 约束模型可见的协议，Pydantic 再完成运行时类型转换。
|
|
Draft202012Validator(registered.definition.parameters).validate(call.arguments)
|
|
arguments = registered.arguments_model.model_validate(call.arguments)
|
|
except (ValidationError, JsonSchemaValidationError) as exc:
|
|
return ToolResult(
|
|
tool_call_id=call.tool_call_id,
|
|
name=call.name,
|
|
success=False,
|
|
error_code="TOOL_ARGUMENT_INVALID",
|
|
error_message=str(exc),
|
|
duration_ms=round((perf_counter() - started) * 1000),
|
|
)
|
|
|
|
try:
|
|
output = registered.executor(arguments, context)
|
|
if inspect.isawaitable(output):
|
|
output = await output
|
|
return ToolResult(
|
|
tool_call_id=call.tool_call_id,
|
|
name=call.name,
|
|
success=True,
|
|
output=output,
|
|
duration_ms=round((perf_counter() - started) * 1000),
|
|
)
|
|
except ToolExecutionError as exc:
|
|
return ToolResult(
|
|
tool_call_id=call.tool_call_id,
|
|
name=call.name,
|
|
success=False,
|
|
error_code=exc.code,
|
|
error_message=exc.message,
|
|
duration_ms=round((perf_counter() - started) * 1000),
|
|
)
|
|
except Exception as exc: # 工具失败转换成结构化结果，由模型决定是否降级或重试。
|
|
return ToolResult(
|
|
tool_call_id=call.tool_call_id,
|
|
name=call.name,
|
|
success=False,
|
|
error_code="TOOL_EXECUTION_FAILED",
|
|
error_message=str(exc),
|
|
duration_ms=round((perf_counter() - started) * 1000),
|
|
)
|