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 (
    DependencyUnavailableError,
    ForbiddenAgentError,
    UpstreamTimeoutError,
    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 UpstreamTimeoutError("工具调用超时") from exc
        except Exception as exc:
            await self._audit(name, intent, context, "failed", type(exc).__name__)
            raise DependencyUnavailableError("工具调用失败") from exc
        record = ToolCallRecord(
            tool_name=name, status="succeeded",
            input_summary={key: "[redacted]" for key in arguments},
            output_summary={"result_type": type(output).__name__},
        )
        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,))

    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),
            ))