Studio: fast RAG path (warm embedder + in-process ingest + FTS5 BM25) behind UNSLOTH_RAG_FAST

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)

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()

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(

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