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
@@ -6,6 +6,8 @@ const DEFAULTS = {
|
||||
think: false, // Qwen3-style reasoning; off = much faster chat/memory
|
||||
temperature: 0.7,
|
||||
num_ctx: 0, // context window in tokens; 0 = model default
|
||||
rag_top_k: 3, // document chunks injected into chat
|
||||
rag_min_score: 0.6, // min cosine similarity for a chunk to count
|
||||
system_prompt: "",
|
||||
timeout: 120,
|
||||
gpu_offload: -1, // -1 = Auto; 0–100 = percent of layers forced onto the GPU
|
||||
@@ -209,6 +211,36 @@ export function Settings() {
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style={{ marginTop: "1.25rem", display: "flex", gap: "1rem" }}>
|
||||
<div style={{ flex: 1 }}>
|
||||
<label style={labelStyle}>Document chunks (top-k)</label>
|
||||
<input
|
||||
type="number" min="0" step="1"
|
||||
value={form.rag_top_k}
|
||||
onChange={e => update("rag_top_k", Math.max(0, parseInt(e.target.value) || 0))}
|
||||
style={{ width: "100%", padding: "0.6rem", background: "#222", color: "#eee", border: "1px solid #333", borderRadius: "8px", boxSizing: "border-box" }}
|
||||
/>
|
||||
</div>
|
||||
<div style={{ flex: 1 }}>
|
||||
<label style={labelStyle}>
|
||||
Min relevance
|
||||
<span style={{ float: "right", color: "#7aa", fontWeight: 600, textTransform: "none", letterSpacing: 0 }}>
|
||||
{Number(form.rag_min_score).toFixed(2)}
|
||||
</span>
|
||||
</label>
|
||||
<input
|
||||
type="range" min="0" max="1" step="0.05"
|
||||
value={form.rag_min_score}
|
||||
onChange={e => update("rag_min_score", parseFloat(e.target.value))}
|
||||
style={{ width: "100%", accentColor: "#007acc" }}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<div style={{ fontSize: "0.72rem", color: "#555", marginTop: "0.2rem" }}>
|
||||
How many uploaded-document chunks to pull into chat, and the minimum cosine
|
||||
similarity each must clear. Higher relevance = fewer, tighter matches.
|
||||
</div>
|
||||
|
||||
<div style={{ marginTop: "1.25rem" }}>
|
||||
<label style={labelStyle}>
|
||||
GPU Offload
|
||||
|
||||
@@ -40,6 +40,9 @@ PyYAML
|
||||
# Document ingest (RAG): pure-Python text extraction, no native deps
|
||||
pypdf
|
||||
python-docx
|
||||
# Vector search: loadable SQLite extension (prebuilt wheels). Brute-force cosine
|
||||
# stays as the fallback when the host Python can't load extensions.
|
||||
sqlite-vec
|
||||
|
||||
# Documentation Support
|
||||
markdown-it-py
|
||||
|
||||
@@ -12,6 +12,9 @@ psutil
|
||||
# Document ingest (RAG): pure-Python text extraction, no native deps
|
||||
pypdf
|
||||
python-docx
|
||||
# Vector search: loadable SQLite extension (prebuilt wheels). Brute-force cosine
|
||||
# stays as the fallback when the host Python can't load extensions.
|
||||
sqlite-vec
|
||||
|
||||
# Native desktop window for the UI (WebView2 on Windows; pulls pythonnet).
|
||||
# Used by bin/nexus_window.py, launched from launch_nexus.ps1.
|
||||
|
||||
+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.
|
||||
|
||||
@@ -57,6 +57,46 @@ def test_search_empty_query_returns_nothing():
|
||||
assert asyncio.run(s.search_documents("", _fake_embed)) == []
|
||||
|
||||
|
||||
def test_chunker_overlap_and_hard_split():
|
||||
s = _store()
|
||||
# a single oversized paragraph (no blank lines, as in PDF text) is split
|
||||
big = "x" * 2000
|
||||
parts = s._chunk_text(big, size=800, overlap=120)
|
||||
# each chunk is one <=size unit, plus at most an overlap tail (+separator)
|
||||
assert len(parts) >= 3 and all(len(p) <= 800 + 120 + 2 for p in parts)
|
||||
# consecutive chunks share an overlap tail
|
||||
two = s._chunk_text("A" * 700 + "\n\n" + "B" * 700, size=800, overlap=120)
|
||||
assert len(two) == 2 and two[1].startswith("A" * 120)
|
||||
|
||||
|
||||
def test_vec_index_used_and_matches_brute_force():
|
||||
# On a host that can load sqlite-vec, the fast path must be exercised (not a
|
||||
# silent fallback) and agree with brute force on the top hit.
|
||||
s = _store()
|
||||
if not s.vec_enabled:
|
||||
import pytest
|
||||
pytest.skip("sqlite-vec not loadable on this host")
|
||||
|
||||
async def run():
|
||||
await s.add_document("Lego", "lego star wars boss fight tips", _fake_embed)
|
||||
await s.add_document("GPU", "gpu vega vram notes", _fake_embed)
|
||||
# vec table populated by the dual-write
|
||||
conn = s._connect()
|
||||
n = conn.execute("SELECT COUNT(*) FROM vec_documents").fetchone()[0]
|
||||
conn.close()
|
||||
assert n == 2
|
||||
|
||||
vec_hits = await s.search_documents("lego star wars", _fake_embed, limit=2, min_score=0.1)
|
||||
assert vec_hits and vec_hits[0]["title"] == "Lego"
|
||||
|
||||
# force the brute-force path and compare the top title
|
||||
s.vec_enabled = False
|
||||
bf_hits = await s.search_documents("lego star wars", _fake_embed, limit=2, min_score=0.1)
|
||||
assert bf_hits[0]["title"] == vec_hits[0]["title"]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_extract_text_by_type():
|
||||
from synapse.main import _extract_text
|
||||
# plain text / markdown -> UTF-8 decode
|
||||
|
||||
Reference in New Issue
Block a user
GitNexus · part of the Nexus stack