引言
工具调用(Function Calling)是现代 AI Agent 的核心能力。CountBot 实现了一套完整的工具系统,包含 12+ 内置工具,支持参数验证、审计日志和动态注册。本文将深入分析其设计与实现。
工具抽象基类
python
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 验证器:
python
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
这个验证器在工具执行前自动运行,防止无效参数导致的运行时错误。
工具注册表
python
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 提供了统一的工具注册函数:
python
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_file | ReadFileTool | 读取文件内容 |
| write_file | WriteFileTool | 写入文件 |
| edit_file | EditFileTool | 编辑文件(按文本或行号) |
| list_dir | ListDirTool | 列出目录内容 |
| exec | ExecTool | 执行 Shell 命令 |
| web_fetch | WebFetchTool | 抓取网页内容 |
| spawn | SpawnTool | 生成子代理 |
| screenshot | ScreenshotTool | 屏幕截图 |
| file_search | FileSearchTool | 文件内容搜索 |
| memory_write/search/read | Memory*Tool | 记忆读写搜索 |
| send_media | SendMediaTool | 发送媒体到渠道 |
工作空间沙箱
文件系统工具通过 WorkspaceValidator 实现路径沙箱:
python
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 兼容的函数定义:
python
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 描述参数、注册表管理生命周期、沙箱保障安全,形成了一套完整的工具管理方案。