徐涛
32d886f870
- 新增 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 两行调用 - 零新增外部依赖
70 lines
2.8 KiB
Rust
70 lines
2.8 KiB
Rust
//! document_demo —— Document + RecursiveCharacterSplitter + MockEmbedding + RagPipeline 完整衔接示例。
|
|
//!
|
|
//! 演示 RAG 管线：
|
|
//! 1. 创建多段落 Document
|
|
//! 2. RecursiveCharacterSplitter 分割为 chunk
|
|
//! 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::{InMemoryVectorStore, RagPipeline, VectorStore};
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
agcore::init_tracing();
|
|
|
|
// 1. 创建多段落 Document（含中英文混合）
|
|
let doc = Document::new(
|
|
"rust-intro",
|
|
"Rust 是一门系统编程语言，注重安全、并发和性能。\n\n\
|
|
Rust 通过所有权系统管理内存，无需垃圾回收器。\
|
|
所有权规则让内存安全在编译期就能得到保证。\n\n\
|
|
Rust 的并发模型通过类型系统区分线程间共享与独占数据，\
|
|
避免数据竞争。Send 和 Sync 两个 trait 标记了类型的线程安全性。\n\n\
|
|
Rust 的性能与 C/C++ 相当，但提供了更现代的开发体验。\
|
|
Cargo 是官方的构建系统和包管理器，使用简单直观。",
|
|
"text/markdown",
|
|
);
|
|
|
|
println!("输入文档: {} 字符", doc.content.chars().count());
|
|
|
|
// 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 pipeline = RagPipeline::new(
|
|
Arc::clone(&embedder),
|
|
Arc::clone(&store),
|
|
Some(splitter),
|
|
);
|
|
|
|
// 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!(
|
|
" id={}, score={:.4}, content={}",
|
|
doc.id, score, doc.content
|
|
);
|
|
}
|
|
|
|
assert!(!results.is_empty(), "至少应返回 1 条检索结果");
|
|
assert!(
|
|
results[0].0.id.starts_with("rust-intro:chunk:0000"),
|
|
"Top 1 应为 chunk 0 自身"
|
|
);
|
|
|
|
println!("\n✓ document_demo 完成");
|
|
} |