AI Agent 的工具调用机制:从 Function Calling 到 Rust 实现的 Tool Registry
AI Agent 的工具调用机制:从 Function Calling 到 Rust 实现的 Tool Registry

一、工具调用是 AI Agent 的手和脚:没有工具的 Agent 只是聊天机器人
大模型的文本生成能力很强,但它无法执行任何实际操作——不能查数据库、不能调 API、不能读写文件。AI Agent 的核心价值在于"能做事",而"做事"的前提是"有工具可用"。Function Calling(函数调用)是连接大模型与外部工具的桥梁:大模型根据用户意图选择合适的工具和参数,应用程序执行工具调用并将结果返回给大模型,大模型再根据结果生成最终回答。
但 Function Calling 只是工具调用的"协议层",真正让 Agent 可靠运行的是工具注册与管理机制(Tool Registry)。一个生产级 Agent 可能需要注册几十个工具,每个工具有不同的参数 Schema、权限要求和执行策略。如果工具注册是硬编码的 if-else,新增一个工具就要改代码、重新编译、重新部署——这在快速迭代的 AI 应用中是不可接受的。
用 Rust 实现一个类型安全、可扩展的 Tool Registry,是构建可靠 AI Agent 的基础设施。Rust 的类型系统可以在编译期保证工具参数的类型安全,trait 机制可以实现工具的多态注册,serde 的序列化能力可以自动处理 JSON Schema 与 Rust 结构体的转换。
二、工具调用机制的底层原理
2.1 Function Calling 的交互流程
Function Calling 的完整流程包含四个步骤:
sequenceDiagram
participant U as 用户
participant A as Agent 运行时
participant L as 大模型 API
participant T as 工具执行器
U->>A: "帮我查北京明天的天气"
A->>L: 发送消息 + 工具定义列表
L-->>A: 返回 tool_call: get_weather(city="北京", date="明天")
A->>T: 执行 get_weather
T-->>A: 返回 {"temp": 28, "condition": "晴"}
A->>L: 发送工具执行结果
L-->>A: 生成最终回答:"北京明天天气晴,气温约28°C"
A->>U: 返回回答
关键细节:大模型不直接执行工具,它只输出"应该调用哪个工具、传什么参数"的 JSON。应用程序负责解析这个 JSON、执行工具调用、将结果返回给大模型。这种设计的优势是安全性——大模型无法直接执行任何有副作用的操作。
2.2 工具定义的 JSON Schema
每个工具需要定义名称、描述和参数 Schema。大模型根据工具描述判断是否需要调用该工具,根据参数 Schema 生成符合格式的参数。例如:
{
"name": "get_weather",
"description": "查询指定城市的天气预报",
"parameters": {
"type": "object",
"properties": {
"city": {"type": "string", "description": "城市名称"},
"date": {"type": "string", "description": "日期,格式 YYYY-MM-DD"}
},
"required": ["city"]
}
}
2.3 Tool Registry 的设计原则
Tool Registry 需要满足三个设计原则:
- 类型安全:工具参数在编译期检查类型,避免运行时类型错误。
- 可扩展:新增工具只需实现 trait 并注册,不需要修改 Registry 的核心代码。
- 自描述:工具的 JSON Schema 自动从 Rust 结构体定义生成,无需手写。
三、Rust 生产级代码实现
3.1 工具 trait 定义与参数自动派生
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
/// 工具参数 trait:自动生成 JSON Schema
pub trait ToolParams: Serialize + for<'de> Deserialize<'de> {
/// 从 JSON Schema 自动生成工具的参数定义
fn json_schema() -> Value;
}
/// 工具 trait:所有工具必须实现此接口
#[async_trait::async_trait]
pub trait Tool: Send + Sync {
/// 工具名称(全局唯一标识)
fn name(&self) -> &str;
/// 工具描述(大模型根据此描述决定是否调用)
fn description(&self) -> &str;
/// 工具参数的 JSON Schema
fn parameters_schema(&self) -> Value;
/// 执行工具调用
async fn execute(&self, params: Value) -> Result<Value, ToolError>;
}
/// 工具执行错误
#[derive(Debug, thiserror::Error)]
pub enum ToolError {
#[error("参数解析失败: {0}")]
InvalidParams(String),
#[error("工具执行失败: {0}")]
ExecutionFailed(String),
#[error("权限不足: {0}")]
PermissionDenied(String),
#[error("超时: {0}")]
Timeout(String),
}
3.2 具体工具实现示例
use schemars::JsonSchema;
/// 天气查询工具的参数
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
pub struct WeatherParams {
/// 城市名称
pub city: String,
/// 日期,格式 YYYY-MM-DD(可选,默认今天)
pub date: Option<String>,
}
impl ToolParams for WeatherParams {
fn json_schema() -> Value {
// schemars 自动从结构体定义生成 JSON Schema
let schema = schemars::schema_for!(WeatherParams);
serde_json::to_value(&schema).unwrap_or_default()
}
}
/// 天气查询工具
pub struct WeatherTool {
api_key: String,
base_url: String,
}
impl WeatherTool {
pub fn new(api_key: String) -> Self {
Self {
api_key,
base_url: "https://api.weather.example.com".to_string(),
}
}
}
#[async_trait::async_trait]
impl Tool for WeatherTool {
fn name(&self) -> &str {
"get_weather"
}
fn description(&self) -> &str {
"查询指定城市的天气预报,返回温度、天气状况和湿度信息"
}
fn parameters_schema(&self) -> Value {
WeatherParams::json_schema()
}
async fn execute(&self, params: Value) -> Result<Value, ToolError> {
// 解析参数(类型安全:serde 自动校验字段类型)
let p: WeatherParams = serde_json::from_value(params)
.map_err(|e| ToolError::InvalidParams(e.to_string()))?;
// 调用天气 API
let url = format!(
"{}/forecast?city={}&date={}&key={}",
self.base_url,
urlencoding::encode(&p.city),
p.date.as_deref().unwrap_or("today"),
self.api_key,
);
let response = reqwest::get(&url)
.await
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
if !response.status().is_success() {
return Err(ToolError::ExecutionFailed(
format!("API 返回状态码: {}", response.status()),
));
}
let result: Value = response
.json()
.await
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
Ok(result)
}
}
3.3 Tool Registry 实现
use std::sync::Arc;
/// 工具注册表:管理所有可用工具
pub struct ToolRegistry {
tools: HashMap<String, Arc<dyn Tool>>,
}
impl ToolRegistry {
pub fn new() -> Self {
Self {
tools: HashMap::new(),
}
}
/// 注册工具
pub fn register(&mut self, tool: impl Tool + 'static) {
let name = tool.name().to_string();
self.tools.insert(name, Arc::new(tool));
}
/// 执行工具调用
pub async fn execute(
&self,
tool_name: &str,
params: Value,
) -> Result<Value, ToolError> {
let tool = self.tools.get(tool_name)
.ok_or_else(|| ToolError::InvalidParams(
format!("未知工具: {}", tool_name),
))?;
tool.execute(params).await
}
/// 生成所有工具的定义列表(发送给大模型)
pub fn tool_definitions(&self) -> Vec<Value> {
self.tools.values().map(|tool| {
serde_json::json!({
"type": "function",
"function": {
"name": tool.name(),
"description": tool.description(),
"parameters": tool.parameters_schema(),
}
})
}).collect()
}
/// 获取工具名称列表
pub fn tool_names(&self) -> Vec<&str> {
self.tools.keys().map(|s| s.as_str()).collect()
}
}
/// 构建器:简化 ToolRegistry 的初始化
pub struct ToolRegistryBuilder {
registry: ToolRegistry,
}
impl ToolRegistryBuilder {
pub fn new() -> Self {
Self {
registry: ToolRegistry::new(),
}
}
pub fn with_tool(mut self, tool: impl Tool + 'static) -> Self {
self.registry.register(tool);
self
}
pub fn build(self) -> ToolRegistry {
self.registry
}
}
3.4 Agent 运行时:串联大模型与工具
use tokio::time::{timeout, Duration};
/// Agent 运行时配置
pub struct AgentConfig {
pub model: String,
pub max_tool_calls: usize,
pub tool_timeout_secs: u64,
}
/// Agent 运行时
pub struct AgentRuntime {
registry: ToolRegistry,
config: AgentConfig,
client: reqwest::Client,
}
impl AgentRuntime {
pub fn new(registry: ToolRegistry, config: AgentConfig) -> Self {
Self {
registry,
config,
client: reqwest::Client::new(),
}
}
/// 运行 Agent:循环调用大模型 → 执行工具 → 返回结果
pub async fn run(&self, user_message: &str) -> Result<String, Box<dyn std::error::Error>> {
let mut messages = vec![
serde_json::json!({
"role": "user",
"content": user_message,
}),
];
for _ in 0..self.config.max_tool_calls {
// 调用大模型
let response = self.call_llm(&messages).await?;
// 检查是否有工具调用
let tool_calls = response.get("tool_calls")
.and_then(|v| v.as_array());
match tool_calls {
Some(calls) if !calls.is_empty() => {
// 执行每个工具调用
for call in calls {
let tool_name = call["function"]["name"].as_str().unwrap_or("");
let params_str = call["function"]["arguments"].as_str().unwrap_or("{}");
let params: Value = serde_json::from_str(params_str)
.unwrap_or_default();
// 带超时的工具执行
let result = timeout(
Duration::from_secs(self.config.tool_timeout_secs),
self.registry.execute(tool_name, params),
)
.await
.unwrap_or_else(|_| {
Err(ToolError::Timeout(
format!("工具 {} 执行超时", tool_name),
))
});
// 将工具结果添加到消息列表
let tool_result = match result {
Ok(v) => v.to_string(),
Err(e) => format!("错误: {}", e),
};
messages.push(serde_json::json!({
"role": "tool",
"tool_call_id": call["id"],
"content": tool_result,
}));
}
}
_ => {
// 大模型直接返回文本回答
let content = response["content"].as_str().unwrap_or("");
return Ok(content.to_string());
}
}
}
Ok("达到最大工具调用次数限制".to_string())
}
async fn call_llm(&self, messages: &[Value]) -> Result<Value, Box<dyn std::error::Error>> {
// 简化实现:调用 OpenAI 兼容 API
let body = serde_json::json!({
"model": self.config.model,
"messages": messages,
"tools": self.registry.tool_definitions(),
});
let resp = self.client
.post("https://api.openai.com/v1/chat/completions")
.json(&body)
.send()
.await?;
let result: Value = resp.json().await?;
// 提取 assistant 的回复
Ok(result["choices"][0]["message"].clone())
}
}
四、Trade-offs:Rust 实现 Tool Registry 的代价
4.1 开发效率与类型安全的权衡
Rust 的类型系统保证了编译期安全,但增加了开发复杂度。每个工具需要定义参数结构体、实现 Tool trait、派生 Serialize/Deserialize/JsonSchema。相比之下,Python 的实现可以用装饰器一行搞定工具注册。但 Rust 的优势在于:参数类型错误在编译期就能发现,而非运行时崩溃。
4.2 动态注册的限制
当前的 ToolRegistry 在编译期确定所有工具类型(通过 impl Tool + 'static)。如果需要运行时动态加载工具(如从插件目录加载 .wasm 文件),需要引入 trait object 和动态分发,这会增加运行时开销和代码复杂度。
4.3 适用边界
Rust Tool Registry 适用于以下场景:Agent 需要高并发处理大量工具调用、对类型安全有严格要求、工具参数结构复杂需要编译期校验。不适用于:快速原型验证(Python 更快)、工具数量少且简单(过度工程)、需要运行时动态加载插件的场景。
五、总结
用 Rust 实现 Tool Registry,将 AI Agent 的工具调用从"硬编码 if-else"升级为"类型安全、可扩展的注册机制"。核心落地步骤如下:
- 定义 Tool trait:统一工具的接口规范,包含名称、描述、参数 Schema 和执行方法。
- 自动生成 JSON Schema:使用 schemars 从 Rust 结构体自动生成工具参数定义,避免手写 Schema 的错误。
- 实现 ToolRegistry:基于 HashMap 的工具注册表,支持动态查询和批量生成工具定义。
- 串联 Agent 运行时:循环调用大模型 → 解析工具调用 → 执行工具 → 返回结果,带超时和错误处理。
- 渐进式扩展:新增工具只需实现 Tool trait 并注册,不需要修改核心代码。
工具调用是 Agent 的基础设施,Registry 是基础设施的基础设施。把这部分做扎实,后续的工具扩展才能稳定可靠。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐



所有评论(0)