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:
+5
-1
@@ -286,7 +286,11 @@ async def chat_stream_endpoint(payload: Dict[str, Any]):
|
||||
system_prompt = (system_prompt + separator + memory_block) if system_prompt else memory_block
|
||||
|
||||
# Retrieve relevant uploaded documents (RAG) and inject the top chunks.
|
||||
doc_hits = await store.search_documents(message, get_ollama_manager().embed, limit=3)
|
||||
doc_hits = await store.search_documents(
|
||||
message, get_ollama_manager().embed,
|
||||
limit=app_settings.get("rag_top_k", 3),
|
||||
min_score=app_settings.get("rag_min_score", 0.6),
|
||||
)
|
||||
doc_titles: list = []
|
||||
if doc_hits:
|
||||
doc_block = "\n\n".join(f"[{d['title']}]\n{d['text']}" for d in doc_hits)
|
||||
|
||||
+133
-10
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user