RAG: replace bm25s with SQLite FTS5 incremental lexical index
This commit is contained in:
Roland Tannous
parent
45207c2bf1
commit
f5673e9bb0
7 changed files with 255 additions and 198 deletions
|
|
@ -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)
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue