feat(rag): overlapping chunker, retrieval knobs, sqlite-vec index

- Chunker: char overlap across boundaries + hard-split of oversized paragraphs.
- Retrieval knobs: rag_top_k / rag_min_score in settings + Settings UI.
- Vector index: sqlite-vec ANN over document embeddings, dual-written and
  backfilled, with brute-force cosine as the guaranteed fallback.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
jon
2026-07-23 13:58:05 -05:00
co-authored by Claude Opus 4.8
parent ac0eb1e5f6
commit f4aea78b55
6 changed files with 216 additions and 11 deletions
+133 -10
View File
@@ -63,16 +63,41 @@ class PersistentMemoryStore:
def __init__(self, db_path: Path):
self.db_path = db_path
os.makedirs(self.db_path.parent, exist_ok=True)
self.vec_enabled = self._probe_vec()
self._ensure_tables()
self._cache: Dict[str, MemoryItem] = self._load_all_memory()
# -----------------------------
# Internal helpers
# -----------------------------
@staticmethod
def _probe_vec() -> bool:
"""True if this host can load the sqlite-vec extension. Some Python
builds ship SQLite with loadable extensions disabled — those fall back
to the brute-force cosine scan, so recall never depends on this."""
try:
import sqlite_vec
c = sqlite3.connect(":memory:")
c.enable_load_extension(True)
sqlite_vec.load(c)
c.execute("SELECT vec_version()")
c.close()
return True
except Exception:
return False
def _connect(self):
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL;")
if self.vec_enabled:
try:
import sqlite_vec
conn.enable_load_extension(True)
sqlite_vec.load(conn)
conn.enable_load_extension(False)
except Exception:
pass
return conn
def _ensure_tables(self):
@@ -622,20 +647,32 @@ class PersistentMemoryStore:
# Documents (RAG)
# -----------------------------
@staticmethod
def _chunk_text(text: str, size: int = 800) -> List[str]:
"""Split on blank lines, then pack paragraphs into ~`size`-char chunks.
ponytail: naive char-based packing, no token counting or overlap — good
enough for local recall; add overlap if retrieval misses boundaries."""
chunks: List[str] = []
buf = ""
def _chunk_text(text: str, size: int = 800, overlap: int = 120) -> List[str]:
"""Pack paragraphs into ~`size`-char chunks with a char `overlap` carried
across boundaries, so a passage spanning two chunks still matches. Any
single paragraph larger than `size` (common in PDFs with few blank lines)
is hard-split into overlapping windows first.
ponytail: char-based, not token-based — fine for local recall; move to a
token splitter only if chunk sizes start hurting the context budget."""
units: List[str] = []
for para in (p.strip() for p in text.split("\n\n")):
if not para:
continue
if buf and len(buf) + len(para) + 2 > size:
chunks.append(buf)
buf = para
if len(para) <= size:
units.append(para)
else:
buf = f"{buf}\n\n{para}" if buf else para
step = max(1, size - overlap)
units.extend(para[i:i + size] for i in range(0, len(para), step))
chunks: List[str] = []
buf = ""
for u in units:
if buf and len(buf) + len(u) + 2 > size:
chunks.append(buf)
tail = buf[-overlap:] if overlap else "" # overlap seed for the next chunk
buf = f"{tail}\n\n{u}" if tail else u
else:
buf = f"{buf}\n\n{u}" if buf else u
if buf:
chunks.append(buf)
return chunks
@@ -656,10 +693,81 @@ class PersistentMemoryStore:
(str(_uuid.uuid4()), doc_id, title, i, piece,
json.dumps(vec) if vec else None, now),
)
self._vec_upsert(conn, cur.lastrowid, vec) # mirror into the ANN index
conn.commit()
conn.close()
return {"doc_id": doc_id, "title": title, "chunks": len(pieces)}
# --- vector index (sqlite-vec) — accelerator over the JSON embedding column,
# --- with brute-force cosine below as the guaranteed fallback ---------------
def _ensure_vec_docs(self, conn, dim: int) -> None:
conn.execute(
f"CREATE VIRTUAL TABLE IF NOT EXISTS vec_documents "
f"USING vec0(embedding float[{dim}] distance_metric=cosine)"
)
def _vec_upsert(self, conn, rowid: int, vec: list) -> None:
"""Best-effort mirror of one chunk's vector into the vec index."""
if not (self.vec_enabled and vec):
return
try:
import sqlite_vec
self._ensure_vec_docs(conn, len(vec))
conn.execute("DELETE FROM vec_documents WHERE rowid = ?", (rowid,))
conn.execute(
"INSERT INTO vec_documents(rowid, embedding) VALUES (?, ?)",
(rowid, sqlite_vec.serialize_float32(vec)),
)
except Exception:
pass # the index is an accelerator, never a requirement
def _backfill_vec(self, conn, dim: int) -> None:
"""Index any document chunks missing from vec_documents (older rows, or
rows written while the extension was unavailable)."""
try:
import sqlite_vec
self._ensure_vec_docs(conn, dim)
rows = conn.execute(
"SELECT d.rowid AS rid, d.embedding AS emb FROM documents d "
"LEFT JOIN vec_documents v ON v.rowid = d.rowid "
"WHERE v.rowid IS NULL AND d.embedding IS NOT NULL"
).fetchall()
for r in rows:
try:
vec = json.loads(r["emb"])
if len(vec) == dim:
conn.execute(
"INSERT INTO vec_documents(rowid, embedding) VALUES (?, ?)",
(r["rid"], sqlite_vec.serialize_float32(vec)),
)
except Exception:
pass
conn.commit()
except Exception:
pass
def _vec_search(self, query_vec: list, limit: int, min_score: float):
"""KNN over the vec index. Returns hits, or None to signal fall back to
the brute-force scan (e.g. extension error or dimension mismatch)."""
try:
import sqlite_vec
conn = self._connect()
self._backfill_vec(conn, len(query_vec))
rows = conn.execute(
"SELECT d.title AS title, d.text AS text, v.distance AS distance "
"FROM vec_documents v JOIN documents d ON d.rowid = v.rowid "
"WHERE v.embedding MATCH ? ORDER BY v.distance LIMIT ?",
(sqlite_vec.serialize_float32(query_vec), limit),
).fetchall()
conn.close()
# sqlite-vec cosine distance = 1 - cosine similarity
return [
{"title": r["title"], "text": r["text"], "score": 1.0 - r["distance"]}
for r in rows if (1.0 - r["distance"]) >= min_score
]
except Exception:
return None
def list_documents(self) -> List[dict]:
conn = self._connect()
cur = conn.cursor()
@@ -686,6 +794,12 @@ class PersistentMemoryStore:
def delete_document(self, doc_id: str) -> bool:
conn = self._connect()
cur = conn.cursor()
if self.vec_enabled:
try:
for r in cur.execute("SELECT rowid FROM documents WHERE doc_id = ?", (doc_id,)).fetchall():
conn.execute("DELETE FROM vec_documents WHERE rowid = ?", (r["rowid"],))
except Exception:
pass
cur.execute("DELETE FROM documents WHERE doc_id = ?", (doc_id,))
deleted = cur.rowcount
conn.commit()
@@ -702,6 +816,11 @@ class PersistentMemoryStore:
query_vec = await embed_fn(self._EMBED_QUERY_PREFIX + query.strip())
if not query_vec:
return []
# Fast path: the sqlite-vec ANN index. None => fall through to brute force.
if self.vec_enabled:
hits = self._vec_search(query_vec, limit, min_score)
if hits is not None:
return hits
conn = self._connect()
cur = conn.cursor()
cur.execute("SELECT title, text, embedding FROM documents WHERE embedding IS NOT NULL")
@@ -729,6 +848,10 @@ class PersistentMemoryStore:
"temperature": 0.7,
# Context window (tokens Ollama keeps in view). 0 → Ollama's model default.
"num_ctx": 0,
# RAG retrieval: how many document chunks to inject, and the minimum
# cosine similarity (0-1) a chunk must clear to count as relevant.
"rag_top_k": 3,
"rag_min_score": 0.6,
"system_prompt": "",
"timeout": 120,
# How long Ollama keeps the model resident in VRAM between messages.