diff --git a/src-tauri/src/commands/rag_commands.rs b/src-tauri/src/commands/rag_commands.rs index 89bd6e3..6def79a 100644 --- a/src-tauri/src/commands/rag_commands.rs +++ b/src-tauri/src/commands/rag_commands.rs @@ -18,6 +18,15 @@ pub struct RagHit { pub struct RagSource { pub source: String, pub chunks: i64, + /// Embedding model the source was indexed with ("" = legacy/unknown). + /// Lets the UI flag sources that need a Reindex after a model change. + pub model: String, +} + +#[derive(Debug, Serialize)] +pub struct RagReindexReport { + pub sources: usize, + pub chunks: usize, } #[tauri::command] @@ -48,10 +57,20 @@ pub fn rag_list(state: tauri::State<'_, AppState>) -> Result, Str .list_sources() .map_err(|e| e.to_string())? .into_iter() - .map(|(source, chunks)| Ok(RagSource { source, chunks })) + .map(|(source, chunks, model)| Ok(RagSource { source, chunks, model })) .collect() } +/// Re-embed every source with the currently configured embed model. +/// The fix-it button when the DM changes models (or for legacy chunks +/// indexed before per-source model stamping). +#[tauri::command] +pub async fn rag_reindex(state: tauri::State<'_, AppState>) -> Result { + let config = state.config.lock().map_err(|e| e.to_string())?.clone(); + let (sources, chunks) = state.rag.reindex(&config).await.map_err(|e| e.to_string())?; + Ok(RagReindexReport { sources, chunks }) +} + #[tauri::command] pub fn rag_clear(state: tauri::State<'_, AppState>, source: Option) -> Result<(), String> { state.rag.clear(source.as_deref()).map_err(|e| e.to_string()) diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 6a5e879..536ddd6 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -82,6 +82,7 @@ pub fn run() { commands::rag_commands::rag_clear, commands::rag_commands::rag_chunks, commands::rag_commands::rag_add_directory, + commands::rag_commands::rag_reindex, commands::generation_commands::generation_add, commands::generation_commands::generation_list, commands::generation_commands::generation_get, diff --git a/src-tauri/src/rag/mod.rs b/src-tauri/src/rag/mod.rs index 96d717b..31fbbd2 100644 --- a/src-tauri/src/rag/mod.rs +++ b/src-tauri/src/rag/mod.rs @@ -9,7 +9,13 @@ use std::sync::Mutex; // that and scan time shows up in profiles. Avoids native extension loading. /// One embedding vector, stored as little-endian f32 bytes. -const MAX_CHUNK_CHARS: usize = 1000; +/// Max chunk size in BYTES (chunk() splits on byte boundaries safely). +const MAX_CHUNK_BYTES: usize = 1000; + +/// Legacy chunks indexed before the embed_model column existed. Unknown +/// model — assumed compatible so search keeps working; a one-click +/// Reindex stamps everything with the current model. +const EMBED_MODEL_UNKNOWN: &str = ""; pub struct RagStore { db: Mutex, @@ -27,13 +33,20 @@ impl RagStore { let db = Connection::open(path)?; db.execute_batch( "CREATE TABLE IF NOT EXISTS chunks ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - source TEXT NOT NULL, - text TEXT NOT NULL, - emb BLOB NOT NULL + id INTEGER PRIMARY KEY AUTOINCREMENT, + source TEXT NOT NULL, + text TEXT NOT NULL, + emb BLOB NOT NULL, + embed_model TEXT NOT NULL DEFAULT '' ); CREATE INDEX IF NOT EXISTS idx_chunks_source ON chunks(source);", )?; + // Migration for pre-embed_model DBs: add the column if it's missing. + // Duplicate-column error is expected on already-migrated DBs — ignore. + let _ = db.execute( + "ALTER TABLE chunks ADD COLUMN embed_model TEXT NOT NULL DEFAULT ''", + [], + ); Ok(Self { db: Mutex::new(db) }) } @@ -47,14 +60,20 @@ impl RagStore { if para.is_empty() { continue; } - if para.len() <= MAX_CHUNK_CHARS { + if para.len() <= MAX_CHUNK_BYTES { out.push(para.to_string()); } else { - // Hard-cap long paragraphs on char boundaries. - for chunk in para.as_bytes().chunks(MAX_CHUNK_CHARS) { - if let Ok(s) = std::str::from_utf8(chunk) { - out.push(s.trim().to_string()); + // Hard-cap long paragraphs on char boundaries. A naive byte + // split can land mid multibyte char — back the boundary off + // so accented/CJK text is never silently dropped. + let mut start = 0; + while start < para.len() { + let mut end = (start + MAX_CHUNK_BYTES).min(para.len()); + while end < para.len() && !para.is_char_boundary(end) { + end -= 1; } + out.push(para[start..end].trim().to_string()); + start = end; } } } @@ -107,9 +126,13 @@ impl RagStore { anyhow::bail!("embedding count mismatch"); } let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?; - let mut stmt = db.prepare("INSERT INTO chunks (source, text, emb) VALUES (?, ?, ?)")?; + // Re-adding a source replaces it (no duplicate chunks skewing scores). + db.execute("DELETE FROM chunks WHERE source = ?", params![source])?; + let mut stmt = db.prepare( + "INSERT INTO chunks (source, text, emb, embed_model) VALUES (?, ?, ?, ?)", + )?; for (text, emb) in chunks.iter().zip(embeddings.iter()) { - stmt.execute(params![source, text, Self::vec_to_blob(emb)])?; + stmt.execute(params![source, text, Self::vec_to_blob(emb), config.embed_model])?; } Ok(chunks.len()) } @@ -128,14 +151,16 @@ impl RagStore { .next() .ok_or_else(|| anyhow::anyhow!("no query embedding"))?; - let rows: Vec<(String, String, Vec)> = { + let rows: Vec<(String, String, Vec, String)> = { let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?; - let mut stmt = db.prepare("SELECT text, source, emb FROM chunks")?; + let mut stmt = + db.prepare("SELECT text, source, emb, embed_model FROM chunks")?; let rows = stmt.query_map([], |r| { Ok(( r.get::<_, String>(0)?, r.get::<_, String>(1)?, r.get::<_, Vec>(2)?, + r.get::<_, String>(3)?, )) })?; rows.filter_map(|r| r.ok()).collect() @@ -145,9 +170,16 @@ impl RagStore { } // ponytail: naive O(n) scan. Fine to ~10k chunks; see module note. + // Skip chunks embedded with a DIFFERENT model — their dot products + // against this query vector are garbage. '' rows are legacy/unknown + // (indexed before per-chunk stamping) and stay in play; Reindex + // re-stamps them. let mut scored: Vec<(String, String, f32)> = rows .into_iter() - .map(|(text, source, blob)| { + .filter(|(_, _, _, model)| { + model.as_str() == EMBED_MODEL_UNKNOWN || model == &config.embed_model + }) + .map(|(text, source, blob, _)| { let v = Self::blob_to_vec(&blob); let dot: f32 = v.iter().zip(q_emb.iter()).map(|(a, b)| a * b).sum(); (text, source, dot) @@ -158,11 +190,15 @@ impl RagStore { Ok(scored) } - /// (source, chunk_count) for every distinct source. - pub fn list_sources(&self) -> anyhow::Result> { + /// (source, chunk_count, embed_model) for every distinct source. + pub fn list_sources(&self) -> anyhow::Result> { let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?; - let mut stmt = db.prepare("SELECT source, COUNT(*) FROM chunks GROUP BY source ORDER BY source")?; - let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)))?; + let mut stmt = db.prepare( + "SELECT source, COUNT(*), MAX(embed_model) FROM chunks GROUP BY source ORDER BY source", + )?; + let rows = stmt.query_map([], |r| { + Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?, r.get::<_, String>(2)?)) + })?; let mut out = Vec::new(); for r in rows { out.push(r?); @@ -170,6 +206,31 @@ impl RagStore { Ok(out) } + /// Re-embed every source with the CURRENT embed model: for each source, + /// join its stored chunk texts and run them through add_document (which + /// replaces the old rows and stamps the model). Also the fix-it button + /// when the DM changes embedding models. Returns (sources, chunks). + pub async fn reindex(&self, config: &LlmConfig) -> anyhow::Result<(usize, usize)> { + let sources: Vec = { + let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?; + let mut stmt = db.prepare("SELECT DISTINCT source FROM chunks ORDER BY source")?; + let rows = stmt.query_map([], |r| r.get::<_, String>(0))?; + rows.filter_map(|r| r.ok()).collect() + }; + let mut chunks = 0; + for s in &sources { + let texts: Vec = { + let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?; + let mut stmt = + db.prepare("SELECT text FROM chunks WHERE source = ? ORDER BY id")?; + let rows = stmt.query_map(params![s], |r| r.get::<_, String>(0))?; + rows.filter_map(|r| r.ok()).collect() + }; + chunks += self.add_document(config, &s, &texts.join("\n\n")).await?; + } + Ok((sources.len(), chunks)) + } + /// Clear all chunks, or just one source if given. pub fn clear(&self, source: Option<&str>) -> anyhow::Result<()> { let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?; @@ -211,7 +272,7 @@ mod tests { #[test] fn chunk_splits_paragraphs_and_caps_long_ones() { - let long = "a".repeat(MAX_CHUNK_CHARS * 2 + 50); + let long = "a".repeat(MAX_CHUNK_BYTES * 2 + 50); let text = format!("short para\n\n{long}\n\nanother"); let chunks = RagStore::chunk(&text); // "short para", (2 or 3 long pieces), "another" @@ -219,10 +280,29 @@ mod tests { assert_eq!(chunks.first().unwrap(), "short para"); assert!(chunks.last().unwrap() == "another"); for c in &chunks { - assert!(c.len() <= MAX_CHUNK_CHARS); + assert!(c.len() <= MAX_CHUNK_BYTES); } } + #[test] + fn chunk_never_drops_multibyte_text() { + // é is 2 bytes: a 1002-byte paragraph forces a split that used to land + // wherever the byte counter said — including mid-character, silently + // dropping the whole chunk. Every byte must survive now. + let para = format!("é{}", "a".repeat(MAX_CHUNK_BYTES)); + let chunks = RagStore::chunk(¶); + let rejoined = chunks.concat(); + assert_eq!(rejoined, para, "chunking must not drop multibyte text"); + for c in &chunks { + assert!(c.len() <= MAX_CHUNK_BYTES); + } + // Force a mid-character boundary: 998 a's then 3-byte chars, so the + // 1000-byte cut lands inside the first 日. + let para2 = format!("{}{}", "a".repeat(MAX_CHUNK_BYTES - 2), "日日日"); + let chunks2 = RagStore::chunk(¶2); + assert_eq!(chunks2.concat(), para2); + } + #[test] fn blob_roundtrip() { let v = vec![0.0, 1.5, -2.25, 3.33]; diff --git a/src/components/LorePanel.tsx b/src/components/LorePanel.tsx index 6602d67..ea58449 100644 --- a/src/components/LorePanel.tsx +++ b/src/components/LorePanel.tsx @@ -8,6 +8,7 @@ import { addToLore } from "../lib/lore"; interface RagSource { source: string; chunks: number; + model: string; } interface RagHit { @@ -39,8 +40,15 @@ export function LorePanel() { const [expanded, setExpanded] = useState(null); const [chunks, setChunks] = useState([]); const [loadingChunks, setLoadingChunks] = useState(false); + // ponytail: current embed model, to flag sources indexed with a different + // one (search skips those — garbage cosine scores). "" = still loading. + const [embedModel, setEmbedModel] = useState(""); + const [reindexing, setReindexing] = useState(false); const { addToast } = useToast(); + const stale = + embedModel !== "" && sources.some((s) => s.model !== "" && s.model !== embedModel); + async function loadSources() { try { setSources(await invoke("rag_list")); @@ -51,6 +59,9 @@ export function LorePanel() { useEffect(() => { loadSources(); + invoke<{ embed_model: string }>("get_llm_config") + .then((c) => setEmbedModel(c.embed_model)) + .catch(() => {}); }, []); async function loadChunks(s: string) { @@ -101,6 +112,18 @@ export function LorePanel() { loadSources(); } + async function reindex() { + setReindexing(true); + try { + const r = await invoke<{ sources: number; chunks: number }>("rag_reindex"); + addToast(`Reindexed ${r.sources} sources · ${r.chunks} chunks`, "success"); + await loadSources(); + } catch (e) { + addToast(`Reindex failed: ${e}`, "error"); + } + setReindexing(false); + } + // ponytail: file upload via native HTML input — no tauri-plugin-dialog // needed. Reads .md/.txt contents in the webview and indexes each file as // its own lore source. Multiple files supported. @@ -214,6 +237,20 @@ export function LorePanel() {

Indexed Sources

+ {/* ponytail: embed-model mismatch banner — search silently skips + these chunks, so surface it with the one-click fix. */} + {stale && ( +
+ Indexed with an old embedding model + +
+ )} {sources.length > 0 && ( confirmClear ? (