ai-video-fullstack/backend/services/tool_executor.py

"""Reusable deterministic tool execution shared by Prompt, Agent, and Action."""

from __future__ import annotations

from copy import deepcopy
from typing import Any
from urllib.parse import quote

import httpx
from models import RuntimeTool, RuntimeToolBinding

from services.runtime_variables import (
    DynamicVariableError,
    DynamicVariableStore,
    value_at_path,
)
from services.client_tools import ClientToolError, ClientToolPort
from services.tool_policy import policy_for_tool
from services.tools import McpClientError, McpToolClient


class ToolExecutionError(RuntimeError):
    pass


class ToolExecutor:
    """Execute supported tools and apply their results to session variables."""

    def __init__(
        self,
        store: DynamicVariableStore,
        *,
        mcp_client: McpToolClient | None = None,
        client_tools: ClientToolPort | None = None,
    ):
        self.store = store
        self._mcp_client = mcp_client or McpToolClient()
        self._client_tools = client_tools

    def set_client_tools(self, client_tools: ClientToolPort | None) -> None:
        self._client_tools = client_tools

    @staticmethod
    def input_schema(tool: RuntimeTool) -> tuple[dict[str, Any], list[str]]:
        config = (tool.definition or {}).get("config") or {}
        if tool.type == "mcp":
            schema = config.get("input_schema") or {}
            return (
                dict(schema.get("properties") or {}),
                [str(name) for name in schema.get("required") or []],
            )
        parameters = list(config.get("parameters") or [])
        properties = {
            str(parameter.get("name")): {
                "type": str(parameter.get("type") or "string"),
                "description": str(parameter.get("description") or ""),
            }
            for parameter in parameters
            if parameter.get("name")
        }
        required = [
            str(parameter["name"])
            for parameter in parameters
            if parameter.get("name") and parameter.get("required", True)
        ]
        return properties, required

    @classmethod
    def schema_parts(
        cls,
        tool: RuntimeTool,
        binding: RuntimeToolBinding | None = None,
    ) -> tuple[dict[str, Any], list[str]]:
        properties, required = cls.input_schema(tool)
        if binding is None:
            return properties, required
        visible = {
            name
            for name, source in binding.argument_bindings.items()
            if str((source or {}).get("source") or "") == "model"
        }
        missing_required = [
            name for name in required if name not in binding.argument_bindings
        ]
        if missing_required:
            raise ToolExecutionError(
                "缺少必填工具参数绑定: " + ", ".join(missing_required)
            )
        return (
            {name: spec for name, spec in properties.items() if name in visible},
            [name for name in required if name in visible],
        )

    async def execute(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any] | None = None,
        *,
        binding: RuntimeToolBinding | None = None,
        result_assignments: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        normalized_arguments = self.resolve_arguments(
            tool,
            binding,
            dict(arguments or {}),
        )
        template_values = self.resolve_template_bindings(binding)
        if tool.type == "http":
            result = await self._execute_http(
                tool,
                normalized_arguments,
                template_values=template_values,
            )
        elif tool.type == "mcp":
            result = await self._execute_mcp(
                tool,
                normalized_arguments,
                template_values=template_values,
            )
        elif tool.type == "client":
            result = await self._execute_client(tool, normalized_arguments)
        else:
            raise ToolExecutionError(f"不支持工具类型: {tool.type}")
        if result.get("status") != "ok":
            return {**result, "updated_variables": []}
        return self._apply_result_assignments(
            tool,
            result,
            result_assignments=(
                result_assignments
                if result_assignments is not None
                else (binding.result_assignments if binding is not None else {})
            ),
        )

    def resolve_arguments(
        self,
        tool: RuntimeTool,
        binding: RuntimeToolBinding | None,
        model_arguments: dict[str, Any],
    ) -> dict[str, Any]:
        if binding is None:
            return model_arguments
        properties, required = self.input_schema(tool)
        visible_model_parameters = {
            name
            for name, source in binding.argument_bindings.items()
            if str((source or {}).get("source") or "") == "model"
        }
        resolved: dict[str, Any] = {}
        for name, source in binding.argument_bindings.items():
            if name not in properties:
                raise ToolExecutionError(f"工具参数不存在: {name}")
            resolved[name] = self._resolve_argument_source(
                name,
                source,
                model_arguments,
            )
        unknown_model_arguments = sorted(
            set(model_arguments) - visible_model_parameters
        )
        if unknown_model_arguments:
            raise ToolExecutionError(
                "模型提供了未授权工具参数: "
                + ", ".join(unknown_model_arguments)
            )
        missing = [name for name in required if name not in resolved]
        if missing:
            raise ToolExecutionError("缺少必填工具参数: " + ", ".join(missing))
        return resolved

    def resolve_template_bindings(
        self,
        binding: RuntimeToolBinding | None,
    ) -> dict[str, Any]:
        if binding is None:
            return {}
        return {
            name: self._resolve_source(source, {})
            for name, source in binding.template_bindings.items()
        }

    def _resolve_source(
        self,
        source: dict[str, Any],
        model_arguments: dict[str, Any],
    ) -> Any:
        kind = str((source or {}).get("source") or "")
        if kind == "model":
            raise ToolExecutionError("模型参数来源缺少参数上下文")
        if kind == "variable":
            name = str(
                source.get("variable_name") or source.get("variableName") or ""
            )
            return self.store.value(name)
        if kind == "literal":
            return deepcopy(source.get("value"))
        if kind == "template":
            return self.store.render_data(deepcopy(source.get("value")))
        raise ToolExecutionError(f"不支持的参数来源: {kind or 'empty'}")

    def _resolve_argument_source(
        self,
        name: str,
        source: dict[str, Any],
        model_arguments: dict[str, Any],
    ) -> Any:
        if str((source or {}).get("source") or "") == "model":
            if name not in model_arguments:
                raise ToolExecutionError(f"模型未提供工具参数: {name}")
            return model_arguments[name]
        return self._resolve_source(source, model_arguments)

    async def _execute_http(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any],
        *,
        template_values: dict[str, Any],
    ) -> dict[str, Any]:
        config = (tool.definition or {}).get("config") or {}
        parameters = list(config.get("parameters") or [])
        renderer = self._tool_renderer(tool, template_values)
        url = renderer.render(str(config.get("url") or ""))
        configured_headers = renderer.render_data(
            deepcopy(config.get("headers") or {}), allow_secrets=True
        )
        secret_headers = renderer.render_data(
            deepcopy((tool.secrets or {}).get("headers") or {}), allow_secrets=True
        )
        headers: dict[str, str] = {}
        query: dict[str, object] = {}
        body = renderer.render_data(deepcopy(config.get("body") or {}))

        for parameter in parameters:
            name = str(parameter.get("name") or "")
            if not name or name not in arguments:
                continue
            value = arguments[name]
            location = str(parameter.get("location") or "body")
            if location == "path":
                url = url.replace(f"{{{name}}}", quote(str(value), safe=""))
            elif location == "query":
                query[name] = value
            elif location == "header":
                headers[name] = str(value)
            else:
                body[name] = value

        headers.update({str(key): str(value) for key, value in configured_headers.items()})
        headers.update({str(key): str(value) for key, value in secret_headers.items()})
        try:
            async with httpx.AsyncClient(
                timeout=float(config.get("timeout_seconds") or 15),
                follow_redirects=False,
            ) as client:
                response = await client.request(
                    str(config.get("method") or "GET"),
                    url,
                    headers=headers,
                    params=query,
                    json=body if body else None,
                )
            response.raise_for_status()
        except httpx.TimeoutException as exc:
            raise ToolExecutionError("HTTP 工具调用超时") from exc
        except httpx.HTTPStatusError as exc:
            raise ToolExecutionError(f"HTTP 工具返回错误状态:{exc.response.status_code}") from exc
        except httpx.RequestError as exc:
            raise ToolExecutionError(f"HTTP 工具调用失败:{exc}") from exc

        if len(response.content) > 1_000_000:
            raise ToolExecutionError("HTTP 工具响应超过 1 MB 限制")
        try:
            payload: Any = response.json()
        except ValueError:
            payload = {"text": response.text[:8000]}

        return {
            "status": "ok",
            "status_code": response.status_code,
            "data": payload,
        }

    async def _execute_mcp(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any],
        *,
        template_values: dict[str, Any],
    ) -> dict[str, Any]:
        server = tool.mcp_server
        if server is None:
            raise ToolExecutionError("MCP 工具缺少可用的 Server 连接")
        config = (tool.definition or {}).get("config") or {}
        remote_tool_name = str(config.get("remote_tool_name") or "")
        if not remote_tool_name:
            raise ToolExecutionError("MCP 工具缺少远端工具名称")

        renderer = self._tool_renderer(tool, template_values)
        resolved_server = server.model_copy(
            update={
                "url": renderer.render(server.url),
                "headers": renderer.render_data(server.headers),
                "secret_headers": dict(server.secret_headers),
            }
        )
        try:
            payload = await self._mcp_client.call_tool(
                resolved_server,
                remote_tool_name,
                arguments,
            )
        except McpClientError as exc:
            raise ToolExecutionError(str(exc)) from exc
        return {"status": "ok", "data": payload}

    def _tool_renderer(
        self,
        tool: RuntimeTool,
        template_values: dict[str, Any],
    ) -> DynamicVariableStore:
        values = dict(self.store.values)
        values.update(
            {f"input__{name}": value for name, value in template_values.items()}
        )
        secrets = {
            str(name): str(value)
            for name, value in (
                (tool.secrets or {}).get("dynamic_variables") or {}
            ).items()
        }
        invalid = [name for name in secrets if not name.startswith("secret__")]
        if invalid:
            raise DynamicVariableError(
                "工具密钥变量必须以 secret__ 开头: " + ", ".join(invalid)
            )
        return DynamicVariableStore(
            values,
            secrets,
            optional_names=self.store.optional_names,
            variable_types=self.store.variable_types,
        )

    async def _execute_client(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any],
    ) -> dict[str, Any]:
        if self._client_tools is None:
            raise ToolExecutionError("当前运行模式不支持客户端工具")
        supports = getattr(self._client_tools, "supports", None)
        if supports is not None and not supports(tool.function_name):
            raise ToolExecutionError(f"客户端未注册工具: {tool.function_name}")
        config = (tool.definition or {}).get("config") or {}
        policy = policy_for_tool(tool)
        try:
            return await self._client_tools.call(
                tool.function_name,
                arguments,
                timeout_seconds=float(config.get("timeout_seconds") or 3),
                wait_for_response=policy.wait_for_response,
                response_wait_mode=policy.response_wait_mode,
            )
        except ClientToolError as exc:
            raise ToolExecutionError(str(exc)) from exc

    def _apply_result_assignments(
        self,
        tool: RuntimeTool,
        result: dict[str, Any],
        *,
        result_assignments: dict[str, str] | None,
    ) -> dict[str, Any]:
        assignments = result_assignments or {}
        payload = result.get("data")
        staged: dict[str, Any] = {}
        for variable_name, path in assignments.items():
            try:
                value = value_at_path(payload, str(path))
            except KeyError:
                try:
                    value = value_at_path({"response": payload}, str(path))
                except KeyError:
                    continue
            staged[str(variable_name)] = value
        updated = self.store.assign_declared_many(staged) if staged else []
        return {**result, "updated_variables": updated}