diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 085235e8fb..21ecb9edf4 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -503,19 +503,13 @@ TERMINAL_TOOL = { } -# Lazy import — keeps studio.db init lazy so tools.py doesn't pull in -# the whole rag stack on inference paths that never see RAG. +# Lazy import: don't pull rag stack on inference paths that never see RAG. def _get_rag_tool_spec(): from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL return SEARCH_KNOWLEDGE_BASE_TOOL -# RAG_SEARCH_TOOL is included in ALL_TOOLS; routes/inference.py filters -# the list against payload.enabled_tools so each request only sees the -# tools the frontend explicitly enabled. When the RAG button is off -# the tool name won't be in enabled_tools and the LLM will never see -# the spec. ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()] diff --git a/studio/backend/core/rag/bm25.py b/studio/backend/core/rag/bm25.py index 508d75129d..0d45ba2420 100644 --- a/studio/backend/core/rag/bm25.py +++ b/studio/backend/core/rag/bm25.py @@ -1,17 +1,9 @@ # 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 lexical index using the ``bm25s`` library. +"""Per-scope BM25 index (rebuild on change; bm25s has no cheap incremental insert). -bm25s does not support incremental insertion cheaply, so we rebuild the -full per-scope index whenever its document set changes. At studio -scale (a few hundred to a few tens of thousands of chunks per KB) this -is fast enough; the upside is that deletes are trivial. - -A scope is ``kb_`` or ``thread_``. Each scope stores: - - ``/`` directory holding the bm25s index files - - ``/ids.json`` mapping the index's positional ids back to - chunk-id strings (bm25s returns row indices, not our ids). +Each scope dir holds the bm25s files + ids.json mapping row index → chunk_id. """ from __future__ import annotations @@ -48,11 +40,7 @@ def _evict(scope: str) -> None: def rebuild_index(scope: str, chunks: list[dict]) -> None: - """Rebuild the BM25 index for ``scope`` from the full chunk list. - - Each chunk dict must contain ``id`` and ``text``. Passing an empty - list deletes the scope's index files. - """ + """Rebuild scope's BM25 from full chunk list. Empty list deletes the index.""" import bm25s base = _scope_dir(scope) @@ -64,9 +52,7 @@ def rebuild_index(scope: str, chunks: list[dict]) -> None: tokens = bm25s.tokenize(texts, show_progress = False) retriever = bm25s.BM25() retriever.index(tokens, show_progress = False) - # Drop stale files from any previous build/library version before - # writing the new index so we never mix old + new artifacts in the - # scope dir (bm25s.BM25.save does not unlink files it does not write). + # bm25s.BM25.save does not unlink stale files; clear the dir first. delete_scope(scope) ensure_dir(base) retriever.save(str(base)) @@ -87,10 +73,7 @@ def _load(scope: str) -> tuple[Any, list[str]] | None: 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 or partially-written index files: treat as a - # missing index so search returns empty and a future - # re-ingest can rebuild cleanly. Log so the failure - # is visible without crashing the request. + # 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, diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index 3d07a7ce24..0f645a957e 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -20,7 +20,6 @@ TokenCounter = Callable[[str], int] def _char_token_estimate(text: str) -> int: - # Rough 4 chars / token heuristic — only used if no real tokenizer is provided. return max(1, (len(text) + 3) // 4) @@ -60,7 +59,6 @@ def _atomic_split( tail = separators[separators.index(sep) + 1 :] out.extend(_atomic_split(piece, tail, max_tokens, count)) return out - # No separator made progress — hard-slice by characters. approx_chars = max(1, max_tokens * 4) return [text[i : i + approx_chars] for i in range(0, len(text), approx_chars)] @@ -71,7 +69,7 @@ def _merge( overlap_tokens: int, count: TokenCounter, ) -> list[str]: - """Greedy-merge atomic pieces into chunks <= max_tokens with overlap between adjacent chunks.""" + """Greedy-merge into <= max_tokens chunks with overlap.""" chunks: list[str] = [] buffer: list[str] = [] buffer_tokens = 0 @@ -100,12 +98,8 @@ def _merge( return [c.strip() for c in chunks if c.strip()] +# Markdown headings first so layout-aware parser output splits at sections. DEFAULT_SEPARATORS: tuple[str, ...] = ( - # Markdown heading boundaries first — when the parser emits - # layout-aware Markdown (PDF via pymupdf4llm, DOCX via mammoth, - # HTML via markdownify) chunks split at section breaks rather - # than mid-paragraph. Falls back to the original separators on - # plain text input where headings are absent. "\n# ", "\n## ", "\n### ", @@ -126,11 +120,7 @@ def chunk_pages( token_counter: TokenCounter | None = None, separators: tuple[str, ...] = DEFAULT_SEPARATORS, ) -> list[Chunk]: - """Split parsed pages into overlapping chunks. - - Each page is split independently so page_number stays meaningful for - PDFs — cross-page chunks would lose source attribution. - """ + """Split pages independently so page_number stays attached to chunks.""" count = token_counter or _char_token_estimate out: list[Chunk] = [] for page in pages: @@ -158,23 +148,11 @@ def chunk_pages_with_spans( token_counter: TokenCounter | None = None, separators: tuple[str, ...] = DEFAULT_SEPARATORS, ) -> tuple[str, list[Chunk], list[tuple[int, int]]]: - """Late-chunking-friendly variant of :func:`chunk_pages`. + """Late-chunking variant: joins pages so the embedder sees the whole doc. - Joins all pages into a single document so the embedder sees the - whole text in one pass (that's the point of late chunking — chunk - vectors that carry full-document context via the model's - bidirectional attention). - - Returns ``(full_doc, chunks, char_spans)`` where - ``char_spans[i] = (start, end)`` are byte-character offsets of - ``chunks[i].text`` inside ``full_doc``. The embedder layer maps - char spans → token spans via the tokenizer's offsets_mapping and - mean-pools per chunk. - - Page-number metadata on each :class:`Chunk` is recovered from the - chunk's char span — the first page whose range overlaps the chunk - wins. PDFs keep useful citations even though chunking ignores page - boundaries here. + Returns ``(full_doc, chunks, char_spans)``; ``char_spans[i]`` is the + (start, end) char offset of ``chunks[i].text`` inside ``full_doc``. + Page numbers are recovered by overlap with the original page ranges. """ count = token_counter or _char_token_estimate @@ -203,13 +181,9 @@ def chunk_pages_with_spans( continue idx = full_doc.find(text, search_cursor) if idx < 0: - # Overlap can push the search cursor past a chunk's true - # start — restart from the document head as a fallback. + # Overlap can push past a chunk's true start; restart from head. idx = full_doc.find(text) if idx < 0: - # Chunker output diverged from the source (rare — happens - # if a separator-splice mangled the text). Skip the chunk - # rather than corrupt the vector store with a wrong span. continue end_idx = idx + len(text) page_number = _page_for_span(idx, end_idx, page_ranges) @@ -221,8 +195,7 @@ def chunk_pages_with_spans( ) ) char_spans.append((idx, end_idx)) - # Advance past the *start* of this chunk so an overlapping - # next chunk can still be found. + # Advance past start (not end) so overlapping next chunk is findable. search_cursor = idx + 1 return full_doc, chunks, char_spans diff --git a/studio/backend/core/rag/db.py b/studio/backend/core/rag/db.py index 1bb05c5fbd..954bffb0ad 100644 --- a/studio/backend/core/rag/db.py +++ b/studio/backend/core/rag/db.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""sqlite-vec backed connection helper for RAG vectors. - -Single process-wide connection, opened lazily on first use. The -extension-load step runs once at open time. studio.db (chat history, -RAG metadata) stays untouched, so the extension-load surface is scoped -to the RAG code path only — chat code keeps its plain sqlite handle. -""" +"""Lazy process-wide rag.db connection with sqlite-vec loaded.""" from __future__ import annotations @@ -29,14 +23,6 @@ def rag_db_path() -> Path: def _load_sqlite_vec(conn: sqlite3.Connection) -> None: - """Enable extension loading and pull in sqlite-vec. - - install.sh creates the studio venv via `uv venv --python `, - which uses uv's managed python-build-standalone build. That CPython - is compiled with --enable-loadable-sqlite-extensions, so this path - succeeds on standard installs. The actionable error message is - here for the rare custom-interpreter case. - """ try: conn.enable_load_extension(True) except AttributeError as exc: @@ -77,12 +63,6 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: def get_rag_connection() -> sqlite3.Connection: - """Lazy process-wide sqlite connection to rag.db with sqlite-vec loaded. - - Returns the cached connection on subsequent calls. FastAPI's thread - pool plus check_same_thread=False + WAL mode handles concurrent - reads; writes are serialized by SQLite itself. - """ global _conn with _conn_lock: if _conn is None: @@ -101,7 +81,6 @@ def get_rag_connection() -> sqlite3.Connection: def _reset_for_tests() -> None: - """Drop the cached connection. Test-only — production never calls.""" global _conn with _conn_lock: if _conn is not None: diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 299190fab3..0de28eef81 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -1,15 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Embedding model singleton for RAG. - -Loads the configured embedder via Unsloth's ``FastSentenceTransformer`` -wrapper with ``for_inference=True`` (which returns a plain -``sentence_transformers.SentenceTransformer`` instance with proper dtype -and device handling). Lifecycle is fully independent of the chat -``InferenceBackend`` so loading an embedder cannot evict the active -chat model. -""" +"""RAG embedder singleton. Independent of the chat InferenceBackend.""" from __future__ import annotations @@ -30,19 +22,13 @@ _embedding_dim: int | None = None def _load(model_name: str) -> Any: logger.info("Loading RAG embedder: %s", model_name) - # BGE-VL ships a sentence-transformers shim that's tightly coupled - # to a specific ST internal API and breaks across ST version bumps. - # Bypass ST entirely and load via the canonical transformers - # AutoModel path, wrapped to match the SentenceTransformer API - # slice the RAG ingester uses. + # BGE-VL's ST shim breaks across ST versions; load via AutoModel. if model_name.startswith("BAAI/BGE-VL"): return _BGEVLAdapter(model_name) from unsloth import FastSentenceTransformer - # trust_remote_code is required for nomic-embed-text-v1.5 (custom - # modeling for 8K context). Safe to enable because the embedder - # matrix is config-pinned — users don't supply arbitrary names. + # trust_remote_code: nomic-embed-text-v1.5 needs custom modeling for 8K ctx. return FastSentenceTransformer.from_pretrained( model_name, for_inference = True, @@ -51,18 +37,7 @@ def _load(model_name: str) -> Any: class _BGEVLAdapter: - """Adapter exposing the slice of SentenceTransformer API the RAG - ingester depends on, backed by BGE-VL's transformers AutoModel. - - Supports: - - ``encode(list_of_strings, ...)`` → text embeddings - - ``encode(list_of_PIL_images, ...)`` → image embeddings - - ``get_sentence_embedding_dimension()`` - - ``tokenize([text])`` for token-aware chunking (best-effort) - - Auto-detects image vs text inputs from the first element. Returns - L2-normalized numpy arrays when ``normalize_embeddings=True``. - """ + """SentenceTransformer-shaped adapter over BGE-VL's AutoModel.""" def __init__(self, hf_model_name: str): from transformers import AutoModel @@ -72,9 +47,7 @@ class _BGEVLAdapter: hf_model_name, trust_remote_code = True, ) - # BGE-VL's encode() requires set_processor to install the - # tokenizer / image processor on the model. Without it, the - # first encode() raises with a missing-processor error. + # Required: BGE-VL's encode() raises without an installed processor. self._model.set_processor(hf_model_name) device = "cuda" if torch.cuda.is_available() else "cpu" self._model.to(device).eval() @@ -86,10 +59,7 @@ class _BGEVLAdapter: return F.normalize(tensor, p = 2.0, dim = -1) - # CLIP-family text encoder context cap. BGE-VL inherits CLIP's - # 77-token text positional embedding table — exceeding it triggers - # a shape-mismatch in the embedding layer. Pre-truncate any text - # chunk to this length before calling the model. + # CLIP positional embedding cap; longer text triggers shape mismatch. _CLIP_TEXT_MAX_TOKENS = 77 def encode( @@ -140,13 +110,7 @@ class _BGEVLAdapter: return out.numpy() if convert_to_numpy else out def _encode_text_truncated(self, texts: list[str]): - """Tokenize with explicit truncation to CLIP's 77-token limit, then - call get_text_features directly. BGE-VL's high-level encode() does - not truncate, so longer chunks overflow the position-embedding - table and crash inside the text model. Text chunks beyond the cap - are silently truncated — the multimodal mode is primarily about - the image side; large text chunks should land in text-mode RAG. - """ + """Truncate to CLIP's 77-token limit; long text in multimodal mode is lossy.""" import torch tokenizer = self._get_text_tokenizer() @@ -158,14 +122,7 @@ class _BGEVLAdapter: max_length = self._CLIP_TEXT_MAX_TOKENS, ) inputs = {k: v.to(self._device) for k, v in inputs.items()} - # Some chunks were almost certainly truncated; flag it once per - # batch so users running multimodal on long-form text know the - # text channel is lossy by design. - overflowed = any( - len(t.split()) > 30 # ~rough proxy; tokens vary by lang - for t in texts - ) - if overflowed: + if any(len(t.split()) > 30 for t in texts): logger.info( "BGE-VL text encode: truncating chunks to %d tokens (CLIP cap)", self._CLIP_TEXT_MAX_TOKENS, @@ -174,7 +131,6 @@ class _BGEVLAdapter: return self._model.get_text_features(**inputs) def _get_text_tokenizer(self): - """Locate the tokenizer set up by ``set_processor`` for text input.""" processor = getattr(self._model, "processor", None) if processor is not None: tok = getattr(processor, "tokenizer", None) @@ -192,11 +148,6 @@ class _BGEVLAdapter: return self._dim def tokenize(self, texts): - """Best-effort tokenize for the token_counter chunking path. - - Falls back gracefully — the caller already handles exceptions - by approximating tokens as ``len(text) // 4`` when this raises. - """ return self._get_text_tokenizer()( texts, return_tensors = "pt", @@ -205,7 +156,6 @@ class _BGEVLAdapter: def get_embedder(model_name: str | None = None) -> Any: - """Return the cached SentenceTransformer, loading it on first use.""" global _model, _model_name, _embedding_dim target = model_name or RAG_EMBEDDING_MODEL with _lock: @@ -255,14 +205,7 @@ def encode_images( batch_size: int | None = None, normalize: bool = True, ): - """Embed raw image bytes via a multimodal SentenceTransformer. - - Works with CLIP-family models (BGE-VL, openai/clip-*) whose - `encode` accepts PIL.Image objects in the same call as text. The - returned vectors live in the same 512-d (or model-specific) space - as text vectors from this model, so a single scope's vector rows - hold both kinds. - """ + """Embed image bytes via a CLIP-family multimodal encoder.""" from io import BytesIO from PIL import Image @@ -281,11 +224,7 @@ def encode_images( def token_counter(model_name: str | None = None): - """Return a ``len(tokenize(text))`` callable using the embedder's tokenizer. - - Avoid loading the model just for chunking by reaching through the - SentenceTransformer's ``tokenize`` API. - """ + """Return a token-count callable backed by the embedder's tokenizer.""" model = get_embedder(model_name) def _count(text: str) -> int: @@ -301,9 +240,7 @@ def token_counter(model_name: str | None = None): return _count -# ------------------------------------------------------------------ -# Late chunking (Phase 3B-late) -# ------------------------------------------------------------------ +# --- Late chunking (Jina technique) --- _LATE_WINDOW_OVERLAP_TOKENS = 512 @@ -315,18 +252,7 @@ def late_chunk_encode( model_name: str | None = None, normalize: bool = True, ): - """Embed each chunk via late-chunking pooling. - - Single forward pass over the full document, then mean-pool the - token embeddings whose offset ranges fall inside each chunk's - char span. Chunks therefore carry full-document context via the - encoder's bidirectional attention — Jina's published technique, - works with any encoder that exposes per-token outputs. - - When the doc exceeds the embedder's context, falls back to - windowed late chunking with a 512-token overlap between windows - so cross-window context is partially preserved. - """ + """Single forward pass over the doc, mean-pool token embeddings per chunk span.""" import numpy as np if not char_spans: @@ -373,7 +299,6 @@ def late_chunk_encode( def _encode_tokens(model, encoded): - """Run the embedder's underlying transformer to get per-token last_hidden_state.""" import torch transformer = model[0].auto_model @@ -395,15 +320,11 @@ def _pool_spans( doc_text: str, token_index_offset: int = 0, ): - """Mean-pool token embeddings per (char_start, char_end) span. - - `token_index_offset` shifts char_span-derived token indices into - a sub-window's local frame (used by the windowed code path). - """ + """Mean-pool token embeddings per (char_start, char_end) span.""" vectors = [] n_rows = token_embeddings.shape[0] for char_start, char_end in char_spans: - # Special tokens (CLS / SEP) report offsets (0, 0) — exclude them. + # Skip special tokens whose offsets are (0, 0). indices = [ i - token_index_offset for i, (ts, te) in enumerate(offsets) @@ -411,9 +332,6 @@ def _pool_spans( ] indices = [i for i in indices if 0 <= i < n_rows] if not indices: - # Fall back to a standalone encode of the chunk text — rare - # (would mean tokenizer produced zero non-special tokens for - # the span), but keeps the pipeline alive. vec = model.encode( doc_text[char_start:char_end], normalize_embeddings = normalize, @@ -440,12 +358,7 @@ def _windowed_late_chunk_encode( normalize: bool, np_module, ): - """Doc exceeds context window — slice into overlapping windows. - - Each chunk is pooled against the window that contains the most of - its tokens. The 512-token window overlap means chunks near a - boundary still see context from both sides. - """ + """Doc > ctx window: pool each chunk against the window containing most of its tokens.""" import torch tokenizer = model.tokenizer @@ -464,7 +377,6 @@ def _windowed_late_chunk_encode( n_tokens = int(all_input_ids.shape[0]) stride = max(1, max_length - _LATE_WINDOW_OVERLAP_TOKENS) - # Build (start_token, end_token) windows. windows: list[tuple[int, int]] = [] pos = 0 while pos < n_tokens: @@ -474,7 +386,6 @@ def _windowed_late_chunk_encode( break pos += stride - # Cache window → token embeddings (only encode when needed). window_embeddings: dict[int, "np_module.ndarray"] = {} def _window_embeddings(window_index: int): @@ -491,7 +402,6 @@ def _windowed_late_chunk_encode( vectors = [] for char_start, char_end in char_spans: - # Collect global token indices in the chunk. chunk_token_indices = [ i for i, (ts, te) in enumerate(all_offsets) @@ -506,7 +416,6 @@ def _windowed_late_chunk_encode( ) vectors.append(vec) continue - # Pick the window covering the most of this chunk's tokens. best_window = 0 best_overlap = 0 for wi, (ws, we) in enumerate(windows): diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index b43fc849de..8520c2c96d 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -1,19 +1,10 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Document ingestion pipeline. +"""RAG ingestion pipeline. -Follows the studio's existing job pattern (`core/data_recipe/jobs/manager.py`): -spawn a fresh subprocess per job with ``mp.get_context("spawn")`` and stream -progress events back over a queue. The subprocess does the heavy work -(parse → chunk → load embedder → embed in batches) and ships -``(chunks, vectors)`` batches back. The parent persists everything: -sqlite rows, vector_store rows in rag.db, and (at job completion) a -rebuilt BM25 index. - -Only the parent process owns the rag.db connection (sqlite-vec loaded -there); subprocesses never open it directly. This keeps search -available throughout the lifetime of an ingestion job. +Spawn-subprocess per job (parse/chunk/embed); parent persists chunks, +vectors, and rebuilds BM25 on completion. Only the parent opens rag.db. """ from __future__ import annotations @@ -44,9 +35,7 @@ _CTX = mp.get_context("spawn") _QUEUE_TIMEOUT_SECONDS = 300 -# ------------------------------------------------------------------ -# Subprocess worker -# ------------------------------------------------------------------ +# --- Subprocess worker --- _MIME_TO_EXT = { "image/png": ".png", @@ -108,7 +97,6 @@ def _subprocess_worker( ) return - # Standard chunking path (text + optional images for multimodal mode). text_count = _run_standard_chunking( pages = pages, chunk_size = chunk_size, @@ -147,12 +135,7 @@ def _run_standard_chunking( out_queue, send_complete: bool = True, ) -> int: - """Stream text chunks back to the parent. Returns the number streamed. - - When called as part of the multimodal pipeline, `send_complete` is - False because image chunks still need to be streamed before the - document is marked complete. - """ + """Stream text chunks; returns count. send_complete=False when images follow.""" out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) chunks = chunk_pages( pages, @@ -207,17 +190,7 @@ def _stream_image_chunks( out_queue, first_index: int, ) -> int: - """Save extracted images to disk and stream image+caption chunks. - - Each image becomes an `image`-kind chunk; if the parser found an - adjacent caption, a paired `caption`-kind chunk is also emitted. - Pairs share a ``pair_group`` so the parent can link them via - rag_chunks.linked_chunk_id. - - Images are written to ``rag_uploads_root() / 'images' / - / img-NNN.`` so the parent can serve them via - the static-image route without holding bytes in memory. - """ + """Persist images, emit image+caption chunks; pairs share pair_group.""" from core.rag.embeddings import encode, encode_images from utils.paths.storage_roots import ensure_dir, rag_uploads_root @@ -228,7 +201,6 @@ def _stream_image_chunks( img_dir = ensure_dir(rag_uploads_root() / "images" / document_id) - # Persist bytes, build parallel lists for encoding. paths: list[str] = [] bytes_for_encoding: list[bytes] = [] captions: list[str] = [] @@ -251,7 +223,6 @@ def _stream_image_chunks( image_vectors = encode_images(bytes_for_encoding, model_name = model_name) - # Embed only the non-empty captions; track which images they map to. caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()] if caption_to_image: caption_vectors_arr = encode( @@ -280,8 +251,6 @@ def _stream_image_chunks( } ) out_vectors.append(image_vectors[idx].tolist()) - # Emit the caption chunk right after its image so the parent - # sees them adjacent (simplifies pair linking). if next_cap is not None and next_cap[0] == idx: _cap_index, cap_vec = next_cap out_chunks.append( @@ -319,13 +288,7 @@ def _run_late_chunking( late_chunk_encode, out_queue, ) -> None: - """Late chunking: chunk once over the whole doc, embed in a single pass. - - There's no per-batch streaming here — the whole doc is encoded in - one forward pass (or one per window for long docs). We ship all - chunks back to the parent in one message; the parent's pump still - handles them via the same chunks_batch handler. - """ + """Chunk once, embed in one pass, ship all chunks in one chunks_batch.""" from core.rag.chunking import chunk_pages_with_spans out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) @@ -366,9 +329,7 @@ def _run_late_chunking( out_queue.put({"type": "complete", "num_chunks": len(chunks)}) -# ------------------------------------------------------------------ -# Job manager (parent side) -# ------------------------------------------------------------------ +# --- Job manager (parent side) --- class _JobState: @@ -450,13 +411,7 @@ def _insert_chunks_and_collect_for_bm25( chunks_meta: list[dict], vectors: list[list[float]], ) -> list[dict]: - """Insert chunks into sqlite + vector_store; return [{id, text}] for BM25. - - Image-kind chunks ship a stable image_path and skip BM25 (no text - body to tokenise). Paired image/caption chunks share a pair_group - field — the second pass links them via rag_chunks.linked_chunk_id - so retrieval can dereference an image hit to its caption. - """ + """Insert chunks into sqlite + vector_store; return [{id, text}] for BM25.""" rows: list[tuple] = [] points: list[dict] = [] bm25_rows: list[dict] = [] @@ -496,9 +451,6 @@ def _insert_chunks_and_collect_for_bm25( }, } ) - # BM25 indexes text + caption chunks. Image chunks have no - # tokenisable body — their caption (if any) is in a separate - # caption-kind chunk that BM25 will index. if kind in ("text", "caption") and meta["text"]: bm25_rows.append({"id": chunk_id, "text": meta["text"]}) with get_connection() as conn: @@ -511,9 +463,7 @@ def _insert_chunks_and_collect_for_bm25( """, rows, ) - # Link image ↔ caption pairs. We only set linked_chunk_id when - # a pair_group has exactly two members; lone images stay - # unlinked (no caption was paired). + # Link only when exactly two members in a pair_group. for ids in pair_groups.values(): if len(ids) != 2: continue @@ -560,7 +510,7 @@ def _pump( proc: Any, out_queue: Any, ) -> None: - """Drain queue messages until the subprocess signals complete/error or dies.""" + """Drain queue until subprocess completes/errors/dies.""" bm25_buffer: list[dict] = [] embedding_dim: int | None = None final_status = "failed" @@ -597,7 +547,6 @@ def _pump( vector_store.ensure_collection(state.scope, embedding_dim) elif mtype == "chunks_batch": if embedding_dim is None: - # defensive: subprocess should always emit "dim" first embedding_dim = len(msg["vectors"][0]) if msg["vectors"] else None if embedding_dim is not None: vector_store.ensure_collection(state.scope, embedding_dim) @@ -675,15 +624,7 @@ def enqueue_ingestion( chunking_strategy: str = "standard", mode: str = "text", ) -> str: - """Create the job row, spawn the subprocess, and start the pump thread. - - Returns the job_id. The caller can poll via ``GET /api/rag/jobs/{job_id}/events`` - or read the ``rag_ingestion_jobs`` table directly. - - chunking_strategy / mode default to today's behaviour. KB-scoped - uploads should pass the KB's stored values; per-thread uploads - default unless an override is set in chat_settings. - """ + """Create the job row, spawn the subprocess, start the pump; return job_id.""" from utils.rag.config import resolve_embedder scope = _scope_for(kb_id, thread_id) @@ -736,11 +677,7 @@ def enqueue_ingestion( def delete_document_artifacts(document_id: str, scope: str) -> None: - """Remove a document's vectors, then rebuild BM25 for the scope. - - The caller is responsible for the sqlite cascade (deleting the - rag_documents row triggers ON DELETE CASCADE on rag_chunks). - """ + """Drop the doc's vectors, rebuild BM25. Caller deletes the rag_documents row.""" vector_store.delete_document(scope, document_id) remaining = _all_scope_chunks(scope) if remaining: @@ -755,11 +692,7 @@ def delete_scope_artifacts(scope: str) -> None: def purge_thread_documents(thread_ids: list[str]) -> None: - """Remove all RAG artifacts owned by the given chat thread ids. - - Used by the chat-thread DELETE handlers because rag_documents has - no FK cascade to chat_threads (see schema comment). - """ + """Drop RAG artifacts for the given thread ids (no FK cascade to chat_threads).""" if not thread_ids: return import os @@ -791,7 +724,7 @@ def purge_thread_documents(thread_ids: list[str]) -> None: def purge_all_thread_documents() -> None: - """Drop every per-thread RAG artifact. Used by clear-all-history.""" + """Drop every per-thread RAG artifact.""" with get_connection() as conn: rows = conn.execute( "SELECT DISTINCT thread_id FROM rag_documents WHERE thread_id IS NOT NULL" diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py index b3d910ad0e..040e2b0861 100644 --- a/studio/backend/core/rag/parsers/__init__.py +++ b/studio/backend/core/rag/parsers/__init__.py @@ -9,15 +9,7 @@ from pathlib import Path @dataclass(frozen = True) class ParsedPage: - """One page (or page-equivalent) of Markdown-rendered text from a source document. - - For PDFs `page_number` is the 1-indexed physical page. For DOCX / HTML / - TXT / MD the whole document is one ParsedPage with `page_number = None`. - - Text is expected to be Markdown — heading markers (`#`, `##`, …), - pipe-tables, and list bullets survive extraction so the chunker can - split on them. Parsers MUST emit Markdown, not bare plain text. - """ + """Markdown text from one page (PDF) or whole doc (others).""" text: str page_number: int | None = None @@ -25,14 +17,7 @@ class ParsedPage: @dataclass(frozen = True) class ParsedImage: - """One image extracted from a source document. - - Captured only when `parse(..., want_images=True)` is set — the - multimodal ingestion path in Phase 3B-multimodal consumes these. - `nearest_caption` is best-effort paragraph-adjacency; can be empty - when no caption could be paired (the image still ingests, just - without the paired-caption chunk). - """ + """Image extracted with want_images=True; nearest_caption may be empty.""" image_bytes: bytes mime_type: str @@ -42,13 +27,6 @@ class ParsedImage: @dataclass(frozen = True) class ParseResult: - """Result of parsing a single source document. - - `pages` is always populated; `images` is empty unless the caller - passed `want_images=True`. Iteration aliases for `pages` so legacy - code that did `for page in parse(path)` keeps working. - """ - pages: list[ParsedPage] = field(default_factory = list) images: list[ParsedImage] = field(default_factory = list) diff --git a/studio/backend/core/rag/parsers/docx.py b/studio/backend/core/rag/parsers/docx.py index bf16a9052c..c3ad127140 100644 --- a/studio/backend/core/rag/parsers/docx.py +++ b/studio/backend/core/rag/parsers/docx.py @@ -1,17 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""DOCX parsing via mammoth. - -Mammoth converts Word documents to Markdown while preserving Heading -styles (`# `, `## `, ...), bullet/numbered lists, tables, and basic -emphasis. This replaces the previous python-docx paragraph-concat -approach that lost all heading metadata. - -Images are captured via the `convert_image` handler when -`want_images=True`, falling back to python-docx for inline image bytes -if mammoth's docx adapter can't reach them. -""" +"""DOCX → Markdown via mammoth (preserves headings, lists, tables).""" from __future__ import annotations @@ -24,9 +14,7 @@ from . import ParsedImage, ParsedPage, ParseResult logger = logging.getLogger(__name__) -# Mammoth uses some default DOCX-style-name → Markdown mappings, but a -# few common variants ship with non-default names. Map them explicitly -# so we don't lose headings. +# Force common DOCX heading style names to Markdown headings. _STYLE_MAP = """ p[style-name='Title'] => h1.title:fresh p[style-name='Subtitle'] => h2.subtitle:fresh @@ -40,11 +28,9 @@ p[style-name='Heading 6'] => h6:fresh def _html_to_markdown(html: str) -> str: - """Convert mammoth's HTML output to Markdown via markdownify.""" from markdownify import markdownify md = markdownify(html, heading_style = "ATX", strip = ["script", "style"]) - # markdownify can emit excessive blank lines on tables; tighten up. md = re.sub(r"\n{3,}", "\n\n", md) return md.strip() @@ -55,12 +41,7 @@ def extract(path: Path, *, want_images: bool = False) -> ParseResult: images: list[ParsedImage] = [] if want_images: - # mammoth's image converter is called for every inline image. - # We capture the bytes here and substitute a stable placeholder - # in the rendered Markdown so the chunker doesn't trip over - # base64 blobs. Caption-pairing is approximate — we use the - # full document text as the caption pool (better than nothing - # for DOCX where heading→figure adjacency isn't reliable). + # Capture bytes; suppress src so base64 doesn't land in Markdown. def _convert(image): with image.open() as image_bytes: blob = image_bytes.read() @@ -77,8 +58,6 @@ def extract(path: Path, *, want_images: bool = False) -> ParseResult: convert_image = mammoth.images.img_element(_convert) else: - # Drop image elements entirely — cheaper and avoids embedding - # base64 in Markdown when the caller doesn't want images. convert_image = mammoth.images.img_element(lambda _image: {"src": ""}) with open(path, "rb") as fp: @@ -92,9 +71,7 @@ def extract(path: Path, *, want_images: bool = False) -> ParseResult: markdown = _html_to_markdown(result.value) if want_images and images: - # Best-effort: every image inherits the whole doc text as a - # caption pool. Phase 3B-multimodal will improve this once - # multimodal embedders consume captions directly. + # Approximate caption: first 1500 chars of doc. caption_pool = markdown[:1500] images = [ ParsedImage( diff --git a/studio/backend/core/rag/parsers/html.py b/studio/backend/core/rag/parsers/html.py index 721e280736..8050dee770 100644 --- a/studio/backend/core/rag/parsers/html.py +++ b/studio/backend/core/rag/parsers/html.py @@ -1,18 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""HTML parsing via markdownify. - -Converts HTML to Markdown so headings (`

`…`

`), tables, and lists -arrive at the chunker as Markdown structure. The previous -`BeautifulSoup.get_text()` approach stripped all tags and made every -heading indistinguishable from body text. - -Image extraction (when `want_images=True`) only handles local file -references — remote URLs are skipped to avoid network calls during -ingestion. Phase 3B-multimodal can revisit this if HTML inputs with -remote images become a common pattern. -""" +"""HTML → Markdown via markdownify. Images only resolve local file refs.""" from __future__ import annotations @@ -35,13 +24,12 @@ def _collect_local_images(soup, html_path: Path) -> list[ParsedImage]: src = tag.get("src") or "" parsed = urlparse(src) if parsed.scheme and parsed.scheme not in ("file", ""): - # Remote / data URLs — skip; we don't fetch over network. continue local_path = (base_dir / unquote(parsed.path or src)).resolve() try: local_path.relative_to(base_dir.resolve()) except ValueError: - # Refuse to read outside the source's own directory. + # Path traversal: refuse to read outside the source's directory. continue if not local_path.is_file(): continue diff --git a/studio/backend/core/rag/parsers/pdf.py b/studio/backend/core/rag/parsers/pdf.py index 0d312072f1..eeb46e2ce7 100644 --- a/studio/backend/core/rag/parsers/pdf.py +++ b/studio/backend/core/rag/parsers/pdf.py @@ -1,16 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Layout-aware PDF parsing via pymupdf + pymupdf4llm. - -Produces Markdown per page (headings, pipe-tables, lists survive) so the -recursive chunker can split on heading boundaries. Falls back to pypdf -text-extraction only when pymupdf fails to open the file — keeps the -pipeline alive for malformed PDFs. - -Image extraction is gated behind `want_images=True` so text-only KBs -pay zero cost for images they don't index. -""" +"""PDF → Markdown via pymupdf4llm; pypdf fallback for malformed files.""" from __future__ import annotations @@ -39,8 +30,7 @@ def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult: show_progress = False, ) except Exception: - # pymupdf4llm can choke on individual pages (rare). Fall - # back to plain text extraction for just that page. + # pymupdf4llm can choke on a single page; fall back to plain text. md = doc[page_index].get_text("text") or "" md = md.strip() if md: @@ -55,7 +45,6 @@ def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult: def _extract_images_pymupdf(doc, pages: list[ParsedPage]) -> list[ParsedImage]: - """Pull embedded images and pair each with the nearest text on the same page.""" captions_by_page: dict[int, str] = { p.page_number: p.text for p in pages if p.page_number } diff --git a/studio/backend/core/rag/parsers/text.py b/studio/backend/core/rag/parsers/text.py index 439a07440b..78b3cb610a 100644 --- a/studio/backend/core/rag/parsers/text.py +++ b/studio/backend/core/rag/parsers/text.py @@ -9,7 +9,6 @@ from . import ParsedPage, ParseResult def extract(path: Path, *, want_images: bool = False) -> ParseResult: - # want_images is ignored — plain text / Markdown have no embedded images. raw = path.read_bytes() try: text = raw.decode("utf-8") diff --git a/studio/backend/core/rag/reranker.py b/studio/backend/core/rag/reranker.py index 1496b2c271..d234bf00d6 100644 --- a/studio/backend/core/rag/reranker.py +++ b/studio/backend/core/rag/reranker.py @@ -1,18 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Optional cross-encoder reranking stage. - -Off by default. Callers opt in per request via ``enable_rerank`` on -``SearchRequest``. The reranker model is lazy-loaded on first opt-in -query and competes with the active chat model for GPU memory — keeping -it opt-in protects chat latency on smaller GPUs. - -``sentence_transformers.CrossEncoder`` has no unsloth wrapper today; -this module loads it directly. A future ``FastCrossEncoder`` addition -to ``unsloth/models/sentence_transformer.py`` would slot in here by -replacing the import in ``_load``. -""" +"""Opt-in CrossEncoder reranker (off by default; shares GPU with chat model).""" from __future__ import annotations @@ -51,12 +40,7 @@ def get_reranker(model_name: str | None = None) -> Any: def unload() -> None: - """Drop the reranker reference and trigger a GC pass. - - Useful when memory pressure is high — callers can free the - reranker without restarting the studio process. Next ``rerank`` - call lazy-loads it again. - """ + """Drop the reranker; next call lazy-loads again.""" global _model, _model_name with _lock: if _model is not None: @@ -79,16 +63,9 @@ def rerank( model_name: str | None = None, top_k: int | None = None, ) -> list[Hit]: - """Re-order ``pairs`` by CrossEncoder relevance to ``query``. - - Each pair is ``(Hit, chunk_text)``. Returns Hits with the new - cross-encoder scores attached. If ``top_k`` is given, truncates. - """ + """Re-order (Hit, text) pairs by CrossEncoder score; image hits are appended last.""" if not pairs: return [] - # CrossEncoder is text-only; image-kind hits get appended at the end - # in their original relative order so they're never dropped, just - # never reranked. Caption-kind hits are eligible (they carry text). text_pairs = [(h, t) for h, t in pairs if h.kind != "image"] image_hits = [h for h, _t in pairs if h.kind == "image"] diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index 4c48519151..e6eb45c276 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -1,17 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""High-level retrieval surface for RAG: BM25, dense, and RRF hybrid. - -Reciprocal Rank Fusion is parameter-light: each candidate's fused score -is the sum of ``1 / (rrf_k + rank)`` across rankers. It avoids the -need to calibrate score scales between BM25 (raw, unbounded) and cosine -similarity (-1..1). - -Hits also carry the raw dense cosine score when available so callers -can apply a meaningful similarity threshold (e.g., "drop hits below -0.3 cosine") — the fused RRF score isn't on a comparable scale. -""" +"""RAG retrieval: BM25, dense, and RRF hybrid. Hits carry dense_score for thresholding.""" from __future__ import annotations @@ -34,10 +24,7 @@ class Hit: document_id: str | None = None chunk_index: int | None = None kind: str = "text" - # Raw cosine similarity from the dense retriever (0..1 for - # normalized embeddings). None when this chunk wasn't returned by - # the dense pass (BM25-only hit) — callers applying a similarity - # floor should treat None as "no signal" and exclude it. + # Raw cosine; None for BM25-only hits. dense_score: float | None = None @@ -54,12 +41,7 @@ def retrieve_dense( document_ids: list[str] | None = None, embedder_model: str | None = None, ) -> list[Hit]: - """Dense retrieval. `embedder_model` MUST match the model that - populated this scope's vectors — using a different one yields a - dim mismatch (shape (N, scope_dim) vs (query_dim,)) at distance - compute time. Callers should resolve from the KB / thread settings - before passing. - """ + """Dense retrieval. embedder_model MUST match the model that populated this scope.""" limit = k or RAG_TOP_K_DENSE vector = embeddings.encode( [query], @@ -96,8 +78,7 @@ def _rrf_fuse( ) -> list[Hit]: fused: dict[str, float] = {} seen: dict[str, Hit] = {} - # Track the dense cosine score per chunk so it survives fusion — - # callers downstream filter on this, not the RRF score. + # Preserve dense_score through fusion for downstream thresholding. dense_scores: dict[str, float] = {} for ranking in rankings: for rank, hit in enumerate(ranking): @@ -148,12 +129,7 @@ def retrieve_hybrid( def filter_by_min_score(hits: list[Hit], min_score: float) -> list[Hit]: - """Drop hits whose dense cosine score is below ``min_score``. - - Hits without a dense score (BM25-only) are dropped too — there's - no comparable signal to evaluate them against the similarity floor. - Use ``min_score = 0.0`` (or negative) to disable the filter. - """ + """Drop hits whose dense_score < min_score; BM25-only hits dropped too.""" if min_score <= 0.0: return hits return [h for h in hits if h.dense_score is not None and h.dense_score >= min_score] diff --git a/studio/backend/core/rag/scope.py b/studio/backend/core/rag/scope.py index 270693faba..054915ce24 100644 --- a/studio/backend/core/rag/scope.py +++ b/studio/backend/core/rag/scope.py @@ -1,25 +1,13 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Scope identifiers + per-scope embedder resolution. - -A "scope" is the namespace key carried by vector rows and BM25 -indexes — ``kb_`` for stand-alone Knowledge Bases or -``thread_`` for per-thread document sets. - -`resolve_scope_embedder` looks up which embedder populated a scope's -vectors so the query side can re-use the same model. Lives in -``core/rag`` rather than ``routes`` so the inference-side tool -handler can call it without a route-→-core import cycle. -""" +"""Scope identifiers (kb_ / thread_) + per-scope embedder resolver.""" from __future__ import annotations from storage.studio_db import get_connection, list_chat_settings from utils.rag.config import resolve_embedder -# Persisted chat-settings keys. Defined here so the resolver (core) -# and the route handlers (routes/rag.py) share one source of truth. RAG_DEFAULTS_KEY = "rag.defaults" @@ -28,16 +16,7 @@ def thread_settings_key(thread_id: str) -> str: def resolve_scope_embedder(scope: str) -> str | None: - """Return the embedder used to populate ``scope``'s vector rows. - - Resolution order: - - ``kb_`` → ``rag_knowledge_bases.embedding_model`` column. - - ``thread_`` → per-thread override → app-level defaults - override → ``RAG_EMBEDDER_MATRIX[(mode, chunking)]``. - - Returns ``None`` for unrecognised scope strings; callers treat - ``None`` as "fall back to the configured default embedder". - """ + """KB → kb.embedding_model; thread → per-thread/defaults/matrix. None = use default.""" if scope.startswith("kb_"): kb_id = scope[len("kb_") :] with get_connection() as conn: diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 178da09991..9db3fa7fd2 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -3,15 +3,8 @@ """`search_knowledge_base` tool — RAG retrieval surfaced to the LLM. -Invoked from `core/inference/tools.execute_tool` when the local model -emits a `search_knowledge_base` call. The handler runs the existing -hybrid retrieval, hydrates chunk text + filename + page number from -sqlite, and returns a Markdown-with-numbered-citations string that -the LLM consumes as the tool-result message. - -Scope (`kb_id` / `thread_id`) is not exposed as a tool argument — it -comes from the chat-completions request body (`rag_scope`) so the -LLM doesn't need to know about KB UUIDs. +Scope comes from the request body (`rag_scope`), not from the tool args, +so the model never sees KB UUIDs. """ from __future__ import annotations @@ -61,12 +54,7 @@ SEARCH_KNOWLEDGE_BASE_TOOL = { def _format_hits_for_llm(hits: list[Any]) -> str: - """Render hits as numbered Markdown citations for the LLM. - - Empty results produce a one-line message rather than an empty - string — the model needs to know the search ran but found nothing - so it can fall back to its own knowledge or ask the user. - """ + """Render hits as numbered Markdown citations; empty results return a message, not ''.""" if not hits: return ( "No matching chunks were found in the attached documents. " @@ -95,13 +83,7 @@ def search_knowledge_base( min_score: float = 0.0, mode: Literal["bm25", "dense", "hybrid"] = "hybrid", ) -> str: - """Execute the RAG search and return a tool-result string. - - `kb_id` takes precedence over `thread_id` when both are set — - matches the create/upload contract that a document belongs to one - or the other, never both. ``min_score`` is a cosine-similarity - floor on dense hits; chunks below it are dropped. - """ + """Run RAG and return a tool-result string. kb_id takes precedence over thread_id.""" if not query or not query.strip(): return "Error: empty query." @@ -215,8 +197,7 @@ def search_knowledge_base( else: hits = hits[:k] - # Image-kind hits don't carry LLM-friendly text — skip them. The - # paired caption (linked_chunk_id) usually surfaces separately. + # Skip image-kind hits; the paired caption surfaces separately. formatted = [ lookup[hit.chunk_id] for hit in hits diff --git a/studio/backend/core/rag/vector_store.py b/studio/backend/core/rag/vector_store.py index 40536ff688..bcc07d3c72 100644 --- a/studio/backend/core/rag/vector_store.py +++ b/studio/backend/core/rag/vector_store.py @@ -1,21 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""SQLite-backed vector store for RAG, distance math via sqlite-vec. - -Replaces the previous Qdrant local-mode store. Vectors are stored as -BLOBs in a single `rag_vectors` table inside rag.db, keyed by -chunk_id; cosine distance is computed by sqlite-vec's -`vec_distance_cosine(blob, blob)` scalar function. - -A "scope" is `kb_` for standalone knowledge bases or -`thread_` for per-thread document sets. The scope column is -indexed; queries filter by scope before computing distances so -different scopes can hold vectors of different dimensions without -breaking the cosine math. The per-scope embedder resolver -(core/rag/scope.py:resolve_scope_embedder) guarantees one embedder -per scope, so dims within a scope are always consistent. -""" +"""sqlite-vec vector store. Scope filter (kb_/thread_) keeps mixed-dim +scopes safe — per-scope embedder resolver guarantees one dim per scope.""" from __future__ import annotations @@ -36,11 +23,6 @@ def thread_scope(thread_id: str) -> str: def collection_exists(scope: str) -> bool: - """Whether the scope has any indexed vectors. - - Used by callers (notably retrieval.retrieve_dense) to short-circuit - when the scope was never populated. Cheap — a covering index hit. - """ from core.rag.db import get_rag_connection conn = get_rag_connection() @@ -52,21 +34,12 @@ def collection_exists(scope: str) -> bool: def ensure_collection(scope: str, dim: int) -> None: - """No-op for sqlite-vec — vectors get inserted directly into the - shared table. Kept for API parity with the previous Qdrant store - so ingestion callers don't need conditional logic. - """ - _ = scope, dim # unused; signature preserved + """No-op; kept for API parity. Vectors go straight into the shared table.""" + _ = scope, dim def upsert_chunks(scope: str, points: Iterable[dict]) -> None: - """Insert/update vectors. Each point: {id, vector, payload}. - - Conflict resolution is per chunk_id (the primary key): re-ingesting - overwrites in place. Payload is round-tripped as JSON so the - Qdrant-shaped {filename, page_number, kind, ...} dicts callers - already build can be reused unchanged. - """ + """Insert/update vectors. Each point: {id, vector, payload}.""" import sqlite_vec from core.rag.db import get_rag_connection @@ -116,14 +89,7 @@ def search( top_k: int, document_ids: list[str] | None = None, ) -> list[dict]: - """Cosine-distance search filtered to a single scope. - - Returns rows in the same shape as the old Qdrant path — - {chunk_id, score, payload} — where ``score`` is cosine similarity - in [0, 1] (vec_distance_cosine returns 1 - similarity, so we - invert). Filtered-by-document_ids variant for the per-thread - "search only these uploads" case. - """ + """Cosine search; returns {chunk_id, score, payload} with score = 1 - distance.""" import sqlite_vec from core.rag.db import get_rag_connection diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index f80a3b6fc4..a45fb0d128 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -1,21 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""RAG API routes. - -Surface: - - Knowledge-base CRUD - - Document upload (KB-scoped and per-thread) - - Document list/delete - - Ingestion-job SSE stream - - Search (BM25 / dense / hybrid) - -Per-thread document uploads are scoped to a single chat thread and -share the same chunk/embed/index pipeline as KB documents — they only -differ in the scope key (``thread_`` vs ``kb_``) and lifecycle -(per-thread docs are dropped when the thread is deleted, via the -ON DELETE CASCADE on rag_documents.thread_id). -""" +"""RAG API: KB CRUD, document upload (KB + per-thread), ingestion SSE, search.""" from __future__ import annotations @@ -69,9 +55,7 @@ router = APIRouter() logger = get_logger(__name__) -# ------------------------------------------------------------------ -# Pydantic schemas -# ------------------------------------------------------------------ +# --- Pydantic schemas --- ChunkingStrategy = Literal["standard", "late"] KBMode = Literal["text", "multimodal"] @@ -81,9 +65,6 @@ class CreateKBRequest(BaseModel): name: str = Field(min_length = 1, max_length = 200) description: str | None = None embedding_model: str | None = None - # Phase 3 introduces two per-KB knobs. Both default to today's - # behaviour so existing API clients are unaffected. The (multimodal, - # late) combination is rejected at create time — see _validate_mode_combo. chunking_strategy: ChunkingStrategy = "standard" mode: KBMode = "text" @@ -145,9 +126,6 @@ class SearchRequest(BaseModel): document_ids: list[str] | None = None enable_rerank: bool = False reranker_model: str | None = None - # Cosine-similarity floor on the dense retrieval score. Hits whose - # dense_score is below this (or absent — BM25-only hits) are - # dropped before the response is sent. 0.0 disables the filter. min_score: float = Field(default = 0.0, ge = 0.0, le = 1.0) @@ -167,9 +145,7 @@ class SearchResponse(BaseModel): hits: list[SearchHit] -# ------------------------------------------------------------------ -# Helpers -# ------------------------------------------------------------------ +# --- Helpers --- def _sanitize_filename(filename: str) -> str: @@ -181,15 +157,10 @@ def _now_ms() -> int: return int(time.time()) -# Per-scope embedder lookup lives in core/rag/scope.py so the -# inference-side tool handler can use it without a routes-→-core -# import cycle. from core.rag.scope import resolve_scope_embedder as _resolve_scope_embedder # noqa: E402 def _row_to_kb(row: Any) -> KBResponse: - # chunking_strategy / mode may be absent on rows fetched through a - # pre-Phase-3 connection in tests; fall back to the schema defaults. keys = row.keys() if hasattr(row, "keys") else () chunking_strategy = ( row["chunking_strategy"] if "chunking_strategy" in keys else "standard" @@ -207,12 +178,7 @@ def _row_to_kb(row: Any) -> KBResponse: def _validate_mode_combo(mode: KBMode, chunking_strategy: ChunkingStrategy) -> None: - """Reject the one illegal (mode, strategy) combination. - - No public open-weight embedder supports both late-chunking pooling - and shared text/image embedding. Surface the constraint as a 400 - rather than failing silently during ingestion. - """ + """Reject (multimodal, late) — no embedder supports both at once.""" if mode == "multimodal" and chunking_strategy == "late": raise HTTPException( status_code = 400, @@ -287,11 +253,9 @@ async def _save_upload(file: UploadFile) -> tuple[Path, str, int]: stored_path = upload_dir / stored_name max_bytes = RAG_MAX_UPLOAD_MB * 1024 * 1024 written = 0 - # anyio.open_file routes each write through a worker thread so a - # multi-MB upload doesn't stall concurrent requests on the event - # loop. The async-with handles close on both happy and error - # paths; the outer try/except cleans up the partial file after - # the file handle is closed (Windows refuses unlink on an open fd). + # Route writes through anyio worker thread so the event loop stays free. + # Outer try/except cleans up partial files after async-with closes the fd + # (Windows refuses unlink on an open fd). try: async with await anyio.open_file(stored_path, "wb") as f: while True: @@ -369,9 +333,7 @@ def _unlink_if_under_uploads(path: Path) -> None: real.unlink(missing_ok = True) -# ------------------------------------------------------------------ -# Knowledge bases -# ------------------------------------------------------------------ +# --- Knowledge bases --- @router.post("/knowledge-bases", response_model = KBResponse) @@ -384,9 +346,7 @@ def create_knowledge_base( _validate_mode_combo(payload.mode, payload.chunking_strategy) kb_id = str(uuid4()) - # If the caller didn't override embedding_model, resolve from the - # Phase-3 matrix using their (mode, strategy) selection. Unknown - # combos fall back to the legacy default — see resolve_embedder. + # No override: resolve from (mode, strategy) matrix. embedding_model = payload.embedding_model or resolve_embedder( payload.mode, payload.chunking_strategy ) @@ -483,8 +443,7 @@ def set_rag_defaults( current = _load_rag_defaults() new_strategy = payload.chunking_strategy or current.chunking_strategy new_mode = payload.mode or current.mode - # PATCH-style — passing an empty string clears the override; a - # null/missing field keeps the current value. + # PATCH-style: empty string clears, null/missing keeps current. if payload.embedding_model is None: new_embedder = current.embedding_model elif payload.embedding_model.strip() == "": @@ -526,11 +485,7 @@ def _thread_settings_key(thread_id: str) -> str: def _load_thread_settings(thread_id: str) -> ThreadRagSettings: - """Per-thread RAG settings, falling back to app-level defaults. - - Stored in chat_settings under "thread::rag" as a nested JSON - dict — same shape as RagDefaults. - """ + """Per-thread RAG settings (chat_settings['thread::rag']) with defaults fallback.""" settings = list_chat_settings() raw = settings.get(_thread_settings_key(thread_id)) or {} if not isinstance(raw, dict): @@ -611,11 +566,7 @@ def _reingest_scope( mode: str, embedding_model: str, ) -> ReingestResponse: - """Wipe scope artifacts and re-enqueue every stored document. - - The chat_settings / per-thread defaults aren't touched — caller is - responsible for updating any associated metadata before calling. - """ + """Wipe scope artifacts and re-enqueue every document; metadata untouched.""" scope = kb_scope(kb_id) if kb_id else thread_scope(thread_id) # type: ignore[arg-type] with get_connection() as conn: if kb_id: @@ -628,9 +579,7 @@ def _reingest_scope( "SELECT id, stored_path FROM rag_documents WHERE thread_id = ?", (thread_id,), ).fetchall() - # Delete the rag_documents rows (cascade drops chunks); the - # uploaded file on disk is preserved so we can re-ingest from - # it. We re-INSERT a fresh row per stored_path below. + # Drop rag_documents (chunks cascade); files on disk are reused below. doc_ids = [r["id"] for r in rows] if doc_ids: placeholders = ",".join("?" for _ in doc_ids) @@ -648,8 +597,7 @@ def _reingest_scope( if not stored_path.is_file(): continue filename = stored_path.name - # Strip the UUID prefix we attached at upload time so the - # re-inserted document carries the original name. + # Strip the upload-time UUID prefix; keep the original filename. if "_" in filename: _uuid_prefix, _, original = filename.partition("_") if original: @@ -728,14 +676,7 @@ def reingest_thread_documents( payload: UpdateThreadRagSettingsRequest | None = None, current_subject: str = Depends(get_current_subject), ) -> ReingestResponse: - """Rebuild a thread's RAG index. - - Optional body lets the caller change the thread's chunking - strategy / mode / embedder at the same time — persisted into - chat_settings before re-ingestion so subsequent uploads pick up - the new values too. With an empty body, current settings are - reused. - """ + """Rebuild a thread's RAG index; optional body updates settings before reingest.""" from utils.rag.config import resolve_embedder if payload is None: @@ -745,7 +686,6 @@ def reingest_thread_documents( or payload.mode is not None or payload.embedding_model is not None ): - # set_thread_rag_settings handles validation + persistence. settings = set_thread_rag_settings( thread_id, payload, @@ -786,9 +726,7 @@ def delete_knowledge_base( return {"ok": True} -# ------------------------------------------------------------------ -# Document upload (KB and per-thread) -# ------------------------------------------------------------------ +# --- Document upload (KB and per-thread) --- @router.post("/knowledge-bases/{kb_id}/documents", response_model = UploadResponse) @@ -799,9 +737,7 @@ async def upload_kb_document( ) -> UploadResponse: kb_row = _kb_or_404(kb_id) stored_path, filename, byte_size = await _save_upload(file) - # Defensive .get() — rows fetched through a connection that pre-dates - # the Phase 3 schema (e.g. in tests) lack chunking_strategy/mode; - # fall back to the same defaults as the column. + # Tolerate pre-Phase-3 rows missing chunking_strategy/mode. kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else () chunking_strategy = ( kb_row["chunking_strategy"] if "chunking_strategy" in kb_keys else "standard" @@ -828,12 +764,8 @@ async def upload_thread_document( ) -> UploadResponse: from utils.rag.config import resolve_embedder - # Don't validate against chat_threads — a brand-new chat won't be - # persisted there until after the first runStart/runEnd. Users who - # attach a document on a fresh thread would otherwise hit a 404. + # No chat_threads check — fresh threads aren't persisted until first run. stored_path, filename, byte_size = await _save_upload(file) - # Per-thread settings fall back to app-level defaults inside the - # helper, so first-time-uploaded threads inherit user preferences. settings = _load_thread_settings(thread_id) embedder = settings.embedding_model or resolve_embedder( settings.mode, @@ -852,9 +784,7 @@ async def upload_thread_document( ) -# ------------------------------------------------------------------ -# Document list / delete -# ------------------------------------------------------------------ +# --- Document list / delete --- @router.get("/knowledge-bases/{kb_id}/documents", response_model = DocumentListResponse) @@ -890,13 +820,7 @@ def get_rag_image( filename: str, current_subject: str = Depends(get_current_subject), ) -> FileResponse: - """Serve an image extracted during multimodal ingestion. - - Files live under ``rag_uploads_root() / 'images' / ``. - We realpath-check the resolved file against that root to refuse - path-traversal attempts (``..`` segments, symlinks pointing - elsewhere). Filenames are constrained to a single path component. - """ + """Serve an extracted image; realpath-check against the uploads root.""" if "/" in filename or "\\" in filename or filename.startswith("."): raise HTTPException(status_code = 400, detail = "Invalid filename") root = Path(os.path.realpath(rag_uploads_root() / "images")) @@ -930,12 +854,7 @@ def delete_document( def list_thread_indexes( current_subject: str = Depends(get_current_subject), ) -> ThreadIndexListResponse: - """List every chat thread that has at least one RAG document. - - LEFT JOIN to chat_threads so threads that were never persisted - (user attached a file but never sent a message) still show up — - just with a null title. - """ + """List threads with >=1 RAG doc. LEFT JOIN keeps unpersisted threads (null title).""" with get_connection() as conn: rows = conn.execute( """ @@ -969,19 +888,12 @@ def clear_thread_documents( thread_id: str, current_subject: str = Depends(get_current_subject), ) -> dict: - """Purge every RAG document attached to ``thread_id``. - - Removes the per-thread vector rows, the bm25 index, the - rag_documents/rag_chunks rows, and the uploaded files. The chat - thread itself is untouched. - """ + """Drop all RAG artifacts for thread_id; chat thread itself untouched.""" ingestion.purge_thread_documents([thread_id]) return {"ok": True} -# ------------------------------------------------------------------ -# Ingestion job SSE -# ------------------------------------------------------------------ +# --- Ingestion job SSE --- @router.get("/jobs/{job_id}/events") @@ -1050,9 +962,7 @@ async def _replay_terminal_state(row: Any): yield f"data: {json.dumps(payload)}\n\n" -# ------------------------------------------------------------------ -# Search -# ------------------------------------------------------------------ +# --- Search --- @router.post("/search", response_model = SearchResponse) @@ -1071,9 +981,7 @@ def search( else: scope = thread_scope(payload.thread_id) - # Embed the query with the same model that populated this scope's - # vectors. Mixing spaces (e.g. Qwen3-VL 2048-d docs vs bge-small - # 384-d query) crashes inside the cosine-distance compute. + # Query must use the same embedder as the scope (dim must match). scope_embedder = _resolve_scope_embedder(scope) logger.info( "RAG search: scope=%s embedder=%s mode=%s top_k=%d min_score=%.3f rerank=%s query=%r", @@ -1086,8 +994,7 @@ def search( payload.query[:120], ) - # When reranking is opt-in, pull a wider candidate pool so the - # CrossEncoder has more to choose from before truncating to top_k. + # Reranker needs a wider candidate pool than top_k. candidate_k = ( max(payload.top_k, RAG_RERANK_CANDIDATE_K) if payload.enable_rerank @@ -1168,8 +1075,7 @@ def search( image_url: str | None = None if kind == "image" and meta.get("image_path"): image_url = ( - f"/api/rag/images/{meta['document_id']}/" - f"{Path(meta['image_path']).name}" + f"/api/rag/images/{meta['document_id']}/{Path(meta['image_path']).name}" ) out.append( SearchHit( diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index c4541fbf9a..49e66d526d 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -205,13 +205,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) WITHOUT ROWID """ ) - # RAG: knowledge bases, documents, chunks, ingestion jobs. - # rag_documents enforces XOR on (kb_id, thread_id): a document belongs - # to either a standalone KB or a single chat thread, never both. - # chunking_strategy and mode are set at KB-creation time and are - # immutable thereafter — changing either invalidates existing chunks - # because they were ingested through a specific pipeline (different - # chunker, different embedder). See Phase 3 in the plan. + # RAG schema. rag_documents enforces XOR on (kb_id, thread_id). + # chunking_strategy/mode are immutable post-create (invalidates chunks). conn.execute( """ CREATE TABLE IF NOT EXISTS rag_knowledge_bases ( @@ -226,9 +221,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) - # Idempotent ALTER for existing installs that pre-date the columns. - # Mirrors the chat_threads display_name / *_code_exec_container_id - # pattern earlier in this file. + # Idempotent ALTER for pre-existing installs. kb_cols = { row[1] for row in conn.execute("PRAGMA table_info(rag_knowledge_bases)").fetchall() @@ -243,11 +236,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: "ALTER TABLE rag_knowledge_bases " "ADD COLUMN mode TEXT NOT NULL DEFAULT 'text'" ) - # thread_id has no FK to chat_threads. A user can attach a document - # to a thread that hasn't yet been persisted (saveThread only runs - # after the first model exchange — see runtime-provider.tsx). The - # chat-thread DELETE handlers in routes/chat_history.py purge - # matching rag_documents explicitly so lifecycle stays clean. + # thread_id has no FK: docs can attach before the thread is persisted. + # chat_history DELETE handlers purge matching rag_documents explicitly. conn.execute( """ CREATE TABLE IF NOT EXISTS rag_documents ( @@ -272,11 +262,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute( "CREATE INDEX IF NOT EXISTS idx_rag_documents_thread_id ON rag_documents(thread_id)" ) - # kind = 'text' | 'image' | 'caption'. image_path is set when kind = 'image' - # (path to the extracted figure under rag_uploads_root() / 'images/'). - # linked_chunk_id pairs an 'image' chunk with its 'caption' chunk (and vice - # versa) so retrieval can hydrate the matching half. Both default to NULL - # so today's text-only ingestion is unaffected. See Phase 3B-multimodal. + # kind: text|image|caption. linked_chunk_id pairs image↔caption (both null for text). conn.execute( """ CREATE TABLE IF NOT EXISTS rag_chunks ( diff --git a/studio/backend/utils/rag/config.py b/studio/backend/utils/rag/config.py index 73edbf41b6..e394288368 100644 --- a/studio/backend/utils/rag/config.py +++ b/studio/backend/utils/rag/config.py @@ -31,32 +31,17 @@ RAG_EMBEDDING_MODEL: str = ( or "BAAI/bge-small-en-v1.5" ) -# Phase 3: default embedders per (mode, chunking_strategy). Ingestion in -# Phase 3B-late and Phase 3B-multimodal looks the embedder up here at job -# start, falling back to RAG_EMBEDDING_MODEL (above) for legacy KBs that -# pre-date the columns. The (multimodal, late) combo is intentionally -# absent — no public open-weight embedder supports both at once, and -# routes/rag.py rejects the combo with a 400 at KB create time. +# Default embedder per (mode, chunking). (multimodal, late) is unsupported +# and rejected at KB-create time in routes/rag.py. RAG_EMBEDDER_MATRIX: dict[tuple[str, str], str] = { ("text", "standard"): "BAAI/bge-small-en-v1.5", ("text", "late"): "nomic-ai/nomic-embed-text-v1.5", - # Qwen3-VL-Embedding-2B: 2B-param multimodal embedder built on - # Qwen3 (not CLIP), so no 77-token text cap — long chunks embed - # losslessly. Loads through vanilla SentenceTransformer with - # trust_remote_code, no shim. 2048-d shared text/image space. - # Trade-off: ~4 GB download / VRAM vs BGE-VL-base's ~600 MB. ("multimodal", "standard"): "Qwen/Qwen3-VL-Embedding-2B", } def resolve_embedder(mode: str, chunking_strategy: str) -> str: - """Look up the default embedder for a (mode, chunking_strategy) pair. - - Unknown combos fall back to the legacy single default so old KBs - keep working. Callers that explicitly require the new matrix - behaviour (Phase 3B paths) should validate the inputs before - calling. - """ + """Embedder for (mode, chunking); unknown combos fall back to RAG_EMBEDDING_MODEL.""" return RAG_EMBEDDER_MATRIX.get( (mode, chunking_strategy), RAG_EMBEDDING_MODEL, @@ -76,9 +61,6 @@ RAG_MAX_UPLOAD_MB: int = _env_int("UNSLOTH_RAG_MAX_UPLOAD_MB", 50) RAG_EMBED_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_EMBED_BATCH_SIZE", 32) -# Reranking is off by default. The CrossEncoder runs on GPU and competes -# with the active chat model — callers opt in per-request via -# `enable_rerank` on SearchRequest. RAG_RERANKER_MODEL: str = ( os.environ.get("UNSLOTH_RAG_RERANKER_MODEL", "").strip() or "BAAI/bge-reranker-base" ) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 440a6817c7..769b90a781 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -337,8 +337,7 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { event.preventDefault(); return; } - // Send is going through — drop the chips. The docs themselves - // remain in the backend thread KB and stay searchable. + // Drop chips on send; docs stay searchable in the backend. clearDocs(); }, [ @@ -941,9 +940,7 @@ const ImagesToggle: FC = () => { ); }; -// Mirror of shared-composer's RAG pill (the master switch for -// retrieval). Visible on every model; the sidebar Retrieval section -// configures source / top-K / reranker once this is on. +// Master RAG switch (mirrors shared-composer); sidebar configures the rest. const RagToggle: FC = () => { const modelLoaded = useChatRuntimeStore( (s) => !!s.params.checkpoint && !s.modelLoading, @@ -1037,11 +1034,7 @@ const ToolStatusDisplay: FC = () => { ); }; -// Custom + button used when RAG is on: opens a picker accepting -// doc formats the ingester handles and routes selected files to -// the RAG thread-document pipeline. Replaces the stock -// ComposerAddAttachment (which base64-attaches files inline) so -// docs become indexed chunks instead of one-shot model context. +// RAG-aware + button: picks doc formats and routes to ingest pipeline. const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({ onSelect, }) => { @@ -1061,7 +1054,6 @@ const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({ const f = files[i]; if (f && isDocumentFile(f)) onSelect(f); } - // Reset so re-selecting the same file fires onChange. e.target.value = ""; }} /> diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index b743948012..a920b26594 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -992,19 +992,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }); } - // RAG: optionally retrieve context for the last user turn and - // prepend it as a system-role block. Failures are logged but - // don't break the chat — better to answer without context than - // to drop a message the user just sent. - // - // Pre-fetch RAG context unconditionally when the RAG button is - // on and a source is selected. This runs for every provider — - // local-tool, local-no-tool, and external — so users don't have - // to phrase their query as "the document I attached" for - // retrieval to fire. On local tool-capable models the - // `search_knowledge_base` tool is *also* registered below as an - // optional refinement path (the LLM can run a second, narrower - // query if the pre-fetched chunks weren't enough). + // Pre-fetch RAG context for the last user turn; failures don't block chat. + // Runs for all providers; local tool-capable models also get the tool below + // for a narrower follow-up query if needed. const ragSource = runtime.ragSource; const ragToolEnabled = runtime.ragToolEnabled; const ragToolPathTaken = @@ -1647,11 +1637,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(codeToolsEnabled ? ["python", "terminal"] : []), ...(ragToolPathTaken ? ["search_knowledge_base"] : []), ], - // Phase 4: per-request RAG context the backend's - // `search_knowledge_base` handler reads when the LLM - // invokes the tool. Only sent when the tool path - // is taken — external providers fall through to the - // pre-fetch block above. + // Per-request scope for the LLM-invoked tool; tool path only. ...(ragToolPathTaken ? { rag_scope: { diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 3f3bdef19a..5748726bbb 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -464,10 +464,8 @@ export function ChatSettingsPanel({ const updateThreadSettings = useRagStore((s) => s.updateThreadSettings); const ragDefaults = useRagStore((s) => s.defaults); - // Load this thread's RAG settings once per threadId. Ref-guarded so - // `threadSettings` isn't a dep — if it were, the post-load - // store-mutation re-triggers the effect and any failure mode where - // the selector flickers undefined produces an update loop. + // Load thread RAG settings once per threadId. Ref-guarded; keep + // `threadSettings` out of deps to avoid post-load update loops. const threadSettingsLoadedRef = useRef(null); useEffect(() => { if ( @@ -492,7 +490,6 @@ export function ChatSettingsPanel({ ) => { if (!activeThreadId) return; if (threadDocs.length === 0) { - // No existing chunks to invalidate — just persist. void updateThreadSettings(activeThreadId, patch); return; } @@ -503,8 +500,7 @@ export function ChatSettingsPanel({ if (ok) { void reingestThread(activeThreadId, patch); } else { - // User declined — refresh the store so the select snaps back - // to the unchanged settings. + // User declined: refresh so the select snaps back. void loadThreadSettings(activeThreadId); } }; @@ -645,8 +641,6 @@ export function ChatSettingsPanel({ activeExternalProvider.baseUrl, ) && activeExternalProvider.providerType === "openai"; - // (activeThreadId is declared earlier in this component — see the - // RAG retrieval-section block above.) const openAiApiKeyForSection = activeExternalProvider ? getExternalProviderApiKey(activeExternalProvider.id) || null : null; diff --git a/studio/frontend/src/features/chat/components/pending-doc-chips.tsx b/studio/frontend/src/features/chat/components/pending-doc-chips.tsx index d65d02c129..a300c61910 100644 --- a/studio/frontend/src/features/chat/components/pending-doc-chips.tsx +++ b/studio/frontend/src/features/chat/components/pending-doc-chips.tsx @@ -12,9 +12,6 @@ interface PendingDocChipsProps { onRemove: (id: string) => void; } -// Renders the upload/index/ready/error chips for in-flight RAG -// document uploads above the composer textarea. Mirrors the -// inline chip markup SharedComposer uses on the empty state. export const PendingDocChips: FC = ({ docs, onRemove }) => { if (docs.length === 0) return null; return ( diff --git a/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts index 5e77cbbc0d..39ef5302dc 100644 --- a/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts +++ b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts @@ -27,8 +27,6 @@ const DOCUMENT_EXTENSIONS = new Set([ ".htm", ]); -// File-input accept attribute matching DOCUMENT_EXTENSIONS so the -// browser picker filters to formats the RAG ingester actually handles. export const DOCUMENT_ACCEPT = ".pdf,.txt,.md,.markdown,.docx,.html,.htm,application/pdf,text/plain,text/markdown,application/vnd.openxmlformats-officedocument.wordprocessingml.document,text/html"; @@ -47,12 +45,7 @@ export interface UseThreadDocUploadsResult { isIndexing: boolean; } -// Encapsulates the per-thread RAG document upload lifecycle: -// pick file → POST /api/rag/threads/{id}/documents → subscribe to -// ingestion SSE → flip chip status → on send, clear chips (docs -// live in the backend KB and don't need to be re-attached). -// Used by both the empty-state SharedComposer and the in-thread -// assistant-ui Composer so the upload UX is identical in both. +/** Per-thread RAG upload: file → POST → SSE → chip status. */ export function useThreadDocUploads(): UseThreadDocUploadsResult { const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const [pendingDocs, setPendingDocs] = useState([]); @@ -83,9 +76,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { d.id === id ? { ...d, status: "ready" } : d, ), ); - // First successful ingest in an off-source thread should - // flip the source to 'thread' so the tool / pre-fetch - // path has somewhere to search. + // First ingest on an off-source thread → flip to 'thread'. if (useChatRuntimeStore.getState().ragSource.kind === "off") { useChatRuntimeStore .getState() @@ -126,9 +117,7 @@ export function useThreadDocUploads(): UseThreadDocUploadsResult { void useRagStore .getState() .deleteDocument(doc.documentId, `thread:${activeThreadId ?? ""}`) - .catch(() => { - // Best effort — the chip is going away regardless. - }); + .catch(() => {}); } return prev.filter((d) => d.id !== id); }); diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 078b589756..fb08cf0483 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -580,12 +580,10 @@ export function SharedComposer({ }); continue; } - // Handle RAG document files (route to per-thread ingest pipeline). if (isDocumentFile(file)) { addDoc(file); continue; } - // Handle image files if (!file.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue; if (file.size > MAX_IMAGE_SIZE) continue; if (attachUnavailableReason) { @@ -614,9 +612,7 @@ export function SharedComposer({ doc.documentId, `thread:${activeThreadId ?? ""}`, ) - .catch(() => { - // Best effort — the chip is going away regardless. - }); + .catch(() => {}); } return prev.filter((d) => d.id !== id); }); @@ -709,9 +705,7 @@ export function SharedComposer({ setPendingImages([]); setPendingAudio(null); clearPendingAudioStore(); - // Docs remain in the backend (already uploaded); just drop the - // composer-side chips. The settings panel still shows them in the - // thread's document list. + // Docs stay in backend; drop chips only. setPendingDocs([]); textareaRef.current?.focus(); @@ -1277,12 +1271,7 @@ export function SharedComposer({ Images )} - {/* RAG: master switch for retrieval. On local models with - tool-use support, registers `search_knowledge_base` as a - tool the LLM can call. On external providers, falls back - to the pre-fetch path. The sidebar Retrieval section - configures the source / top-K / reranker; the button is - the only on/off control. */} + {/* Master RAG toggle; sidebar Retrieval section configures the rest. */}