216 lines
7.7 KiB
Python
216 lines
7.7 KiB
Python
"""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,
|
|
)
|
|
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,
|
|
):
|
|
self.store = store
|
|
self._mcp_client = mcp_client or McpToolClient()
|
|
|
|
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 {}
|
|
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
|
|
|
|
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)
|
|
normalized_arguments = dict(arguments or {})
|
|
if tool.type == "http":
|
|
result = await self._execute_http(tool, normalized_arguments)
|
|
elif tool.type == "mcp":
|
|
result = await self._execute_mcp(tool, normalized_arguments)
|
|
else:
|
|
raise ToolExecutionError(f"不支持工具类型: {tool.type}")
|
|
return self._apply_result_assignments(
|
|
tool,
|
|
result,
|
|
result_assignments=result_assignments,
|
|
)
|
|
|
|
async def _execute_http(
|
|
self,
|
|
tool: RuntimeTool,
|
|
arguments: dict[str, Any],
|
|
) -> 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]}
|
|
|
|
return {
|
|
"status": "ok",
|
|
"status_code": response.status_code,
|
|
"data": payload,
|
|
}
|
|
|
|
async def _execute_mcp(
|
|
self,
|
|
tool: RuntimeTool,
|
|
arguments: 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 工具缺少远端工具名称")
|
|
|
|
resolved_server = server.model_copy(
|
|
update={
|
|
"url": self.store.render(server.url),
|
|
"headers": self.store.render_data(server.headers),
|
|
"secret_headers": self.store.render_data(
|
|
server.secret_headers,
|
|
allow_secrets=True,
|
|
),
|
|
}
|
|
)
|
|
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 _apply_result_assignments(
|
|
self,
|
|
tool: RuntimeTool,
|
|
result: dict[str, Any],
|
|
*,
|
|
result_assignments: dict[str, str] | None,
|
|
) -> dict[str, Any]:
|
|
config = (tool.definition or {}).get("config") or {}
|
|
assignments = (
|
|
result_assignments
|
|
if result_assignments is not None
|
|
else config.get("dynamic_variable_assignments") or {}
|
|
)
|
|
payload = result.get("data")
|
|
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 {**result, "updated_variables": updated}
|