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

from services.runtime_variables import (
    DynamicVariableError,
    DynamicVariableStore,
    value_at_path,
)


class ToolExecutionError(RuntimeError):
    pass


class ToolExecutor:
    def __init__(self, store: DynamicVariableStore):
        self.store = store

    def register_secrets(self, tool: RuntimeTool) -> None:
        dynamic = (tool.secrets or {}).get("dynamic_variables") or {}
        for name, value in dynamic.items():
            if not str(name).startswith("secret__"):
                raise DynamicVariableError(f"工具密钥变量必须以 secret__ 开头: {name}")
            self.store.secrets[str(name)] = str(value)

    @staticmethod
    def schema_parts(tool: RuntimeTool) -> tuple[dict[str, Any], list[str]]:
        config = (tool.definition or {}).get("config") 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

    async def execute(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any] | None = None,
        *,
        result_assignments: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        self.register_secrets(tool)
        if tool.type != "http":
            raise ToolExecutionError(f"Action 暂不支持工具类型: {tool.type}")
        return await self._execute_http(
            tool,
            dict(arguments or {}),
            result_assignments=result_assignments,
        )

    async def _execute_http(
        self,
        tool: RuntimeTool,
        arguments: dict[str, Any],
        *,
        result_assignments: dict[str, str] | None,
    ) -> dict[str, Any]:
        config = (tool.definition or {}).get("config") or {}
        parameters = list(config.get("parameters") or [])
        url = self.store.render(str(config.get("url") or ""))
        configured_headers = self.store.render_data(
            deepcopy(config.get("headers") or {}), allow_secrets=True
        )
        secret_headers = self.store.render_data(
            deepcopy((tool.secrets or {}).get("headers") or {}), allow_secrets=True
        )
        headers: dict[str, str] = {}
        query: dict[str, object] = {}
        body = self.store.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]}

        assignments = (
            result_assignments
            if result_assignments is not None
            else config.get("dynamic_variable_assignments") or {}
        )
        updated: list[str] = []
        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
            self.store.assign(str(variable_name), value)
            updated.append(str(variable_name))
        return {
            "status": "ok",
            "status_code": response.status_code,
            "data": payload,
            "updated_variables": updated,
        }