上一章
CH.185 小时实战: AI 代码审查助手

AI 应用集成

LLM API · RAG · 流式响应 · Agent

为什么用 Rust 做 AI 后端?

AI 应用通常瓶颈在网络 I/O(调用 LLM API)而非 CPU,这正是 Rust 异步模型的甜区:

  • 并发处理大量 AI 请求:Tokio 可以同时处理数千个 LLM 请求,内存占用极低
  • 流式响应:Rust 的 Stream trait 天然适合 SSE(Server-Sent Events)
  • 向量计算:Rust 的 SIMD 支持让向量相似度计算比 Python 快 10x+
  • 边缘部署:Rust 二进制可以部署到 Cloudflare Workers、Fly.io 边缘节点

调用 LLM API

对比:调用 OpenAI / Claude API
JS
JavaScript(fetch)
const response = await fetch("https://api.openai.com/v1/chat/completions", {
method: "POST",
headers: {
  "Authorization": `Bearer ${process.env.OPENAI_API_KEY}`,
  "Content-Type": "application/json",
},
body: JSON.stringify({
  model: "gpt-4o",
  messages: [{ role: "user", content: "Hello!" }],
}),
});
const data = await response.json();
console.log(data.choices[0].message.content);
Rs
Rust(reqwest + serde)
use reqwest::Client;
use serde::{Deserialize, Serialize};

#[derive(Serialize)]
struct ChatRequest {
  model:    String,
  messages: Vec<Message>,
}

#[derive(Serialize, Deserialize)]
struct Message { role: String, content: String }

#[derive(Deserialize)]
struct ChatResponse { choices: Vec<Choice> }

#[derive(Deserialize)]
struct Choice { message: Message }

async fn chat(prompt: &str) -> anyhow::Result<String> {
  let client = Client::new();
  let res: ChatResponse = client
      .post("https://api.openai.com/v1/chat/completions")
      .bearer_auth(std::env::var("OPENAI_API_KEY")?)
      .json(&ChatRequest {
          model:    "gpt-4o".into(),
          messages: vec![Message { role: "user".into(), content: prompt.into() }],
      })
      .send().await?
      .json().await?;
  Ok(res.choices[0].message.content.clone())
}

Rust 的 serde 让 JSON 序列化/反序列化完全类型安全。如果 API 返回的字段和你定义的结构体不匹配,编译或解析时立即报错——不会在运行时才发现字段拼错了。


流式响应:SSE + Axum

AI 应用最重要的 UX 之一——流式输出让用户看到逐字生成的效果:

[dependencies]
axum          = "0.8"
reqwest       = { version = "0.12", features = ["stream"] }
tokio-stream  = "0.1"
futures       = "0.3"
eventsource-stream = "0.2"
// src/routes/chat.rs — 流式 AI 响应
use axum::{
    extract::Json,
    response::{IntoResponse, Sse},
    response::sse::{Event, KeepAlive},
};
use futures::StreamExt;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tokio_stream::wrappers::ReceiverStream;
use tokio::sync::mpsc;
 
#[derive(Deserialize)]
pub struct ChatInput { message: String }
 
pub async fn stream_chat(
    Json(input): Json<ChatInput>,
) -> impl IntoResponse {
    let (tx, rx) = mpsc::channel::<Result<Event, std::convert::Infallible>>(32);
 
    // 在后台任务中调用 LLM 并流式发送
    tokio::spawn(async move {
        let client = Client::new();
 
        let mut stream = client
            .post("https://api.openai.com/v1/chat/completions")
            .bearer_auth(std::env::var("OPENAI_API_KEY").unwrap())
            .json(&serde_json::json!({
                "model": "gpt-4o",
                "stream": true,
                "messages": [{ "role": "user", "content": input.message }]
            }))
            .send().await.unwrap()
            .bytes_stream();
 
        while let Some(chunk) = stream.next().await {
            if let Ok(bytes) = chunk {
                let text = String::from_utf8_lossy(&bytes);
                // 解析 SSE data 行,提取 delta content
                for line in text.lines() {
                    if let Some(data) = line.strip_prefix("data: ") {
                        if data == "[DONE]" { break; }
                        if let Ok(val) = serde_json::from_str::<serde_json::Value>(data) {
                            if let Some(content) = val
                                .pointer("/choices/0/delta/content")
                                .and_then(|v| v.as_str())
                            {
                                let _ = tx.send(Ok(Event::default().data(content))).await;
                            }
                        }
                    }
                }
            }
        }
    });
 
    // 将 channel 转为 SSE 流
    let stream = ReceiverStream::new(rx);
    Sse::new(stream).keep_alive(KeepAlive::default())
}

前端消费 SSE

// frontend/src/components/ChatStream.tsx
"use client";
 
import { useState } from "react";
 
export function ChatStream() {
  const [response, setResponse] = useState("");
  const [loading, setLoading]   = useState(false);
 
  async function sendMessage(message: string) {
    setResponse("");
    setLoading(true);
 
    const eventSource = new EventSource(
      `/api/chat/stream?message=${encodeURIComponent(message)}`
    );
 
    eventSource.onmessage = (e) => {
      setResponse((prev) => prev + e.data);
    };
 
    eventSource.onerror = () => {
      setLoading(false);
      eventSource.close();
    };
  }
 
  return (
    <div>
      <button onClick={() => sendMessage("解释 Rust 所有权")}>
AI
      </button>
      <pre className="whitespace-pre-wrap">{response}</pre>
    </div>
  );
}

向量数据库与 RAG

检索增强生成(RAG)——让 AI 回答你代码库的专属问题:

[dependencies]
pgvector   = { version = "0.4", features = ["sqlx"] }  # PostgreSQL 向量扩展
async-openai = "0.28"  # OpenAI 官方 Rust 客户端
-- 启用 pgvector 扩展
CREATE EXTENSION IF NOT EXISTS vector;
 
-- 文档向量表
CREATE TABLE documents (
    id        SERIAL PRIMARY KEY,
    content   TEXT    NOT NULL,
    metadata  JSONB,
    embedding vector(1536)  -- OpenAI text-embedding-3-small 维度
);
 
-- 向量索引(加速相似度搜索)
CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100);
// src/rag.rs — RAG 管道
use async_openai::{Client, types::{CreateEmbeddingRequestArgs, CreateChatCompletionRequestArgs}};
use pgvector::Vector;
use sqlx::PgPool;
 
pub struct RagPipeline {
    openai: Client<async_openai::config::OpenAIConfig>,
    pool:   PgPool,
}
 
impl RagPipeline {
    // 1. 将文档向量化并存储
    pub async fn index_document(&self, content: &str, metadata: serde_json::Value) -> anyhow::Result<()> {
        let embedding = self.embed(content).await?;
 
        sqlx::query!(
            "INSERT INTO documents (content, metadata, embedding) VALUES ($1, $2, $3)",
            content,
            metadata,
            embedding as Vector,
        )
        .execute(&self.pool)
        .await?;
 
        Ok(())
    }
 
    // 2. 相似度搜索
    pub async fn search(&self, query: &str, limit: i64) -> anyhow::Result<Vec<String>> {
        let query_embedding = self.embed(query).await?;
 
        // pgvector 余弦相似度搜索(<=> 操作符)
        let rows = sqlx::query!(
            r#"
            SELECT content
            FROM documents
            ORDER BY embedding <=> $1
            LIMIT $2
            "#,
            query_embedding as Vector,
            limit,
        )
        .fetch_all(&self.pool)
        .await?;
 
        Ok(rows.into_iter().map(|r| r.content).collect())
    }
 
    // 3. RAG:检索 + 生成
    pub async fn ask(&self, question: &str) -> anyhow::Result<String> {
        // 检索相关文档
        let contexts = self.search(question, 5).await?;
        let context = contexts.join("\n\n---\n\n");
 
        // 构建带上下文的 prompt
        let prompt = format!(
            "根据以下代码文档回答问题。\n\n文档:\n{context}\n\n问题:{question}"
        );
 
        // 调用 LLM 生成答案
        let request = CreateChatCompletionRequestArgs::default()
            .model("gpt-4o")
            .messages([async_openai::types::ChatCompletionRequestUserMessageArgs::default()
                .content(prompt)
                .build()?
                .into()])
            .build()?;
 
        let response = self.openai.chat().create(request).await?;
        Ok(response.choices[0].message.content.clone().unwrap_or_default())
    }
 
    async fn embed(&self, text: &str) -> anyhow::Result<Vector> {
        let request = CreateEmbeddingRequestArgs::default()
            .model("text-embedding-3-small")
            .input(text)
            .build()?;
 
        let response = self.openai.embeddings().create(request).await?;
        let values: Vec<f32> = response.data[0].embedding.clone();
        Ok(Vector::from(values))
    }
}

AI Agent:工具调用(Function Calling)

让 AI 自主决定调用哪个函数:

// src/agent.rs — 简单 AI Agent
use serde_json::{json, Value};
 
// 定义工具
const TOOLS: &str = r#"[
  {
    "type": "function",
    "function": {
      "name": "run_cargo_check",
      "description": "对 Rust 代码运行 cargo check,返回编译错误",
      "parameters": {
        "type": "object",
        "properties": {
          "code": { "type": "string", "description": "要检查的 Rust 代码" }
        },
        "required": ["code"]
      }
    }
  },
  {
    "type": "function",
    "function": {
      "name": "search_crates_io",
      "description": "在 crates.io 搜索 Rust crate",
      "parameters": {
        "type": "object",
        "properties": {
          "query": { "type": "string" }
        },
        "required": ["query"]
      }
    }
  }
]"#;
 
pub async fn run_agent(user_request: &str) -> anyhow::Result<String> {
    let client = reqwest::Client::new();
    let tools: Value = serde_json::from_str(TOOLS)?;
    let mut messages = vec![json!({ "role": "user", "content": user_request })];
 
    loop {
        let response = client
            .post("https://api.openai.com/v1/chat/completions")
            .bearer_auth(std::env::var("OPENAI_API_KEY")?)
            .json(&json!({
                "model": "gpt-4o",
                "messages": messages,
                "tools": tools,
            }))
            .send().await?
            .json::<Value>().await?;
 
        let choice = &response["choices"][0];
        let message = &choice["message"];
 
        // 无工具调用 → 完成
        if choice["finish_reason"] == "stop" {
            return Ok(message["content"].as_str().unwrap_or_default().to_string());
        }
 
        // 有工具调用 → 执行并返回结果
        if let Some(tool_calls) = message["tool_calls"].as_array() {
            messages.push(message.clone());
 
            for call in tool_calls {
                let fn_name = call["function"]["name"].as_str().unwrap();
                let args: Value = serde_json::from_str(
                    call["function"]["arguments"].as_str().unwrap()
                )?;
 
                let result = match fn_name {
                    "run_cargo_check" => execute_cargo_check(&args).await,
                    "search_crates_io" => search_crates(&args).await,
                    _ => "未知工具".to_string(),
                };
 
                messages.push(json!({
                    "role": "tool",
                    "tool_call_id": call["id"],
                    "content": result,
                }));
            }
        }
    }
}
 
async fn execute_cargo_check(args: &Value) -> String {
    use tokio::process::Command;
    // 将代码写入临时文件并运行 cargo check
    let code = args["code"].as_str().unwrap_or("");
    // ...实际实现略
    format!("cargo check 结果:{code}")
}
 
async fn search_crates(args: &Value) -> String {
    let query = args["query"].as_str().unwrap_or("");
    // 调用 crates.io API
    format!("搜索 '{query}' 的结果:...")
}

实战项目:AI 代码审查助手

实战项目

AI 代码审查助手

90 分钟高级

接收 Rust 代码,自动运行 clippy 检查,调用 LLM 提供改进建议,通过 SSE 流式返回结果。

OpenAI / Claude API流式 SSE 响应Function CallingRAG 文档检索async-openai
cargo new code-reviewer && cd code-reviewer
# 代码审查助手功能
POST /api/review
  body: { code: "fn main() { ... }" }
SSE 流式返回:
    1. 运行 cargo clippy
    2. 检索相关 Rust 最佳实践文档(RAG
    3. AI 综合分析并给出改进建议
 
# 示例输出(流式)
data: 正在运行 cargo clippy...
data: 发现 2 个 warning
data: 检索相关文档...
data: \n\n**改进建议:**\n\n1. `unwrap()` 替换为 `?` 操作符...

本地 LLM:Ollama 集成

无需 API Key,在本地运行开源模型:

# 安装 Ollama 并下载模型
brew install ollama
ollama pull llama3.2
ollama pull codellama  # 代码专用模型
// Ollama API 兼容 OpenAI 格式,只需改 base_url
let client = async_openai::Client::with_config(
    async_openai::config::OpenAIConfig::new()
        .with_api_base("http://localhost:11434/v1")
        .with_api_key("ollama"),  // 本地不需要真实 key
);
 
// 之后和调用 OpenAI 完全一样

成本优化策略:开发阶段用本地 Ollama(免费),测试通过后切换到 Claude / GPT-4o。Rust 的配置注入让切换 LLM 只需改一个环境变量。


实战视频

Rust + AI:LLM API 与 RAG 架构
1 分 4 秒精讲
0:00 / 0:00
CC
OpenAI API、SSE 流式响应、pgvector RAG。

课程完结:你已经掌握了什么

恭喜!走到这里,你已经系统掌握了:

章节核心技能实战项目
1. Rust 基础类型系统、控制流、迭代器TODO CLI
2. 所有权与借用Move/Copy、借用规则、生命周期字符串处理器
3. 异步与并发Tokio、Future、Channel并发爬虫
4. Web 后端Axum、Router、中间件博客 REST API
5. 数据库SQLx、迁移、事务用户认证系统
6. 全栈实战ts-rs、WebSocket、Server Components实时协作任务板
7. DevOpsDocker、GitHub Actions、K8sCI/CD 流水线
8. AI 集成LLM API、SSE、RAG、Agent代码审查助手

下一步

  • 加入 Rust 中文社区 持续学习
  • 贡献开源 Rust 项目
  • 将你学到的用于实际工作项目

感谢你完成 RustForge 全部课程!🦀