跳转至

工具调用与执行

演示 Agent 中工具执行的工程化要点:参数校验、异常捕获、执行超时与结果标准化。 所有逻辑基于标准库,可在离线环境测试。 对应教程:第5卷《Agent 开发》第3章 工具执行。

文件路径: examples/volume-05-agent\tool_execution.py


源代码

"""工具执行引擎示例。

演示 Agent 中工具执行的工程化要点:参数校验、异常捕获、执行超时与结果标准化。
所有逻辑基于标准库,可在离线环境测试。

对应教程:第5卷《Agent 开发》第3章 工具执行。
"""

from __future__ import annotations

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


@dataclass
class ToolSchema:
    """工具参数 schema(简化版 JSON Schema)。

    Attributes:
        name: 工具名。
        parameters: 参数名到 ``(类型, 是否必填)`` 的映射。
    """

    name: str
    parameters: dict[str, tuple[type, bool]] = field(default_factory=dict)


@dataclass
class ExecutionResult:
    """工具执行结果。

    Attributes:
        ok: 是否成功。
        output: 输出内容(成功时)或错误信息(失败时)。
        error: 错误类型名,成功时为空。
    """

    ok: bool
    output: Any = None
    error: str = ""

    def to_dict(self) -> dict[str, Any]:
        """转为字典表示。"""
        return {"ok": self.ok, "output": self.output, "error": self.error}


def validate_arguments(args: dict[str, Any], schema: ToolSchema) -> list[str]:
    """根据 schema 校验参数。

    Args:
        args: 待校验参数。
        schema: 工具 schema。

    Returns:
        错误信息列表,为空表示通过。
    """
    errors: list[str] = []
    for param_name, (expected_type, required) in schema.parameters.items():
        if param_name not in args:
            if required:
                errors.append(f"缺少必填参数: {param_name}")
            continue
        value = args[param_name]
        # bool 是 int 的子类,需单独处理
        if expected_type is int and isinstance(value, bool):
            errors.append(f"参数 {param_name} 应为 int,实际为 bool")
        elif not isinstance(value, expected_type):
            errors.append(
                f"参数 {param_name} 应为 {expected_type.__name__},实际为 {type(value).__name__}"
            )
    return errors


class ToolExecutor:
    """工具执行器,封装校验、执行、超时与异常处理。"""

    def __init__(self, timeout: float = 5.0) -> None:
        """初始化执行器。

        Args:
            timeout: 单次执行超时(秒)。
        """
        self.timeout = timeout
        self._registry: dict[str, tuple[Callable[..., Any], ToolSchema]] = {}

    def register(self, func: Callable[..., Any], schema: ToolSchema) -> None:
        """注册工具函数及其 schema。

        Args:
            func: 工具函数。
            schema: 工具 schema。
        """
        self._registry[schema.name] = (func, schema)

    def execute(self, name: str, args: dict[str, Any]) -> ExecutionResult:
        """执行指定工具。

        Args:
            name: 工具名。
            args: 参数。

        Returns:
            ``ExecutionResult``。
        """
        if name not in self._registry:
            return ExecutionResult(ok=False, error="ToolNotFound", output=f"未注册工具: {name}")

        func, schema = self._registry[name]
        errors = validate_arguments(args, schema)
        if errors:
            return ExecutionResult(ok=False, error="ValidationError", output="; ".join(errors))

        try:
            value = self._invoke_with_timeout(func, args)
        except TimeoutError:
            return ExecutionResult(
                ok=False, error="TimeoutError", output=f"执行超过 {self.timeout}s"
            )
        except Exception as exc:  # noqa: BLE001
            return ExecutionResult(ok=False, error=type(exc).__name__, output=str(exc))
        return ExecutionResult(ok=True, output=value)

    def _invoke_with_timeout(self, func: Callable[..., Any], args: dict[str, Any]) -> Any:
        """在线程中调用函数并施加超时。

        Args:
            func: 待调用函数。
            args: 参数。

        Returns:
            函数返回值。

        Raises:
            TimeoutError: 超时。
        """
        if self.timeout <= 0:
            return func(**args)
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(func, **args)
            return future.result(timeout=self.timeout)


# ===== 示例工具 =====


def add(a: int, b: int) -> int:
    """示例工具:两整数相加。"""
    return a + b


def slow_task(seconds: float) -> str:
    """示例工具:模拟耗时任务,用于演示超时。"""
    import time

    time.sleep(seconds)
    return "done"


def build_default_executor(timeout: float = 5.0) -> ToolExecutor:
    """构造注册了示例工具的执行器。

    Args:
        timeout: 默认超时。

    Returns:
        ``ToolExecutor``。
    """
    executor = ToolExecutor(timeout=timeout)
    executor.register(add, ToolSchema(name="add", parameters={"a": (int, True), "b": (int, True)}))
    executor.register(
        slow_task,
        ToolSchema(name="slow_task", parameters={"seconds": (float, True)}),
    )
    return executor


def main() -> None:
    """演示工具执行引擎。"""
    executor = build_default_executor(timeout=1.0)

    print("正常执行 add:")
    print(" ", executor.execute("add", {"a": 2, "b": 3}).to_dict())

    print("\n参数校验失败:")
    print(" ", executor.execute("add", {"a": 2}).to_dict())
    print(" ", executor.execute("add", {"a": 2, "b": "x"}).to_dict())

    print("\n未知工具:")
    print(" ", executor.execute("nope", {}).to_dict())

    print("\n超时:")
    print(" ", executor.execute("slow_task", {"seconds": 2.0}).to_dict())


if __name__ == "__main__":
    main()