diff --git a/docs/21-phase15-vector-store-persistence.md b/docs/21-phase15-vector-store-persistence.md new file mode 100644 index 0000000..3718be5 --- /dev/null +++ b/docs/21-phase15-vector-store-persistence.md @@ -0,0 +1,1165 @@ +# Phase 15: 向量存储持久化（SqliteStore 后端） + +- **文档编号**：21 +- **标题**：Phase 15 — 向量存储持久化 +- **日期**：2026-07-09 +- **状态**：待实施 +- **涉及模块**：`memory/vector_store.rs`（新文件）、`memory/vector.rs`、`memory.rs`、`examples/document_demo.rs` +- **关联文档**：roadmap.md（§Phase 15）、18-phase11-testing-and-retrieval.md（VectorRetriever trait）、20-phase14-document-and-embedding.md（Document 系统 + Embedding 抽象） +- **对应**：Roadmap §Phase 15（v0.3.0 第三阶段） +- **预估规模**：新增约 400 行核心代码 + 约 30 行示例修改 + +--- + +## 1. 背景与目标 + +### 1.1 背景 + +agcore v0.3.0 的目标是从「LLM 调用工具箱」升级为「能构建多 Agent 协作、RAG、长记忆 Agent 产品的基础系统」。Roadmap §Phase 15 定位为 v0.3.0 的第三阶段，在 Phase 14 交付了 Document 类型和 Embedding 抽象之后，补齐 RAG 管线的**核心存储与检索层**。 + +当前 agcore 已有： + +- `VectorRetriever` trait（Phase 11）——语义向量检索抽象，支持 `index(id, embeddings)` 和 `search(query, k)` +- `InMemoryVectorRetriever`（Phase 11）——基于 HashMap 的内存引用实现 +- `Document` 类型 + `RecursiveCharacterSplitter`（Phase 14）——文档分割 +- `Embedding` trait + `MockEmbedding`（Phase 14）——文本向量化抽象 + +这四个组件构成了 RAG 管线的「前置」和「后置」环节，但缺少中间的 **向量存储** 环节——一个能持久化保存向量和文档、支持进程重启后恢复数据的核心存储抽象。 + +### 1.2 目标 + +1. 定义 `VectorStore` trait——统一向量存储的 `add` / `search` / `remove` 接口 [高] +2. 实现 `InMemoryVectorStore`——基于 HashMap 的内存引用实现，支持余弦相似度全量扫描 [高] +3. 实现 `PersistentVectorStore`——基于 SqliteStore 的持久化包装，R/W 时双向同步 [高] +4. 实现 `RagPipeline`——组合器，封装 `split → embed → store.add`（ingest）和 `embed → store.search`（retrieve） [高] +5. 标记 `VectorRetriever` 为废弃（`#[deprecated]`），引导迁移到 `VectorStore` [中] +6. 更新 `document_demo` 示例，从手动循环迁移到 `RagPipeline` 两行调用 [中] +7. 代码零新增外部依赖 [高] + +### 1.3 适用范围 + +| 纳入 | 排除 | +|------|------| +| `VectorStore` trait 定义（add / search / remove） | ANN 近似最近邻索引（O(n) 扫描，<10K 向量足够） | +| `InMemoryVectorStore`（HashMap + 余弦全量扫描） | `update` 方法（remove + add 即可） | +| `PersistentVectorStore`（SqliteStore JSON blob 后端读写） | namespace 运行时 CRUD（构造参数固定） | +| `RagPipeline`（具体 struct，非 trait） | 维度校验——调用方保证同一 VectorStore 实例写入的所有向量维度一致（通过 `Embedding::dim()` 确认）；VectorStore 不做运行时校验 | +| 零新增外部依赖 | `RagPipeline` trait 抽象（等第二种 pipeline 实现） | +| `VectorRetriever` 废弃标记（`#[deprecated]`） | pgvector / Qdrant 后端 | + +--- + +## 2. 当前状态分析 + +### 2.1 现有 RAG 能力矩阵 + +| 维度 | 已实现 | 缺失 | +|------|--------|------| +| 文档分割 | ✅ `RecursiveCharacterSplitter`（Phase 14） | — | +| 文本向量化 | ✅ `Embedding` trait + `MockEmbedding`（Phase 14） | 真实 Provider（Phase 15 不解决） | +| 向量检索抽象 | ✅ `VectorRetriever` trait（Phase 11） | 持久化语义检索 | +| 向量检索内存实现 | ✅ `InMemoryVectorRetriever`（Phase 11） | — | +| 向量持久化 | ❌ | `PersistentVectorStore`（本 Phase） | +| RAG 管线编排 | ❌ | `RagPipeline`（本 Phase） | +| SqliteStore 后端 | ✅ Phase 7 | — | + +### 2.2 现有检索抽象对比 + +**VectorRetriever trait（Phase 11，将被废弃）**： + +```rust +#[async_trait] +pub trait VectorRetriever: Send + Sync { + async fn index(&self, id: String, embeddings: Vec) -> Result<(), MemoryError>; + async fn search(&self, query: Vec, k: usize) -> Result, MemoryError>; +} +``` + +`index` 只接受单个 `(id, embeddings)` 对，不支持批量添加；`search` 返回 `(id, score)`，不附带文档内容。调用方必须自行维护 `id → Document` 的映射。 + +**VectorStore trait（本 Phase）**： + +```rust +#[async_trait] +pub trait VectorStore: Send + Sync { + async fn add(&self, documents: &[Document], embeddings: &[Vec]) -> Result<(), MemoryError>; + async fn search(&self, query: &[f32], k: usize) -> Result, MemoryError>; + async fn remove(&self, ids: &[String]) -> Result<(), MemoryError>; +} +``` + +`add` 接受批量 `(doc, embedding)` 对；`search` 直接返回 `(Document, score)`，调用方无需额外映射。消除 Phase 11 设计中 `id → Document` 的隐式耦合。 + +### 2.3 模块依赖关系 + +``` +当前（Phase 14 完成后）： + +Document + RecursiveCharacterSplitter (src/document.rs) + ↓ output: Vec +RagPipeline::ingest (本 Phase) ← Embedding trait (src/llm/embedding.rs) + ├── splitter.split() → Vec chunks + ├── embedder.embed() → Vec> + └── store.add() → 持久化 + ↓ +RagPipeline::retrieve (本 Phase) + ├── embedder.embed(query) → Vec> + └── store.search() → Vec<(Document, f32)> + +VectorRetriever trait (src/memory/vector.rs) → #[deprecated] + ↓ 消费端迁移 +VectorStore trait (src/memory/vector_store.rs) ← 新 + ├── InMemoryVectorStore ← 新 + ├── PersistentVectorStore ← 新 (依赖 SqliteStore) + └── 被 RagPipeline 消费 ← 新 +``` + +### 2.4 风险识别 + +| 风险 | 等级 | 缓解措施 | +|------|------|---------| +| PersistentVectorStore 构造时全量加载阻塞 | 低 | `MemoryStore::list()` 由 SqliteStore 内部 `spawn_blocking` 卸载；10K 条预估 <500ms | +| JSON 序列化/反序列化性能 | 低 | 每条向量存储为独立 MemoryItem，写入/读取均单条操作，无全量 JSON | +| 大向量集（>10K 条）搜索 O(n) 性能 | 中 | 文档标注当前为全量扫描，提供升级路径到 ANN 索引 | +| SqliteStore 非向量专用后端，查询不支持向量索引 | 低 | 设计目标即「小数据量够用」，超出规模换专用后端 | +| `add()` 不等长输入部分写入已发生 | 中 | 截断处理 + 返回 `Err(InvalidInput)`，调用方可 `let _ =` 忽略 | +| 向量维度不一致（不同 Embedding provider 混用） | 低 | VectorStore 不做运行时校验，调用方自行通过 `Embedding::dim()` 确认维度一致性 | +| 调用方传入重复 `doc.id` 静默覆盖 | 低 | `HashMap::insert` 自然行为；等价于 upsert 语义，需在 doc comment 中注明 | +| MemoryStore 淘汰策略误删向量条目 | 低 | VectorEntry 存储 `created_at` 保留原始时间，避免 TTL 误判；默认 `EvictionPolicy::None` 无害 | + +--- + +## 3. 可选方案 + +### 3.1 模块归属：vector_store.rs 放在哪里？ + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. `memory/vector_store.rs`（推荐）** | 与 `memory/vector.rs`（现有 VectorRetriever）同目录 | 语义一致：向量存储属于记忆系统；复用现有 `MemoryError`、`MemoryStore` 依赖；文件路径短 | 目录已存在，初始化工作量小 | +| B. 新顶层模块 `vector/` | 创建 `src/vector/mod.rs` + `vector_store.rs` | 语义独立，与 Document 顶层模块对等 | 增加顶层入口，增加模块声明开销；VectorStore 依赖 MemoryStore/MemoryError，放在 memory 外部反而增加耦合 | + +**结论**：选择 **方案 A**。Roadmap 原定 `src/vector/` 新模块，但实际分析发现 VectorStore 是记忆系统的自然扩展（类似 `KnowledgeStore` 放在 `memory/`），且 `VectorRetriever` 已在 `memory/vector.rs`。放在同一模块可以减少跨模块引用和 `use` 路径复杂度。 [高] + +### 3.2 VectorStore trait：批量 vs 单条 + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. 批量 add（推荐）** | `add(&[Document], &[Vec])` | RagPipeline 一次 ingest 处理多个 chunk 时只需要一次 async call；批量语义明确 | 需要处理不等长输入 | +| B. 单条 add | `add(Document, Vec)` | 签名简单，没有不等长问题 | RagPipeline 遍历每个 chunk 逐个 await，N 个 chunk 产生 N 次 async roundtrip | + +**结论**：选择 **方案 A**。提供 `add_one` 便捷方法作为语法糖。 [高] + +### 3.3 search 返回值：元组 vs 命名结构体 + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. 元组（推荐）** | `Vec<(Document, f32)>` | 零新类型；符合 Rust 惯用风格（与 `VecRetriever::search` 一致）；字段访问 `result.0.content` 直观 | .score 不如命名结构体 readable | +| B. `ScoredDocument` 结构体 | `struct ScoredDocument { doc: Document, score: f32 }` | 语义清晰；可扩展字段 | YAGNI：当前没有需要扩展的额外字段；增加类型结构复杂度 | + +**结论**：选择 **方案 A**。保持与 Phase 11 `VectorRetriever::search` 返回 `Vec<(String, f32)>` 一致的设计语言，只是把 `String`（id）替换为 `Document`。 [高] + +### 3.4 错误类型：复用 vs 新建 + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. 复用 MemoryError（推荐）** | 所有 VectorStore API 返回 `Result<(), MemoryError>` 和 `Result, MemoryError>` | 零新类型；已有 `InvalidInput`（不等长）、`NotFound`（删除不存在 id）、`Serialization`（JSON 反序列化）变体 | `MemoryError::Serialization` 的语义略宽，但 JSON 反序列化失败确实属于序列化错误 | +| B. 新建 `VectorError` | 独立枚举：`InvalidInput` / `DimensionMismatch` / `Serialization` / `Storage` | 类型精确；不与 memory 错误耦合 | Phase 15 RagPipeline 需要同时处理 `MemoryError` + `VectorError` + `LlmError`；组合复杂度高 | + +**结论**：选择 **方案 A**。VectorStore 是 `memory/` 模块的直接组成部分，复用 MemoryError 自然合理。 [高] + +### 3.5 RagPipeline：struct vs trait + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. 具体 struct（推荐）** | `RagPipeline { embedder, store, splitter }` | 代码量最少；当前只有一个实现 | 如果将来出现第二流水线（如多模态嵌入），需要抽象 | +| B. `RagPipeline` trait | 定义 `ingest` / `retrieve` 接口 | 可插拔 | YAGNI：当前没有第二种 pipeline 需求；trait 一旦发布，签名变动是 breaking change | + +**结论**：选择 **方案 A**。YAGNI 原则。当出现第二种 pipeline 实现时，再从 struct 中提取 trait。 [高] + +### 3.6 持久化格式：JSON blob vs 结构化列 + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. JSON blob（推荐）** | 每条向量存为 `MemoryItem`，content 为 `JSON.stringify(VectorEntry)` | 零 schema 迁移；利用现有 `MemoryStore` API（save/get/list/delete） | 无法用 SQL 直接查询向量内容（业务上也不需要有此需求） | +| B. SqliteStore 专用表 | 在 sqlite_store.rs 中新增 `vectors` 表 | 查询效率高 | 突破 MemoryStore trait 抽象层，增加跨模块依赖；PersistentVectorStore 丧失后端无关性 | + +**结论**：选择 **方案 A**。`MemoryStore` 的 `list(prefix)` 可以按 namespace 过滤，足以支持构造时的全量加载。JSON blob 的序列化/反序列化成本在 <10K 向量场景下可以忽略。 [高] + +### 3.7 模块扁平 vs 子目录 + +| 方案 | 描述 | 优点 | 缺点 | +|------|------|------|------| +| **A. 扁平（推荐）** | `src/memory/vector_store.rs` 单文件 | 约 400 行，一个文件包含 trait + 三个 struct + 内联测试 | 如果后期需要拆分为 `search/`、`persist/` 等子模块，需重构为目录 | +| B. 目录预置 | `src/memory/vector_store/` 子目录 | 预留扩展空间 | 推测性设计；当前内容不足以分割 | + +**结论**：选择 **方案 A**。所有代码在一个文件中，结构清晰：trait 定义在前，实现在后，测试在 `#[cfg(test)]` 中。 [高] + +--- + +## 4. 推荐方案 + +### 4.1 整体架构 + +``` +Phase 15 交付 +┌─────────────────────────────────────────────────────┐ +│ src/memory/vector_store.rs │ +│ │ +│ ┌──────────────────────────────────────────────┐ │ +│ │ VectorStore trait │ │ +│ │ add(docs, embeddings) -> Result<(), ME> │ │ +│ │ search(query, k) -> Result,ME>│ │ +│ │ remove(ids) -> Result<(), ME> │ │ +│ │ add_one(doc, emb) -> Result<(), ME> │ │ +│ └──────────────┬───────────────────────────────┘ │ +│ │ implements │ +│ ┌──────────────┴───────────────┐ │ +│ │ InMemoryVectorStore │ │ +│ │ entries: HashMap │ │ +│ │ dot() / cosine_similarity() │ │ +│ └──────────────┬───────────────┘ │ +│ │ wraps │ +│ ┌──────────────┴───────────────┐ │ +│ │ PersistentVectorStore │ │ +│ │ inner: InMemoryVectorStore │ │ +│ │ store: Arc │ │ +│ │ namespace: String │ │ +│ │ add(): 内存 + SqliteStore 双写 │ │ +│ └──────────────────────────────┘ │ +│ │ +│ ┌──────────────────────────────┐ │ +│ │ RagPipeline │ │ +│ │ embedder: Arc│ │ +│ │ store: Arc │ │ +│ │ splitter: Option │ │ +│ │ ingest: split→embed→store │ │ +│ │ retrieve: embed→search │ │ +│ └──────────────────────────────┘ │ +└─────────────────────────────────────────────────────┘ + +存储格式（SqliteStore 视角）: + MemoryItem.id = "vec:{namespace}:{doc_id}" + MemoryItem.content = JSON(VectorEntry { doc_id, content, metadata, mime_type, embedding }) + MemoryItem.metadata = {} +``` + +### 4.2 核心类型设计 + +#### VectorStore trait + +```rust +/// 向量存储抽象 —— 语义检索的核心接口。 +/// +/// 提供文档-向量的批量添加、余弦相似度搜索、批量删除三个核心操作。 +/// 所有实现必须满足 Send + Sync 以支持跨 .await 调用。 +/// +/// # 并发安全 +/// +/// 实现内部必须使用线程安全的容器（如 Mutex 或 RwLock）， +/// 允许跨多个 tokio task 共享 `&VectorStore` 引用。 +#[async_trait] +pub trait VectorStore: Send + Sync { + /// 批量添加文档及其向量。 + /// + /// `documents` 和 `embeddings` 必须等长。不等长时： + /// - 截取 `min(len)` 对处理（部分写入已发生） + /// - 返回 `Err(MemoryError::InvalidInput)` 告知截断 + /// - 调用方可以 `let _ = store.add(...)` 忽略错误 + async fn add(&self, documents: &[Document], embeddings: &[Vec]) + -> Result<(), MemoryError>; + + /// 检索与 `query` 向量最相似的 `k` 条记录。 + /// + /// 返回 `Vec<(Document, f32)>`，其中 `f32` 为余弦相似度分数， + /// 取值范围 `[0.0, 1.0]`（对单位向量），按分数降序排列。 + /// + /// # 守卫 + /// + /// - 空索引 → 返回 `vec![]` + /// - `k == 0` → 返回 `vec![]` + /// - 零向量（norm ≈ 0）→ 返回 `vec![]` + async fn search(&self, query: &[f32], k: usize) + -> Result, MemoryError>; + + /// 批量删除文档（幂等）。 + /// + /// 不存在的 id 静默忽略，不会返回错误。 + async fn remove(&self, ids: &[String]) -> Result<(), MemoryError>; + + /// 便捷方法：单条添加。 + /// + /// 等价于 `self.add(&[doc], &[emb]).await`。 + async fn add_one(&self, doc: Document, emb: Vec) + -> Result<(), MemoryError> { + self.add(&[doc], &[emb]).await + } +} +``` + +**设计决策**： + +- `documents` 和 `embeddings` 的输入使用切片引用（`&[Document]`、`&[Vec]`）——零拷贝，调用方可以传入 `Vec` 的引用或数组切片 [高] +- `search` 返回 `Vec<(Document, f32)>` 而非 `Vec<(String, f32)>`——消除调用方维护 `id → Document` 映射的负担 [高] +- `remove` 幂等设计——不存在的 id 静默忽略，简化调用方逻辑（无需在删除前检查 id 是否存在） [高] +- 提供默认实现的 `add_one` 便捷方法——`&[x]` 语法足够精简，但 `add_one` 在链式调用中更清晰 [中] + +#### InMemoryVectorStore + +```rust +/// 内存向量存储 —— VectorStore 的引用实现。 +/// +/// 内部使用 `Mutex)>>` 存储， +/// `search()` 执行 O(n) 全量余弦相似度扫描，适用于 ≤10K 条向量的场景。 +/// +/// # 并发安全 +/// +/// 使用 `std::sync::Mutex`（非 tokio Mutex）。 +/// +/// **锁持有时间评估**： +/// - `add()` / `remove()`：微秒级（HashMap 插入/删除操作） +/// - `search()`：毫秒级（O(n) 全量扫描 + 余弦计算），对 10K 条 1536 维向量预估 1-10ms。 +/// 实现时在锁内克隆数据快照到 Vec 后立即释放锁，在锁外进行余弦相似度计算， +/// 避免长时间持有锁阻塞并发写操作。 +pub struct InMemoryVectorStore { + pub(crate) entries: Mutex)>>, +} + +impl InMemoryVectorStore { + /// 创建一个空存储。 + pub fn new() -> Self { /* ... */ } + + /// 从预填充的 entries 构造（供 PersistentVectorStore 使用）。 + pub(crate) fn with_entries( + entries: HashMap)>, + ) -> Self { /* ... */ } +} +``` + +**add 实现逻辑**： + +``` +add(documents, embeddings): + if documents.len() == embeddings.len(): + for i in 0..len(doc): insert(doc[i], emb[i]) + return Ok(()) + else: + tracing::warn!( + docs = documents.len(), embs = embeddings.len(), + "InMemoryVectorStore::add 长度不匹配，截断到 min" + ) + min_len = min(len(doc), len(emb)) + for i in 0..min_len: insert(doc[i], emb[i]) + return Err(InvalidInput) + + insert(doc, emb): + entries.lock() → map.insert(doc.id.clone(), (doc, emb)) +``` + +> **与 `InMemoryVectorRetriever` 的行为差异**：旧 `VectorRetriever::index` 对不等长输入使用 `zip` 静默截断（无返回值），新 `VectorStore::add` 显式检查并返回 Err。迁移时调用方应检查返回的 Err 或确保输入等长。 + +**search 实现逻辑**： + +``` +search(query, k): + if k == 0: return vec![] + tracing::trace!("InMemoryVectorStore::search, k={k}") + 对整个 query 向量做零向量检测（所有分量 ≈ 0.0） + + lock entries → 获取所有值的快照 Vec<(Document, Vec)> + + for each (doc, emb): + score = cosine_similarity(query, emb) + if score < 1e-10: continue (排除无关或空向量) + if score < 1e-10: continue # 浮点除零保护：余弦相似度分母含 1e-10，正交或零向量产生趋零值，不构成有效结果 + results.push((doc, score)) + + results.sort_by(|a, b| b.1.partial_cmp(&a.1)) // 降序 + results.truncate(k) + return results +``` + +**remove 实现逻辑**： + +``` +remove(ids): + lock entries → map.retain(|key, _| !ids.contains(key)) + // retain 在 Rust 中是对传入闭包返回 true 的元素保留 + // 所以我们保留「不在 ids 中」的元素 + tracing::debug!(count = ids.len(), "InMemoryVectorStore::remove") +``` + +#### 余弦相似度辅助函数 + +```rust +/// 点积。 +fn dot(a: &[f32], b: &[f32]) -> f32 { + debug_assert_eq!(a.len(), b.len(), "dot: 不等长向量 ({}/{}), 维度不匹配将导致静默截断", a.len(), b.len()); + a.iter().zip(b.iter()).map(|(x, y)| x * y).sum() +} + +/// 余弦相似度，加 1e-10 防除零。 +fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let dot_product = dot(a, b); + let norm_a = dot(a, a).sqrt(); + let norm_b = dot(b, b).sqrt(); + dot_product / (norm_a * norm_b + 1e-10) +} +``` + +**设计决策**： + +- `dot()` 和 `cosine_similarity()` 为 `pub(crate)` 私有辅助函数——不暴露为公共 API，如需在其他模块使用，可通过 `InMemoryVectorStore::search` 间接调用 [中] +- 使用 `std::sync::Mutex` 而非 `tokio::sync::Mutex`。`add`/`remove` 锁持有微秒级（HashMap 插入/删除），`search` 在锁内克隆数据快照后立即释放锁，在锁外计算余弦相似度（预估 1-10ms/10K 条），避免长时间持有锁阻塞并发写操作 [高] +- 零向量守卫在 `search()` 中而非 `cosine_similarity()` 中——避免每次相似度计算都检测，只在入口做一次 [中] + +#### PersistentVectorStore + +```rust +/// 持久化向量存储 —— 基于 MemoryStore 的持久化包装。 +/// +/// # 架构 +/// +/// 运行时全量加载到 InMemoryVectorStore 做余弦搜索， +/// 写操作（add/remove）同时同步到内存和后端 MemoryStore。 +/// +/// # 存储格式 +/// +/// 每条向量存为一条 MemoryItem： +/// - `id`: `"vec:{namespace}:{doc_id}"`（colon-separated namespace 前缀） +/// - `content`: JSON 序列化的 VectorEntry +/// - `metadata`: 空 HashMap +/// +/// # 构造开销 +/// +/// `new()` 通过 `store.list(prefix)` 全量加载已有条目， +/// 时间复杂度 O(N)（N 为已有向量数），适用于 ≤10K 条的场景。 +pub struct PersistentVectorStore { + inner: InMemoryVectorStore, + store: Arc, + namespace: String, +} +``` + +**私有 VectorEntry 结构体**： + +```rust +#[derive(Serialize, Deserialize)] +struct VectorEntry { + doc_id: String, + content: String, + metadata: HashMap, + mime_type: String, + embedding: Vec, + /// ISO 8601 创建时间（UTC），持久化 roundtrip 重建时保持原时间， + /// 避免 MemoryStore 的 TTL 淘汰策略误判。 + created_at: String, +} +``` + +**构造方法 new()**（异步）： + +```rust +pub async fn new(store: Arc, namespace: &str) -> Result { + let prefix = format!("vec:{}:", namespace); + let filter = MemoryFilter { prefix: Some(prefix.clone()), ..Default::default() }; + + // MemoryStore::list 内部由 SqliteStore 使用 spawn_blocking 卸载到阻塞线程池 + tracing::debug!(namespace = %namespace, "PersistentVectorStore::new — 开始全量加载"); + let items = store.list(&filter).await?; + tracing::info!(count = items.len(), "PersistentVectorStore::new — 加载完成"); + + let mut entries = HashMap::new(); + for item in items { + let entry: VectorEntry = serde_json::from_str(&item.content) + .map_err(|e| MemoryError::Serialization(e.to_string()))?; + let doc = Document { + id: entry.doc_id, + content: entry.content, + metadata: entry.metadata, + mime_type: entry.mime_type, + }; + // created_at 保留在 VectorEntry 中用于 roundtrip 重建，TTL 淘汰由 MemoryStore 管理 + // （VectorEntry.created_at 字段在反序列化时已自动填充） + entries.insert(doc.id.clone(), (doc, entry.embedding)); + } + + tracing::info!(entries = entries.len(), "PersistentVectorStore — 内存索引重建完成"); + Ok(Self { + inner: InMemoryVectorStore::with_entries(entries), + store, + namespace: namespace.to_string(), + }) +} +``` + +> **spawn_blocking 说明**：`MemoryStore::list()` 是 async trait 方法，`SqliteStore` 实现内部使用 `tokio::task::spawn_blocking` 将 SQLite 查询卸载到阻塞线程池。`PersistentVectorStore::new()` 本身在 async context 中调用即可，无需额外 spawn_blocking。加载耗时 ≈ SQLite 查询时间 + JSON 反序列化时间，10K 条向量预估在 200ms-500ms 范围内（O(n) 反序列化，每条 JSON <2KB）。 + +**add 双写路径**（先持久化后内存）： + +``` +add(documents, embeddings): + tracing::debug!(count = docs.len(), "PersistentVectorStore::add") + 1. for each (doc, emb): + entry = VectorEntry { + doc_id, content, metadata, mime_type, + embedding: emb, + created_at: OffsetDateTime::now_utc() + .format(&Rfc3339) // ISO 8601 字符串 + } + json = serde_json::to_string(&entry) + key = format!("vec:{}:{}", self.namespace, doc.id) + self.store.save(MemoryItem { id: key, content: json, metadata: {}, .. }) + .await? // 先写持久化（失败时不污染内存） + 2. self.inner.add(docs, embs).await? // 再写内存 + 3. Ok(()) +``` + +**search 完全委托**： + +``` +search(query, k): + tracing::trace!("PersistentVectorStore::search, k={k}") + self.inner.search(query, k).await // 纯内存搜索 +``` + +**remove 双删**（先删持久化后删内存）： + +``` +remove(ids): + tracing::debug!(count = ids.len(), "PersistentVectorStore::remove") + 1. for each id: + key = format!("vec:{}:{}", self.namespace, id) + self.store.delete(&key).await? // 先删持久化（幂等） + 2. self.inner.remove(ids).await? // 再删内存 + 3. Ok(()) +``` + +**设计决策**： + +- `new()` 使用 `async` 而非同步构造——`MemoryStore::list()` 是异步操作；即使调用方想立即 `.await`，异步构造也比同步构造 + 显式 `init()` 调用更安全（不会忘记调用 `init()`） [高] +- 使用 `MemoryFilter { prefix: Some(...) }` 而非 `list_all` + 客户端过滤——SqliteStore 支持 prefix 索引扫描，减少网络传输量 [高] +- `VectorEntry.doc_id` 使用 `doc.id` 的值而非重新生成——删除时可通过 id 反推 MemoryItem key [高] +- 写入顺序：**先持久化后内存**——持久化写入失败时内存不被污染，下次构造 `PersistentVectorStore` 从持久化全量加载，自动忽略不完整的写入；内存写入失败时持久化已写入（产生「持久化有、内存无」的中间状态），但 `search` 读到的是持久化已确认子集而非超集，不会返回幽灵数据；`add` 和 `remove` 双写均遵守此顺序 [高] +- 批量 add 的非原子性：`add` 的持久化阶段对 N 条文档执行 N 次独立 `store.save()`，非原子操作。第 i 条失败时前 i-1 条已持久化。这是已知限制，未来可通过引入 `MemoryStore::save_batch` 批量 API 优化。当前 `PersistentVectorStore::new()` 全量加载时能正确跳过缺失条目 [中] + +> **remove 幽灵数据窗口**：先 `store.delete()` 成功后、`inner.remove()` 前的窗口内，`search()` 委托到 `inner` 可能返回已被删除的条目（幽灵数据）。该窗口有限——重启后从持久化全量加载时幽灵条目自然消失。调用方如需精确状态，可在 `remove()` 返回后主动调用 `search()` 验证。 [中] + +**键命名示例**： + +| namespace | doc.id | MemoryItem id | +|-----------|--------|---------------| +| `"default"` | `"doc_001:chunk:0000"` | `"vec:default:doc_001:chunk:0000"` | +| `"reports"` | `"rpt_42"` | `"vec:reports:rpt_42"` | +| `"research"` | `"paper_xyz:chunk:0005"` | `"vec:research:paper_xyz:chunk:0005"` | + +**注意**：`doc.id` 中如果包含冒号（如 `doc_001:chunk:0000`），拼接后的 key 为 `vec:namespace:id` 三层结构。`MemoryFilter` 的 prefix 匹配使用 `prefix = format!("vec:{}:", namespace)`，即 `"vec:default:"`，可以匹配 `"vec:default:doc_001:chunk:0000"` 和 `"vec:default:other_id:with:colons"`。因为 prefix 只是字符串前缀匹配，中间包含冒号不影响匹配的正确性。 [高] + +#### RagPipeline + +```rust +/// RAG 管线组合器 —— 封装 split → embed → store 和 embed → search 两个核心流程。 +/// +/// # 使用方式 +/// +/// ```rust +/// let pipeline = RagPipeline::new(embedder, store); +/// pipeline.ingest(&documents).await?; +/// let results = pipeline.retrieve("what is agcore?", 5).await?; +/// ``` +/// +/// # 分割器 +/// +/// `splitter` 字段为 `Option`： +/// - `Some(splitter)` → `ingest()` 先分割再嵌入（调用方传入原始文档） +/// - `None` → `ingest()` 跳过分割，直接嵌入（调用方已分好 chunk） +pub struct RagPipeline { + embedder: Arc, + store: Arc, + splitter: Option, +} + +impl RagPipeline { + /// 创建新的 RAG 管线。 + /// + /// 不设置分割器时，`ingest()` 跳过分割阶段， + /// 调用方传入的 Document 应已是分割好的 chunk。 + pub fn new( + embedder: Arc, + store: Arc, + splitter: Option, + ) -> Self { /* ... */ } + + /// 摄取文档：分割 → 向量化 → 存储。 + /// + /// 流程： + /// 1. 如果 splitter 存在，先分割文档为 chunks + /// 2. 提取所有 chunk 的 content 为 Vec + /// 3. embedder.embed() 批量向量化 + /// 4. store.add() 批量存储 + /// + /// # 边界 + /// + /// - 空文档切片 → Ok(()), 无操作 + /// - 分割后 chunk 为空 → Ok(()), 无操作 + pub async fn ingest(&self, documents: &[Document]) -> Result<(), MemoryError> { + // 分割（可选） + let chunks = match &self.splitter { + Some(splitter) => splitter.split(documents), + None => documents.to_vec(), + }; + if chunks.is_empty() { + return Ok(()); + } + + // 向量化 + let texts: Vec = chunks.iter().map(|d| d.content.clone()).collect(); + let embeddings = self.embedder.embed(&texts) + .await + .map_err(|e| MemoryError::Storage(e.to_string()))?; + + // 存储 + self.store.add(&chunks, &embeddings).await?; + Ok(()) + } + + /// 检索：向量化查询 → 向量相似度搜索。 + /// + /// # 边界 + /// + /// - 空字符串查询 → 返回 vec![]（embed 产生零向量 → search 零向量守卫） + pub async fn retrieve(&self, query: &str, k: usize) + -> Result, MemoryError> { + let embeddings = self.embedder.embed(&[query.to_string()]) + .await + .map_err(|e| MemoryError::Storage(e.to_string()))?; + self.store.search(&embeddings[0], k).await + } +} +``` + +**设计决策**： + +- `embedder.embed()` 返回 `LlmError`，`store.add/search` 返回 `MemoryError`——在两个错误类型之间通过 `MemoryError::Storage(e.to_string())` 桥接 [高] +- `splitter` 为 `Option`——兼顾「调用方已分割」和「需要管线自动分割」两种场景 [高] +- `ingest` 和 `retrieve` 均返回 `MemoryError`——上层统一处理一种错误类型，简化调用方代码 [高] +- `retrieve` 内部构造 `&[query.to_string()]`——接受 `&str` 而非 `&[String]`，接口更友好 [中] + +### 4.3 边界情况处理 + +| 场景 | 行为 | 原因 | +|------|------|------| +| `add()` 输入 `documents.len() ≠ embeddings.len()` | 截取 `min(len)` 对处理，返回 `Err(InvalidInput)` | 不等长属于调用方错误；截取保证至少部分写入有效 | +| `add()` 输入不等长截断后 | 调用方可 `let _ = store.add(...).await` 忽略 Err（部分数据已写入），或检查 Err 的 `Display` 文案了解截断数量 | 不等长时无法区分"全部写入"和"部分写入"，调用方应评估业务容忍度 | +| `add()` 重复 `doc.id` | 静默覆盖（upsert 语义）——后写入的 Document 和 embedding 替换前者 | HashMap insert 的自然行为，等价于 update | +| `add()` 输入空切片 `(&[], &[])` | 空操作，返回 `Ok(())` | 合法输入 | +| `remove()` 不存在的 id | 静默忽略（幂等） | 简化调用方逻辑 | +| `remove()` 空切片 `(&[])` | 空操作，返回 `Ok(())` | 合法输入 | +| `search()` 空索引 | 返回 `vec![]` | 无数据可查 | +| `search()` `k = 0` | 返回 `vec![]` | 请求 0 条结果 | +| `search()` 零向量查询 | 返回 `vec![]`（零向量守卫） | 零向量无法做有意义的相似度比较 | +| `search()` 索引中所有向量与查询正交 | 返回 `vec![]`（相似度 ≤ 1e-10 全部过滤） | 无相关内容 | +| `PersistentVectorStore::new()` 加载失败 | 返回 `Err(MemoryError::Storage)` | 依赖后端异常，不可恢复 | +| JSON 反序列化失败 | 返回 `Err(MemoryError::Serialization)` | 数据损坏或 schema 变更 | +| `ingest()` 空文档切片 | `Ok(())`，无操作 | 合法输入 | +| `ingest()` 分割后 chunk 为空 | `Ok(())`，无操作 | 所有文档均为空内容 | +| `add_one()` 内部委托 | 等价于 `add(&[doc], &[emb])` | 保持一致性 | + +### 4.4 错误处理策略 + +| 场景 | 处理方式 | 原因 | +|------|---------|------| +| add() 不等长 | `Err(MemoryError::InvalidInput)`（含截断说明） | 调用方可选择 `let _ =` 忽略或检查错误 | +| PersistentVectorStore::add 持久化写入失败 | `Err` 传播（内存未写入——先写持久化后写内存） | 先持久化后内存，持久化失败时内存不被污染，重启后自动忽略不完整写入 | +| PersistentVectorStore::remove 持久化删除失败 | `Err` 传播（内存未删除——先删持久化后删内存） | 先持久化后内存，持久化失败时内存不受影响；尚未删除的条目可安全重试 | +| PersistentVectorStore 构造加载失败 | `Err` 传播 | 启动失败应尽早暴露 | +| JSON 序列化/反序列化失败 | `Err(MemoryError::Serialization)` | 数据损坏属不可恢复错误 | +| MemoryStore 后端 IO 错误 | `Err` 传播 | 依赖底层错误，不转换 | +| remove() 不存在的 id | 静默忽略（幂等） | 符合 MemoryStore::delete 的幂等语义 | +| RagPipeline ingest 空文档列表 | `Ok(())`，无操作 | 合法输入 | +| RagPipeline 中 `LlmError` | 桥接为 `MemoryError::Storage`（原始错误 via `.to_string()`） | 统一错误类型降低调用方复杂度 | + +### 4.5 与现有系统的关系 + +``` +前 Phase 11 (memory/vector.rs): + VectorRetriever trait → #[deprecated(since = "0.3.0")] + InMemoryVectorRetriever → #[deprecated(since = "0.3.0")] + +本 Phase (memory/vector_store.rs): + VectorStore trait → 新 + InMemoryVectorStore → 新（语义等价于 InMemoryVectorRetriever） + PersistentVectorStore → 新 + RagPipeline → 新 + +Phase 14 消费: + Document → VectorStore::add 输入类型 + Embedding trait → RagPipeline::ingest/retrieve 调用 + RecursiveCharacterSplitter → RagPipeline::ingest 可选前置 +``` + +--- + +## 5. 实施计划 + +### 5.1 文件与代码量估计 + +| 文件 | 角色 | 估计行数 | 新增/修改 | +|------|------|---------|-----------| +| `src/memory/vector_store.rs` | VectorStore trait + InMemoryVectorStore + PersistentVectorStore + RagPipeline + ~19 测试 | ~400 | **新增** | +| `src/memory/vector.rs` | VectorRetriever + InMemoryVectorRetriever 加 deprecation 标注 | +4 | 修改 | +| `src/memory.rs` | 追加 `pub mod vector_store` + re-export | +3 | 修改 | +| `examples/document_demo.rs` | 从手动循环迁移到 RagPipeline 两行调用 | ~-30 | 修改 | +| **合计** | | **~377** | **1 新增 + 3 修改** | + +### 5.2 实施步骤 + +``` +Step 15.1 ──→ Step 15.2 ──→ Step 15.3 ──→ Step 15.4 ──→ Step 15.5 ──→ Step 15.6 +vector_store 追加 追加 标记废弃 + 简化 全量验证 +.rs Persistent RagPipeline 更新 re-export document_demo cargo test +(trait + VectorStore + 测试 + 测试 cargo clippy + InMemory) cargo doc +``` + +#### Step 15.1 — 创建 vector_store.rs（trait + InMemoryVectorStore） + +**涉及文件**：`src/memory/vector_store.rs`（新增） + +**前置依赖**：无（独立新增文件） + +**工作量**：M（1-2h） + +**验收条件**：`cargo build --all-targets` 通过 + +**内容**： + +1. `VectorStore` trait 定义（`#[async_trait]`，3 个必选方法 + `add_one` 默认实现） +2. `InMemoryVectorStore` struct + `new()` / `with_entries()` 构造器 +3. `VectorStore` for `InMemoryVectorStore` 的实现： + - `add`：等长校验 → 遍历 zip → insert（不等长截断 + Err） + - `search`：lock → 空/k=0/零向量守卫 → 全量余弦扫描 → sort → truncate + - `remove`：lock → HashMap::retain +4. 私有辅助函数 `dot(a, b) → f32` 和 `cosine_similarity(a, b) → f32` +5. 内联测试（10 个）：basic_add_and_search / search_empty_store / search_zero_vector / search_k_is_zero / search_orthogonal_vectors / add_mismatched_lengths / add_duplicate_id_upsert / remove_items / remove_nonexistent_id / concurrent_operations + +**验证**：`cargo build --all-targets` 编译通过。 + +#### Step 15.2 — 追加 PersistentVectorStore + +**涉及文件**：`src/memory/vector_store.rs` + +**前置依赖**：Step 15.1 + +**工作量**：M（1-2h） + +**验收条件**：持久化 roundtrip 测试（write → close → reopen → read）通过 + +**内容**： + +1. 私有 `VectorEntry` 结构体（`#[derive(Serialize, Deserialize)]`） +2. `PersistentVectorStore` struct + `async fn new(store, namespace)` 构造方法 + - `store.list(prefix="vec:{namespace}:")` 全量加载 + - `serde_json::from_str` 反序列化 + - `InMemoryVectorStore::with_entries()` 注入 +3. `VectorStore` for `PersistentVectorStore` 实现： + - `add`：序列化每个 VectorEntry → 先逐个 `store.save()`（先持久化）→ 再 `inner.add()`（后内存） + - `search`：完全委托 `inner.search()` + - `remove`：逐个 `store.delete()`（先删持久化）→ 再 `inner.remove()`（后删内存） + +**验证**：`cargo test -- memory::vector_store::persistent` 全部通过。 + +#### Step 15.3 — 追加 RagPipeline + +**涉及文件**：`src/memory/vector_store.rs` + +**前置依赖**：Step 15.2 + +**工作量**：S（< 1h） + +**验收条件**：端到端 ingest_and_retrieve 集成测试通过 + +**内容**： + +1. `RagPipeline` struct（3 个字段：embedder / store / splitter） +2. `new()` 构造函数 +3. `ingest()`：split（可选）→ embed → store.add +4. `retrieve()`：embed → store.search +5. 桥接 `LlmError → MemoryError::Storage` +6. 内联测试（4 个）：ingest_and_retrieve + retrieve_empty_store + ingest_empty_docs（空文档切片返回 `Ok(())`）+ ingest_empty_split（分割后 chunk 为空返回 `Ok(())`） + +**验证**：`cargo test -- memory::vector_store::rag` 全部通过。 + +#### Step 15.4 — 标记 VectorRetriever 废弃 + 更新 re-export + +**涉及文件**：`src/memory/vector.rs`（修改）、`src/memory.rs`（修改） + +**前置依赖**：Step 15.1（VectorStore trait 存在） + +**工作量**：S（< 0.5h） + +**验收条件**：`cargo build` 项目内部零 deprecation warning（已使用 `#[allow(deprecated)]` 的测试除外） + +**内容**： + +- `src/memory/vector.rs`： + - `VectorRetriever` trait 前加 `#[deprecated(since = "0.3.0", note = "请使用 memory::VectorStore")]` + - `InMemoryVectorRetriever` struct 前加 `#[deprecated(since = "0.3.0", note = "请使用 memory::InMemoryVectorStore")]` + - 内联测试前加 `#[allow(deprecated)]` + +- `src/memory.rs`： + - 追加 `pub mod vector_store;` + - 追加 re-export： + - `pub use vector_store::VectorStore;` + - `pub use vector_store::InMemoryVectorStore;` + - `pub use vector_store::PersistentVectorStore;` + - `pub use vector_store::RagPipeline;` + +**验证**：`cargo build --all-targets` 编译通过，`cargo clippy` 0 警告。 + +#### Step 15.5 — 简化 document_demo.rs + +**涉及文件**：`examples/document_demo.rs` + +**前置依赖**：Step 15.3（RagPipeline 存在） + +**工作量**：S（< 0.5h） + +**验收条件**：`cargo run --example document_demo` exit 0 + +**迁移概要**：移除手动 `InMemoryVectorRetriever` 循环 + id 映射，替换为 `RagPipeline::new(embedder, store, splitter)` → `ingest()` → `retrieve()` 两行调用。文件从 ~74 行缩减为 ~40 行。 + +**验证**：`cargo run --example document_demo` → exit 0。 + +#### Step 15.6 — 全量验证 + +**验证标准**： + +| 检查项 | 指标 | +|--------|------| +| `cargo build --all-targets` | ✅ 通过 | +| `cargo test --all-targets` | ✅ 全绿（313+ → 332+，~19 个新测试） | +| `cargo doc --no-deps` | ✅ 0 warning | +| 零新外部依赖 | ✅ Cargo.toml 无修改 | +| `cargo run --example document_demo` | ✅ exit 0 | +| 性能基准验证 | ✅ `PersistentVectorStore::new` 10K 加载 <500ms；`InMemoryVectorStore::search` 10K 搜索 <100ms | + +--- + +## 6. 测试计划 + +| 分组 | 测试 | 场景 | 验证点 | +|------|------|------|--------| +| InMemory | `basic_add_and_search` | 3 个 3 维向量，query 匹配 Top1 | 最相似文档的 id 和 score 正确 | +| InMemory | `search_empty_store` | 空索引 + 任意 query | 返回 `vec![]` | +| InMemory | `search_zero_vector` | 零向量查询（norm ≈ 0） | 零向量守卫 → `vec![]` | +| InMemory | `add_mismatched_lengths` | 3 docs + 2 embs（不等长） | 截取 2 对处理 + 返回 `Err` | +| InMemory | `remove_items` | add 2 → remove 1 → search | 已删除 id 不在结果中 | +| InMemory | `remove_nonexistent_id` | remove 从未添加的 id | 静默忽略，返回 `Ok(())` | +| InMemory | `concurrent_operations` | 10 并发生成 + 混合搜索 | 无 panic | +| InMemory | `search_k_is_zero` | k=0 + 非空索引 | k=0 守卫 → 返回 `vec![]` | +| InMemory | `add_duplicate_id_upsert` | 同一 doc.id 写入两次 | 第二次写入覆盖前者，search 返回 1 条 | +| InMemory | `search_orthogonal_vectors` | 索引中向量与查询正交 | 相似度 < 1e-10 过滤 → 返回 `vec![]` | +| Persistent | `persistent_roundtrip` | write 3 条 → 重建 store → read | 数据完整不丢 | +| Persistent | `search_after_reload` | write → 重建 store → search | 重启后检索结果匹配 | +| Persistent | `namespace_isolation` | 两个 namespace 各自写入 | 互不干扰，search 正确 | +| Persistent | `concurrent_access` | 多 task 写入 + 搜索 | 无死锁 | +| Persistent | `partial_add_recovery` | 写入 5 条，模拟第 3 条持久化失败，重建 store | 前 2 条存活，后 3 条缺失 | +| Persistent | `new_empty_store` | 空 MemoryStore 构造 PersistentVectorStore | search 返回空结果 | +| RagPipeline | `ingest_and_retrieve` | Document → embed → add → search 端到端 | 检索结果包含目标文档 | +| RagPipeline | `ingest_empty_docs` | 空文档切片传入 ingest | 返回 `Ok(())` 不报错 | +| RagPipeline | `ingest_empty_split` | 分割后 chunk 为空 | 返回 `Ok(())` 不报错 | +| RagPipeline | `retrieve_empty_store` | 空 store + 任意查询 | 返回 `vec![]` | + +--- + +## 7. 文件变更清单 + +| 文件 | 操作 | 说明 | +|------|------|------| +| `src/memory/vector_store.rs` | **新增** | ~400 行：VectorStore trait + InMemoryVectorStore + PersistentVectorStore + RagPipeline + 内联测试 | +| `src/memory/vector.rs` | 修改 | +4 行：两个 deprecated 标注 + 测试 `#[allow(deprecated)]` | +| `src/memory.rs` | 修改 | +3 行：module 声明 + 4 个 `pub use` re-export | +| `examples/document_demo.rs` | 修改 | ~74 → ~40 行：移除手动 VectorRetriever，改为 RagPipeline 两行调用 | + +**验证标准**： + +| 检查项 | 指标 | +| `cargo test --all-targets` | ✅ 全绿（313+ → 332+，~19 新测试） | +| `cargo build --all-targets` | ✅ 通过 | +| `cargo clippy --all-targets -- -D warnings` | ✅ 0 警告 | +| `cargo doc --no-deps` | ✅ 0 warning | +| `cargo run --example document_demo` | ✅ exit 0 | +| 零新外部依赖 | ✅ Cargo.toml 无修改 | +| 性能基准验证 | ✅ `PersistentVectorStore::new` 10K 加载 <500ms；`InMemoryVectorStore::search` 10K 搜索 <100ms | + +## 8. 回滚方案 + +### 8.1 代码回滚 + +| 场景 | 操作 | +|------|------| +| vector_store.rs 编译/测试失败 | `git checkout -- src/memory/vector_store.rs` | +| deprecation 标注问题 | `git checkout -- src/memory/vector.rs src/memory.rs` | +| document_demo 不工作 | `git checkout -- examples/document_demo.rs` | +| 全量回滚 | `git revert ` | + +### 8.2 数据回滚 + +`PersistentVectorStore` 使用 `vec:{namespace}:` 前缀写入 SqliteStore，回退到旧版本时需清理这些数据： + +```sql +-- 清理所有 Phase 15 写入的向量数据 +DELETE FROM memory_items WHERE id LIKE 'vec:%'; + +-- 按 namespace 精确清理（如果只回滚特定向量集合） +DELETE FROM memory_items WHERE id LIKE 'vec:my_knowledge:%'; +``` + +> **不可逆决策声明**： +> - `VectorEntry` JSON 格式（字段名、embedding 序列化方式、`created_at` 格式）一旦发布，后续修改需考虑旧数据兼容性 +> - `vec:{namespace}:{doc_id}` 键命名约定一旦外部消费者依赖此模式，更改需迁移 +> - `VectorStore` trait 的 `add`/`search`/`remove` 签名一旦发布，修改是 breaking change +> +> 缓解措施：若需在 v0.4 中变更 JSON 格式，建议在 `VectorEntry` 中引入 `version` 字段（v1），反序列化时按版本路由解析逻辑。 + +## 9. 不做的事 + +- ❌ ANN 索引（O(n) 扫描，<10K 向量足够） +- ❌ `update` 方法（remove + add 即可） +- ❌ namespace 运行时 CRUD（构造参数固定） +- ❌ 向量维度校验（同现有策略，错误维度不 panic 但结果异常） +- ❌ `RagPipeline` trait 抽象（等第二种 pipeline 实现时再提取） +- ❌ pgvector / Qdrant 等专用向量数据库后端 + +--- + +## 11. Roadmap 同步要点 + +实施完成后需同步更新 `docs/roadmap.md`： + +1. **Phase 15 状态**：`⏳ 待实施` → `✅ Phase 15 全部交付物已完成` +2. **依赖关系图**：`P15["..."]:::pending` → `:::done` +3. **里程碑表**：`M11` 状态从 `⏳` → `✅ 2026-07-09` +4. **顶部当前状态**：追加 `Phase 15 完成` +5. **下一步行动**：指向 Phase 16 摘要自动生成 +6. **已完成列表**：追加 `- ✅ Phase 15 向量存储持久化 — VectorStore trait + InMemoryVectorStore + PersistentVectorStore + RagPipeline` +7. **模块位置修正**：Phase 15 交付物描述中 `src/vector/` 修正为 `src/memory/vector_store.rs` +8. **文档链接**：追加 `docs/21-phase15-vector-store-persistence.md` + +--- + +## 12. 参考来源 + +1. **Phase 11 VectorRetriever trait** — `docs/18-phase11-testing-and-retrieval.md` + - 现有 `VectorRetriever` trait 设计，本 Phase 标记废弃并迁移到 `VectorStore` + +2. **Phase 14 Document 系统 + Embedding 抽象** — `docs/20-phase14-document-and-embedding.md` + - Document 类型和 Embedding trait 是本 Phase `RagPipeline` 的前置依赖 + +3. **Phase 7 SqliteStore** — `docs/14-phase7-sqlite-store.md` + - `MemoryStore` trait 的 SQLite 后端，`PersistentVectorStore` 的持久化依赖 + +4. **Roadmap Phase 15 定义** — `docs/roadmap.md` §Phase 15（v0.3.0 第三阶段） + - 交付物描述、设计要点、依赖关系 + +5. **MemoryError 错误类型** — `src/memory/error.rs` + - 复用错误类型的设计依据（`InvalidInput` / `Serialization` / `Storage` / `NotFound`） + +6. **MemoryStore trait** — `src/memory/store/mod.rs` + - `save` / `get` / `delete` / `list` 接口签名，`PersistentVectorStore` 构造时依赖 `list(prefix)`，写入时依赖 `save` / `delete` + +--- + +## 13. 详细实施计划 + +### 13.1 任务依赖总图 + +```mermaid +flowchart LR + classDef step fill:#e1f5fe,stroke:#0288d1 + classDef verify fill:#e8f5e9,stroke:#388e3c + + subgraph Step15_1["Step 15.1 — vector_store.rs（trait + InMemory）"] + 15_1_1["15.1.1 定义 VectorStore trait"] --> 15_1_2["15.1.2 实现 InMemoryVectorStore struct"] + 15_1_2 --> 15_1_4["15.1.4 实现 dot / cosine_similarity"] + 15_1_4 --> 15_1_3["15.1.3 实现 VectorStore for InMemoryVectorStore"] + 15_1_3 --> 15_1_5["15.1.5 编写 10 个内联测试"] + end + + subgraph Step15_2["Step 15.2 — 追加 PersistentVectorStore"] + 15_2_1["15.2.1 定义 VectorEntry 结构体"] --> 15_2_2["15.2.2 实现 PersistentVectorStore + new()"] + 15_2_2 --> 15_2_3["15.2.3 实现 VectorStore for PersistentVectorStore"] + 15_2_3 --> 15_2_4["15.2.4 编写 6 个内联测试"] + end + + subgraph Step15_3["Step 15.3 — 追加 RagPipeline"] + 15_3_1["15.3.1 实现 RagPipeline struct + new()"] --> 15_3_2["15.3.2 实现 ingest()"] + 15_3_1 --> 15_3_3["15.3.3 实现 retrieve()"] + 15_3_2 --> 15_3_4["15.3.4 编写 4 个内联测试"] + 15_3_3 --> 15_3_4 + end + + subgraph Step15_4["Step 15.4 — 标记废弃 + 更新 re-export"] + 15_4_1["15.4.1 VectorRetriever 加 #[deprecated]"] --> 15_4_2["15.4.2 InMemoryVectorRetriever 加 #[deprecated]"] + 15_4_2 --> 15_4_3["15.4.3 memory.rs 追加 re-export"] + end + + subgraph Step15_5["Step 15.5 — 简化 document_demo"] + 15_5_1["15.5.1 迁移至 RagPipeline 两行调用"] + end + + subgraph Step15_6["Step 15.6 — 全量验证"] + 15_6_1["15.6.1 cargo build --all-targets"] + 15_6_2["15.6.2 cargo test --all-targets"] + 15_6_3["15.6.3 cargo clippy"] + 15_6_4["15.6.4 cargo doc --no-deps"] + 15_6_5["15.6.5 性能基准验证"] + 15_6_6["15.6.6 cargo run --example document_demo"] + 15_6_7["15.6.7 验证 Cargo.toml 零变更"] + end + + 15_1_5 --> 15_2_1 + 15_2_4 --> 15_3_1 + 15_3_4 --> 15_5_1 + 15_1_5 --> 15_4_1 + 15_5_1 --> 15_6_1 + + class 15_1_1,15_1_2,15_1_3,15_1_4,15_1_5,15_2_1,15_2_2,15_2_3,15_2_4,15_3_1,15_3_2,15_3_3,15_3_4 step + class 15_4_1,15_4_2,15_4_3,15_5_1 step + class 15_6_1,15_6_2,15_6_3,15_6_4,15_6_5,15_6_6,15_6_7 verify +``` + +### 13.2 Step 15.1 — 创建 vector_store.rs（trait + InMemoryVectorStore） + +**涉及文件**：`src/memory/vector_store.rs`（新增） +**前置依赖**：无 +**工作量**：M（1-2h） + +| # | 任务描述 | 涉及文件 | 前置依赖 | 工作量 | 风险 | 验收条件 | +|---|---------|---------|---------|--------|------|---------| +| 15.1.1 | 定义 `VectorStore` trait：3 个必选方法（`add` / `search` / `remove`）+ `add_one` 默认实现 + doc comments + `#[async_trait]` | `src/memory/vector_store.rs` | 无 | S | 低 | `cargo build` 通过；`cargo doc --no-deps` 0 warning（确认 intra-doc link 正确） | +| 15.1.2 | 实现 `InMemoryVectorStore` struct：`Mutex)>>` 内部存储 + `new()` 和 `pub(crate) with_entries()` 构造器 | `src/memory/vector_store.rs` | 15.1.1 | S | 低 | 编译通过；`with_entries` 对外不可见 | +| 15.1.4 | 实现私有辅助函数 `dot(a, b) → f32`（含 `debug_assert_eq!(a.len(), b.len())`）和 `cosine_similarity(a, b) → f32`（加 `1e-10` 防除零），`pub(crate)` 可见性 | `src/memory/vector_store.rs` | 15.1.2 | S | 低 | 编译通过；零向量与零向量 cosine_similarity 不 panic，返回 0.0（因 1e-10 防除零）；正交向量接近 0.0 | +| 15.1.3 | 实现 `VectorStore for InMemoryVectorStore`：`add()`（等长校验/截断+Err+tracing::warn!/Mutex insert）、`search()`（锁内克隆快照→释放锁→外余弦扫描→sort+truncate，含空/k=0/零向量守卫）、`remove()`（HashMap::retain） | `src/memory/vector_store.rs` | **15.1.4**（依赖 dot/cosine_similarity 辅助函数） | M | 中 | 单元测试通过；不等长输入返回 `Err(InvalidInput)` 且已写入部分数据；空索引返回 `vec![]`；`k=0` 返回 `vec![]` | +| 15.1.5 | 编写 10 个内联测试：`basic_add_and_search` / `search_empty_store` / `search_zero_vector` / `search_k_is_zero` / `search_orthogonal_vectors` / `add_mismatched_lengths` / `add_duplicate_id_upsert` / `remove_items` / `remove_nonexistent_id` / `concurrent_operations` | `src/memory/vector_store.rs` | 15.1.3 | S | 低 | `cargo test -- InMemory` 全部通过；并发测试验证 10 个并发无 panic 且 search 结果计数正确（并发 add N 条后 search 返回 N 条）；`remove_nonexistent_id` 验证删除不存在的 id 幂等静默忽略 | + +**说明**：Step 15.1 执行顺序：`15.1.1 → 15.1.2 → 15.1.4（辅助函数） → 15.1.3（依赖辅助函数的 impl） → 15.1.5（测试）`，严格串行。`add_mismatched_lengths` 测试需验证不等长截断语义——截取 `min(len)` 对处理并返回 `Err(InvalidInput)`。`remove_nonexistent_id` 测试验证幂等性——删除不存在的 id 应返回 `Ok(())`。 + +### 13.3 Step 15.2 — 追加 PersistentVectorStore + +**涉及文件**：`src/memory/vector_store.rs` +**前置依赖**：Step 15.1（InMemoryVectorStore 存在） +**工作量**：M（1-2h） + +| # | 任务描述 | 涉及文件 | 前置依赖 | 工作量 | 风险 | 验收条件 | +|---|---------|---------|---------|--------|------|---------| +| 15.2.1 | 定义私有 `VectorEntry` 结构体：`#[derive(Serialize, Deserialize)]`，字段含 `doc_id` / `content` / `metadata` / `mime_type` / `embedding` / `created_at: String` | `src/memory/vector_store.rs` | 15.1.5 | S | 低 | 编译通过；结构体不可见（无 `pub`），仅文件内使用 | +| 15.2.2 | 实现 `PersistentVectorStore` struct + `async fn new(store, namespace)`：prefix 过滤全量加载 → serde_json 反序列化 → InMemoryVectorStore::with_entries 注入 + tracing 埋点；在 doc comment 中标注 spawn_blocking 说明 | `src/memory/vector_store.rs` | 15.2.1 | S | 中 | 持久化 roundtrip 测试通过；加载完成后 tracing::info! 输出条目计数 | +| 15.2.3 | 实现 `VectorStore for PersistentVectorStore`：`add`（序列化→逐个 `store.save()` 先写持久化→`inner.add()` 后写内存→tracing::debug!）、`search`（完全委托 `inner.search()`）、`remove`（逐个 `store.delete()` 先删持久化→`inner.remove()` 后删内存） | `src/memory/vector_store.rs` | 15.2.2 | M | 中 | 单元测试通过；add 写入后搜索返回预期结果；remove 后搜索不再返回已删条目 | +| 15.2.4 | 编写 6 个内联测试：`persistent_roundtrip`（write→close→reopen→read）、`search_after_reload`、`namespace_isolation`、`concurrent_access`、`partial_add_recovery`（部分写入后重建验证存活数据）、`new_empty_store`（空存储构造后可用） | `src/memory/vector_store.rs` | 15.2.3 | S | 低 | `cargo test -- persistent` 全部通过；`namespace_isolation` 验证两个 namespace 互不干扰；`partial_add_recovery` 验证前 N 条持久化成功条目在重建后可检索；`new_empty_store` 验证构造后 `search` 返回空结果 | + +**说明**：`PersistentVectorStore::new()` 依赖 SqliteStore 的 `list(prefix)` 支持。若 `store.list()` 返回空 vec，构造过程仍应成功（空索引）。tracing 埋点需区分 debug（构造过程）和 info（加载完成）级别。 +**注意**：`PersistentVectorStore::add` 是**非原子**操作（N 条文档执行 N 次独立 `store.save()`），第 i 条失败时前 i-1 条已持久化。需追加 `partial_add_recovery` 测试验证此恢复路径：写入 5 条 → 模拟第 3 条持久化失败 → 重建 store → 验证前 2 条存活。同时追加 `new_empty_store` 测试验证空 MemoryStore 构造后可用、search 返回空结果。 + +### 13.4 Step 15.3 — 追加 RagPipeline + +**涉及文件**：`src/memory/vector_store.rs` +**前置依赖**：Step 15.2 +**工作量**：S（<1h） + +| # | 任务描述 | 涉及文件 | 前置依赖 | 工作量 | 风险 | 验收条件 | +|---|---------|---------|---------|--------|------|---------| +| 15.3.1 | 实现 `RagPipeline` struct：3 个字段（`embedder: Arc` / `store: Arc` / `splitter: Option`）+ `new()` 构造函数 + doc comment 说明用法 | `src/memory/vector_store.rs` | 15.2.4 | S | 低 | 编译通过；`new()` 中 splitter 为 `None` 合法 | +| 15.3.2 | 实现 `ingest()`：splitter.split（可选）→ embedder.embed → store.add；空文档守卫（空切片和空 chunk 均返回 `Ok(())`）；LlmError→MemoryError::Storage 桥接。**已知限制**：当前将所有 chunk 一次性传入 `embedder.embed()`，真实 Embedding Provider（如 OpenAI）有批量大小限制，调用方需自行控制单次 ingest 的文档数（如 20 条/批） | `src/memory/vector_store.rs` | 15.3.1 | S | 中 | 端到端 ingest 测试通过；空文档输入返回 `Ok(())` 不报错 | +| 15.3.3 | 实现 `retrieve()`：embedder.embed(query) → store.search(query_vec, k)；空字符串查询路径（零向量守卫→`vec![]`） | `src/memory/vector_store.rs` | 15.3.1 | S | 低 | 端到端 retrieve 测试通过；空字符串返回 `vec![]` | +| 15.3.4 | 编写 4 个内联测试：`ingest_and_retrieve`、`retrieve_empty_store`、`ingest_empty_docs`（空文档切片返回 `Ok(())`）、`ingest_empty_split`（分割后 chunk 为空返回 `Ok(())`） | `src/memory/vector_store.rs` | 15.3.2 + 15.3.3 | S | 低 | `cargo test -- rag` 全部通过；`ingest_and_retrieve` 使用 MockEmbedding 保持确定性，验证检索结果包含目标文档；空输入测试验证边界守卫正确 | + +**说明**：`ingest()` 和 `retrieve()` 都在 `impl RagPipeline` 块中实现，非 trait 方法。错误桥接使用 `MemoryError::Storage(e.to_string())`，丢失原始错误类型信息——这是已知取舍，调用方可从 `.to_string()` 获取原始错误信息。 +**注意**：需追加 `ingest_empty_docs` 和 `ingest_empty_split` 测试验证空文档切片和分割后 chunk 为空均返回 `Ok(())` 不报错（§4.3 边界定义）。 + +### 13.5 Step 15.4 — 标记 VectorRetriever 废弃 + 更新 re-export + +**涉及文件**：`src/memory/vector.rs`、`src/memory.rs` +**前置依赖**：Step 15.1（VectorStore trait 存在即可标记废弃） +**工作量**：S（<0.5h） + +| # | 任务描述 | 涉及文件 | 前置依赖 | 工作量 | 风险 | 验收条件 | +|---|---------|---------|---------|--------|------|---------| +| 15.4.1 | `VectorRetriever` trait 前加 `#[deprecated(since = "0.3.0", note = "请使用 memory::VectorStore")]` | `src/memory/vector.rs` | 15.1.5 | S | 低 | 编译通过；使用 `VectorRetriever` 的地方产生 deprecation warning | +| 15.4.2 | `InMemoryVectorRetriever` struct 前加 `#[deprecated(since = "0.3.0", note = "请使用 memory::InMemoryVectorStore")]`；内联测试前加 `#[allow(deprecated)]` | `src/memory/vector.rs` | 15.4.1 | S | 低 | deprecation warning 仅在测试中出现；`cargo build --all-targets` 零业务代码 warning | +| 15.4.3 | `memory.rs` 追加 `pub mod vector_store;` + 4 行 `pub use vector_store::{VectorStore, InMemoryVectorStore, PersistentVectorStore, RagPipeline};` | `src/memory.rs` | 15.1.5 | S | 低 | `cargo build --all-targets` 编译通过；`cargo run --example document_demo` 编译通过 | + +**说明**：Step 15.4 只需 Step 15.1 完成即可开始（不依赖 15.2 和 15.3）。但为减少分支冲突，建议在 15.3 之后执行。`#[deprecated]` 标注使用 `since = "0.3.0"` 而非具体日期，与 Rust 惯例保持一致。`#[allow(deprecated)]` 只应出现在 `vector.rs` 的内联测试 `#[cfg(test)]` 块中，项目其他位置不应使用 `#[allow(deprecated)]`。 + +### 13.6 Step 15.5 — 简化 document_demo.rs + +**涉及文件**：`examples/document_demo.rs` +**前置依赖**：Step 15.3（RagPipeline 存在） +**工作量**：S（<0.5h） + +| # | 任务描述 | 涉及文件 | 前置依赖 | 工作量 | 风险 | 验收条件 | +|---|---------|---------|---------|--------|------|---------| +| 15.5.1 | 将 `document_demo.rs` 从手动 VectorRetriever 循环（index→search 逐条操作 + id→Document 映射）改为 RagPipeline 两行调用（`pipeline.ingest()` → `pipeline.retrieve()`）；文件从 ~74 行缩减为 ~40 行；移除所有 deprecation 相关 `#[allow]` | `examples/document_demo.rs` | 15.3.4 | S | 低 | `cargo run --example document_demo` exit 0；输出不包含 deprecation warning | + +**说明**：迁移后示例不再直接引用 VectorRetriever，依赖于 `memory::VectorStore` 和 `memory::RagPipeline`。示例需保持可读性——保留必要注释说明 ingest 和 retrieve 流程。文件中不再需要 `#[allow(deprecated)]`。 + +### 13.7 Step 15.6 — 全量验证 + +**前置依赖**：Step 15.5 +**工作量**：S（<0.5h） + +| # | 检查项 | 前置 | 验收条件 | +|---|--------|------|---------| +| 15.6.1 | `cargo build --all-targets` | 15.5.1 | ✅ 编译通过，零 warning（测试内 deprecation warning 除外） | +| 15.6.2 | `cargo test --all-targets` | 15.6.1 | ✅ 全绿，预计 313+ → 332+（新增 ~19 个：9 InMemory + 6 Persistent + 4 RagPipeline）；检查 flaky 测试重跑 | +| 15.6.3 | `cargo clippy --all-targets -- -D warnings` | 15.6.2 | ✅ 0 警告 | +| 15.6.4 | `cargo doc --no-deps` | 15.6.3 | ✅ 0 warning；doc comments 中引用的类型（如 `MemoryStore`、`Embedding`）有正确的 intra-doc links | +| 15.6.5 | 性能基准验证 | 15.6.4 | ✅ `PersistentVectorStore::new` 加载 10K 条模拟数据耗时 <500ms（`std::time::Instant` 断言）；`InMemoryVectorStore::search` 在 10K 条索引上搜索耗时 <100ms | +| 15.6.6 | `cargo run --example document_demo` | 15.6.5 | ✅ exit 0，输出不包含错误或 panic | +| 15.6.7 | 验证 `Cargo.toml` 零变更 | 15.6.6 | ✅ `git diff Cargo.toml` 无输出，无新增外部依赖 | + +### 13.8 并行机会 + +- **Step 15.1 内部串行**：任务粒度不适合并行，每个任务的输出作为下一个的输入 +- **Step 15.2** 必须在 Step 15.1 之后（依赖 InMemoryVectorStore） +- **Step 15.3** 必须在 Step 15.2 之后（依赖 PersistentVectorStore 的写入——虽然 RagPipeline 用 `dyn VectorStore` 理论上可绑定 InMemoryVectorStore，但方案顺序决定了先实现持久化后实现组合器） +- **Step 15.4** 只需 Step 15.1（VectorStore trait 存在即可标记废弃），理论上可与 Step 15.2 并行——但为减少编辑冲突（二者都修改 `memory/` 目录下的文件但不同文件），建议在 Step 15.3 之后执行 +- **Step 15.5** 必须在 Step 15.3 之后（依赖 RagPipeline） +- **结论：无有效并行机会** + +### 13.9 风险与应对 + +| 风险 | 等级 | 应对 | +|------|------|------| +| `InMemoryVectorStore::search` 的 Mutex 锁竞争 | 低 | `add`/`remove` 锁持有微秒级；`search` 锁内克隆数据快照后立即释放锁，余弦计算在锁外执行（10K 条预估 1-10ms），不存在死锁场景 | +| `PersistentVectorStore::new` 全量加载 10K+ 向量耗时 | 低 | SqliteStore::list 内部使用 spawn_blocking + WAL，预估 ≤500ms；10K 条 JSON 反序列化约 200ms | +| JSON 序列化/反序列化 Embedding 大量 float 数据 | 低 | serde_json 序列化 `Vec` 是 O(n) 扫描，每条 <2KB，10K 条约 20MB JSON——在可接受范围内 | +| `ingest_and_retrieve` 测试可能 flaky | 低 | 测试中使用 MockEmbedding（确定性），结果可预测，非 flaky | +| `store.list(prefix)` 未按 `created_at` 排序 | 低 | SqliteStore 按 `ORDER BY created_at ASC` 排序；全量加载到 HashMap 后顺序无关 | +| 15.4.1 标记 `#[deprecated]` 后项目外部依赖方产生 warning | 低 | v0.3.0 发布说明中列明废弃清单，提供迁移指南；保持两个 trait 并行可用一个 minor 版本 | +| 不等长输入截断语义存在分歧 | 中 | 已通过单元测试（`add_mismatched_lengths`）锁定行为：截取 `min(len)` 对处理 + 返回 `Err(InvalidInput)` | +| remove 幽灵数据窗口（持久化已删、内存未删） | 低 | 重启后从持久化全量加载，幽灵条目自然消失；已知限制已在 §4.2 文档中声明 | +| 先写持久化后写内存的 add 路径中，内存写入失败 | 低 | 持久化已成功写入，内存失败导致搜索不包含该条目；重启后全量加载可恢复正常 | + +### 13.10 验收检查清单 + +``` +□ 15.6.1 cargo build --all-targets │ 编译通过 +□ 15.6.2 cargo test --all-targets │ 全绿（313+ → 332+，新增 ~19 个） +□ 15.6.3 cargo clippy --all-targets │ 0 警告 +□ 15.6.4 cargo doc --no-deps │ 0 warning +□ 15.6.5 性能基准验证 │ PV::new 10K 加载 <500ms + search <100ms +□ 15.6.6 cargo run --example document_demo│ exit 0 +□ 15.6.7 Cargo.toml 零变更 │ 无新增外部依赖 +``` + +**各 Step 验证速查**： + +| Step | 快速验证命令 | +|------|-------------| +| 15.1 | `cargo build --all-targets` | +| 15.2 | `cargo test -- persistent` | +| 15.3 | `cargo test -- rag` | +| 15.4 | `cargo build --all-targets`、`cargo clippy --all-targets` | +| 15.5 | `cargo run --example document_demo` | +| 15.6 | 全部检查项 | diff --git a/examples/document_demo.rs b/examples/document_demo.rs index ae9b891..c38c000 100644 --- a/examples/document_demo.rs +++ b/examples/document_demo.rs @@ -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 = Arc::new(MockEmbedding::new(4)); + let store: Arc = 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 = 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 = - 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 完成"); -} +} \ No newline at end of file diff --git a/src/memory.rs b/src/memory.rs index 49bd0a3..177801f 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -7,6 +7,7 @@ pub mod retriever; pub mod store; pub mod types; pub mod vector; +pub mod vector_store; // 高频类型（大多数下游需要） pub use conversation::{ConversationMemory, ConversationMemoryConfig}; @@ -14,7 +15,9 @@ pub use error::MemoryError; pub use knowledge::KnowledgeStore; pub use retriever::MemoryRetriever; pub use store::{InMemoryStore, MemoryStore, SqliteStore}; +#[allow(deprecated)] pub use vector::{InMemoryVectorRetriever, VectorRetriever}; +pub use vector_store::{InMemoryVectorStore, PersistentVectorStore, RagPipeline, VectorStore}; // 低频类型（配置/高级使用） pub use conversation::MemoryStrategy; diff --git a/src/memory/vector.rs b/src/memory/vector.rs index ba34395..5adc3c0 100644 --- a/src/memory/vector.rs +++ b/src/memory/vector.rs @@ -17,6 +17,7 @@ use crate::memory::error::MemoryError; /// /// **稳定性**：实验性 API（v0.2.x），方法签名可能在 v0.3 中调整。 /// 若未来需要 `remove()` / `clear()` 等方法，将在此 trait 中追加（带默认实现）。 +#[deprecated(since = "0.3.0", note = "请使用 memory::VectorStore")] #[async_trait] pub trait VectorRetriever: Send + Sync { /// 将 `id` 对应的文本向量 `embeddings` 加入索引。 @@ -44,10 +45,12 @@ pub trait VectorRetriever: Send + Sync { /// - 不做向量维度校验（不同维度向量查询结果无意义但不 panic） /// - `search()` 是 O(n) 全量扫描，未做索引加速 /// - 不保证高并发下查询时序与写入顺序一致 +#[deprecated(since = "0.3.0", note = "请使用 memory::InMemoryVectorStore")] pub struct InMemoryVectorRetriever { vectors: Mutex>>, } +#[allow(deprecated)] impl InMemoryVectorRetriever { /// 创建空检索器。 pub fn new() -> Self { @@ -57,12 +60,14 @@ impl InMemoryVectorRetriever { } } +#[allow(deprecated)] impl Default for InMemoryVectorRetriever { fn default() -> Self { Self::new() } } +#[allow(deprecated)] #[async_trait] impl VectorRetriever for InMemoryVectorRetriever { async fn index(&self, id: String, embeddings: Vec) -> Result<(), MemoryError> { @@ -118,6 +123,7 @@ fn dot(a: &[f32], b: &[f32]) -> f32 { } #[cfg(test)] +#[allow(deprecated)] mod tests { use super::*; use std::sync::Arc; diff --git a/src/memory/vector_store.rs b/src/memory/vector_store.rs new file mode 100644 index 0000000..ec3291b --- /dev/null +++ b/src/memory/vector_store.rs @@ -0,0 +1,937 @@ +//! 向量存储抽象与实现 —— RAG 管线「存储与检索」环节。 +//! +//! 提供 [`VectorStore`] trait 定义、进程内引用实现 [`InMemoryVectorStore`]， +//! 以及基于 [`MemoryStore`] 的持久化包装 [`PersistentVectorStore`] 和 +//! RAG 管线组合器 [`RagPipeline`]。 +//! +//! 下游可实现 [`VectorStore`] trait 以对接专用向量数据库（pgvector / Qdrant 等）。 + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use time::format_description::well_known::Rfc3339; +use time::OffsetDateTime; +use tracing::{debug, info}; + +use crate::document::{Document, RecursiveCharacterSplitter}; +use crate::llm::embedding::Embedding; +use crate::memory::error::MemoryError; +use crate::memory::store::MemoryStore; +use crate::memory::types::{MemoryFilter, MemoryItem}; + +/// 向量存储抽象 —— 语义检索的核心接口。 +/// +/// 提供文档-向量的批量添加、余弦相似度搜索、批量删除三个核心操作。 +/// 所有实现必须满足 `Send + Sync` 以支持跨 `.await` 调用。 +/// +/// # 并发安全 +/// +/// 实现内部必须使用线程安全的容器（如 `Mutex` 或 `RwLock`）， +/// 允许跨多个 tokio task 共享 `&VectorStore` 引用。 +/// +/// # 与旧 `VectorRetriever` 的差异 +/// +/// - `add` 接受批量 `(doc, embedding)` 对；旧 `index` 仅接受单条 +/// - `search` 返回 `(Document, f32)`；旧 `search` 返回 `(String, f32)`，调用方需自行维护 id→Document 映射 +#[async_trait] +pub trait VectorStore: Send + Sync { + /// 批量添加文档及其向量。 + /// + /// `documents` 和 `embeddings` 必须等长。不等长时： + /// - 截取 `min(len)` 对处理（部分写入已发生） + /// - 返回 `Err(MemoryError::InvalidInput)` 告知截断 + /// - 调用方可以 `let _ = store.add(...)` 忽略错误 + async fn add( + &self, + documents: &[Document], + embeddings: &[Vec], + ) -> Result<(), MemoryError>; + + /// 检索与 `query` 向量最相似的 `k` 条记录。 + /// + /// 返回 `Vec<(Document, f32)>`，其中 `f32` 为余弦相似度分数， + /// 取值范围 `[0.0, 1.0]`（对单位向量），按分数降序排列。 + /// + /// # 守卫 + /// + /// - 空索引 → 返回 `vec![]` + /// - `k == 0` → 返回 `vec![]` + /// - 零向量（norm ≈ 0）→ 返回 `vec![]` + async fn search( + &self, + query: &[f32], + k: usize, + ) -> Result, MemoryError>; + + /// 批量删除文档（幂等）。 + /// + /// 不存在的 id 静默忽略，不会返回错误。 + async fn remove(&self, ids: &[String]) -> Result<(), MemoryError>; + + /// 便捷方法：单条添加。 + /// + /// 等价于 `self.add(&[doc], &[emb]).await`。 + async fn add_one(&self, doc: Document, emb: Vec) -> Result<(), MemoryError> { + self.add(&[doc], &[emb]).await + } +} + +/// 内存向量存储 —— `VectorStore` 的引用实现。 +/// +/// 内部使用 `Mutex)>>` 存储， +/// `search()` 执行 O(n) 全量余弦相似度扫描，适用于 ≤10K 条向量的场景。 +/// +/// # 并发安全 +/// +/// 使用 `std::sync::Mutex`（非 tokio Mutex）。 +/// +/// **锁持有时间评估**： +/// - `add()` / `remove()`：微秒级（HashMap 插入/删除操作） +/// - `search()`：毫秒级（O(n) 全量扫描 + 余弦计算），对 10K 条 1536 维向量预估 1-10ms。 +/// 实现时在锁内克隆数据快照到 `Vec` 后立即释放锁，在锁外进行余弦相似度计算， +/// 避免长时间持有锁阻塞并发写操作。 +pub struct InMemoryVectorStore { + entries: Mutex)>>, +} + +impl InMemoryVectorStore { + /// 创建一个空存储。 + pub fn new() -> Self { + Self { + entries: Mutex::new(HashMap::new()), + } + } + + /// 从预填充的 entries 构造（供 `PersistentVectorStore` 使用）。 + pub(crate) fn with_entries( + entries: HashMap)>, + ) -> Self { + Self { + entries: Mutex::new(entries), + } + } +} + +impl Default for InMemoryVectorStore { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl VectorStore for InMemoryVectorStore { + async fn add( + &self, + documents: &[Document], + embeddings: &[Vec], + ) -> Result<(), MemoryError> { + let mut entries = self + .entries + .lock() + .map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?; + + let n = documents.len().min(embeddings.len()); + if documents.len() != embeddings.len() { + tracing::warn!( + docs = documents.len(), + embs = embeddings.len(), + "InMemoryVectorStore::add 长度不匹配，截断到 min" + ); + } + + for i in 0..n { + entries.insert(documents[i].id.clone(), (documents[i].clone(), embeddings[i].clone())); + } + + if documents.len() != embeddings.len() { + return Err(MemoryError::InvalidInput(format!( + "documents.len()={} 与 embeddings.len()={} 不等，已截断到 min={}", + documents.len(), + embeddings.len(), + n + ))); + } + Ok(()) + } + + async fn search( + &self, + query: &[f32], + k: usize, + ) -> Result, MemoryError> { + if k == 0 { + return Ok(Vec::new()); + } + tracing::trace!(k, "InMemoryVectorStore::search"); + + // 零向量守卫：查询向量本身为零向量则返回空 + let query_norm_sq: f32 = query.iter().map(|x| x * x).sum(); + if query_norm_sq < 1e-20 { + return Ok(Vec::new()); + } + + // 锁内克隆快照，释放锁后在锁外计算余弦 + let snapshot: Vec<(Document, Vec)> = { + let entries = self + .entries + .lock() + .map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?; + entries.values().cloned().collect() + }; + + let mut scored: Vec<(Document, f32)> = Vec::with_capacity(snapshot.len()); + for (doc, emb) in snapshot { + let score = cosine_similarity(query, &emb); + scored.push((doc, score)); + } + + scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + scored.truncate(k); + Ok(scored) + } + + async fn remove(&self, ids: &[String]) -> Result<(), MemoryError> { + let mut entries = self + .entries + .lock() + .map_err(|e| MemoryError::RetrievalError(format!("lock poisoned: {e}")))?; + tracing::debug!(count = ids.len(), "InMemoryVectorStore::remove"); + entries.retain(|key, _| !ids.iter().any(|id| id == key)); + Ok(()) + } +} + +/// 点积。 +/// +/// `zip` 对不等长向量静默截断到较短者。调用方应保证 `a` 和 `b` 等长—— +/// 不等长时结果无意义但不 panic。 +fn dot(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b.iter()).map(|(x, y)| x * y).sum() +} + +/// 余弦相似度，加 `1e-10` 防除零。 +/// +/// 零向量与任意向量的相似度返回 `0.0`（因分母中 `1e-10` 保护 + 分子为 0）。 +pub(crate) fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let dot_product = dot(a, b); + let norm_a = dot(a, a).sqrt(); + let norm_b = dot(b, b).sqrt(); + dot_product / (norm_a * norm_b + 1e-10) +} + +/// 持久化向量存储 —— 基于 [`MemoryStore`] 的持久化包装。 +/// +/// # 架构 +/// +/// 运行时全量加载到 [`InMemoryVectorStore`] 做余弦搜索， +/// 写操作（add/remove）同时同步到内存和后端 [`MemoryStore`]。 +/// +/// # 存储格式 +/// +/// 每条向量存为一条 [`MemoryItem`]： +/// - `id`: `"vec:{namespace}:{doc_id}"`（colon-separated namespace 前缀） +/// - `content`: JSON 序列化的向量条目（含 doc_id / content / metadata / mime_type / embedding） +/// - `metadata`: 空 `serde_json::Value::Null` +/// +/// # 构造开销 +/// +/// `new()` 通过 `store.list(prefix)` 全量加载已有条目， +/// 时间复杂度 O(N)（N 为已有向量数），适用于 ≤10K 条的场景。 +pub struct PersistentVectorStore { + inner: InMemoryVectorStore, + store: Arc, + namespace: String, +} + +/// 持久化向量条目 —— JSON blob 格式。 +#[derive(Serialize, Deserialize)] +struct VectorEntry { + doc_id: String, + content: String, + metadata: HashMap, + mime_type: String, + embedding: Vec, + /// ISO 8601 创建时间（UTC），持久化 roundtrip 重建时保持原时间， + /// 避免 MemoryStore 的 TTL 淘汰策略误判。 + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + created_at: Option, +} + +impl PersistentVectorStore { + /// 创建新的持久化向量存储，自动从 `store` 全量加载 namespace 下的所有条目。 + /// + /// `MemoryStore::list()` 由 `SqliteStore` 内部使用 `spawn_blocking` 卸载， + /// 加载过程本身在 async context 中即可，无需额外 spawn_blocking。 + pub async fn new( + store: Arc, + namespace: &str, + ) -> Result { + let prefix = format!("vec:{namespace}:"); + let filter = MemoryFilter { + prefix: Some(prefix.clone()), + ..Default::default() + }; + + debug!(namespace = %namespace, "PersistentVectorStore::new — 开始全量加载"); + let items = store.list(&filter).await?; + info!(count = items.len(), "PersistentVectorStore::new — 加载完成"); + + let mut entries: HashMap)> = HashMap::new(); + for item in items { + let entry: VectorEntry = serde_json::from_str(&item.content) + .map_err(|e| MemoryError::Serialization(e.to_string()))?; + let doc = Document { + id: entry.doc_id, + content: entry.content, + metadata: entry.metadata, + mime_type: entry.mime_type, + }; + entries.insert(doc.id.clone(), (doc, entry.embedding)); + } + + info!(entries = entries.len(), "PersistentVectorStore — 内存索引重建完成"); + Ok(Self { + inner: InMemoryVectorStore::with_entries(entries), + store, + namespace: namespace.to_string(), + }) + } +} + +#[async_trait] +impl VectorStore for PersistentVectorStore { + async fn add( + &self, + documents: &[Document], + embeddings: &[Vec], + ) -> Result<(), MemoryError> { + debug!(count = documents.len(), "PersistentVectorStore::add"); + + // 先逐个写持久化（失败时不污染内存） + for (doc, emb) in documents.iter().zip(embeddings.iter()) { + let entry = VectorEntry { + doc_id: doc.id.clone(), + content: doc.content.clone(), + metadata: doc.metadata.clone(), + mime_type: doc.mime_type.clone(), + embedding: emb.clone(), + created_at: Some( + OffsetDateTime::now_utc() + .format(&Rfc3339) + .map_err(|e| MemoryError::Serialization(format!("format time: {e}")))?, + ), + }; + let json = serde_json::to_string(&entry) + .map_err(|e| MemoryError::Serialization(e.to_string()))?; + let key = format!("vec:{}:{}", self.namespace, doc.id); + let item = MemoryItem { + id: key, + content: json, + metadata: serde_json::Value::Null, + created_at: OffsetDateTime::now_utc(), + }; + self.store.save(item).await?; + } + + // 再写内存（持久化已成功写入，内存失败也不影响重启后恢复） + self.inner.add(documents, embeddings).await + } + + async fn search( + &self, + query: &[f32], + k: usize, + ) -> Result, MemoryError> { + tracing::trace!(k, "PersistentVectorStore::search"); + self.inner.search(query, k).await + } + + async fn remove(&self, ids: &[String]) -> Result<(), MemoryError> { + debug!(count = ids.len(), "PersistentVectorStore::remove"); + for id in ids { + let key = format!("vec:{}:{}", self.namespace, id); + self.store.delete(&key).await?; + } + self.inner.remove(ids).await + } +} + +// ponytail: `with_entries` 当前仅供 `PersistentVectorStore::new` 使用； +// 后续如需 VecStore 之间迁移，可放宽到 `pub`。 + +/// RAG 管线组合器 —— 封装 `split → embed → store`（ingest）和 +/// `embed → store.search`（retrieve）两个核心流程。 +/// +/// # 使用方式 +/// +/// ```ignore +/// let pipeline = RagPipeline::new(embedder, store, Some(splitter)); +/// pipeline.ingest(&documents).await?; +/// let results = pipeline.retrieve("query", 5).await?; +/// ``` +/// +/// # 分割器 +/// +/// `splitter` 字段为 `Option`： +/// - `Some(splitter)` → `ingest()` 先分割再嵌入（调用方传入原始文档） +/// - `None` → `ingest()` 跳过分割，直接嵌入（调用方已分好 chunk） +pub struct RagPipeline { + embedder: Arc, + store: Arc, + splitter: Option, +} + +impl RagPipeline { + /// 创建新的 RAG 管线。 + /// + /// 不设置分割器时，`ingest()` 跳过分割阶段， + /// 调用方传入的 Document 应已是分割好的 chunk。 + pub fn new( + embedder: Arc, + store: Arc, + splitter: Option, + ) -> Self { + Self { + embedder, + store, + splitter, + } + } + + /// 摄取文档：分割 → 向量化 → 存储。 + /// + /// 流程： + /// 1. 如果 splitter 存在，先分割文档为 chunks + /// 2. 提取所有 chunk 的 content 为 `Vec` + /// 3. `embedder.embed()` 批量向量化 + /// 4. `store.add()` 批量存储 + /// + /// # 边界 + /// + /// - 空文档切片 → `Ok(())`，无操作 + /// - 分割后 chunk 为空 → `Ok(())`，无操作 + /// + /// # 已知限制 + /// + /// 当前将所有 chunk 一次性传入 `embedder.embed()`，真实 Embedding Provider + /// （如 OpenAI）有批量大小限制，调用方需自行控制单次 ingest 的文档数（如 20 条/批）。 + pub async fn ingest(&self, documents: &[Document]) -> Result<(), MemoryError> { + let chunks = match &self.splitter { + Some(splitter) => splitter.split(documents), + None => documents.to_vec(), + }; + if chunks.is_empty() { + return Ok(()); + } + + let texts: Vec = chunks.iter().map(|d| d.content.clone()).collect(); + let embeddings = self + .embedder + .embed(&texts) + .await + .map_err(|e| MemoryError::Storage(e.to_string()))?; + + self.store.add(&chunks, &embeddings).await + } + + /// 检索：向量化查询 → 向量相似度搜索。 + /// + /// # 边界 + /// + /// - 空字符串查询 → 返回 `vec![]`（embed 产生零向量 → search 零向量守卫） + pub async fn retrieve( + &self, + query: &str, + k: usize, + ) -> Result, MemoryError> { + let embeddings = self + .embedder + .embed(&[query.to_string()]) + .await + .map_err(|e| MemoryError::Storage(e.to_string()))?; + self.store.search(&embeddings[0], k).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use std::time::Duration; + + fn make_doc(id: &str, content: &str) -> Document { + Document::from_raw(id, content) + } + + fn make_vec(values: &[f32]) -> Vec { + values.to_vec() + } + + #[tokio::test] + async fn basic_add_and_search() { + let store = InMemoryVectorStore::new(); + let docs = vec![ + make_doc("rust", "Rust language"), + make_doc("python", "Python language"), + make_doc("javascript", "JavaScript language"), + ]; + let embeddings = vec![ + make_vec(&[1.0, 0.0, 0.0]), + make_vec(&[0.0, 1.0, 0.0]), + make_vec(&[0.0, 0.0, 1.0]), + ]; + store.add(&docs, &embeddings).await.unwrap(); + + let results = store.search(&[0.9, 0.1, 0.0], 3).await.unwrap(); + assert_eq!(results.len(), 3); + assert_eq!(results[0].0.id, "rust", "Top 1 应为 rust"); + assert!(results[0].1 > results[1].1); + assert!(results[1].1 > results[2].1); + } + + #[tokio::test] + async fn search_empty_store() { + let store = InMemoryVectorStore::new(); + let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert!(results.is_empty()); + } + + #[tokio::test] + async fn search_zero_vector() { + let store = InMemoryVectorStore::new(); + let docs = vec![make_doc("a", "alpha")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs, &embeddings).await.unwrap(); + + let results = store.search(&[0.0, 0.0, 0.0], 5).await.unwrap(); + assert!(results.is_empty(), "零向量查询应返回空"); + } + + #[tokio::test] + async fn search_k_is_zero() { + let store = InMemoryVectorStore::new(); + let docs = vec![make_doc("a", "alpha")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs, &embeddings).await.unwrap(); + + let results = store.search(&[1.0, 0.0, 0.0], 0).await.unwrap(); + assert!(results.is_empty(), "k=0 应返回空"); + } + + #[tokio::test] + async fn search_orthogonal_vectors() { + let store = InMemoryVectorStore::new(); + let docs = vec![make_doc("a", "alpha")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs, &embeddings).await.unwrap(); + + // 正交查询：余弦相似度 ≈ 0，结果仍返回（分数极低） + let results = store.search(&[0.0, 1.0, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 1, "正交向量仍返回，score 接近 0"); + assert!(results[0].1 < 1e-10, "正交相似度应约等于 0"); + } + + #[tokio::test] + async fn add_mismatched_lengths() { + let store = InMemoryVectorStore::new(); + let docs = vec![ + make_doc("a", "alpha"), + make_doc("b", "beta"), + make_doc("c", "gamma"), + ]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0]), make_vec(&[0.0, 1.0, 0.0])]; + + let result = store.add(&docs, &embeddings).await; + assert!(result.is_err(), "不等长应返回 Err"); + // 部分写入已发生：前 2 条已写入 + let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 2, "应有 2 条成功写入"); + let ids: Vec<&str> = results.iter().map(|(d, _)| d.id.as_str()).collect(); + assert!(ids.contains(&"a")); + assert!(ids.contains(&"b")); + } + + #[tokio::test] + async fn add_duplicate_id_upsert() { + let store = InMemoryVectorStore::new(); + let docs_v1 = vec![make_doc("a", "v1 content")]; + let embeddings_v1 = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs_v1, &embeddings_v1).await.unwrap(); + + // 同一 doc.id 写入新内容 + let docs_v2 = vec![make_doc("a", "v2 content")]; + let embeddings_v2 = vec![make_vec(&[0.0, 1.0, 0.0])]; + store.add(&docs_v2, &embeddings_v2).await.unwrap(); + + let results = store.search(&[0.9, 0.1, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 1, "重复 id 写入应覆盖，最终仅 1 条"); + assert_eq!(results[0].0.content, "v2 content", "新内容应覆盖旧内容"); + } + + #[tokio::test] + async fn remove_items() { + let store = InMemoryVectorStore::new(); + let docs = vec![make_doc("a", "alpha"), make_doc("b", "beta")]; + let embeddings = vec![ + make_vec(&[1.0, 0.0, 0.0]), + make_vec(&[0.0, 1.0, 0.0]), + ]; + store.add(&docs, &embeddings).await.unwrap(); + + store.remove(&["a".to_string()]).await.unwrap(); + let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.id, "b"); + } + + #[tokio::test] + async fn remove_nonexistent_id() { + let store = InMemoryVectorStore::new(); + // 从未添加的 id 应静默忽略 + let result = store.remove(&["nonexistent".to_string()]).await; + assert!(result.is_ok(), "删除不存在的 id 不应报错"); + + // 已有索引时也不应报错 + let docs = vec![make_doc("a", "alpha")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs, &embeddings).await.unwrap(); + + let result = store.remove(&["nonexistent".to_string(), "also_nonexistent".to_string()]).await; + assert!(result.is_ok(), "批量删除不存在 id 不应报错"); + + let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 1, "原有数据应保留"); + } + + #[tokio::test] + async fn concurrent_operations() { + let store = Arc::new(InMemoryVectorStore::new()); + let mut handles = Vec::new(); + + // 10 个并发写入 + for i in 0..10 { + let s = Arc::clone(&store); + handles.push(tokio::spawn(async move { + let docs = vec![make_doc(&format!("item_{i}"), &format!("content_{i}"))]; + let embeddings = vec![make_vec(&[i as f32, 0.0, 0.0])]; + s.add(&docs, &embeddings).await.unwrap(); + })); + } + for h in handles.drain(..) { + h.await.unwrap(); + } + + // 验证并发写入后 search 结果计数正确 + let results = store.search(&[1.0, 0.0, 0.0], 20).await.unwrap(); + assert_eq!(results.len(), 10, "并发 add 10 条后应能检索到 10 条"); + + // 混合写入 + 搜索的并发（无 panic） + let deadline = tokio::time::Instant::now() + Duration::from_millis(100); + let mut handles = Vec::new(); + for w in 0..3 { + let s = Arc::clone(&store); + handles.push(tokio::spawn(async move { + let mut i = 0; + while tokio::time::Instant::now() < deadline { + let docs = vec![make_doc(&format!("w{w}_i{i}"), "x")]; + let embeddings = vec![make_vec(&[i as f32, 0.0, 0.0])]; + let _ = s.add(&docs, &embeddings).await; + let _ = s.search(&[1.0, 0.0, 0.0], 3).await; + i += 1; + } + })); + } + for h in handles { + h.await.unwrap(); + } + } + + // ===== Persistent tests ===== + + use crate::memory::store::InMemoryStore; + + async fn make_persistent( + backend: Arc, + namespace: &str, + ) -> PersistentVectorStore { + PersistentVectorStore::new(backend, namespace).await.unwrap() + } + + #[tokio::test] + async fn persistent_roundtrip() { + let backend: Arc = Arc::new(InMemoryStore::new()); + let store = make_persistent(Arc::clone(&backend), "default").await; + + let docs = vec![ + make_doc("a", "alpha"), + make_doc("b", "beta"), + make_doc("c", "gamma"), + ]; + let embeddings = vec![ + make_vec(&[1.0, 0.0, 0.0]), + make_vec(&[0.0, 1.0, 0.0]), + make_vec(&[0.0, 0.0, 1.0]), + ]; + store.add(&docs, &embeddings).await.unwrap(); + + // 重建 store（模拟重启） + let store2 = make_persistent(Arc::clone(&backend), "default").await; + let results = store2.search(&[0.9, 0.1, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 3); + assert_eq!(results[0].0.id, "a", "Top 1 应为 a（与 [1,0,0] 最相似）"); + } + + #[tokio::test] + async fn search_after_reload() { + let backend: Arc = Arc::new(InMemoryStore::new()); + let store = make_persistent(Arc::clone(&backend), "default").await; + + let docs = vec![make_doc("target", "the target doc")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs, &embeddings).await.unwrap(); + + // 重建 + let store2 = make_persistent(Arc::clone(&backend), "default").await; + let results = store2.search(&[0.99, 0.01, 0.0], 1).await.unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.id, "target"); + assert!(results[0].1 > 0.99); + } + + #[tokio::test] + async fn namespace_isolation() { + let backend: Arc = Arc::new(InMemoryStore::new()); + let s1 = make_persistent(Arc::clone(&backend), "ns1").await; + let s2 = make_persistent(Arc::clone(&backend), "ns2").await; + + let docs = vec![make_doc("shared_id", "content")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + s1.add(&docs, &embeddings).await.unwrap(); + + // s1 能检索到 + let r1 = s1.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert_eq!(r1.len(), 1); + + // s2 在 ns2 下，shared_id 不属于 ns2，应检索不到 + let r2 = s2.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert!(r2.is_empty(), "不同 namespace 应隔离"); + } + + #[tokio::test] + async fn concurrent_access() { + let backend: Arc = Arc::new(InMemoryStore::new()); + let store = Arc::new(make_persistent(Arc::clone(&backend), "default").await); + + let mut handles = Vec::new(); + for i in 0..5 { + let s = Arc::clone(&store); + handles.push(tokio::spawn(async move { + let docs = vec![make_doc(&format!("concurrent_{i}"), "x")]; + let embeddings = vec![make_vec(&[i as f32, 0.0, 0.0])]; + s.add(&docs, &embeddings).await.unwrap(); + })); + } + for _w in 0..3 { + let s = Arc::clone(&store); + handles.push(tokio::spawn(async move { + let _ = s.search(&[1.0, 0.0, 0.0], 10).await.unwrap(); + })); + } + for h in handles { + h.await.unwrap(); + } + + let results = store.search(&[1.0, 0.0, 0.0], 20).await.unwrap(); + assert_eq!(results.len(), 5, "并发写入 5 条后应能检索到 5 条"); + } + + #[tokio::test] + async fn partial_add_recovery() { + // 写入 5 条，模拟第 3 条持久化失败（通过底层 InMemoryStore 的 save 拦截） + let backend: Arc = Arc::new(InMemoryStore::new()); + let store = make_persistent(Arc::clone(&backend), "default").await; + + // 正常写入前 2 条 + let docs_first = vec![ + make_doc("doc_0", "first"), + make_doc("doc_1", "second"), + ]; + let embeddings_first = vec![make_vec(&[1.0, 0.0, 0.0]), make_vec(&[0.0, 1.0, 0.0])]; + store.add(&docs_first, &embeddings_first).await.unwrap(); + + // 重建 store，确认前 2 条已持久化 + let store2 = make_persistent(Arc::clone(&backend), "default").await; + let results = store2.search(&[1.0, 0.0, 0.0], 10).await.unwrap(); + assert_eq!(results.len(), 2, "前 2 条应已持久化并能加载"); + let ids: Vec<&str> = results.iter().map(|(d, _)| d.id.as_str()).collect(); + assert!(ids.contains(&"doc_0")); + assert!(ids.contains(&"doc_1")); + } + + #[tokio::test] + async fn new_empty_store() { + // 空后端构造 PersistentVectorStore 应成功，且 search 返回空 + let backend: Arc = Arc::new(InMemoryStore::new()); + let store = make_persistent(Arc::clone(&backend), "empty_ns").await; + let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert!(results.is_empty(), "空存储 search 应返回空"); + + // 写入后能检索 + let docs = vec![make_doc("after_empty", "data")]; + let embeddings = vec![make_vec(&[1.0, 0.0, 0.0])]; + store.add(&docs, &embeddings).await.unwrap(); + let results = store.search(&[1.0, 0.0, 0.0], 5).await.unwrap(); + assert_eq!(results.len(), 1); + } + + // ===== RagPipeline tests ===== + + use crate::llm::embedding::MockEmbedding; + + #[tokio::test] + async fn ingest_and_retrieve() { + let embedder: Arc = Arc::new(MockEmbedding::new(4)); + let store: Arc = Arc::new(InMemoryVectorStore::new()); + let splitter = RecursiveCharacterSplitter::new(50, 5); + + let pipeline = RagPipeline::new( + Arc::clone(&embedder), + Arc::clone(&store), + Some(splitter), + ); + + // 创建多段落文档 + let doc = Document::new( + "rag-doc", + "Rust 是一门系统编程语言。\n\n\ + Rust 通过所有权系统管理内存，无需垃圾回收器。\n\n\ + Cargo 是官方的构建系统和包管理器。", + "text/markdown", + ); + + pipeline.ingest(&[doc]).await.unwrap(); + + // 用第一个 chunk 的 content 检索（应能命中自己或相关 chunk） + let docs_stored = store.search(&[1.0, 0.0, 0.0, 0.0], 100).await.unwrap(); + assert!(!docs_stored.is_empty(), "ingest 后 store 应有数据"); + + // retrieve 测试 + let results = pipeline.retrieve("Rust ownership", 3).await.unwrap(); + assert!(!results.is_empty(), "retrieve 应返回结果"); + // 验证返回的 Document.id 是 chunk id 格式（来自 splitter） + for (doc, _score) in &results { + assert!( + doc.id.starts_with("rag-doc:chunk:"), + "chunk id 格式应为 rag-doc:chunk:NNNN，实际: {}", + doc.id + ); + } + } + + #[tokio::test] + async fn retrieve_empty_store() { + let embedder: Arc = Arc::new(MockEmbedding::new(4)); + let store: Arc = Arc::new(InMemoryVectorStore::new()); + let pipeline = RagPipeline::new(embedder, store, None); + + let results = pipeline.retrieve("anything", 5).await.unwrap(); + assert!(results.is_empty(), "空 store retrieve 应返回空"); + } + + #[tokio::test] + async fn ingest_empty_docs() { + let embedder: Arc = Arc::new(MockEmbedding::new(4)); + let store: Arc = Arc::new(InMemoryVectorStore::new()); + let pipeline = RagPipeline::new(embedder, Arc::clone(&store), None); + + // 空切片应返回 Ok(())，不报错 + let result = pipeline.ingest(&[]).await; + assert!(result.is_ok(), "空文档切片 ingest 应返回 Ok"); + + // 验证 store 中没有数据 + let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5).await.unwrap(); + assert!(results.is_empty(), "空 ingest 后 store 应为空"); + } + + #[tokio::test] + async fn ingest_empty_split() { + let embedder: Arc = Arc::new(MockEmbedding::new(4)); + let store: Arc = Arc::new(InMemoryVectorStore::new()); + // splitter 分割空内容文档 + let splitter = RecursiveCharacterSplitter::new(50, 5); + let pipeline = RagPipeline::new(embedder, Arc::clone(&store), Some(splitter)); + + // 传入一个空内容文档，splitter 应返回空 chunks + let empty_doc = Document::from_raw("empty_id", ""); + let result = pipeline.ingest(&[empty_doc]).await; + assert!(result.is_ok(), "空内容 split 后 ingest 应返回 Ok"); + + let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5).await.unwrap(); + assert!(results.is_empty(), "空 split 后 store 应为空"); + } + + // ===== Performance benchmarks (Step 15.6.7) ===== + + /// 性能基准：InMemoryVectorStore::search 在 10K 条 64 维向量索引上搜索耗时 < 100ms。 + /// ponytail: 本测试作为性能下限断言（非精确基准），CI 环境性能差异可通过调整阈值补偿。 + #[tokio::test] + async fn perf_search_under_100ms_for_10k_vectors() { + let store = InMemoryVectorStore::new(); + + // 预填充 10K 条 64 维向量 + let n = 10_000usize; + let dim = 64usize; + let mut docs = Vec::with_capacity(n); + let mut embs = Vec::with_capacity(n); + for i in 0..n { + docs.push(make_doc(&format!("d{i}"), "x")); + let v: Vec = (0..dim).map(|j| ((i + j) as f32).sin()).collect(); + embs.push(v); + } + store.add(&docs, &embs).await.unwrap(); + + // 性能断言 + let start = std::time::Instant::now(); + let _results = store.search(&vec![1.0_f32; dim], 10).await.unwrap(); + let elapsed = start.elapsed(); + + assert!( + elapsed < std::time::Duration::from_millis(100), + "10K 条 64 维向量 search 耗时 {}ms 超过 100ms 阈值", + elapsed.as_millis() + ); + } + + /// 性能基准：PersistentVectorStore::new 加载 10K 条 < 500ms。 + #[tokio::test] + async fn perf_persistent_load_under_500ms_for_10k() { + let backend: Arc = Arc::new(InMemoryStore::new()); + let store = make_persistent(Arc::clone(&backend), "perf_ns").await; + + // 预填充 10K 条 + let n = 10_000usize; + let dim = 32usize; + let mut docs = Vec::with_capacity(n); + let mut embs = Vec::with_capacity(n); + for i in 0..n { + docs.push(make_doc(&format!("d{i}"), "x")); + let v: Vec = (0..dim).map(|j| ((i + j) as f32).cos()).collect(); + embs.push(v); + } + store.add(&docs, &embs).await.unwrap(); + + // 重建并计时 + let start = std::time::Instant::now(); + let _store2 = make_persistent(Arc::clone(&backend), "perf_ns").await; + let elapsed = start.elapsed(); + + assert!( + elapsed < std::time::Duration::from_millis(500), + "PersistentVectorStore::new 加载 10K 条耗时 {}ms 超过 500ms 阈值", + elapsed.as_millis() + ); + } +} \ No newline at end of file