|
| 1 | +//! kb 向量回填存储（REQ-259，v0.19.5）。 |
| 2 | +//! |
| 3 | +//! @ai-context: 派生索引铁律：kb_* 全部可由 reindex_all 重建——向量列亦然。 |
| 4 | +//! 全量重建（kb_reindex_all 命令）成功后由本模块做后置回填： |
| 5 | +//! 拉取全部 chunk 文本 → 引擎批嵌入 → 单事务写回 embedding 列 + |
| 6 | +//! kb_meta 元数据（model/dim/format，检索合流的 dim 校验数据源）。 |
| 7 | +//! @ai-context: 软重建钩子（保存/删除）不增量补向量——语义召回一致性依赖 |
| 8 | +//! 全量重建（REQ-262 UI 按钮）；未回填行 embedding=NULL，语义 |
| 9 | +//! 合流自然缺席、FTS 精度不受影响（诚实降级）。 |
| 10 | +
|
| 11 | +use crate::db::Db; |
| 12 | +use crate::error::Result; |
| 13 | +use crate::kb_embed::{ |
| 14 | + EmbeddingEngine, FORMAT_F32LE, META_DIM, META_FORMAT, META_MODEL, encode_embedding, |
| 15 | +}; |
| 16 | +use crate::kb_index::meta_set; |
| 17 | + |
| 18 | +impl Db { |
| 19 | + /// 全量向量回填：全部已切块文本 → 引擎嵌入 → 写列 + 元数据（幂等—— |
| 20 | + /// reindex 后 embedding 全 NULL，重跑即覆盖；引擎失败 → Err 保持可诊断）。 |
| 21 | + pub fn kb_fill_embeddings(&self, engine: &dyn EmbeddingEngine) -> Result<usize> { |
| 22 | + let dim = engine |
| 23 | + .dims() |
| 24 | + .ok_or_else(|| crate::error::AppError::Db("引擎不可用（无 dim）".to_string()))?; |
| 25 | + let chunks: Vec<(i64, String)> = self.with_conn(|conn| { |
| 26 | + let mut stmt = conn.prepare("SELECT id, text FROM kb_chunks ORDER BY id")?; |
| 27 | + let rows = stmt.query_map([], |r| Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?)))?; |
| 28 | + rows.collect::<rusqlite::Result<Vec<_>>>().map_err(Into::into) |
| 29 | + })?; |
| 30 | + if chunks.is_empty() { |
| 31 | + // 无块也写元数据（重建后空库——状态如实：引擎就绪但无向量） |
| 32 | + self.with_conn(|conn| { |
| 33 | + meta_set(conn, META_MODEL, "onnx:bge-small-zh-v1.5")?; |
| 34 | + meta_set(conn, META_DIM, &dim.to_string())?; |
| 35 | + meta_set(conn, META_FORMAT, FORMAT_F32LE)?; |
| 36 | + Ok(()) |
| 37 | + })?; |
| 38 | + return Ok(0); |
| 39 | + } |
| 40 | + let texts: Vec<String> = chunks.iter().map(|(_, t)| t.clone()).collect(); |
| 41 | + let vectors = engine |
| 42 | + .embed(&texts) |
| 43 | + .map_err(|e| crate::error::AppError::Db(format!("嵌入失败: {e}")))?; |
| 44 | + let paired: Vec<(i64, Vec<f32>)> = chunks |
| 45 | + .into_iter() |
| 46 | + .zip(vectors) |
| 47 | + .map(|((id, _), v)| (id, v)) |
| 48 | + .collect(); |
| 49 | + self.with_conn(|conn| { |
| 50 | + // 回填幂等可重跑（失败可整轮重试）——逐条自动提交，无需事务； |
| 51 | + // 单条失败中断并报错（kb_meta 最后写——元数据即"完成标记"， |
| 52 | + // 未写完=未完成，检索合流按缺 dim 自动 FTS-only） |
| 53 | + { |
| 54 | + let mut stmt = |
| 55 | + conn.prepare("UPDATE kb_chunks SET embedding = ?1 WHERE id = ?2")?; |
| 56 | + for (id, vec) in &paired { |
| 57 | + stmt.execute(rusqlite::params![encode_embedding(vec), id])?; |
| 58 | + } |
| 59 | + } |
| 60 | + meta_set(conn, META_MODEL, "onnx:bge-small-zh-v1.5")?; |
| 61 | + meta_set(conn, META_DIM, &dim.to_string())?; |
| 62 | + meta_set(conn, META_FORMAT, FORMAT_F32LE)?; |
| 63 | + Ok(()) |
| 64 | + })?; |
| 65 | + Ok(paired.len()) |
| 66 | + } |
| 67 | +} |
| 68 | + |
| 69 | +#[cfg(test)] |
| 70 | +mod tests { |
| 71 | + use crate::db::Db; |
| 72 | + use crate::kb_embed::EmbeddingEngine; |
| 73 | + |
| 74 | + /// 假引擎：dim=2，每文本返回确定性向量（[字符数×0.1, 1.0]） |
| 75 | + struct FakeEngine; |
| 76 | + impl EmbeddingEngine for FakeEngine { |
| 77 | + fn dims(&self) -> Option<usize> { |
| 78 | + Some(2) |
| 79 | + } |
| 80 | + fn embed(&self, texts: &[String]) -> std::result::Result<Vec<Vec<f32>>, String> { |
| 81 | + Ok(texts |
| 82 | + .iter() |
| 83 | + .map(|t| vec![t.chars().count() as f32 * 0.1, 1.0]) |
| 84 | + .collect()) |
| 85 | + } |
| 86 | + } |
| 87 | + |
| 88 | + fn seed_chunk(db: &Db, id: i64, text: &str) { |
| 89 | + db.with_conn(|c| { |
| 90 | + // 先建 notes 事实行（kb_chunks.note_id 有 FK——外键约束先满足） |
| 91 | + c.execute( |
| 92 | + "INSERT INTO notes (id, title, content, source, created_at, updated_at) |
| 93 | + VALUES (?1, ?2, ?3, 'manual', 1, 1)", |
| 94 | + rusqlite::params![id, format!("测试笔记 {id}"), text], |
| 95 | + )?; |
| 96 | + c.execute( |
| 97 | + "INSERT INTO kb_chunks (id, source_kind, note_id, ord, char_start, char_end, text) |
| 98 | + VALUES (?1, 'note', ?1, 1, 0, ?2, ?3)", |
| 99 | + rusqlite::params![id, text.chars().count() as i64, text], |
| 100 | + )?; |
| 101 | + Ok(()) |
| 102 | + }) |
| 103 | + .unwrap(); |
| 104 | + } |
| 105 | + |
| 106 | + #[test] |
| 107 | + fn fill_writes_blobs_and_meta_idempotently() { |
| 108 | + let db = Db::open(":memory:").unwrap(); |
| 109 | + seed_chunk(&db, 1, "学习"); |
| 110 | + seed_chunk(&db, 2, "配色与晕染"); |
| 111 | + let n = db.kb_fill_embeddings(&FakeEngine).unwrap(); |
| 112 | + assert_eq!(n, 2); |
| 113 | + // 列内 BLOB = dim×4；元数据三键齐备 |
| 114 | + db.with_conn(|c| { |
| 115 | + let bytes: i64 = c.query_row( |
| 116 | + "SELECT SUM(length(embedding)) FROM kb_chunks", |
| 117 | + [], |
| 118 | + |r| r.get(0), |
| 119 | + ).unwrap(); |
| 120 | + assert_eq!(bytes, 2 * 2 * 4); |
| 121 | + let dim: String = c.query_row( |
| 122 | + "SELECT value FROM kb_meta WHERE key='embedding_dim'", |
| 123 | + [], |
| 124 | + |r| r.get(0), |
| 125 | + ).unwrap(); |
| 126 | + assert_eq!(dim, "2"); |
| 127 | + Ok(()) |
| 128 | + }) |
| 129 | + .unwrap(); |
| 130 | + // 幂等重跑：全量覆盖不报错 |
| 131 | + let again = db.kb_fill_embeddings(&FakeEngine).unwrap(); |
| 132 | + assert_eq!(again, 2); |
| 133 | + } |
| 134 | + |
| 135 | + #[test] |
| 136 | + fn fill_empty_library_still_records_meta() { |
| 137 | + let db = Db::open(":memory:").unwrap(); |
| 138 | + assert_eq!(db.kb_fill_embeddings(&FakeEngine).unwrap(), 0); |
| 139 | + db.with_conn(|c| { |
| 140 | + let format: String = c.query_row( |
| 141 | + "SELECT value FROM kb_meta WHERE key='embedding_format'", |
| 142 | + [], |
| 143 | + |r| r.get(0), |
| 144 | + ).unwrap(); |
| 145 | + assert_eq!(format, "f32le"); |
| 146 | + Ok(()) |
| 147 | + }) |
| 148 | + .unwrap(); |
| 149 | + } |
| 150 | + |
| 151 | + #[test] |
| 152 | + fn fill_requires_engine_with_dim() { |
| 153 | + let db = Db::open(":memory:").unwrap(); |
| 154 | + seed_chunk(&db, 1, "x"); |
| 155 | + assert!(db.kb_fill_embeddings(&crate::kb_embed::NoopEmbedding).is_err()); |
| 156 | + } |
| 157 | +} |
0 commit comments