harness_zbt/backend/app/harness/tooling.py

"""Tool 注册、上下文注入与执行审计。"""

from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

from langchain_core.tools import BaseTool, StructuredTool
from pydantic import BaseModel

from app.core.errors import AppError
from app.domains.identity.models import AdminUser, H5User
from app.harness.policy import HarnessPolicyEngine, ToolPolicy
from app.harness.schemas import HarnessToolResult

ToolHandler = Callable[["ToolExecutionContext", BaseModel], HarnessToolResult]


@dataclass(frozen=True)
class ToolDefinition:
    name: str
    description: str
    arguments: type[BaseModel]
    policy: ToolPolicy
    handler: ToolHandler


@dataclass(frozen=True)
class ToolExecution:
    name: str
    result: HarnessToolResult


@dataclass
class ToolExecutionContext:
    persona: str
    h5_user: H5User | None = None
    admin_user: AdminUser | None = None
    executions: list[ToolExecution] = field(default_factory=list)

    @property
    def principal_id(self) -> str:
        principal = self.h5_user or self.admin_user
        if principal is None:
            raise AppError("AGENT_PRINCIPAL_REQUIRED", "Agent缺少登录身份", 401)
        return principal.id

    @property
    def permissions(self) -> tuple[str, ...]:
        return self.admin_user.permissions if self.admin_user else ()


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

    def register(self, definition: ToolDefinition) -> None:
        if definition.name in self._items:
            raise AppError(
                "HARNESS_TOOL_DUPLICATED",
                f"Tool标识重复：{definition.name}",
                500,
            )
        self._items[definition.name] = definition

    def names(self) -> tuple[str, ...]:
        return tuple(sorted(self._items))

    def build(
        self,
        *,
        names: tuple[str, ...],
        context: ToolExecutionContext,
        policy_engine: HarnessPolicyEngine,
    ) -> list[BaseTool]:
        tools: list[BaseTool] = []
        for name in names:
            definition = self._items.get(name)
            if definition is None:
                raise AppError(
                    "HARNESS_TOOL_NOT_FOUND",
                    f"未找到Tool：{name}",
                    500,
                )
            if not policy_engine.can_use(
                persona=context.persona,
                allowed_by_persona=names,
                tool_name=name,
                policy=definition.policy,
                permissions=context.permissions,
            ):
                continue
            tools.append(
                self._to_langchain_tool(
                    definition,
                    names,
                    context,
                    policy_engine,
                )
            )
        return tools

    @staticmethod
    def _to_langchain_tool(
        definition: ToolDefinition,
        allowed_names: tuple[str, ...],
        context: ToolExecutionContext,
        policy_engine: HarnessPolicyEngine,
    ) -> BaseTool:
        def invoke(**kwargs: Any) -> str:
            policy_engine.authorize(
                persona=context.persona,
                allowed_by_persona=allowed_names,
                tool_name=definition.name,
                policy=definition.policy,
                permissions=context.permissions,
            )
            arguments = definition.arguments.model_validate(kwargs)
            result = definition.handler(context, arguments)
            context.executions.append(ToolExecution(name=definition.name, result=result))
            return result.model_dump_json(exclude_none=True)

        return StructuredTool.from_function(
            func=invoke,
            name=definition.name,
            description=definition.description,
            args_schema=definition.arguments,
        )