301 lines
11 KiB
Python
301 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Awaitable, Callable, Iterable
|
|
|
|
from fastapi import FastAPI
|
|
|
|
from .errors import (
|
|
ArgumentValidationError,
|
|
ExecutionError,
|
|
ToolNotFoundError,
|
|
)
|
|
from .openapi_parser import extract_operations, validate_operations
|
|
from .tool import CompositeTool, PilotTool
|
|
from .types import ExecuteResponse, OperationInfo, ToolDefinition, ToolRecord
|
|
from .validation import validate_arguments
|
|
|
|
|
|
UserT = Any
|
|
GetUserFn = Callable[[Any], Awaitable[UserT] | UserT] | None
|
|
AllowToolFn = Callable[[UserT, ToolRecord], Awaitable[bool] | bool] | None
|
|
|
|
|
|
@dataclass
|
|
class ApiPilotConfig:
|
|
default_timeout: float = 30.0
|
|
default_max_response_size: int = 1_000_000
|
|
expose_http_endpoints: bool = False
|
|
http_prefix: str = "/apipilot"
|
|
require_auth_for_http: bool = True
|
|
app: FastAPI | None = None
|
|
get_current_user: GetUserFn = None
|
|
allow_tool: AllowToolFn = None
|
|
|
|
|
|
class ApiPilot:
|
|
def __init__(self, config: ApiPilotConfig | None = None) -> None:
|
|
self.config = config or ApiPilotConfig()
|
|
self._tools: dict[str, ToolRecord] = {}
|
|
self._composites: dict[str, CompositeTool] = {}
|
|
self._user_dependent: bool = False
|
|
|
|
def register_fastapi_app(self, app: FastAPI) -> None:
|
|
self.config.app = app
|
|
openapi_root = app.openapi()
|
|
operations = extract_operations(app)
|
|
results = validate_operations(operations, openapi_root)
|
|
for definition, op, specs in results:
|
|
extras = op.apipilot
|
|
timeout = float(extras.get("timeout", self.config.default_timeout))
|
|
max_size = int(extras.get("max_response_size", self.config.default_max_response_size))
|
|
record = ToolRecord(
|
|
name=definition["name"],
|
|
description=definition["description"],
|
|
tool_type="endpoint",
|
|
definition=definition,
|
|
operation=op,
|
|
param_specs=specs,
|
|
timeout=timeout,
|
|
max_response_size=max_size,
|
|
)
|
|
self._tools[record.name] = record
|
|
|
|
def add_tool(self, tool: CompositeTool | PilotTool) -> None:
|
|
name = tool.name
|
|
if name in self._tools or name in self._composites:
|
|
from .errors import DuplicateToolError
|
|
|
|
raise DuplicateToolError(f"Nome tool duplicato in registrazione manuale: {name}")
|
|
params_schema = getattr(tool, "parameters_schema", None)
|
|
if params_schema is None:
|
|
params_schema = {
|
|
"type": "object",
|
|
"properties": {},
|
|
"required": [],
|
|
}
|
|
definition: ToolDefinition = {
|
|
"name": name,
|
|
"description": tool.description,
|
|
"parameters": params_schema,
|
|
}
|
|
timeout = float(getattr(tool, "timeout", self.config.default_timeout))
|
|
max_size = int(getattr(tool, "max_response_size", self.config.default_max_response_size))
|
|
run_fn: Callable[..., Awaitable[Any]] | None = getattr(tool, "run_fn", None)
|
|
record = ToolRecord(
|
|
name=name,
|
|
description=tool.description,
|
|
tool_type="composite",
|
|
definition=definition,
|
|
composite_run=run_fn or (lambda args, user: tool.run(args, user)),
|
|
timeout=timeout,
|
|
max_response_size=max_size,
|
|
)
|
|
self._tools[name] = record
|
|
if isinstance(tool, CompositeTool):
|
|
self._composites[name] = tool
|
|
|
|
def list_tool_names(self) -> list[str]:
|
|
return list(self._tools.keys())
|
|
|
|
def get_tool(self, name: str) -> ToolRecord | None:
|
|
return self._tools.get(name)
|
|
|
|
async def _is_tool_allowed(self, user: UserT, record: ToolRecord) -> bool:
|
|
allow_fn = self.config.allow_tool
|
|
if allow_fn is None:
|
|
return True
|
|
result = allow_fn(user, record)
|
|
if isinstance(result, Awaitable):
|
|
result = await result
|
|
return bool(result)
|
|
|
|
async def tools_for(self, user: UserT = None) -> list[ToolDefinition]:
|
|
out: list[ToolDefinition] = []
|
|
for record in self._tools.values():
|
|
if await self._is_tool_allowed(user, record):
|
|
out.append(record.definition)
|
|
return out
|
|
|
|
async def execute(
|
|
self,
|
|
*,
|
|
user: UserT = None,
|
|
name: str,
|
|
arguments: dict[str, Any] | None = None,
|
|
) -> ExecuteResponse:
|
|
record = self._tools.get(name)
|
|
if record is None:
|
|
raise ToolNotFoundError(f"Tool non trovato: {name}")
|
|
if not await self._is_tool_allowed(user, record):
|
|
raise ToolNotFoundError(f"Tool non autorizzato per l'utente corrente: {name}")
|
|
|
|
try:
|
|
validated = validate_arguments(arguments, record.definition["parameters"])
|
|
except ArgumentValidationError as e:
|
|
return {
|
|
"success": False,
|
|
"data": None,
|
|
"error": {
|
|
"code": "INVALID_ARGUMENTS",
|
|
"message": str(e),
|
|
"details": e.errors,
|
|
},
|
|
"meta": {"tool": name},
|
|
}
|
|
|
|
try:
|
|
if record.tool_type == "composite":
|
|
result = await self._run_composite(record, validated, user)
|
|
else:
|
|
result = await self._run_endpoint(record, validated, user)
|
|
except ExecutionError as e:
|
|
return {
|
|
"success": False,
|
|
"data": None,
|
|
"error": {
|
|
"code": "EXECUTION_ERROR",
|
|
"message": str(e),
|
|
"status_code": e.status_code,
|
|
"endpoint": e.endpoint,
|
|
},
|
|
"meta": {"tool": name},
|
|
}
|
|
except asyncio.TimeoutError as e:
|
|
return {
|
|
"success": False,
|
|
"data": None,
|
|
"error": {
|
|
"code": "TIMEOUT",
|
|
"message": f"Timeout dopo {record.timeout}s",
|
|
},
|
|
"meta": {"tool": name},
|
|
}
|
|
except Exception as e: # pragma: no cover - defensive
|
|
return {
|
|
"success": False,
|
|
"data": None,
|
|
"error": {
|
|
"code": "UNEXPECTED_EXECUTION_ERROR",
|
|
"message": str(e),
|
|
"type": type(e).__name__,
|
|
},
|
|
"meta": {"tool": name},
|
|
}
|
|
|
|
serialized, truncated = self._enforce_size(result, record.max_response_size)
|
|
meta: dict[str, Any] = {"tool": name}
|
|
if truncated:
|
|
meta["truncated"] = True
|
|
meta["max_response_size"] = record.max_response_size
|
|
return {
|
|
"success": True,
|
|
"data": serialized,
|
|
"error": None,
|
|
"meta": meta,
|
|
}
|
|
|
|
def _enforce_size(self, result: Any, max_size: int) -> tuple[Any, bool]:
|
|
if result is None:
|
|
return None, False
|
|
try:
|
|
serialized = json.dumps(result, ensure_ascii=False, default=str)
|
|
except Exception:
|
|
serialized = str(result)
|
|
if len(serialized) <= max_size:
|
|
return result, False
|
|
if isinstance(result, str):
|
|
return result[:max_size] + "... [truncated]", True
|
|
if isinstance(result, dict) and "items" in result and isinstance(result["items"], list):
|
|
items = result["items"]
|
|
sample: list[Any] = []
|
|
current = len(json.dumps({**result, "items": sample}, ensure_ascii=False, default=str))
|
|
for item in items:
|
|
s = len(json.dumps(item, ensure_ascii=False, default=str))
|
|
if current + s + 32 > max_size:
|
|
break
|
|
sample.append(item)
|
|
current += s
|
|
trimmed = {**result, "items": sample, "_truncated": True, "_total": len(items)}
|
|
return trimmed, True
|
|
compact = serialized[:max_size] + "... [truncated]"
|
|
return compact, True
|
|
|
|
async def _run_composite(
|
|
self, record: ToolRecord, arguments: dict[str, Any], user: UserT
|
|
) -> Any:
|
|
fn = record.composite_run
|
|
if fn is None:
|
|
raise ExecutionError("Tool composto senza funzione run", endpoint=record.name)
|
|
try:
|
|
async with asyncio.timeout(record.timeout):
|
|
return await fn(arguments, user)
|
|
except asyncio.TimeoutError:
|
|
raise
|
|
except ExecutionError:
|
|
raise
|
|
except Exception as e:
|
|
raise ExecutionError(str(e), cause=e, endpoint=record.name) from e
|
|
|
|
async def _run_endpoint(
|
|
self, record: ToolRecord, arguments: dict[str, Any], user: UserT
|
|
) -> Any:
|
|
op = record.operation
|
|
if op is None:
|
|
raise ExecutionError("Operazione endpoint mancante", endpoint=record.name)
|
|
endpoint_fn = op.endpoint_fn
|
|
if endpoint_fn is None:
|
|
raise ExecutionError(
|
|
f"Impossibile trovare l'endpoint {op.method} {op.path}",
|
|
endpoint=op.path,
|
|
)
|
|
|
|
kwargs: dict[str, Any] = {}
|
|
sig_params = self._extract_signature_params(endpoint_fn)
|
|
for spec in record.param_specs:
|
|
if spec.name in arguments:
|
|
kwargs[spec.name] = arguments[spec.name]
|
|
elif spec.default is not None:
|
|
kwargs[spec.name] = spec.default
|
|
elif not spec.required:
|
|
kwargs[spec.name] = None
|
|
|
|
if "user" in sig_params and "user" not in kwargs:
|
|
kwargs["user"] = user
|
|
for extra in ("request", "db", "session"):
|
|
if extra in sig_params and extra not in kwargs and hasattr(self, "_" + extra):
|
|
kwargs[extra] = getattr(self, "_" + extra)
|
|
|
|
try:
|
|
async with asyncio.timeout(record.timeout):
|
|
if asyncio.iscoroutinefunction(endpoint_fn):
|
|
result = await endpoint_fn(**kwargs)
|
|
else:
|
|
result = await asyncio.to_thread(endpoint_fn, **kwargs)
|
|
return result
|
|
except asyncio.TimeoutError:
|
|
raise
|
|
except ExecutionError:
|
|
raise
|
|
except Exception as e:
|
|
status = getattr(e, "status_code", None)
|
|
raise ExecutionError(
|
|
str(e),
|
|
status_code=status,
|
|
endpoint=f"{op.method} {op.path}",
|
|
cause=e,
|
|
) from e
|
|
|
|
@staticmethod
|
|
def _extract_signature_params(fn: Callable[..., Any]) -> set[str]:
|
|
import inspect
|
|
|
|
try:
|
|
sig = inspect.signature(fn)
|
|
return set(sig.parameters.keys())
|
|
except (ValueError, TypeError):
|
|
return set()
|