jameshtwose/dm-pal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

fix: RAG — UTF-8 chunk boundary, source dedupe, embed-model guard + reindex

Browse Source 
- chunk(): byte splits that landed mid multibyte char silently dropped
  the whole chunk (accented/CJK lore). Back the boundary off to a char
  boundary; regression test included.
- add_document(): DELETE the source first — re-adding a file no longer
  doubles its chunks.
- chunks now record their embed_model (ALTER TABLE migration for old
  DBs); search skips chunks from a different model; rag_list reports it;
  new rag_reindex re-embeds everything, surfaced in the Lore panel as a
  mismatch banner with a one-click Reindex.
This commit is contained in:
itsamejms
2026-09-06 23:25:22 +01:00
parent 913345a05e
commit a69d7caefb
4 changed files with 159 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/rag/mod.rs
+101 -21
View File
@@ -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<Connection>,
@@ -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<u8>)> = {
let rows: Vec<(String, String, Vec<u8>, 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<u8>>(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<Vec<(String, i64)>> {
/// (source, chunk_count, embed_model) for every distinct source.
pub fn list_sources(&self) -> anyhow::Result<Vec<(String, i64, String)>> {
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<String> = {
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<String> = {
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(&para);
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(&para2);
assert_eq!(chunks2.concat(), para2);
}
#[test]
fn blob_roundtrip() {
let v = vec![0.0, 1.5, -2.25, 3.33];