[Rust] 基于Rust的AI Agent SDK设计与实现：从多模型适配到上下文管理

在大模型能力飞速迭代的今天，Agent 已成为 AI 应用的新范式。但 Agent 真正落地生产，需要的不只是模型，更需要一套扎实的基础设施。

当我们动手开发一个真正的 AI Agent 应用时，会立刻撞上几堵现实的墙：

• 多模型适配：OpenAI、Anthropic、Google 各家 API 风格迥异，参数结构都不一样。
• 工具调用：Function Calling 的实现细节繁琐，并发、依赖、生命周期管理全是坑。
• 上下文管理：长对话和复杂任务会让 token 消耗像无底洞，成本难以控制。
• 成本追踪：不同模型定价策略天差地别，没有精确统计，预算就是一笔糊涂账。

这些问题让我们清醒地认识到一个事实：Agent 时代，稳固的基础设施不是锦上添花，而是不可或缺的刚需。

为此，我们开源了 agent-io —— 一个用 Rust 编写的 AI Agent SDK，希望能为开发者铺下一块坚实的地基。

为什么需要一个专用的 Agent SDK？

痛点一：多模型适配的复杂性

每个大模型提供商都有自己独特的 API 风格，看看下面这些代码你就明白了：

# OpenAI 风格
response = client.chat.completions.create(
    model="gpt-5.2",
    messages=[{"role": "user", "content": "..."}],
    tools=[...]
)

# Anthropic 风格  
response = client.messages.create(
    model="claude-opus-4-6",
    messages=[{"role": "user", "content": "..."}],
    tools=[...]
)

# Google 风格
response = client.generate_content(
    contents=[{"role": "user", "parts": [...]}],
    tools=[...]
)

参数结构、响应格式、流式处理方式全都不一样。如果你的业务需要灵活切换或同时使用多个模型，光是适配工作就足以让人头疼。

痛点二：工具调用的繁琐细节

Tool Calling 听起来简单，真做起来处处是细节：

• 如何定义严谨且易用的工具 JSON Schema？
• 多个工具调用并发时怎么处理？
• 工具返回的结果如何管理其生命周期？
• 工具之间如何实现依赖注入？

痛点三：上下文爆炸与成本黑盒

Agent 执行复杂任务时会进行多轮迭代，每轮对话都在累积消息。如果不加以控制，token 消耗会指数级增长，最终可能直接超出模型的上下文窗口。

与此同时，不同模型的定价差异巨大，每次调用的 token 数量也难以精确预估。没有清晰的成本追踪，优化和预算控制就无从谈起。

agent-io 的核心设计理念

面对这些挑战，agent-io 的设计围绕几个核心原则展开：

1. 抽象统一，实现灵活

我们通过定义统一的 Trait 来抽象所有 LLM 的能力，让业务代码与具体模型解耦。

/// 所有 LLM 都实现这个统一接口
#[async_trait]
pub trait BaseChatModel: Send + Sync {
    /// 获取模型名称
    fn model(&self) -> &str;

    /// 获取提供商名称
    fn provider(&self) -> &str;

    /// 同步调用
    async fn invoke(
        &self,
        messages: Vec<Message>,
        tools: Option<Vec<ToolDefinition>>,
        tool_choice: Option<ToolChoice>,
    ) -> Result<ChatCompletion, LlmError>;

    /// 流式调用
    async fn invoke_stream(
        &self,
        messages: Vec<Message>,
        tools: Option<Vec<ToolDefinition>>,
        tool_choice: Option<ToolChoice>,
    ) -> Result<ChatStream, LlmError>;
}

这样，无论底层用的是 GPT-5 还是 Claude，上层的业务逻辑都能保持一致。

2. 工具即函数

利用 Rust 强大的类型系统和闭包，让工具的定义和使用变得直观且安全。

// 定义一个天气工具，类型安全
let weather_tool = FunctionTool::new(
    "get_weather",
    "获取指定城市的天气信息",
    json!({
        "type": "object",
        "properties": {
            "location": {
                "type": "string",
                "description": "城市名称，如 ‘北京’"
            }
        },
        "required": ["location"]
    })
    .as_object()
    .unwrap()
    .clone(),
    |args: WeatherArgs| {
        Box::pin(async move {
            // 实际业务逻辑
            Ok(format!("{} 当前天气：晴，25°C", args.location))
        })
    },
);

3. 事件驱动架构

Agent 的执行本质上是流式的。我们用事件来描述执行过程中的每一个阶段，方便开发者进行精细化的控制和状态展示。

pub enum AgentEvent {
    /// 文本输出
    Text(TextEvent),
    /// 思考过程（Claude 支持）
    Thinking(ThinkingEvent),
    /// 工具调用
    ToolCall(ToolCallEvent),
    /// 工具结果
    ToolResult(ToolResultEvent),
    /// 最终响应
    FinalResponse(FinalResponseEvent),
    /// 错误
    Error(ErrorEvent),
    /// 步骤开始
    StepStart(StepStartEvent),
    /// 步骤完成
    StepComplete(StepCompleteEvent),
}

4. 智能上下文压缩

为了解决上下文爆炸的问题，我们引入了 Ephemeral Tool（临时工具） 的概念。这类工具的执行结果在下一轮迭代前会被自动清理，从而显著降低 token 消耗。

// 工具结果用完即销毁，不占用上下文
let search_tool = FunctionTool::new(...)
    .with_ephemeral(EphemeralConfig::Single);

从零打造 Agent SDK：核心实现解析

如果你也对自研 Agent 框架感兴趣，下面这部分关于 agent-io 核心实现的解析，或许能给你带来一些启发。

第一步：定义统一的消息模型

不同 LLM 的消息格式各异，我们需要一个统一的内部表示作为“中间语言”。

/// 统一的消息类型
pub enum Message {
    /// 用户消息
    User(UserMessage),
    /// 助手消息
    Assistant(AssistantMessage),
    /// 系统消息
    System(SystemMessage),
    /// 工具消息
    Tool(ToolMessage),
}

/// 用户消息
pub struct UserMessage {
    pub role: String,
    pub content: Vec<ContentPart>,
}

/// 内容部分（支持文本和多模态）
pub enum ContentPart {
    Text { text: String },
    Image { image_url: String },
    // ... 其他类型
}

这样，无论调用哪个 LLM，流程都变为：外部格式 -> 统一内部格式 -> 各 Provider 特定格式。

第二步：实现 Agent 核心执行循环

Agent 的核心是一个循环：调用 LLM -> 处理工具调用 -> 再调用 LLM -> ... 直到得到最终答案。

impl Agent {
    fn execute_loop(&self) -> impl Stream<Item = AgentEvent> + '_ {
        async_stream::stream! {
            let mut step = 0;

            loop {
                // 1. 检查迭代次数
                if step >= self.config.max_iterations {
                    yield AgentEvent::Error(...);
                    break;
                }

                // 2. 清理 ephemeral 消息
                self.destroy_ephemeral_messages();

                // 3. 构建请求消息
                let messages = self.build_messages();

                // 4. 调用 LLM
                let completion = self.llm.invoke(messages, ...).await?;

                // 5. 处理响应
                if completion.has_tool_calls() {
                    // 执行工具，继续循环
                    for tool_call in &completion.tool_calls {
                        let result = self.execute_tool(tool_call).await;
                        yield AgentEvent::ToolResult(result);
                    }
                    step += 1;
                    continue;
                }

                // 6. 返回最终结果
                yield AgentEvent::FinalResponse(...);
                break;
            }
        }
    }
}

第三步：优雅的错误处理与重试机制

LLM API 调用常会遇到限流等问题，一个健壮的 SDK 必须内置重试逻辑。

async fn call_llm_with_retry(
    llm: &dyn BaseChatModel,
    messages: Vec<Message>,
    tools: Option<Vec<ToolDefinition>>,
    tool_choice: Option<ToolChoice>,
) -> Result<ChatCompletion> {
    let max_retries = 3;
    let mut delay = Duration::from_millis(100);

    for attempt in 0..=max_retries {
        match llm.invoke(messages.clone(), ...).await {
            Ok(completion) => return Ok(completion),
            Err(LlmError::RateLimit) if attempt < max_retries => {
                tokio::time::sleep(delay).await;
                delay *= 2; // 指数退避
            }
            Err(e) => return Err(e.into()),
        }
    }

    Err(Error::Agent("Max retries exceeded".into()))
}

第四步：精确的 Token 消耗与成本追踪

我们统计每次调用的详细信息，让成本变得透明可控。

/// 使用量汇总
pub struct UsageSummary {
    /// 总 token 数
    pub total_tokens: u64,
    /// 按模型分组的使用量
    pub by_model: HashMap<String, ModelUsage>,
    /// 估算成本（美元）
    pub estimated_cost: f64,
}

impl Agent {
    /// 获取当前会话的使用量统计
    pub async fn get_usage(&self) -> UsageSummary {
        self.usage.read().await.clone()
    }
}

第五步：Builder 模式提升易用性

一个优秀的 SDK 应该让常见场景的代码简洁优雅。

let agent = Agent::builder()
    .with_llm(Arc::new(llm))
    .tool(weather_tool)
    .tool(search_tool)
    .system_prompt("你是一个专业的助手")
    .max_iterations(100)
    .build()?;

// 流式查询
let stream = agent.query_stream("今天北京天气如何？").await?;
pin_mut!(stream);

while let Some(event) = stream.next().await {
    match event {
        AgentEvent::Text(e) => print!("{}", e.content),
        AgentEvent::ToolCall(e) => println!("\n[调用工具: {}]", e.name),
        AgentEvent::FinalResponse(e) => println!("\n{}", e.content),
        _ => {}
    }
}

为什么选择 Rust 作为实现语言？

你可能会问，现在 AI Agent 生态的主流是 Python，为什么我们要用 Rust？

1. 性能与安全的完美平衡：零成本抽象让高级特性在运行时无额外开销；所有权系统保障内存安全，无 GC 停顿，非常适合需要长时间稳定运行的服务。
2. 成熟的异步生态：Tokio 提供了生产级的异步运行时，处理高并发请求游刃有余。
3. 强大的工具链与库生态：reqwest 用于 HTTP 请求、serde 处理序列化、tracing 做结构化日志、tower 构建中间件，生态已经非常完善。
4. 卓越的跨语言调用能力：通过 FFI，Rust 编写的核心 SDK 可以轻松被 Python、Node.js、Go 等语言调用，实现“一次编写，多处集成”。

当前已支持的核心能力

目前，agent-io 已经实现了以下关键特性：

能力	状态
多模型支持	OpenAI, Anthropic, Google, Groq, Mistral, DeepSeek, Ollama, OpenRouter
工具调用	Function Calling，支持依赖注入
流式响应	事件驱动的实时输出
上下文压缩	Ephemeral Tool 自动清理
Token 追踪	使用量和成本统计
重试机制	指数退避，自动处理限流

结语：基础设施决定能力上限

在 AI 应用日新月异的今天，人们常问：Agent 的未来到底什么样？我们认为，Agent 的能力上限，不仅取决于模型本身的强弱，更取决于其底层基础设施的成熟度。就像操作系统决定了应用软件的发挥空间，一个健壮、灵活的 Agent SDK 也直接决定了 AI 应用所能达到的复杂度和可靠性。

agent-io 是我们朝着这个方向迈出的一小步。它当然还不完美：需要更丰富的工具生态库、更智能的上下文优化策略、以及更多真实生产场景的锤炼。

但我们坚信，开源社区的力量能让它走得更远、变得更好。

欢迎一起贡献 🤝

agent-io 是一个完全开源的项目，我们诚挚邀请所有开发者一起参与构建。如果你对 AI Agent 基础架构感兴趣，这里是你实践想法的绝佳舞台。

GitHub: https://github.com/lispking/agent-io

你可以：

• 🐛 提交 Issue：报告 Bug 或提议新功能。
• 🔧 贡献代码：实现新的 LLM 适配器或核心特性。
• 📖 完善文档：帮助改进使用指南和代码示例。
• 💬 分享实践：在 云栈社区 等平台交流你的使用场景和最佳实践。

让我们共同努力，为 AI Agent 的蓬勃生态打造更坚实、更高效的基础设施！