04:工具系统设计:从抽象基类到 JSON Schema 的完整实现

10 阅读3分钟

引言

工具调用(Function Calling)是现代 AI Agent 的核心能力。CountBot 实现了一套完整的工具系统,包含 12+ 内置工具,支持参数验证、审计日志和动态注册。本文将深入分析其设计与实现。

工具抽象基类

class Tool(ABC):
    _TYPE_MAP = {
        "string": str, "integer": int, "number": (int, float),
        "boolean": bool, "array": list, "object": dict,
    }

    @property
    @abstractmethod
    def name(self) -> str: ...

    @property
    @abstractmethod
    def description(self) -> str: ...

    @property
    @abstractmethod
    def parameters(self) -> dict[str, Any]: ...

    @abstractmethod
    async def execute(self, **kwargs: Any) -> str: ...

设计要点:

  • 使用 @property + @abstractmethod 确保子类必须提供元数据
  • parameters 返回标准 JSON Schema,与 OpenAI Function Calling 格式完全兼容
  • execute 统一返回 str,简化结果处理

参数验证系统

Tool 基类内置了递归的 JSON Schema 验证器:

def validate_params(self, params: dict[str, Any]) -> list[str]:
    schema = self.parameters or {}
    return self._validate(params, {**schema, "type": "object"}, "")

def _validate(self, val, schema, path) -> list[str]:
    t = schema.get("type")
    errors = []
    
    # 类型检查
    if t in self._TYPE_MAP and not isinstance(val, self._TYPE_MAP[t]):
        return [f"{label} should be {t}"]
    
    # 枚举检查
    if "enum" in schema and val not in schema["enum"]:
        errors.append(f"{label} must be one of {schema['enum']}")
    
    # 数值范围
    if t in ("integer", "number"):
        if "minimum" in schema and val < schema["minimum"]:
            errors.append(f"{label} must be >= {schema['minimum']}")
    
    # 递归验证嵌套对象和数组
    if t == "object":
        for k in schema.get("required", []):
            if k not in val:
                errors.append(f"missing required {k}")
    
    return errors

这个验证器在工具执行前自动运行,防止无效参数导致的运行时错误。

工具注册表

class ToolRegistry:
    def __init__(self):
        self._tools: dict[str, Tool] = {}
        self._audit_enabled: bool = True
        self._session_id: str | None = None

    def register(self, tool: Tool) -> None:
        if tool.name in self._tools:
            raise ValueError(f"Tool '{tool.name}' is already registered")
        self._tools[tool.name] = tool

    def get_definitions(self) -> list[dict]:
        """获取所有工具的 OpenAI 格式定义"""
        return [tool.get_definition() for tool in self._tools.values()]

注册表负责:

  • 工具的注册/注销/查询
  • 生成 OpenAI 格式的工具定义列表
  • 管理审计日志和会话上下文

统一注册入口

setup.py 提供了统一的工具注册函数:

def register_all_tools(workspace, command_timeout=30, ...) -> ToolRegistry:
    tools = ToolRegistry()
    
    # 1. 文件系统工具
    tools.register(ReadFileTool(workspace))
    tools.register(WriteFileTool(workspace))
    tools.register(EditFileTool(workspace))
    tools.register(ListDirTool(workspace))
    
    # 2. Shell 工具
    tools.register(ExecTool(workspace, timeout=command_timeout))
    
    # 3. Web 工具(条件注册)
    if brave_api_key:
        tools.register(WebSearchTool(api_key=brave_api_key))
    tools.register(WebFetchTool())
    
    # 4. 子代理工具(条件注册)
    if subagent_manager:
        tools.register(SpawnTool(subagent_manager))
    
    # 5. 记忆工具、截图工具、文件搜索工具...
    
    return tools

条件注册模式确保只有配置了必要依赖的工具才会被注册。

内置工具一览

工具名功能
read_fileReadFileTool读取文件内容
write_fileWriteFileTool写入文件
edit_fileEditFileTool编辑文件(按文本或行号)
list_dirListDirTool列出目录内容
execExecTool执行 Shell 命令
web_fetchWebFetchTool抓取网页内容
spawnSpawnTool生成子代理
screenshotScreenshotTool屏幕截图
file_searchFileSearchTool文件内容搜索
memory_write/search/readMemory*Tool记忆读写搜索
send_mediaSendMediaTool发送媒体到渠道

工作空间沙箱

文件系统工具通过 WorkspaceValidator 实现路径沙箱:

class WorkspaceValidator:
    def __init__(self, workspace: Path, restrict_to_workspace: bool = True):
        self._workspace = workspace.resolve()
        self._restrict = restrict_to_workspace

    def validate_path(self, path: str) -> Path:
        resolved = (self._workspace / path).resolve()
        if self._restrict and not str(resolved).startswith(str(self._workspace)):
            raise ValueError(f"路径 {path} 超出工作空间范围")
        return resolved

这确保 AI Agent 无法访问工作空间之外的文件,是安全设计的关键一环。

工具定义生成

每个工具通过 get_definition() 生成 OpenAI 兼容的函数定义:

def get_definition(self) -> dict[str, Any]:
    return {
        "type": "function",
        "function": {
            "name": self.name,
            "description": self.description,
            "parameters": self.parameters,
        },
    }

这些定义会被传递给 LLM,让 LLM 了解可用的工具及其参数格式。

总结

CountBot 的工具系统展示了如何构建一个类型安全、可扩展、安全可控的 AI 工具框架。通过 ABC 定义接口、JSON Schema 描述参数、注册表管理生命周期、沙箱保障安全,形成了一套完整的工具管理方案。