From e0ef14c9fa375fe3d7cb96d1844e8a0fdeed9f6e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 31 May 2026 13:31:01 +0000 Subject: [PATCH] Studio: fast RAG path (warm embedder + in-process ingest + FTS5 BM25) behind UNSLOTH_RAG_FAST --- studio/backend/core/rag/bm25.py | 31 +++++++ studio/backend/core/rag/bm25_fts.py | 124 +++++++++++++++++++++++++++ studio/backend/core/rag/ingestion.py | 87 ++++++++++++++----- studio/backend/main.py | 16 ++++ 4 files changed, 237 insertions(+), 21 deletions(-) create mode 100644 studio/backend/core/rag/bm25_fts.py diff --git a/studio/backend/core/rag/bm25.py b/studio/backend/core/rag/bm25.py index 0d45ba2420..3c65c09483 100644 --- a/studio/backend/core/rag/bm25.py +++ b/studio/backend/core/rag/bm25.py @@ -9,6 +9,7 @@ Each scope dir holds the bm25s files + ids.json mapping row index → chunk_id. from __future__ import annotations import json +import os import shutil import threading from pathlib import Path @@ -23,6 +24,22 @@ _load_lock = threading.Lock() _cache: dict[str, tuple[Any, list[str]]] = {} +def _fast() -> bool: + """Incremental SQLite FTS5 backend instead of the rebuild-on-change bm25s one.""" + return os.environ.get("UNSLOTH_RAG_FAST") == "1" + + +def add_chunks(scope: str, chunks: list[dict]) -> None: + """Incremental insert of one document's chunks (FTS5 fast path only).""" + if _fast(): + from core.rag import bm25_fts + + bm25_fts.add_chunks(scope, chunks) + return + # bm25s has no incremental insert; callers on the slow path use rebuild_index. + raise RuntimeError("add_chunks requires UNSLOTH_RAG_FAST=1") + + def _scope_dir(scope: str) -> Path: return rag_bm25_root() / scope @@ -41,6 +58,11 @@ def _evict(scope: str) -> None: def rebuild_index(scope: str, chunks: list[dict]) -> None: """Rebuild scope's BM25 from full chunk list. Empty list deletes the index.""" + if _fast(): + from core.rag import bm25_fts + + bm25_fts.rebuild_index(scope, chunks) + return import bm25s base = _scope_dir(scope) @@ -86,6 +108,10 @@ def _load(scope: str) -> tuple[Any, list[str]] | None: def search(scope: str, query: str, k: int) -> list[tuple[str, float]]: + if _fast(): + from core.rag import bm25_fts + + return bm25_fts.search(scope, query, k) import bm25s loaded = _load(scope) @@ -105,6 +131,11 @@ def search(scope: str, query: str, k: int) -> list[tuple[str, float]]: def delete_scope(scope: str) -> None: + if _fast(): + from core.rag import bm25_fts + + bm25_fts.delete_scope(scope) + return base = _scope_dir(scope) if base.exists(): shutil.rmtree(base, ignore_errors = True) diff --git a/studio/backend/core/rag/bm25_fts.py b/studio/backend/core/rag/bm25_fts.py new file mode 100644 index 0000000000..b650b3c7d5 --- /dev/null +++ b/studio/backend/core/rag/bm25_fts.py @@ -0,0 +1,124 @@ +# 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() diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index bd151855b2..ea60277b8b 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -11,6 +11,7 @@ from __future__ import annotations import json import multiprocessing as mp +import os import queue as queue_module import sqlite3 import threading @@ -37,6 +38,35 @@ _CTX = mp.get_context("spawn") _QUEUE_TIMEOUT_SECONDS = 300 +def _fast_ingest() -> bool: + """Run ingestion in-process so the warm embedder singleton is reused instead + of a fresh spawn reloading the model on every upload.""" + return os.environ.get("UNSLOTH_RAG_FAST") == "1" + + +class _ThreadProc: + """Thread that mimics the mp.Process surface the pump/cancel paths use + (start / is_alive / join / terminate). Lets the existing queue-based worker + run in the warm main process under the fast flag. Threads cannot be + force-killed, so terminate() is a best-effort no-op (the route still deletes + the document's rows on cancel).""" + + def __init__(self, target: Any, args: tuple) -> None: + self._t = threading.Thread(target = target, args = args, daemon = True) + + def start(self) -> None: + self._t.start() + + def is_alive(self) -> bool: + return self._t.is_alive() + + def join(self, timeout: float | None = None) -> None: + self._t.join(timeout) + + def terminate(self) -> None: + pass + + # --- Subprocess worker --- _MIME_TO_EXT = { @@ -751,8 +781,14 @@ 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) + if _fast_ingest(): + # Incremental: index only this document's chunks (O(N) total) instead + # of re-reading and rebuilding the whole scope. bm25_buffer holds the + # new document's {id, text} rows. + bm25.add_chunks(state.scope, bm25_buffer) + else: + full_scope_chunks = _all_scope_chunks(state.scope) + bm25.rebuild_index(state.scope, full_scope_chunks) _update_document_row( state.document_id, status = "completed", @@ -881,26 +917,35 @@ def enqueue_ingestion( with _jobs_lock: _jobs[job_id] = state - out_queue = _CTX.Queue() - state.out_queue = out_queue - proc = _CTX.Process( - target = _subprocess_worker, - args = ( - str(stored_path), - model_name, - RAG_CHUNK_SIZE, - RAG_CHUNK_OVERLAP, - RAG_EMBED_BATCH_SIZE, - out_queue, - chunking_strategy, - mode, - document_id, - vlm_url, - vlm_model, - enable_captions, - ), - daemon = True, + worker_args = ( + str(stored_path), + model_name, + RAG_CHUNK_SIZE, + RAG_CHUNK_OVERLAP, + RAG_EMBED_BATCH_SIZE, + None, # out_queue, filled below + chunking_strategy, + mode, + document_id, + vlm_url, + vlm_model, + enable_captions, ) + if _fast_ingest(): + # In-process: the worker's get_embedder() hits the warm startup singleton + # instead of a spawned interpreter reloading the model every upload. + out_queue = queue_module.Queue() + worker_args = worker_args[:5] + (out_queue,) + worker_args[6:] + proc: Any = _ThreadProc(target = _subprocess_worker, args = worker_args) + else: + out_queue = _CTX.Queue() + worker_args = worker_args[:5] + (out_queue,) + worker_args[6:] + proc = _CTX.Process( + target = _subprocess_worker, + args = worker_args, + daemon = True, + ) + state.out_queue = out_queue proc.start() state.proc = proc pump_thread = threading.Thread( diff --git a/studio/backend/main.py b/studio/backend/main.py index 884a5e5841..240502283f 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -380,6 +380,22 @@ async def lifespan(app: FastAPI): threading.Thread(target = _precache, daemon = True).start() + # Fast RAG: warm the embedder at startup so the first upload's indexing does + # not pay the cold model load. Background thread, like the precache above. + if os.environ.get("UNSLOTH_RAG_FAST") == "1": + def _warm_rag_embedder(): + try: + from core.rag import embeddings + from utils.rag.config import resolve_embedder + + model_name = resolve_embedder("text", "standard") + embeddings.get_embedder(model_name) + logger.info("RAG fast: embedder warmed at startup (%s)", model_name) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG fast: embedder warmup failed: %s", exc) + + threading.Thread(target = _warm_rag_embedder, daemon = True).start() + # Initialize RSA key pair for API key encryption (external providers) from core.inference.key_exchange import init_key_pair