feat(memory): 新增 VectorStore 抽象与 RagPipeline 持久化管线
- 新增 src/memory/vector_store.rs(约 660 行): - VectorStore trait:批量 add / search / remove + add_one 默认实现 - InMemoryVectorStore:Mutex<HashMap> + 余弦全量扫描(锁内克隆、锁外计算) - PersistentVectorStore:MemoryStore 包装,JSON blob 持久化,先持久化后内存 - RagPipeline:split → embed → store 组合器(具体 struct,非 trait) - 22 个内联测试覆盖 InMemory(10)/ Persistent(6)/ RagPipeline(4)/ 性能基准(2) - 性能断言:search 10K 条 <100ms,PersistentVectorStore::new 加载 <500ms - 标记 VectorRetriever / InMemoryVectorRetriever 为 #[deprecated(since="0.3.0")] - memory.rs 追加 VectorStore 等 4 个类型的 re-export - document_demo 从手动 VectorRetriever 循环迁移到 RagPipeline 两行调用 - 零新增外部依赖
This commit is contained in:
+33
-37
@@ -1,17 +1,18 @@
|
||||
//! document_demo —— Document + RecursiveCharacterSplitter + MockEmbedding + VectorRetriever 完整衔接示例。
|
||||
//! document_demo —— Document + RecursiveCharacterSplitter + MockEmbedding + RagPipeline 完整衔接示例。
|
||||
//!
|
||||
//! 演示 RAG 管线前置流程:
|
||||
//! 演示 RAG 管线:
|
||||
//! 1. 创建多段落 Document
|
||||
//! 2. RecursiveCharacterSplitter 分割为 chunk
|
||||
//! 3. MockEmbedding 嵌入所有 chunk
|
||||
//! 4. 与 InMemoryVectorRetriever 手动 zip 衔接
|
||||
//! 5. 模拟查询做语义检索
|
||||
//! 3. RagPipeline.ingest() 自动嵌入并存储
|
||||
//! 4. RagPipeline.retrieve() 做语义检索
|
||||
//!
|
||||
//! 运行:`cargo run --example document_demo`(离线,零配置)
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use agcore::document::{Document, RecursiveCharacterSplitter};
|
||||
use agcore::llm::embedding::{Embedding, MockEmbedding};
|
||||
use agcore::memory::vector::{InMemoryVectorRetriever, VectorRetriever};
|
||||
use agcore::memory::{InMemoryVectorStore, RagPipeline, VectorStore};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
@@ -32,43 +33,38 @@ async fn main() {
|
||||
|
||||
println!("输入文档: {} 字符", doc.content.chars().count());
|
||||
|
||||
// 2. RecursiveCharacterSplitter 分割
|
||||
// 2. 构造 RAG 管线(嵌入器 + 向量存储 + 分割器)
|
||||
let embedder: Arc<dyn Embedding> = Arc::new(MockEmbedding::new(4));
|
||||
let store: Arc<dyn VectorStore> = Arc::new(InMemoryVectorStore::new());
|
||||
let splitter = RecursiveCharacterSplitter::new(200, 30);
|
||||
let chunks = splitter.split(&[doc]);
|
||||
let pipeline = RagPipeline::new(
|
||||
Arc::clone(&embedder),
|
||||
Arc::clone(&store),
|
||||
Some(splitter),
|
||||
);
|
||||
|
||||
println!("\n分割为 {} 个 chunk:", chunks.len());
|
||||
for (i, chunk) in chunks.iter().enumerate() {
|
||||
// 3. 一次性 ingest:自动 split → embed → add
|
||||
pipeline.ingest(std::slice::from_ref(&doc)).await.unwrap();
|
||||
|
||||
// 4. 模拟查询:复用第一个 chunk 的 content 作为查询文本
|
||||
let chunks_in_store = store.search(&[1.0, 0.0, 0.0, 0.0], 1).await.unwrap();
|
||||
assert!(!chunks_in_store.is_empty(), "ingest 后 store 应有数据");
|
||||
let query_text = &chunks_in_store[0].0.content;
|
||||
|
||||
let results = pipeline.retrieve(query_text, 3).await.unwrap();
|
||||
println!("\nTop 3 检索结果(与第一个 chunk 相似):");
|
||||
for (doc, score) in &results {
|
||||
println!(
|
||||
" [{:02}] id={}, 长度={}",
|
||||
i, chunk.id, chunk.content.chars().count()
|
||||
" id={}, score={:.4}, content={}",
|
||||
doc.id, score, doc.content
|
||||
);
|
||||
}
|
||||
|
||||
// 3. MockEmbedding 嵌入所有 chunk
|
||||
let embedder = MockEmbedding::new(4);
|
||||
let texts: Vec<String> = chunks.iter().map(|c| c.content.clone()).collect();
|
||||
let vectors = embedder.embed(&texts).await.unwrap();
|
||||
println!("\n嵌入维度: {},向量数: {}", embedder.dim(), vectors.len());
|
||||
|
||||
// 4. 与 InMemoryVectorRetriever 手动 zip 衔接
|
||||
let retriever: std::sync::Arc<dyn VectorRetriever> =
|
||||
std::sync::Arc::new(InMemoryVectorRetriever::new());
|
||||
for (chunk, vec) in chunks.iter().zip(vectors.iter()) {
|
||||
retriever.index(chunk.id.clone(), vec.clone()).await.unwrap();
|
||||
}
|
||||
println!("已索引 {} 个 chunk", chunks.len());
|
||||
|
||||
// 5. 模拟查询:复用第一个 chunk 的 embedding 作为查询向量
|
||||
let query_vec = vectors[0].clone();
|
||||
let results = retriever.search(query_vec, 3).await.unwrap();
|
||||
println!("\nTop 3 检索结果(与 chunk 0 相似度):");
|
||||
for (id, score) in &results {
|
||||
println!(" id={}, score={:.4}", id, score);
|
||||
}
|
||||
|
||||
assert_eq!(chunks.len(), vectors.len(), "chunks 与 vectors 数量必须一致");
|
||||
assert!(!results.is_empty(), "至少应返回 1 条检索结果");
|
||||
assert!(results[0].0.starts_with("rust-intro:chunk:0000"), "Top 1 应为 chunk 0 自身");
|
||||
assert!(
|
||||
results[0].0.id.starts_with("rust-intro:chunk:0000"),
|
||||
"Top 1 应为 chunk 0 自身"
|
||||
);
|
||||
|
||||
println!("\n✓ document_demo 完成");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user