RAG(检索增强生成)
为您的 agent 提供知识库,使其能够使用您自己的数据回答问题。
什么是 RAG?
RAG 允许您的 agent 在回答问题前从文档中查找相关信息。agent 不再只依赖 LLM 训练时使用的内容,而是搜索您的数据,并将结果用作上下文。
流程如下:
- 摄取 — 文档被拆分为多个块,转换为向量嵌入并存储
- 查询 — 问题被嵌入,并根据相似度与已存储的文本块进行匹配
- 生成 — 最相关的文本块作为上下文传递给 LLM,用于生成回答
这意味着您的 agent 可以回答有关产品文档、公司政策、代码库或您提供的任何文本的问题。
主要亮点:
- 📄 摄取任意文本 — 产品文档、Markdown、代码、政策
- 🔍 语义搜索 — 根据含义而不仅仅是关键词查找相关内容
- 🤖 Agentic 检索 — agent 通过
RagTool决定何时进行搜索- 🔌 可插拔后端 — 无需更改代码即可替换嵌入提供程序和向量存储
安装
[dependencies]
# Core only (in-memory store, all chunkers, no external deps)
adk-rag = "2.1.0"
# With Gemini embeddings (recommended for getting started)
adk-rag = { version = "2.1.0", features = ["gemini"] }
步骤 1:构建管道
RagPipeline 将四个组件连接起来:分块器、嵌入提供程序、向量存储以及可选的重排序器。
use std::collections::HashMap;
use std::sync::Arc;
use adk_rag::{
Document, FixedSizeChunker, InMemoryVectorStore,
RagConfig, RagPipeline, EmbeddingProvider,
};
// Mock embedder for demos — no API key needed.
// In production, use GeminiEmbeddingProvider or OpenAIEmbeddingProvider.
struct MockEmbedder;
#[async_trait::async_trait]
impl EmbeddingProvider for MockEmbedder {
async fn embed(&self, text: &str) -> adk_rag::Result<Vec<f32>> {
let hash = text.bytes().fold(0u64, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u64));
let mut v = vec![0.0f32; 64];
for (i, x) in v.iter_mut().enumerate() {
*x = ((hash.wrapping_add(i as u64)) as f32).sin();
}
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 { v.iter_mut().for_each(|x| *x /= norm); }
Ok(v)
}
fn dimensions(&self) -> usize { 64 }
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let pipeline = RagPipeline::builder()
.config(RagConfig::builder()
.chunk_size(256)
.chunk_overlap(50)
.top_k(3)
.build()?)
.embedding_provider(Arc::new(MockEmbedder))
.vector_store(Arc::new(InMemoryVectorStore::new()))
.chunker(Arc::new(FixedSizeChunker::new(256, 50)))
.build()?;
// Create a collection and ingest a document
pipeline.create_collection("docs").await?;
pipeline.ingest("docs", &Document {
id: "intro".into(),
text: "Rust is a systems programming language focused on safety and speed.".into(),
metadata: HashMap::from([("topic".into(), "rust".into())]),
source_uri: None,
}).await?;
// Query
let results = pipeline.query("docs", "safe programming").await?;
for r in &results {
println!("[{:.3}] {}", r.score, r.chunk.text);
}
Ok(())
}
工作原理:
FixedSizeChunker将文档拆分为每块 256 个字符、重叠 50 个字符的文本块MockEmbedder将每个文本块转换为 64 维向量InMemoryVectorStore存储向量,并通过余弦相似度进行搜索query()对问题进行嵌入,查找最接近的文本块,并按分数排序后返回
步骤 2:将 RAG 添加到 Agent
RAG 的真正强大之处在于,agent 将其用作工具时。RagTool 将该流程封装为 adk_core::Tool — agent 在需要信息时调用 rag_search。
当你将 RagTool 与 Gemini 支持的 agent 一起使用时,ADK 会自动将工具结果规范化为兼容 Gemini 的函数响应。这一点很重要,因为 rag_search 通常会返回一个分块列表,而 Gemini 要求 functionResponse.response 在传输过程中是一个 JSON 对象。
use std::sync::Arc;
use adk_agent::LlmAgentBuilder;
use adk_model::gemini::GeminiModel;
use adk_rag::{
Document, GeminiEmbeddingProvider, InMemoryVectorStore,
RagConfig, RagPipeline, RagTool, RecursiveChunker,
};
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let api_key = std::env::var("GOOGLE_API_KEY")?;
// Build pipeline with real embeddings
let pipeline = Arc::new(
RagPipeline::builder()
.config(RagConfig::builder().chunk_size(300).chunk_overlap(50).top_k(3).build()?)
.embedding_provider(Arc::new(GeminiEmbeddingProvider::new(&api_key)?))
.vector_store(Arc::new(InMemoryVectorStore::new()))
.chunker(Arc::new(RecursiveChunker::new(300, 50)))
.build()?,
);
// Ingest documents
pipeline.create_collection("kb").await?;
pipeline.ingest("kb", &Document {
id: "returns".into(),
text: "Our return policy allows returns within 30 days with a receipt.".into(),
metadata: Default::default(),
source_uri: None,
}).await?;
// Wrap pipeline as a tool and attach to an agent
let agent = LlmAgentBuilder::new("support")
.instruction("Answer questions using the rag_search tool. Cite your sources.")
.model(Arc::new(GeminiModel::new(&api_key, "gemini-3.7-flash")?))
.tool(Arc::new(RagTool::new(pipeline, "kb")))
.build()?;
// The agent now calls rag_search automatically when it needs knowledge base info
adk_cli::console::run_console(Arc::new(agent), "app".into(), "user1".into()).await?;
Ok(())
}
当用户询问“你们的退货政策是什么?”时,agent 会:
- 判断需要搜索知识库
- 使用
{"query": "return policy"}调用rag_search - 获取带有分数的相关分块
- 使用这些分块作为上下文来生成自然的回答
第 3 步:选择分块策略
拆分文档的方式会影响检索质量。adk-rag 提供了三种分块器:
| 分块器 | 最适合 | 拆分方式 |
|---|---|---|
FixedSizeChunker | 一般文本、日志 | 每 N 个字符,并保留重叠部分 |
RecursiveChunker | 文章、文档、代码注释 | 段落 → 句子 → 单词 |
MarkdownChunker | Markdown 文件、README | 按标题,保留章节层级 |
use adk_rag::{FixedSizeChunker, RecursiveChunker, MarkdownChunker};
// Fixed: simple, predictable chunks
let chunker = FixedSizeChunker::new(512, 100);
// Recursive: respects natural text boundaries
let chunker = RecursiveChunker::new(512, 100);
// Markdown: each section becomes a chunk with header_path metadata
let chunker = MarkdownChunker::new(512, 100);
RecursiveChunker 是最佳默认选择 — 它会优先尝试按段落分隔,然后按句子边界,最后按单词边界进行分割,相比固定大小的分割方式生成更自然的文本块。
MarkdownChunker 会为每个文本块添加 header_path 元数据字段(例如 "Getting Started > Installation"),帮助代理引用特定章节。
配置
use adk_rag::RagConfig;
let config = RagConfig::builder()
.chunk_size(256) // max characters per chunk (default: 512)
.chunk_overlap(50) // overlap between chunks (default: 100)
.top_k(5) // results to return (default: 10)
.similarity_threshold(0.5) // minimum score to include (default: 0.0)
.build()?;
| 参数 | 控制内容 | 指导 |
|---|---|---|
chunk_size | 每个分块的最大字符数 | 对于大多数使用场景,建议为 200–500。更小 = 更精确,更大 = 更多上下文 |
chunk_overlap | 相邻分块之间的共享字符数 | chunk_size 的 10–20% 可防止丢失边界处的信息 |
top_k | 返回的结果数量 | 结果越多,为 LLM 提供的上下文越多,但令牌用量也越高 |
similarity_threshold | 纳入的最低分数 | 0.0 返回所有结果;0.3–0.7 可过滤较弱的匹配项 |
嵌入提供商
| 提供商 | 功能标志 | 模型 | 要求 |
|---|---|---|---|
GeminiEmbeddingProvider | gemini | gemini-embedding-2 | GOOGLE_API_KEY |
OpenAIEmbeddingProvider | openai | text-embedding-3-small | OPENAI_API_KEY |
// Gemini
use adk_rag::GeminiEmbeddingProvider;
let embedder = GeminiEmbeddingProvider::new(&api_key)?;
// OpenAI
use adk_rag::OpenAIEmbeddingProvider;
let embedder = OpenAIEmbeddingProvider::new(&api_key, "text-embedding-3-small");
您还可以为任何自定义嵌入服务实现 EmbeddingProvider。
向量存储后端
| 后端 | 功能标志 | 最适合 |
|---|---|---|
InMemoryVectorStore | (默认) | 开发、测试、小型数据集 |
QdrantVectorStore | qdrant | 使用专用向量数据库的生产环境 |
LanceDBVectorStore | lancedb | 嵌入式向量数据库(无需服务器) |
PgVectorStore | pgvector | 当你已经使用 PostgreSQL 时 |
// In-memory (no setup needed)
use adk_rag::InMemoryVectorStore;
let store = InMemoryVectorStore::new();
// Qdrant (requires running Qdrant server)
use adk_rag::QdrantVectorStore;
let store = QdrantVectorStore::new("http://localhost:6334").await?;
// pgvector (requires PostgreSQL with pgvector extension)
use adk_rag::PgVectorStore;
let store = PgVectorStore::new("postgres://user:pass@localhost/db").await?;
自定义重排序器
默认的 NoOpReranker 原样传递结果。编写自定义重排序器以提高精确率:
use adk_rag::{Reranker, SearchResult};
struct KeywordBoostReranker { boost: f32 }
#[async_trait::async_trait]
impl Reranker for KeywordBoostReranker {
async fn rerank(
&self,
query: &str,
mut results: Vec<SearchResult>,
) -> adk_rag::Result<Vec<SearchResult>> {
let keywords: Vec<String> = query.split_whitespace()
.filter(|w| w.len() > 3)
.map(|w| w.to_lowercase())
.collect();
for r in &mut results {
let text = r.chunk.text.to_lowercase();
let hits = keywords.iter().filter(|kw| text.contains(kw.as_str())).count();
r.score += hits as f32 * self.boost;
}
results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
Ok(results)
}
}
// Add to pipeline
let pipeline = RagPipeline::builder()
.config(config)
.embedding_provider(embedder)
.vector_store(store)
.chunker(chunker)
.reranker(Arc::new(KeywordBoostReranker { boost: 0.1 }))
.build()?;
多个集合
针对不同的知识领域使用不同的集合。代理可以搜索特定集合,或者您可以创建多个 RagTool 实例:
// Create collections for different content types
pipeline.create_collection("docs").await?;
pipeline.create_collection("faq").await?;
pipeline.create_collection("changelog").await?;
// Ingest into each
pipeline.ingest("docs", &setup_doc).await?;
pipeline.ingest("faq", &faq_doc).await?;
pipeline.ingest("changelog", &release_doc).await?;
// One tool per collection — the agent picks which to search
let docs_tool = RagTool::new(pipeline.clone(), "docs");
let faq_tool = RagTool::new(pipeline.clone(), "faq");
let agent = LlmAgentBuilder::new("support")
.instruction("Search 'docs' for how-to questions, 'faq' for common questions.")
.model(Arc::new(model))
.tool(Arc::new(docs_tool))
.tool(Arc::new(faq_tool))
.build()?;
代理还可以通过在工具调用中传递 "collection": "faq",在查询时覆盖集合。
功能标志
仅引入所需的依赖:
| 功能 | 启用内容 | 额外依赖 |
|---|---|---|
| (默认) | 核心 trait、InMemoryVectorStore、所有分块器 | 无 |
gemini | GeminiEmbeddingProvider | adk-gemini |
openai | OpenAIEmbeddingProvider | reqwest |
qdrant | QdrantVectorStore | qdrant-client |
lancedb | LanceDBVectorStore | lancedb, arrow |
pgvector | PgVectorStore | sqlx |
full | 以上全部 | 全部 |
# Just core
adk-rag = "2.1.0"
# With Gemini embeddings
adk-rag = { version = "2.1.0", features = ["gemini"] }
# Everything
adk-rag = { version = "2.1.0", features = ["full"] }
注意:
lancedb功能需要安装protoc。请使用brew install protobuf(macOS)或apt install protobuf-compiler(Ubuntu)进行安装。
架构
Ingestion
Documents ──→ [Chunker] ──→ [EmbeddingProvider] ──→ [VectorStore]
Query
Question ──→ [EmbeddingProvider] ──→ [VectorStore search] ──→ [Reranker] ──→ Results
│
▼
Agent uses as context
RagPipeline 协调这两个流程。RagTool 将管道封装为 adk_core::Tool,以便代理按需调用它。
运行示例
cargo adk new rag_agent --template rag
cd rag_agent
cargo run
最佳实践
| 实践 | 原因 |
|---|---|
默认使用 RecursiveChunker | 生成自然的分块边界 |
| 将分块保持在 200–500 个字符 | 平衡精确性与上下文 |
| 在生产环境中使用真实嵌入 | 模拟嵌入器仅用于测试 |
将 similarity_threshold 设置为大于 0 | 过滤无关噪声 |
| 按领域分离集合 | 提高精确率,并让代理能够针对性地进行搜索 |
仅将 InMemoryVectorStore 用于开发 | 在生产环境中切换到 Qdrant/pgvector |