124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Incremental BM25 on SQLite FTS5, living in the shared rag.db.
|
|
|
|
Drop-in for the bm25s-backed `bm25.py` (same `rebuild_index` / `search` /
|
|
`delete_scope` surface) plus an incremental `add_chunks`. FTS5 supports row-level
|
|
INSERT/DELETE, so a new document inserts only its own rows -- no scope-wide
|
|
rebuild. `MATCH` returns only rows that contain a query term (no zero-score
|
|
padding), and the `porter` tokenizer adds stemming. Scope is an UNINDEXED column
|
|
filtered in the WHERE clause; the dense leg stays in sqlite-vec untouched.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import threading
|
|
|
|
from core.rag.db import get_rag_connection
|
|
from loggers import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_schema_lock = threading.Lock()
|
|
_schema_ready = False
|
|
|
|
# FTS5 query terms: alphanumeric runs only. Anything else (quotes, hyphens,
|
|
# operators) is dropped so a raw user query can never form invalid MATCH syntax.
|
|
_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
|
|
|
|
|
|
def _ensure_schema() -> None:
|
|
global _schema_ready
|
|
if _schema_ready:
|
|
return
|
|
with _schema_lock:
|
|
if _schema_ready:
|
|
return
|
|
conn = get_rag_connection()
|
|
conn.executescript(
|
|
"""
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS rag_fts USING fts5(
|
|
chunk_id UNINDEXED,
|
|
scope UNINDEXED,
|
|
text,
|
|
tokenize = 'porter unicode61'
|
|
);
|
|
"""
|
|
)
|
|
conn.commit()
|
|
_schema_ready = True
|
|
|
|
|
|
def add_chunks(scope: str, chunks: list[dict]) -> None:
|
|
"""Insert only these chunks' rows. Each chunk: {id, text}. Incremental O(len)."""
|
|
if not chunks:
|
|
return
|
|
_ensure_schema()
|
|
conn = get_rag_connection()
|
|
conn.executemany(
|
|
"INSERT INTO rag_fts (chunk_id, scope, text) VALUES (?, ?, ?)",
|
|
[(c["id"], scope, c["text"]) for c in chunks],
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def rebuild_index(scope: str, chunks: list[dict]) -> None:
|
|
"""Compat path: replace a scope's rows. Empty list clears the scope."""
|
|
_ensure_schema()
|
|
conn = get_rag_connection()
|
|
conn.execute("DELETE FROM rag_fts WHERE scope = ?", (scope,))
|
|
conn.commit()
|
|
add_chunks(scope, chunks)
|
|
|
|
|
|
def _match_query(query: str) -> str | None:
|
|
terms = _TOKEN_RE.findall(query.lower())
|
|
if not terms:
|
|
return None
|
|
# OR the terms (recall-oriented, like bm25s default); quote each so FTS5
|
|
# treats it as a bare token, never an operator.
|
|
return " OR ".join(f'"{t}"' for t in terms)
|
|
|
|
|
|
def search(scope: str, query: str, k: int) -> list[tuple[str, float]]:
|
|
"""Best-first (chunk_id, score). score = -bm25() so higher is better."""
|
|
_ensure_schema()
|
|
match = _match_query(query)
|
|
if match is None:
|
|
return []
|
|
conn = get_rag_connection()
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT chunk_id, bm25(rag_fts) AS score
|
|
FROM rag_fts
|
|
WHERE scope = ? AND rag_fts MATCH ?
|
|
ORDER BY score
|
|
LIMIT ?
|
|
""",
|
|
(scope, match, k),
|
|
).fetchall()
|
|
# FTS5 bm25() is negative with more-negative = better; negate so callers see
|
|
# higher = better, list already best-first from ORDER BY score ASC.
|
|
return [(row["chunk_id"], -float(row["score"])) for row in rows]
|
|
|
|
|
|
def delete_scope(scope: str) -> None:
|
|
_ensure_schema()
|
|
conn = get_rag_connection()
|
|
conn.execute("DELETE FROM rag_fts WHERE scope = ?", (scope,))
|
|
conn.commit()
|
|
|
|
|
|
def delete_document(document_id: str, chunk_ids: list[str]) -> None:
|
|
"""Incremental per-document delete by chunk ids (FTS has no document_id col)."""
|
|
if not chunk_ids:
|
|
return
|
|
_ensure_schema()
|
|
conn = get_rag_connection()
|
|
conn.executemany(
|
|
"DELETE FROM rag_fts WHERE chunk_id = ?",
|
|
[(cid,) for cid in chunk_ids],
|
|
)
|
|
conn.commit()
|