fix: RAG — UTF-8 chunk boundary, source dedupe, embed-model guard + reindex
- 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
@@ -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<Vec<RagSource>, 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<RagReindexReport, String> {
|
||||
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<String>) -> Result<(), String> {
|
||||
state.rag.clear(source.as_deref()).map_err(|e| e.to_string())
|
||||
|
||||
@@ -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,
|
||||
|
||||
+101
-21
@@ -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(¶);
|
||||
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];
|
||||
|
||||
@@ -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<string | null>(null);
|
||||
const [chunks, setChunks] = useState<RagChunk[]>([]);
|
||||
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<RagSource[]>("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() {
|
||||
<div className="flex flex-col gap-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<h3 className="font-heading text-[var(--color-gold-bright)] text-sm">Indexed Sources</h3>
|
||||
{/* ponytail: embed-model mismatch banner — search silently skips
|
||||
these chunks, so surface it with the one-click fix. */}
|
||||
{stale && (
|
||||
<div className="flex items-center gap-1.5">
|
||||
<span className="text-[10px] text-[var(--color-gold-bright)]">Indexed with an old embedding model</span>
|
||||
<button
|
||||
onClick={reindex}
|
||||
disabled={reindexing}
|
||||
className="rounded bg-[var(--color-gold-bright)] text-[var(--color-bg-deep)] px-2 py-0.5 text-[10px] font-semibold cursor-pointer disabled:opacity-50"
|
||||
>
|
||||
{reindexing ? "Reindexing…" : "Reindex all"}
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
{sources.length > 0 && (
|
||||
confirmClear ? (
|
||||
<div className="flex items-center gap-1">
|
||||
|
||||
Reference in New Issue
Block a user