From f5673e9bb07ef2aaef165ad2cfd91fc016f8431d Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 2 Jun 2026 21:04:23 +0400 Subject: [PATCH] RAG: replace bm25s with SQLite FTS5 incremental lexical index --- studio/backend/core/rag/bm25.py | 111 --------------- studio/backend/core/rag/db.py | 31 ++++ studio/backend/core/rag/ingestion.py | 55 ++----- studio/backend/core/rag/retrieval.py | 7 +- studio/backend/core/rag/vector_store.py | 60 +++++++- studio/backend/requirements/rag.txt | 5 +- tests/python/test_rag_bm25.py | 182 +++++++++++++++++++----- 7 files changed, 254 insertions(+), 197 deletions(-) delete mode 100644 studio/backend/core/rag/bm25.py diff --git a/studio/backend/core/rag/bm25.py b/studio/backend/core/rag/bm25.py deleted file mode 100644 index 0d45ba2420..0000000000 --- a/studio/backend/core/rag/bm25.py +++ /dev/null @@ -1,111 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""Per-scope BM25 index (rebuild on change; bm25s has no cheap incremental insert). - -Each scope dir holds the bm25s files + ids.json mapping row index → chunk_id. -""" - -from __future__ import annotations - -import json -import shutil -import threading -from pathlib import Path -from typing import Any - -from loggers import get_logger -from utils.paths.storage_roots import ensure_dir, rag_bm25_root - -logger = get_logger(__name__) - -_load_lock = threading.Lock() -_cache: dict[str, tuple[Any, list[str]]] = {} - - -def _scope_dir(scope: str) -> Path: - return rag_bm25_root() / scope - - -def _ids_path(scope: str) -> Path: - return _scope_dir(scope) / "ids.json" - - -def _has_index(scope: str) -> bool: - return _ids_path(scope).is_file() - - -def _evict(scope: str) -> None: - _cache.pop(scope, None) - - -def rebuild_index(scope: str, chunks: list[dict]) -> None: - """Rebuild scope's BM25 from full chunk list. Empty list deletes the index.""" - import bm25s - - base = _scope_dir(scope) - if not chunks: - delete_scope(scope) - return - texts = [c["text"] for c in chunks] - ids = [c["id"] for c in chunks] - tokens = bm25s.tokenize(texts, show_progress = False) - retriever = bm25s.BM25() - retriever.index(tokens, show_progress = False) - # bm25s.BM25.save does not unlink stale files; clear the dir first. - delete_scope(scope) - ensure_dir(base) - retriever.save(str(base)) - _ids_path(scope).write_text(json.dumps(ids)) - with _load_lock: - _cache[scope] = (retriever, ids) - - -def _load(scope: str) -> tuple[Any, list[str]] | None: - if not _has_index(scope): - return None - with _load_lock: - if scope in _cache: - return _cache[scope] - import bm25s - - try: - retriever = bm25s.BM25.load(str(_scope_dir(scope)), load_corpus = False) - ids = json.loads(_ids_path(scope).read_text()) - except (FileNotFoundError, OSError, json.JSONDecodeError, ValueError) as exc: - # Corrupt/partial index: treat as missing so re-ingest rebuilds cleanly. - logger.warning( - "bm25 index unreadable for scope %s (%s: %s); treating as missing", - scope, - type(exc).__name__, - exc, - ) - return None - _cache[scope] = (retriever, ids) - return _cache[scope] - - -def search(scope: str, query: str, k: int) -> list[tuple[str, float]]: - import bm25s - - loaded = _load(scope) - if loaded is None: - return [] - retriever, ids = loaded - if not ids: - return [] - k_actual = min(k, len(ids)) - q_tokens = bm25s.tokenize([query], show_progress = False) - indices, scores = retriever.retrieve(q_tokens, k = k_actual, show_progress = False) - out: list[tuple[str, float]] = [] - for pos in range(indices.shape[1]): - idx = int(indices[0][pos]) - out.append((ids[idx], float(scores[0][pos]))) - return out - - -def delete_scope(scope: str) -> None: - base = _scope_dir(scope) - if base.exists(): - shutil.rmtree(base, ignore_errors = True) - _evict(scope) diff --git a/studio/backend/core/rag/db.py b/studio/backend/core/rag/db.py index 954bffb0ad..f5fc7ccb98 100644 --- a/studio/backend/core/rag/db.py +++ b/studio/backend/core/rag/db.py @@ -57,11 +57,42 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ON rag_vectors(scope); CREATE INDEX IF NOT EXISTS idx_rag_vectors_scope_doc ON rag_vectors(scope, document_id); + CREATE VIRTUAL TABLE IF NOT EXISTS rag_chunks_fts USING fts5( + text, + chunk_id UNINDEXED, + scope UNINDEXED, + tokenize = 'porter unicode61' + ); """ ) + _backfill_fts(conn) conn.commit() +def _backfill_fts(conn: sqlite3.Connection) -> None: + """One-time: seed FTS from existing vectors for rag.db files created before + FTS5 replaced the on-disk bm25s index. Runs only when FTS is empty but + vectors exist; idempotent thereafter.""" + fts_seeded = conn.execute( + "SELECT 1 FROM rag_chunks_fts LIMIT 1" + ).fetchone() + if fts_seeded is not None: + return + has_vectors = conn.execute("SELECT 1 FROM rag_vectors LIMIT 1").fetchone() + if has_vectors is None: + return + conn.execute( + """ + INSERT INTO rag_chunks_fts (text, chunk_id, scope) + SELECT json_extract(payload_json, '$.text'), chunk_id, scope + FROM rag_vectors + WHERE kind IN ('text', 'caption') + AND json_extract(payload_json, '$.text') IS NOT NULL + """ + ) + logger.info("RAG FTS5 backfilled from existing vectors") + + def get_rag_connection() -> sqlite3.Connection: global _conn with _conn_lock: diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 37178e84c0..ff2c43a2dd 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -28,7 +28,7 @@ from utils.rag.config import ( RAG_EMBEDDING_MODEL, ) -from . import bm25, embeddings, vector_store +from . import embeddings, vector_store from .vector_store import kb_scope, thread_scope logger = get_logger(__name__) @@ -297,17 +297,18 @@ def _update_document_row(document_id: str, **fields: Any) -> None: conn.commit() -def _insert_chunks_and_collect_for_bm25( +def _insert_chunks( document_id: str, scope: str, first_index: int, chunks_meta: list[dict], vectors: list[list[float]], ) -> list[dict]: - """Insert chunks into sqlite + vector_store; return [{id, text}] for BM25.""" + """Insert chunks into sqlite + vector_store (which also indexes them into + FTS5); return [{id, text}] of the text chunks for the completion count.""" rows: list[tuple] = [] points: list[dict] = [] - bm25_rows: list[dict] = [] + text_rows: list[dict] = [] pair_groups: dict[str, list[str]] = {} for offset, (meta, vec) in enumerate(zip(chunks_meta, vectors)): @@ -359,7 +360,7 @@ def _insert_chunks_and_collect_for_bm25( } ) if kind in ("text", "caption") and meta["text"]: - bm25_rows.append({"id": chunk_id, "text": meta["text"]}) + text_rows.append({"id": chunk_id, "text": meta["text"]}) with closing_connection() as conn: conn.executemany( """ @@ -386,7 +387,7 @@ def _insert_chunks_and_collect_for_bm25( ) conn.commit() vector_store.upsert_chunks(scope, points) - return bm25_rows + return text_rows def _replace_document_pages(document_id: str, pages: list[dict]) -> None: @@ -426,37 +427,13 @@ def _replace_document_pages(document_id: str, pages: list[dict]) -> None: conn.commit() -def _all_scope_chunks(scope: str) -> list[dict]: - if scope.startswith("kb_"): - kb_id = scope[len("kb_") :] - sql = ( - "SELECT c.id, c.text FROM rag_chunks c " - "JOIN rag_documents d ON d.id = c.document_id " - "WHERE d.kb_id = ?" - ) - bind = (kb_id,) - elif scope.startswith("thread_"): - thread_id = scope[len("thread_") :] - sql = ( - "SELECT c.id, c.text FROM rag_chunks c " - "JOIN rag_documents d ON d.id = c.document_id " - "WHERE d.thread_id = ?" - ) - bind = (thread_id,) - else: - return [] - with closing_connection() as conn: - rows = conn.execute(sql, bind).fetchall() - return [{"id": r["id"], "text": r["text"]} for r in rows] - - def _pump( state: _JobState, proc: Any, out_queue: Any, ) -> None: """Drain queue until subprocess completes/errors/dies.""" - bm25_buffer: list[dict] = [] + text_buffer: list[dict] = [] embedding_dim: int | None = None final_status = "failed" final_error: str | None = None @@ -512,7 +489,7 @@ def _pump( if embedding_dim is not None: vector_store.ensure_collection(state.scope, embedding_dim) try: - bm25_rows = _insert_chunks_and_collect_for_bm25( + text_rows = _insert_chunks( state.document_id, state.scope, int(msg["first_index"]), @@ -526,10 +503,10 @@ def _pump( f"document was removed before ingestion finished ({exc})" ) break - bm25_buffer.extend(bm25_rows) + text_buffer.extend(text_rows) elif mtype == "complete": final_status = "completed" - final_num_chunks = int(msg.get("num_chunks", len(bm25_buffer))) + final_num_chunks = int(msg.get("num_chunks", len(text_buffer))) break elif mtype == "error": final_error = str(msg.get("error", "unknown error")) @@ -557,8 +534,6 @@ def _pump( state.push_event({"type": "cancelled"}) return if final_status == "completed": - full_scope_chunks = _all_scope_chunks(state.scope) - bm25.rebuild_index(state.scope, full_scope_chunks) _update_document_row( state.document_id, status = "completed", @@ -732,18 +707,12 @@ def cancel_ingestion(job_id: str) -> bool: def delete_document_artifacts(document_id: str, scope: str) -> None: - """Drop the doc's vectors, rebuild BM25. Caller deletes the rag_documents row.""" + """Drop the doc's vectors + FTS rows. Caller deletes the rag_documents row.""" vector_store.delete_document(scope, document_id) - remaining = _all_scope_chunks(scope) - if remaining: - bm25.rebuild_index(scope, remaining) - else: - bm25.delete_scope(scope) def delete_scope_artifacts(scope: str) -> None: vector_store.delete_scope(scope) - bm25.delete_scope(scope) def purge_thread_documents(thread_ids: list[str]) -> None: diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index 242446aac9..bfe65ad6cb 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -15,7 +15,7 @@ from utils.rag.config import ( RAG_TOP_K_HYBRID, ) -from . import bm25, embeddings, vector_store +from . import embeddings, vector_store # Match "Figure 1", "Figure 1.2", "Figure B.1", "Table 4", "Fig. 5" anywhere in # the query. Feeds a third retrieval source that looks up chunks anchored by these @@ -59,7 +59,10 @@ class Hit: def retrieve_bm25(scope: str, query: str, k: int | None = None) -> list[Hit]: limit = k or RAG_TOP_K_BM25 - return [Hit(chunk_id = cid, score = s) for cid, s in bm25.search(scope, query, limit)] + return [ + Hit(chunk_id = cid, score = s) + for cid, s in vector_store.search_lexical(scope, query, limit) + ] def retrieve_figure_refs( diff --git a/studio/backend/core/rag/vector_store.py b/studio/backend/core/rag/vector_store.py index 7da99a2274..1a8c660b09 100644 --- a/studio/backend/core/rag/vector_store.py +++ b/studio/backend/core/rag/vector_store.py @@ -7,6 +7,7 @@ scopes safe — per-scope embedder resolver guarantees one dim per scope.""" from __future__ import annotations import json +import re from typing import Iterable from loggers import get_logger @@ -39,27 +40,35 @@ def ensure_collection(scope: str, dim: int) -> None: def upsert_chunks(scope: str, points: Iterable[dict]) -> None: - """Insert/update vectors. Each point: {id, vector, payload}.""" + """Insert/update vectors + the lexical FTS5 index. Each point: {id, vector, payload}.""" import sqlite_vec from core.rag.db import get_rag_connection rows = [] + fts_rows: list[tuple[str, str, str]] = [] # (text, chunk_id, scope) + fts_delete: list[tuple[str]] = [] for p in points: payload = p.get("payload") or {} vec = list(p["vector"]) + kind = str(payload.get("kind") or "text") rows.append( ( p["id"], scope, str(payload.get("document_id") or ""), int(payload.get("chunk_index") or 0), - str(payload.get("kind") or "text"), + kind, len(vec), sqlite_vec.serialize_float32(vec), json.dumps(payload, default = str), ) ) + text = payload.get("text") + if kind in ("text", "caption") and text: + # FTS5 has no UPSERT; delete-then-insert keeps re-ingest idempotent. + fts_delete.append((p["id"],)) + fts_rows.append((text, p["id"], scope)) if not rows: return conn = get_rag_connection() @@ -79,6 +88,14 @@ def upsert_chunks(scope: str, points: Iterable[dict]) -> None: """, rows, ) + if fts_rows: + conn.executemany( + "DELETE FROM rag_chunks_fts WHERE chunk_id = ?", fts_delete + ) + conn.executemany( + "INSERT INTO rag_chunks_fts (text, chunk_id, scope) VALUES (?, ?, ?)", + fts_rows, + ) conn.commit() @@ -175,6 +192,7 @@ def delete_scope(scope: str) -> None: conn = get_rag_connection() conn.execute("DELETE FROM rag_vectors WHERE scope = ?", (scope,)) + conn.execute("DELETE FROM rag_chunks_fts WHERE scope = ?", (scope,)) conn.commit() @@ -182,8 +200,46 @@ def delete_document(scope: str, document_id: str) -> None: from core.rag.db import get_rag_connection conn = get_rag_connection() + chunk_ids = [ + row["chunk_id"] + for row in conn.execute( + "SELECT chunk_id FROM rag_vectors WHERE scope = ? AND document_id = ?", + (scope, document_id), + ).fetchall() + ] conn.execute( "DELETE FROM rag_vectors WHERE scope = ? AND document_id = ?", (scope, document_id), ) + if chunk_ids: + conn.executemany( + "DELETE FROM rag_chunks_fts WHERE chunk_id = ?", + [(cid,) for cid in chunk_ids], + ) conn.commit() + + +_FTS_TOKEN = re.compile(r"\w+", re.UNICODE) + + +def _match_query(query: str) -> str: + """User text → safe FTS5 OR-of-quoted-terms; quoting defuses FTS5 operators.""" + tokens = _FTS_TOKEN.findall(query.lower()) + return " OR ".join(f'"{t}"' for t in tokens) + + +def search_lexical(scope: str, query: str, k: int) -> list[tuple[str, float]]: + """BM25 lexical search via SQLite FTS5. Returns [(chunk_id, score)], higher = better.""" + from core.rag.db import get_rag_connection + + match = _match_query(query) + if not match: + return [] + conn = get_rag_connection() + rows = conn.execute( + "SELECT chunk_id, bm25(rag_chunks_fts) AS s FROM rag_chunks_fts " + "WHERE rag_chunks_fts MATCH ? AND scope = ? ORDER BY s LIMIT ?", + (match, scope, int(k)), + ).fetchall() + # bm25() is negative (more negative = better); flip to higher-is-better. + return [(row["chunk_id"], -float(row["s"])) for row in rows] diff --git a/studio/backend/requirements/rag.txt b/studio/backend/requirements/rag.txt index 5ac1840b41..6d77f45ae5 100644 --- a/studio/backend/requirements/rag.txt +++ b/studio/backend/requirements/rag.txt @@ -9,10 +9,9 @@ # adds vector functions (vec_distance_cosine, serialize_float32, vec0 # virtual tables). The studio loads it into a dedicated rag.db file and # stores vectors as BLOB columns alongside the RAG metadata — no separate -# vector server, no second client library. bm25s persists per-scope -# lexical indexes to disk. +# vector server, no second client library. The lexical index lives in the +# same rag.db as a SQLite FTS5 virtual table (built-in, no extra dependency). sqlite-vec>=0.1.5 -bm25s>=0.2 # Image preprocessing helpers required by Qwen3-VL-Embedding-2B (the # multimodal embedder). Not used in text-only mode. diff --git a/tests/python/test_rag_bm25.py b/tests/python/test_rag_bm25.py index b592a397fa..303b47b9a0 100644 --- a/tests/python/test_rag_bm25.py +++ b/tests/python/test_rag_bm25.py @@ -1,4 +1,4 @@ -"""BM25 index lifecycle tests (skipped if bm25s is unavailable).""" +"""FTS5 lexical index lifecycle tests (sqlite-vec backed rag.db).""" import sys from pathlib import Path @@ -10,60 +10,170 @@ STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" if str(STUDIO_BACKEND) not in sys.path: sys.path.insert(0, str(STUDIO_BACKEND)) -pytest.importorskip("bm25s") +pytest.importorskip("sqlite_vec") @pytest.fixture -def isolated_bm25_root(tmp_path, monkeypatch): - from utils.paths import storage_roots - +def isolated_rag_db(tmp_path, monkeypatch): + """Point rag.db at tmp_path and reset the cached connection so each + test gets a fresh database.""" monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) - # Reset module cache between tests. - from core.rag import bm25 + from core.rag import db as rag_db - bm25._cache.clear() - return tmp_path + rag_db._reset_for_tests() + yield tmp_path + rag_db._reset_for_tests() -def test_bm25_index_search_roundtrip(isolated_bm25_root): - from core.rag import bm25 +def _chunk(cid: str, text: str, document_id: str = "doc1", index: int = 0) -> dict: + return { + "id": cid, + "vector": [1.0, 0.0, 0.0, 0.0], + "payload": { + "document_id": document_id, + "chunk_index": index, + "kind": "text", + "text": text, + }, + } + + +def test_lexical_search_roundtrip(isolated_rag_db): + from core.rag import vector_store scope = "kb_test" - chunks = [ - {"id": "c1", "text": "the quick brown fox jumps over the lazy dog"}, - {"id": "c2", "text": "machine learning models predict outputs from inputs"}, - {"id": "c3", "text": "fox terriers are small dogs"}, - ] - bm25.rebuild_index(scope, chunks) - results = bm25.search(scope, "fox", k = 3) + vector_store.upsert_chunks( + scope, + [ + _chunk("c1", "the quick brown fox jumps over the lazy dog", index = 0), + _chunk("c2", "machine learning models predict outputs from inputs", index = 1), + _chunk("c3", "fox terriers are small dogs", index = 2), + ], + ) + results = vector_store.search_lexical(scope, "fox", k = 3) ids = [cid for cid, _ in results] assert "c1" in ids assert "c3" in ids + assert "c2" not in ids + # Scores are flipped to higher-is-better. + assert all(score >= 0 for _, score in results) -def test_bm25_empty_returns_empty(isolated_bm25_root): - from core.rag import bm25 +def test_lexical_empty_returns_empty(isolated_rag_db): + from core.rag import vector_store - assert bm25.search("kb_nonexistent", "anything", k = 5) == [] + assert vector_store.search_lexical("kb_nonexistent", "anything", k = 5) == [] -def test_bm25_delete_scope(isolated_bm25_root): - from core.rag import bm25 +def test_lexical_blank_query_returns_empty(isolated_rag_db): + from core.rag import vector_store + + scope = "kb_blank" + vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta gamma")]) + # No word tokens → no MATCH expression → empty (not a syntax error). + assert vector_store.search_lexical(scope, "!!! ???", k = 5) == [] + + +def test_lexical_query_operators_are_defused(isolated_rag_db): + from core.rag import vector_store + + scope = "kb_ops" + vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta gamma")]) + # Bare FTS5 operators would raise without sanitization. + results = vector_store.search_lexical(scope, "alpha OR NOT (beta)", k = 5) + assert [cid for cid, _ in results] == ["a"] + + +def test_lexical_scope_isolation(isolated_rag_db): + from core.rag import vector_store + + vector_store.upsert_chunks("kb_one", [_chunk("x", "shared keyword here")]) + vector_store.upsert_chunks("kb_two", [_chunk("y", "shared keyword here")]) + results = vector_store.search_lexical("kb_one", "keyword", k = 5) + assert [cid for cid, _ in results] == ["x"] + + +def test_lexical_delete_scope(isolated_rag_db): + from core.rag import vector_store scope = "kb_del" - chunks = [{"id": "a", "text": "alpha beta gamma"}] - bm25.rebuild_index(scope, chunks) - assert bm25.search(scope, "alpha", k = 1) - bm25.delete_scope(scope) - assert bm25.search(scope, "alpha", k = 1) == [] + vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta gamma")]) + assert vector_store.search_lexical(scope, "alpha", k = 1) + vector_store.delete_scope(scope) + assert vector_store.search_lexical(scope, "alpha", k = 1) == [] -def test_bm25_rebuild_replaces_old_corpus(isolated_bm25_root): - from core.rag import bm25 +def test_lexical_delete_document(isolated_rag_db): + from core.rag import vector_store - scope = "kb_replace" - bm25.rebuild_index(scope, [{"id": "old", "text": "alpha beta"}]) - bm25.rebuild_index(scope, [{"id": "new", "text": "gamma delta"}]) - results = bm25.search(scope, "alpha", k = 5) - ids = [cid for cid, _ in results] - assert "old" not in ids + scope = "kb_doc_del" + vector_store.upsert_chunks( + scope, + [ + _chunk("keep1", "alpha keyword", document_id = "keep"), + _chunk("drop1", "beta keyword", document_id = "drop"), + ], + ) + vector_store.delete_document(scope, "drop") + ids = [cid for cid, _ in vector_store.search_lexical(scope, "keyword", k = 5)] + assert ids == ["keep1"] + + +def test_lexical_reingest_is_idempotent(isolated_rag_db): + from core.rag import vector_store + + scope = "kb_reingest" + vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta")]) + vector_store.upsert_chunks(scope, [_chunk("a", "gamma delta")]) + # No duplicate FTS row; old text no longer matches. + assert vector_store.search_lexical(scope, "alpha", k = 5) == [] + ids = [cid for cid, _ in vector_store.search_lexical(scope, "gamma", k = 5)] + assert ids == ["a"] + + +def test_lexical_caption_kind_is_indexed(isolated_rag_db): + from core.rag import vector_store + + scope = "kb_caption" + point = _chunk("cap", "Figure 1: a diagram of the pipeline") + point["payload"]["kind"] = "caption" + vector_store.upsert_chunks(scope, [point]) + ids = [cid for cid, _ in vector_store.search_lexical(scope, "diagram", k = 5)] + assert ids == ["cap"] + + +def test_fts_backfill_from_existing_vectors(isolated_rag_db): + """rag.db files created before FTS5 have vectors but an empty FTS table; + opening the connection backfills it once.""" + import sqlite_vec + + from core.rag import db as rag_db + + # Seed vectors directly, bypassing upsert_chunks' incremental FTS insert. + conn = rag_db.get_rag_connection() + conn.execute( + """ + INSERT INTO rag_vectors + (chunk_id, scope, document_id, chunk_index, kind, dim, vector, payload_json) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + "legacy", + "kb_legacy", + "doc", + 0, + "text", + 4, + sqlite_vec.serialize_float32([1.0, 0.0, 0.0, 0.0]), + '{"text": "backfilled lexical content", "kind": "text"}', + ), + ) + conn.execute("DELETE FROM rag_chunks_fts") + conn.commit() + # Force a fresh connection so _ensure_schema runs the guarded backfill. + rag_db._reset_for_tests() + + from core.rag import vector_store + + ids = [cid for cid, _ in vector_store.search_lexical("kb_legacy", "lexical", k = 5)] + assert ids == ["legacy"]