group_fqcd_jr/app/service/tool_executor.py

import asyncio
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, Protocol

from pydantic import BaseModel, ValidationError

from app.core.contracts import RequestContext, SourceReference, ToolCallRecord
from app.core.errors import ForbiddenAgentError, RecoverableAgentError, ValidationAgentError
from app.infrastructure.db import SessionFactory
from app.model.audit import InteractionAudit


class ToolHandler(Protocol):
    async def __call__(self, arguments: BaseModel, context: RequestContext) -> Any: ...


@dataclass(frozen=True)
class ToolDefinition:
    name: str
    input_model: type[BaseModel]
    handler: ToolHandler
    required_permission: str
    allowed_roles: tuple[str, ...]
    read_only: bool = True
    timeout_seconds: float = 5


@dataclass(frozen=True)
class ToolExecution:
    output: Any
    record: ToolCallRecord
    references: tuple[SourceReference, ...] = ()


class ToolRegistry:
    def __init__(self) -> None:
        self._tools: dict[str, ToolDefinition] = {}

    def register(self, definition: ToolDefinition) -> None:
        if definition.name in self._tools:
            raise ValidationAgentError("工具名称重复")
        if not definition.read_only:
            raise ValidationAgentError("Agent 公共工具仅允许只读")
        self._tools[definition.name] = definition

    def get(self, name: str) -> ToolDefinition:
        definition = self._tools.get(name)
        if definition is None:
            raise ForbiddenAgentError("工具未注册")
        return definition


class ToolExecutor:
    def __init__(self, registry: ToolRegistry) -> None:
        self.registry = registry

    async def execute(
        self, *, name: str, arguments: dict[str, Any], intent: str,
        configured_tools: dict[str, tuple[str, ...]], context: RequestContext,
    ) -> ToolExecution:
        definition = self.registry.get(name)
        allowed = configured_tools.get(intent, ())
        reason = None
        if name not in allowed:
            reason = "工具不在当前意图白名单"
        elif definition.required_permission not in context.permissions:
            reason = "缺少工具权限"
        elif not set(definition.allowed_roles).intersection(context.roles):
            reason = "角色不能使用工具"
        if reason:
            await self._audit(name, intent, context, "denied", reason)
            raise ForbiddenAgentError(reason)
        try:
            validated = definition.input_model.model_validate(arguments)
        except ValidationError as exc:
            await self._audit(name, intent, context, "denied", "参数校验失败")
            raise ValidationAgentError("工具参数校验失败") from exc
        try:
            async with asyncio.timeout(definition.timeout_seconds):
                output = await definition.handler(validated, context)
        except TimeoutError as exc:
            await self._audit(name, intent, context, "failed", "timeout")
            raise RecoverableAgentError("工具调用超时") from exc
        except Exception as exc:
            await self._audit(name, intent, context, "failed", type(exc).__name__)
            raise RecoverableAgentError("工具调用失败") from exc
        record = ToolCallRecord(
            tool_name=name, status="succeeded",
            input_summary={key: "[redacted]" for key in arguments},
            output_summary=self._output_summary(output),
        )
        await self._audit(name, intent, context, "succeeded", "ok")
        reference = SourceReference(source_type="tool", source_id=f"{context.trace_id}:{name}",
                                    title=name)
        return ToolExecution(output=output, record=record, references=(reference,))

    @staticmethod
    def _output_summary(output: Any) -> dict[str, Any]:
        summary: dict[str, Any] = {"result_type": type(output).__name__}
        if not isinstance(output, dict):
            return summary
        summary["status"] = output.get("status")
        audit = output.get("audit")
        if isinstance(audit, dict):
            summary["query_plan"] = audit.get("query_plan")
            summary["generated_sql"] = audit.get("generated_sql")
            summary["permission_check"] = audit.get("permission_check")
            summary["execution"] = audit.get("execution")
        return summary

    async def _audit(
        self, name: str, intent: str, context: RequestContext, status: str, reason: str
    ) -> None:
        async with SessionFactory() as session, session.begin():
            session.add(InteractionAudit(
                actor_type="agent", actor_id=int(context.user_id), portal=context.portal,
                action_type="agent.tool_executed", detail={
                    "tool_name": name, "intent": intent, "status": status,
                    "reason": reason, "trace_id": context.trace_id,
                }, created_at=datetime.now(UTC).replace(tzinfo=None),
            ))