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

cover

一、工具调用是 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"升级为"类型安全、可扩展的注册机制"。核心落地步骤如下:

  1. 定义 Tool trait:统一工具的接口规范,包含名称、描述、参数 Schema 和执行方法。
  2. 自动生成 JSON Schema:使用 schemars 从 Rust 结构体自动生成工具参数定义,避免手写 Schema 的错误。
  3. 实现 ToolRegistry:基于 HashMap 的工具注册表,支持动态查询和批量生成工具定义。
  4. 串联 Agent 运行时:循环调用大模型 → 解析工具调用 → 执行工具 → 返回结果,带超时和错误处理。
  5. 渐进式扩展:新增工具只需实现 Tool trait 并注册,不需要修改核心代码。

工具调用是 Agent 的基础设施,Registry 是基础设施的基础设施。把这部分做扎实,后续的工具扩展才能稳定可靠。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐