From 92994e8b83fad307db7fa6a0954ba74bb97604ed Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sat, 23 May 2026 18:39:58 +0400 Subject: [PATCH 001/122] Studio: add RAG with hybrid search, reranker, chat integration Backend (studio/backend/): - core/rag/: parsers (PDF/TXT/MD/DOCX/HTML via pypdf/python-docx/bs4), recursive token-aware chunker, embeddings singleton via FastSentenceTransformer.from_pretrained(for_inference=True), Qdrant local vector store, bm25s lexical index, RRF hybrid retrieval, spawn-subprocess ingestion job with SSE progress, optional CrossEncoder reranker (off-by-default). - routes/rag.py: KB CRUD, doc upload (KB + per-thread), doc list/delete, ingestion SSE, hybrid+rerank search, thread-index list/clear. - routes/chat_history.py: purge thread RAG artifacts on thread delete and clear-all (rag_documents has no FK cascade to chat_threads so uploads work on un-persisted threads). - studio.db gains 4 RAG tables; storage_roots gains rag_*() helpers. - auth/authentication.py: get_current_subject_sse accepts ?token=... so EventSource can stream ingestion progress. Frontend (studio/frontend/): - features/rag/: api client, Zustand store, hooks, dropzone, KB list, doc rows, ingestion-progress, thread-index list components. - Settings dialog gains a Knowledge Bases tab (master/detail + thread documents list); /knowledge-bases deep-links to it. - features/chat/: per-thread ragSource/enableRerank/ragTopK state in chat-runtime-store; Retrieval section in chat-settings-sheet with KB DropdownMenu (active highlight + per-row trash), thread doc list with Clear-thread-index button, RAG Top K slider, reranker toggle; chat-adapter retrieves before /v1/chat/completions and injects hits as a system block; shared-composer + button routes documents into pendingDocs (auto-uploads, send blocked while indexing). --- studio/backend/auth/authentication.py | 34 + studio/backend/core/rag/__init__.py | 2 + studio/backend/core/rag/bm25.py | 111 +++ studio/backend/core/rag/chunking.py | 129 ++++ studio/backend/core/rag/embeddings.py | 102 +++ studio/backend/core/rag/ingestion.py | 498 +++++++++++++ studio/backend/core/rag/parsers/__init__.py | 32 + studio/backend/core/rag/parsers/docx.py | 27 + studio/backend/core/rag/parsers/html.py | 26 + studio/backend/core/rag/parsers/pdf.py | 24 + studio/backend/core/rag/parsers/text.py | 27 + studio/backend/core/rag/reranker.py | 113 +++ studio/backend/core/rag/retrieval.py | 113 +++ studio/backend/core/rag/vector_store.py | 156 ++++ studio/backend/main.py | 2 + .../backend/requirements/no-torch-runtime.txt | 12 + studio/backend/routes/__init__.py | 2 + studio/backend/routes/chat_history.py | 5 + studio/backend/routes/rag.py | 681 ++++++++++++++++++ studio/backend/storage/studio_db.py | 77 ++ studio/backend/utils/paths/storage_roots.py | 20 + studio/backend/utils/rag/__init__.py | 2 + studio/backend/utils/rag/config.py | 59 ++ studio/frontend/src/app/router.tsx | 2 + .../src/app/routes/knowledge-bases.tsx | 22 + .../src/features/chat/api/chat-adapter.ts | 99 +++ .../features/chat/api/chat-settings-api.ts | 8 + .../src/features/chat/chat-settings-sheet.tsx | 207 ++++++ .../src/features/chat/shared-composer.tsx | 179 ++++- .../chat/stores/chat-runtime-store.ts | 37 +- .../frontend/src/features/rag/api/rag-api.ts | 256 +++++++ .../features/rag/components/document-row.tsx | 78 ++ .../components/document-upload-dropzone.tsx | 83 +++ .../rag/components/ingestion-progress.tsx | 66 ++ .../rag/components/kb-create-dialog.tsx | 131 ++++ .../rag/components/kb-detail-panel.tsx | 87 +++ .../src/features/rag/components/kb-list.tsx | 101 +++ .../rag/components/thread-index-list.tsx | 79 ++ .../rag/hooks/use-ingestion-events.ts | 22 + .../features/rag/hooks/use-kb-documents.ts | 68 ++ .../features/rag/hooks/use-knowledge-bases.ts | 24 + .../src/features/rag/stores/rag-store.ts | 245 +++++++ .../src/features/settings/settings-dialog.tsx | 11 + .../settings/stores/settings-dialog-store.ts | 2 + .../settings/tabs/knowledge-bases-tab.tsx | 43 ++ tests/python/test_rag_bm25.py | 69 ++ tests/python/test_rag_chunking.py | 84 +++ tests/python/test_rag_parsers.py | 89 +++ tests/python/test_rag_reranker.py | 57 ++ tests/python/test_rag_retrieval.py | 46 ++ tests/python/test_rag_vector_store.py | 84 +++ 51 files changed, 4527 insertions(+), 6 deletions(-) create mode 100644 studio/backend/core/rag/__init__.py create mode 100644 studio/backend/core/rag/bm25.py create mode 100644 studio/backend/core/rag/chunking.py create mode 100644 studio/backend/core/rag/embeddings.py create mode 100644 studio/backend/core/rag/ingestion.py create mode 100644 studio/backend/core/rag/parsers/__init__.py create mode 100644 studio/backend/core/rag/parsers/docx.py create mode 100644 studio/backend/core/rag/parsers/html.py create mode 100644 studio/backend/core/rag/parsers/pdf.py create mode 100644 studio/backend/core/rag/parsers/text.py create mode 100644 studio/backend/core/rag/reranker.py create mode 100644 studio/backend/core/rag/retrieval.py create mode 100644 studio/backend/core/rag/vector_store.py create mode 100644 studio/backend/routes/rag.py create mode 100644 studio/backend/utils/rag/__init__.py create mode 100644 studio/backend/utils/rag/config.py create mode 100644 studio/frontend/src/app/routes/knowledge-bases.tsx create mode 100644 studio/frontend/src/features/rag/api/rag-api.ts create mode 100644 studio/frontend/src/features/rag/components/document-row.tsx create mode 100644 studio/frontend/src/features/rag/components/document-upload-dropzone.tsx create mode 100644 studio/frontend/src/features/rag/components/ingestion-progress.tsx create mode 100644 studio/frontend/src/features/rag/components/kb-create-dialog.tsx create mode 100644 studio/frontend/src/features/rag/components/kb-detail-panel.tsx create mode 100644 studio/frontend/src/features/rag/components/kb-list.tsx create mode 100644 studio/frontend/src/features/rag/components/thread-index-list.tsx create mode 100644 studio/frontend/src/features/rag/hooks/use-ingestion-events.ts create mode 100644 studio/frontend/src/features/rag/hooks/use-kb-documents.ts create mode 100644 studio/frontend/src/features/rag/hooks/use-knowledge-bases.ts create mode 100644 studio/frontend/src/features/rag/stores/rag-store.ts create mode 100644 studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx create mode 100644 tests/python/test_rag_bm25.py create mode 100644 tests/python/test_rag_chunking.py create mode 100644 tests/python/test_rag_parsers.py create mode 100644 tests/python/test_rag_reranker.py create mode 100644 tests/python/test_rag_retrieval.py create mode 100644 tests/python/test_rag_vector_store.py diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 6ddcbc8e0b..ec53467189 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -147,6 +147,40 @@ async def get_current_subject( ) +async def get_current_subject_sse( + token: Optional[str] = None, + authorization: Optional[str] = None, +) -> str: + """Auth dep for SSE endpoints. + + EventSource cannot send custom headers, so callers pass the bearer + as a ``?token=…`` query param. Falls back to the Authorization + header so curl / API clients keep working. + + Wire with ``Query(None)`` and ``Header(None)`` at the route layer: + + async def stream( + current_subject: str = Depends( + lambda token = Query(None), authorization = Header(None): + get_current_subject_sse(token, authorization) + ), + ): ... + """ + raw = token + if not raw and authorization and authorization.lower().startswith("bearer "): + raw = authorization[len("bearer "):].strip() + if not raw: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Missing token", + ) + credentials = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = raw) + return await _get_current_subject( + credentials, + allow_password_change = False, + ) + + async def get_current_subject_allow_password_change( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: diff --git a/studio/backend/core/rag/__init__.py b/studio/backend/core/rag/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/core/rag/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/core/rag/bm25.py b/studio/backend/core/rag/bm25.py new file mode 100644 index 0000000000..c3da6806d6 --- /dev/null +++ b/studio/backend/core/rag/bm25.py @@ -0,0 +1,111 @@ +# 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. + +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). +""" + +from __future__ import annotations + +import json +import logging +import shutil +import threading +from pathlib import Path +from typing import Any + +from utils.paths.storage_roots import ensure_dir, rag_bm25_root + +logger = logging.getLogger(__name__) + +_load_lock = threading.Lock() +_cache: dict[str, tuple[Any, list[str]]] = {} + + +def _scope_dir(scope: str) -> Path: + return rag_bm25_root() / scope + + +def _ids_path(scope: str) -> Path: + return _scope_dir(scope) / "ids.json" + + +def _has_index(scope: str) -> bool: + return _ids_path(scope).is_file() + + +def _evict(scope: str) -> None: + _cache.pop(scope, None) + + +def rebuild_index(scope: str, chunks: list[dict]) -> None: + """Rebuild 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. + """ + import bm25s + + base = _scope_dir(scope) + if not chunks: + delete_scope(scope) + return + texts = [c["text"] for c in chunks] + ids = [c["id"] for c in chunks] + tokens = bm25s.tokenize(texts, show_progress = False) + retriever = bm25s.BM25() + retriever.index(tokens, show_progress = False) + ensure_dir(base) + retriever.save(str(base)) + _ids_path(scope).write_text(json.dumps(ids)) + with _load_lock: + _cache[scope] = (retriever, ids) + + +def _load(scope: str) -> tuple[Any, list[str]] | None: + if not _has_index(scope): + return None + with _load_lock: + if scope in _cache: + return _cache[scope] + import bm25s + + retriever = bm25s.BM25.load(str(_scope_dir(scope)), load_corpus = False) + ids = json.loads(_ids_path(scope).read_text()) + _cache[scope] = (retriever, ids) + return _cache[scope] + + +def search(scope: str, query: str, k: int) -> list[tuple[str, float]]: + import bm25s + + loaded = _load(scope) + if loaded is None: + return [] + retriever, ids = loaded + if not ids: + return [] + k_actual = min(k, len(ids)) + q_tokens = bm25s.tokenize([query], show_progress = False) + indices, scores = retriever.retrieve(q_tokens, k = k_actual, show_progress = False) + out: list[tuple[str, float]] = [] + for pos in range(indices.shape[1]): + idx = int(indices[0][pos]) + out.append((ids[idx], float(scores[0][pos]))) + return out + + +def delete_scope(scope: str) -> None: + base = _scope_dir(scope) + if base.exists(): + shutil.rmtree(base, ignore_errors = True) + _evict(scope) diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py new file mode 100644 index 0000000000..e2a5c06493 --- /dev/null +++ b/studio/backend/core/rag/chunking.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +from .parsers import ParsedPage + + +@dataclass(frozen = True) +class Chunk: + text: str + token_count: int + page_number: int | None = None + + +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) + + +def _split_on(text: str, separator: str) -> list[str]: + if separator == "": + return list(text) + parts = text.split(separator) + if len(parts) == 1: + return parts + glued: list[str] = [] + for i, part in enumerate(parts): + if i < len(parts) - 1: + glued.append(part + separator) + else: + if part: + glued.append(part) + return [p for p in glued if p] + + +def _atomic_split( + text: str, + separators: tuple[str, ...], + max_tokens: int, + count: TokenCounter, +) -> list[str]: + if count(text) <= max_tokens: + return [text] + for sep in separators: + pieces = _split_on(text, sep) + if len(pieces) <= 1: + continue + out: list[str] = [] + for piece in pieces: + if count(piece) <= max_tokens: + out.append(piece) + else: + 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)] + + +def _merge( + pieces: list[str], + max_tokens: int, + overlap_tokens: int, + count: TokenCounter, +) -> list[str]: + """Greedy-merge atomic pieces into chunks <= max_tokens with overlap between adjacent chunks.""" + chunks: list[str] = [] + buffer: list[str] = [] + buffer_tokens = 0 + for piece in pieces: + piece_tokens = count(piece) + if buffer and buffer_tokens + piece_tokens > max_tokens: + chunks.append("".join(buffer)) + if overlap_tokens > 0: + overlap: list[str] = [] + running = 0 + for prev in reversed(buffer): + prev_tokens = count(prev) + if running + prev_tokens > overlap_tokens: + break + overlap.insert(0, prev) + running += prev_tokens + buffer = list(overlap) + buffer_tokens = running + else: + buffer = [] + buffer_tokens = 0 + buffer.append(piece) + buffer_tokens += piece_tokens + if buffer: + chunks.append("".join(buffer)) + return [c.strip() for c in chunks if c.strip()] + + +def chunk_pages( + pages: list[ParsedPage], + *, + max_tokens: int, + overlap_tokens: int, + token_counter: TokenCounter | None = None, + separators: tuple[str, ...] = ("\n\n", "\n", ". ", " ", ""), +) -> 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. + """ + count = token_counter or _char_token_estimate + out: list[Chunk] = [] + for page in pages: + atomic = _atomic_split(page.text, separators, max_tokens, count) + merged = _merge(atomic, max_tokens, overlap_tokens, count) + for piece in merged: + out.append( + Chunk( + text = piece, + token_count = count(piece), + page_number = page.page_number, + ) + ) + return out diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py new file mode 100644 index 0000000000..4953d48ffe --- /dev/null +++ b/studio/backend/core/rag/embeddings.py @@ -0,0 +1,102 @@ +# 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. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +from utils.rag.config import RAG_EMBED_BATCH_SIZE, RAG_EMBEDDING_MODEL + +logger = logging.getLogger(__name__) + +_lock = threading.Lock() +_model: Any | None = None +_model_name: str | None = None +_embedding_dim: int | None = None + + +def _load(model_name: str) -> Any: + from unsloth import FastSentenceTransformer + + logger.info("Loading RAG embedder: %s", model_name) + return FastSentenceTransformer.from_pretrained( + model_name, + for_inference = True, + ) + + +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: + if _model is None or _model_name != target: + _model = _load(target) + _model_name = target + try: + _embedding_dim = int(_model.get_sentence_embedding_dimension()) + except Exception: + _embedding_dim = None + return _model + + +def get_embedding_dim(model_name: str | None = None) -> int: + model = get_embedder(model_name) + global _embedding_dim + if _embedding_dim is None: + _embedding_dim = int(model.get_sentence_embedding_dimension()) + return _embedding_dim + + +def get_active_model_name() -> str | None: + return _model_name + + +def encode( + texts: list[str], + *, + model_name: str | None = None, + batch_size: int | None = None, + normalize: bool = True, +): + model = get_embedder(model_name) + return model.encode( + texts, + batch_size = batch_size or RAG_EMBED_BATCH_SIZE, + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + + +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. + """ + model = get_embedder(model_name) + + def _count(text: str) -> int: + try: + tokens = model.tokenize([text]) + ids = tokens.get("input_ids") + if ids is None: + return max(1, len(text) // 4) + return int(ids.shape[1]) + except Exception: + return max(1, len(text) // 4) + + return _count diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py new file mode 100644 index 0000000000..a6aa7c720c --- /dev/null +++ b/studio/backend/core/rag/ingestion.py @@ -0,0 +1,498 @@ +# 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. + +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, Qdrant points, and (at job completion) a rebuilt BM25 index. + +Only the parent process holds the Qdrant local-mode file lock — the +subprocess never opens it directly. This keeps search available +throughout the lifetime of an ingestion job. +""" + +from __future__ import annotations + +import logging +import multiprocessing as mp +import queue as queue_module +import threading +import time +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from storage.studio_db import get_connection +from utils.rag.config import ( + RAG_CHUNK_OVERLAP, + RAG_CHUNK_SIZE, + RAG_EMBED_BATCH_SIZE, + RAG_EMBEDDING_MODEL, +) + +from . import bm25, embeddings, vector_store +from .vector_store import kb_scope, thread_scope + +logger = logging.getLogger(__name__) + +_CTX = mp.get_context("spawn") +_QUEUE_TIMEOUT_SECONDS = 300 + + +# ------------------------------------------------------------------ +# Subprocess worker +# ------------------------------------------------------------------ + +def _subprocess_worker( + stored_path: str, + model_name: str, + chunk_size: int, + overlap: int, + batch_size: int, + out_queue: Any, +) -> None: + try: + from core.rag.chunking import chunk_pages + from core.rag.parsers import parse + + out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) + pages = parse(Path(stored_path)) + if not pages: + out_queue.put({"type": "error", "error": "no extractable text in document"}) + return + + out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1}) + from core.rag.embeddings import get_embedder, token_counter + + model = get_embedder(model_name) + counter = token_counter(model_name) + dim = int(model.get_sentence_embedding_dimension()) + out_queue.put({"type": "dim", "dim": dim}) + + out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) + chunks = chunk_pages( + pages, + max_tokens = chunk_size, + overlap_tokens = overlap, + token_counter = counter, + ) + if not chunks: + out_queue.put({"type": "error", "error": "chunker produced no chunks"}) + return + + total = len(chunks) + for i in range(0, total, batch_size): + batch = chunks[i : i + batch_size] + vectors = model.encode( + [c.text for c in batch], + batch_size = batch_size, + normalize_embeddings = True, + convert_to_numpy = True, + show_progress_bar = False, + ) + out_queue.put( + { + "type": "chunks_batch", + "first_index": i, + "chunks": [ + { + "text": c.text, + "token_count": c.token_count, + "page_number": c.page_number, + } + for c in batch + ], + "vectors": vectors.tolist(), + } + ) + progress = 0.3 + 0.65 * min(1.0, (i + len(batch)) / total) + out_queue.put({"type": "progress", "stage": "embed", "progress": progress}) + + out_queue.put({"type": "complete", "num_chunks": total}) + except Exception as exc: # noqa: BLE001 + logger.exception("ingestion subprocess failed") + out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"}) + + +# ------------------------------------------------------------------ +# Job manager (parent side) +# ------------------------------------------------------------------ + +class _JobState: + def __init__(self, job_id: str, document_id: str, scope: str) -> None: + self.job_id = job_id + self.document_id = document_id + self.scope = scope + self.status = "pending" + self.stage: str | None = None + self.progress: float = 0.0 + self.error: str | None = None + self.subscribers: list[queue_module.Queue[dict]] = [] + self.lock = threading.Lock() + + def push_event(self, event: dict) -> None: + with self.lock: + subs = list(self.subscribers) + for q in subs: + try: + q.put_nowait(event) + except queue_module.Full: + pass + + def subscribe(self) -> queue_module.Queue[dict]: + q: queue_module.Queue[dict] = queue_module.Queue(maxsize = 256) + with self.lock: + self.subscribers.append(q) + return q + + def unsubscribe(self, q: queue_module.Queue[dict]) -> None: + with self.lock: + if q in self.subscribers: + self.subscribers.remove(q) + + +_jobs: dict[str, _JobState] = {} +_jobs_lock = threading.Lock() + + +def get_job_state(job_id: str) -> _JobState | None: + with _jobs_lock: + return _jobs.get(job_id) + + +def _scope_for(kb_id: str | None, thread_id: str | None) -> str: + if kb_id: + return kb_scope(kb_id) + if thread_id: + return thread_scope(thread_id) + raise ValueError("must supply kb_id or thread_id") + + +def _update_job_row(job_id: str, **fields: Any) -> None: + if not fields: + return + keys = list(fields.keys()) + set_clause = ", ".join(f"{k} = ?" for k in keys) + values = list(fields.values()) + [job_id] + with get_connection() as conn: + conn.execute(f"UPDATE rag_ingestion_jobs SET {set_clause} WHERE id = ?", values) + conn.commit() + + +def _update_document_row(document_id: str, **fields: Any) -> None: + if not fields: + return + keys = list(fields.keys()) + set_clause = ", ".join(f"{k} = ?" for k in keys) + values = list(fields.values()) + [document_id] + with get_connection() as conn: + conn.execute(f"UPDATE rag_documents SET {set_clause} WHERE id = ?", values) + conn.commit() + + +def _insert_chunks_and_collect_for_bm25( + document_id: str, + scope: str, + first_index: int, + chunks_meta: list[dict], + vectors: list[list[float]], +) -> list[dict]: + """Insert chunks into sqlite + Qdrant; return [{id, text}] for BM25.""" + rows: list[tuple] = [] + points: list[dict] = [] + bm25_rows: list[dict] = [] + for offset, (meta, vec) in enumerate(zip(chunks_meta, vectors)): + chunk_index = first_index + offset + chunk_id = str(uuid4()) + rows.append( + ( + chunk_id, + document_id, + chunk_index, + meta["text"], + meta["token_count"], + meta["page_number"], + ) + ) + points.append( + { + "id": chunk_id, + "vector": vec, + "payload": { + "document_id": document_id, + "chunk_index": chunk_index, + "text": meta["text"], + "page_number": meta["page_number"], + }, + } + ) + bm25_rows.append({"id": chunk_id, "text": meta["text"]}) + with get_connection() as conn: + conn.executemany( + """ + INSERT INTO rag_chunks + (id, document_id, chunk_index, text, token_count, page_number) + VALUES (?, ?, ?, ?, ?, ?) + """, + rows, + ) + conn.commit() + vector_store.upsert_chunks(scope, points) + return bm25_rows + + +def _all_scope_chunks(scope: str) -> list[dict]: + if scope.startswith("kb_"): + kb_id = scope[len("kb_"):] + sql = ( + "SELECT c.id, c.text FROM rag_chunks c " + "JOIN rag_documents d ON d.id = c.document_id " + "WHERE d.kb_id = ?" + ) + bind = (kb_id,) + elif scope.startswith("thread_"): + thread_id = scope[len("thread_"):] + sql = ( + "SELECT c.id, c.text FROM rag_chunks c " + "JOIN rag_documents d ON d.id = c.document_id " + "WHERE d.thread_id = ?" + ) + bind = (thread_id,) + else: + return [] + with get_connection() as conn: + rows = conn.execute(sql, bind).fetchall() + return [{"id": r["id"], "text": r["text"]} for r in rows] + + +def _pump( + state: _JobState, + proc: Any, + out_queue: Any, +) -> None: + """Drain queue messages until the subprocess signals complete/error or dies.""" + bm25_buffer: list[dict] = [] + embedding_dim: int | None = None + final_status = "failed" + final_error: str | None = None + final_num_chunks = 0 + + started_at = int(time.time()) + state.status = "running" + _update_job_row(state.job_id, status = "running", started_at = started_at) + _update_document_row(state.document_id, status = "running") + state.push_event({"type": "status", "status": "running"}) + + try: + while True: + try: + msg = out_queue.get(timeout = _QUEUE_TIMEOUT_SECONDS) + except queue_module.Empty: + if not proc.is_alive(): + final_error = "subprocess exited without completion message" + break + continue + mtype = msg.get("type") + if mtype == "progress": + state.stage = msg.get("stage") + state.progress = float(msg.get("progress", 0.0)) + _update_job_row( + state.job_id, + stage = state.stage, + progress = state.progress, + ) + state.push_event(msg) + elif mtype == "dim": + embedding_dim = int(msg["dim"]) + 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) + bm25_rows = _insert_chunks_and_collect_for_bm25( + state.document_id, + state.scope, + int(msg["first_index"]), + msg["chunks"], + msg["vectors"], + ) + bm25_buffer.extend(bm25_rows) + elif mtype == "complete": + final_status = "completed" + final_num_chunks = int(msg.get("num_chunks", len(bm25_buffer))) + break + elif mtype == "error": + final_error = str(msg.get("error", "unknown error")) + break + else: + logger.warning("ingestion: unknown message type %r", mtype) + finally: + proc.join(timeout = 30) + if proc.is_alive(): + proc.terminate() + proc.join(timeout = 5) + + finished_at = int(time.time()) + if final_status == "completed": + full_scope_chunks = _all_scope_chunks(state.scope) + bm25.rebuild_index(state.scope, full_scope_chunks) + _update_document_row( + state.document_id, + status = "completed", + num_chunks = final_num_chunks, + ) + _update_job_row( + state.job_id, + status = "completed", + progress = 1.0, + stage = "done", + finished_at = finished_at, + ) + state.status = "completed" + state.progress = 1.0 + state.push_event( + { + "type": "complete", + "num_chunks": final_num_chunks, + } + ) + else: + _update_document_row( + state.document_id, + status = "failed", + error = final_error, + ) + _update_job_row( + state.job_id, + status = "failed", + error = final_error, + finished_at = finished_at, + ) + state.status = "failed" + state.error = final_error + state.push_event({"type": "error", "error": final_error}) + + +def enqueue_ingestion( + document_id: str, + stored_path: Path, + *, + kb_id: str | None = None, + thread_id: str | None = None, + embedding_model: str | None = None, +) -> 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. + """ + scope = _scope_for(kb_id, thread_id) + model_name = embedding_model or RAG_EMBEDDING_MODEL + job_id = str(uuid4()) + with get_connection() as conn: + conn.execute( + """ + INSERT INTO rag_ingestion_jobs + (id, document_id, status, progress, stage) + VALUES (?, ?, 'pending', 0.0, 'queued') + """, + (job_id, document_id), + ) + conn.commit() + + state = _JobState(job_id = job_id, document_id = document_id, scope = scope) + with _jobs_lock: + _jobs[job_id] = state + + out_queue = _CTX.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, + ), + daemon = True, + ) + proc.start() + pump_thread = threading.Thread( + target = _pump, + args = (state, proc, out_queue), + name = f"rag-ingest-pump-{job_id[:8]}", + daemon = True, + ) + pump_thread.start() + return job_id + + +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). + """ + vector_store.delete_document(scope, document_id) + remaining = _all_scope_chunks(scope) + if remaining: + bm25.rebuild_index(scope, remaining) + else: + bm25.delete_scope(scope) + + +def delete_scope_artifacts(scope: str) -> None: + vector_store.delete_scope(scope) + bm25.delete_scope(scope) + + +def purge_thread_documents(thread_ids: list[str]) -> None: + """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). + """ + if not thread_ids: + return + import os + from pathlib import Path + + from utils.paths.storage_roots import rag_uploads_root + + placeholders = ",".join("?" for _ in thread_ids) + uploads_root = Path(os.path.realpath(rag_uploads_root())) + with get_connection() as conn: + rows = conn.execute( + f"SELECT stored_path FROM rag_documents WHERE thread_id IN ({placeholders})", + thread_ids, + ).fetchall() + conn.execute( + f"DELETE FROM rag_documents WHERE thread_id IN ({placeholders})", + thread_ids, + ) + conn.commit() + for row in rows: + try: + real = Path(os.path.realpath(row["stored_path"])) + real.relative_to(uploads_root) + except (OSError, ValueError): + continue + real.unlink(missing_ok = True) + for thread_id in thread_ids: + delete_scope_artifacts(thread_scope(thread_id)) + + +def purge_all_thread_documents() -> None: + """Drop every per-thread RAG artifact. Used by clear-all-history.""" + with get_connection() as conn: + rows = conn.execute( + "SELECT DISTINCT thread_id FROM rag_documents WHERE thread_id IS NOT NULL" + ).fetchall() + purge_thread_documents([r["thread_id"] for r in rows]) diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py new file mode 100644 index 0000000000..105e2a89d8 --- /dev/null +++ b/studio/backend/core/rag/parsers/__init__.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + + +@dataclass(frozen = True) +class ParsedPage: + text: str + page_number: int | None = None + + +class UnsupportedFormatError(ValueError): + pass + + +def parse(path: Path) -> list[ParsedPage]: + suffix = path.suffix.lower() + if suffix == ".pdf": + from .pdf import extract + elif suffix in (".txt", ".md", ".markdown"): + from .text import extract + elif suffix == ".docx": + from .docx import extract + elif suffix in (".html", ".htm"): + from .html import extract + else: + raise UnsupportedFormatError(f"Unsupported file type: {suffix}") + return extract(path) diff --git a/studio/backend/core/rag/parsers/docx.py b/studio/backend/core/rag/parsers/docx.py new file mode 100644 index 0000000000..fedc401773 --- /dev/null +++ b/studio/backend/core/rag/parsers/docx.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from pathlib import Path + +from . import ParsedPage + + +def extract(path: Path) -> list[ParsedPage]: + from docx import Document + + document = Document(str(path)) + parts: list[str] = [] + for paragraph in document.paragraphs: + if paragraph.text and paragraph.text.strip(): + parts.append(paragraph.text) + for table in document.tables: + for row in table.rows: + cells = [cell.text.strip() for cell in row.cells if cell.text.strip()] + if cells: + parts.append(" | ".join(cells)) + text = "\n\n".join(parts).strip() + if not text: + return [] + return [ParsedPage(text = text, page_number = None)] diff --git a/studio/backend/core/rag/parsers/html.py b/studio/backend/core/rag/parsers/html.py new file mode 100644 index 0000000000..121cb6faf3 --- /dev/null +++ b/studio/backend/core/rag/parsers/html.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from pathlib import Path + +from . import ParsedPage + + +_SKIP_TAGS = {"script", "style", "noscript", "template"} + + +def extract(path: Path) -> list[ParsedPage]: + from bs4 import BeautifulSoup + + raw = path.read_bytes() + soup = BeautifulSoup(raw, "lxml") + for tag in soup(_SKIP_TAGS): + tag.decompose() + text = soup.get_text(separator = "\n").strip() + lines = [line.strip() for line in text.splitlines() if line.strip()] + cleaned = "\n".join(lines) + if not cleaned: + return [] + return [ParsedPage(text = cleaned, page_number = None)] diff --git a/studio/backend/core/rag/parsers/pdf.py b/studio/backend/core/rag/parsers/pdf.py new file mode 100644 index 0000000000..2a3f3fbda1 --- /dev/null +++ b/studio/backend/core/rag/parsers/pdf.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from pathlib import Path + +from . import ParsedPage + + +def extract(path: Path) -> list[ParsedPage]: + from pypdf import PdfReader + + reader = PdfReader(str(path)) + pages: list[ParsedPage] = [] + for index, page in enumerate(reader.pages): + try: + text = page.extract_text() or "" + except Exception: + text = "" + text = text.strip() + if text: + pages.append(ParsedPage(text = text, page_number = index + 1)) + return pages diff --git a/studio/backend/core/rag/parsers/text.py b/studio/backend/core/rag/parsers/text.py new file mode 100644 index 0000000000..402b5f32f9 --- /dev/null +++ b/studio/backend/core/rag/parsers/text.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from pathlib import Path + +from . import ParsedPage + + +def extract(path: Path) -> list[ParsedPage]: + raw = path.read_bytes() + try: + text = raw.decode("utf-8") + except UnicodeDecodeError: + try: + import chardet + + detected = chardet.detect(raw) + encoding = detected.get("encoding") or "latin-1" + except ImportError: + encoding = "latin-1" + text = raw.decode(encoding, errors = "replace") + text = text.strip() + if not text: + return [] + return [ParsedPage(text = text, page_number = None)] diff --git a/studio/backend/core/rag/reranker.py b/studio/backend/core/rag/reranker.py new file mode 100644 index 0000000000..a6ac81b12e --- /dev/null +++ b/studio/backend/core/rag/reranker.py @@ -0,0 +1,113 @@ +# 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``. +""" + +from __future__ import annotations + +import gc +import logging +import threading +from typing import Any + +from utils.rag.config import RAG_RERANK_BATCH_SIZE, RAG_RERANKER_MODEL + +from .retrieval import Hit + +logger = logging.getLogger(__name__) + +_lock = threading.Lock() +_model: Any | None = None +_model_name: str | None = None + + +def _load(model_name: str) -> Any: + from sentence_transformers import CrossEncoder + + logger.info("Loading RAG reranker: %s", model_name) + return CrossEncoder(model_name) + + +def get_reranker(model_name: str | None = None) -> Any: + global _model, _model_name + target = model_name or RAG_RERANKER_MODEL + with _lock: + if _model is None or _model_name != target: + unload() + _model = _load(target) + _model_name = target + return _model + + +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. + """ + global _model, _model_name + with _lock: + if _model is not None: + _model = None + _model_name = None + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except ImportError: + pass + + +def rerank( + query: str, + pairs: list[tuple[Hit, str]], + *, + 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. + """ + if not pairs: + return [] + model = get_reranker(model_name) + inputs = [(query, text) for _, text in pairs] + scores = model.predict( + inputs, + batch_size = RAG_RERANK_BATCH_SIZE, + show_progress_bar = False, + ) + ranked = sorted( + zip(pairs, scores), + key = lambda item: float(item[1]), + reverse = True, + ) + out: list[Hit] = [] + for (hit, _text), score in ranked: + out.append( + Hit( + chunk_id = hit.chunk_id, + score = float(score), + document_id = hit.document_id, + chunk_index = hit.chunk_index, + ) + ) + if top_k is not None: + out = out[:top_k] + return out diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py new file mode 100644 index 0000000000..c75ab04347 --- /dev/null +++ b/studio/backend/core/rag/retrieval.py @@ -0,0 +1,113 @@ +# 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). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from utils.rag.config import ( + RAG_RRF_K, + RAG_TOP_K_BM25, + RAG_TOP_K_DENSE, + RAG_TOP_K_HYBRID, +) + +from . import bm25, embeddings, vector_store + + +@dataclass(frozen = True) +class Hit: + chunk_id: str + score: float + document_id: str | None = None + chunk_index: int | None = None + + +def retrieve_bm25(scope: str, query: str, k: int | None = None) -> list[Hit]: + limit = k or RAG_TOP_K_BM25 + return [Hit(chunk_id = cid, score = s) for cid, s in bm25.search(scope, query, limit)] + + +def retrieve_dense( + scope: str, + query: str, + k: int | None = None, + *, + document_ids: list[str] | None = None, +) -> list[Hit]: + limit = k or RAG_TOP_K_DENSE + vector = embeddings.encode([query], normalize = True)[0].tolist() + raw = vector_store.search( + scope, + query_vector = vector, + top_k = limit, + document_ids = document_ids, + ) + out: list[Hit] = [] + for r in raw: + payload = r["payload"] + out.append( + Hit( + chunk_id = r["chunk_id"], + score = r["score"], + document_id = payload.get("document_id"), + chunk_index = payload.get("chunk_index"), + ) + ) + return out + + +def _rrf_fuse( + rankings: list[list[Hit]], + *, + rrf_k: int, + top_k: int, +) -> list[Hit]: + fused: dict[str, float] = {} + seen: dict[str, Hit] = {} + for ranking in rankings: + for rank, hit in enumerate(ranking): + fused[hit.chunk_id] = fused.get(hit.chunk_id, 0.0) + 1.0 / (rrf_k + rank + 1) + if hit.chunk_id not in seen: + seen[hit.chunk_id] = hit + ordered = sorted(fused.items(), key = lambda kv: kv[1], reverse = True)[:top_k] + return [ + Hit( + chunk_id = cid, + score = score, + document_id = seen[cid].document_id, + chunk_index = seen[cid].chunk_index, + ) + for cid, score in ordered + ] + + +def retrieve_hybrid( + scope: str, + query: str, + *, + k: int | None = None, + k_bm25: int | None = None, + k_dense: int | None = None, + document_ids: list[str] | None = None, +) -> list[Hit]: + bm25_hits = retrieve_bm25(scope, query, k_bm25 or RAG_TOP_K_BM25) + dense_hits = retrieve_dense( + scope, + query, + k_dense or RAG_TOP_K_DENSE, + document_ids = document_ids, + ) + return _rrf_fuse( + [bm25_hits, dense_hits], + rrf_k = RAG_RRF_K, + top_k = k or RAG_TOP_K_HYBRID, + ) diff --git a/studio/backend/core/rag/vector_store.py b/studio/backend/core/rag/vector_store.py new file mode 100644 index 0000000000..44e7581c6d --- /dev/null +++ b/studio/backend/core/rag/vector_store.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Qdrant local-mode vector store. + +Qdrant's local mode (``QdrantClient(path=...)``) acquires a file lock on +the storage directory, so only one process at a time can hold the +client. That means *all* vector reads/writes funnel through this module +in the FastAPI parent process. Ingestion subprocesses do not open +Qdrant directly — they compute vectors and send them back over a queue +for the parent to persist. + +A "scope" is the collection name: ``kb_`` for standalone +knowledge bases and ``thread_`` for per-thread document sets. +""" + +from __future__ import annotations + +import logging +import threading +from typing import Any, Iterable + +from utils.paths.storage_roots import ensure_dir, rag_vectordb_root + +logger = logging.getLogger(__name__) + +_client: Any | None = None +_client_lock = threading.Lock() + + +def get_qdrant() -> Any: + """Lazy singleton; parent process only.""" + global _client + with _client_lock: + if _client is None: + from qdrant_client import QdrantClient + + path = ensure_dir(rag_vectordb_root()) + _client = QdrantClient(path = str(path)) + return _client + + +def kb_scope(kb_id: str) -> str: + return f"kb_{kb_id}" + + +def thread_scope(thread_id: str) -> str: + return f"thread_{thread_id}" + + +def collection_exists(scope: str) -> bool: + client = get_qdrant() + try: + client.get_collection(collection_name = scope) + return True + except Exception: + return False + + +def ensure_collection(scope: str, dim: int) -> None: + client = get_qdrant() + if collection_exists(scope): + return + from qdrant_client.models import Distance, VectorParams + + client.create_collection( + collection_name = scope, + vectors_config = VectorParams(size = dim, distance = Distance.COSINE), + ) + + +def upsert_chunks( + scope: str, + points: Iterable[dict], +) -> None: + """Insert/update chunk vectors. + + Each point must have keys ``id`` (str), ``vector`` (list[float]) and + ``payload`` (dict with at least ``document_id`` and ``chunk_index``). + """ + from qdrant_client.models import PointStruct + + client = get_qdrant() + structured = [ + PointStruct(id = p["id"], vector = p["vector"], payload = p["payload"]) + for p in points + ] + if not structured: + return + client.upsert(collection_name = scope, points = structured) + + +def search( + scope: str, + query_vector: list[float], + *, + top_k: int, + document_ids: list[str] | None = None, +) -> list[dict]: + from qdrant_client.models import FieldCondition, Filter, MatchAny + + client = get_qdrant() + query_filter = None + if document_ids: + query_filter = Filter( + must = [ + FieldCondition( + key = "document_id", + match = MatchAny(any = document_ids), + ) + ] + ) + if not collection_exists(scope): + return [] + results = client.search( + collection_name = scope, + query_vector = query_vector, + limit = top_k, + query_filter = query_filter, + ) + return [ + { + "chunk_id": str(r.id), + "score": float(r.score), + "payload": dict(r.payload or {}), + } + for r in results + ] + + +def delete_scope(scope: str) -> None: + client = get_qdrant() + if not collection_exists(scope): + return + client.delete_collection(collection_name = scope) + + +def delete_document(scope: str, document_id: str) -> None: + from qdrant_client.models import FieldCondition, Filter, FilterSelector, MatchValue + + if not collection_exists(scope): + return + client = get_qdrant() + client.delete( + collection_name = scope, + points_selector = FilterSelector( + filter = Filter( + must = [ + FieldCondition( + key = "document_id", + match = MatchValue(value = document_id), + ) + ] + ) + ), + ) diff --git a/studio/backend/main.py b/studio/backend/main.py index 004ae404cd..d1a61d4417 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -122,6 +122,7 @@ from routes import ( inference_studio_router, models_router, providers_router, + rag_router, training_history_router, training_router, ) @@ -528,6 +529,7 @@ app.include_router(export_router, prefix = "/api/export", tags = ["export"]) app.include_router( training_history_router, prefix = "/api/train", tags = ["training-history"] ) +app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"]) # ============ Health and System Endpoints ============ diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index c33ebf4d94..36d1e0f5c4 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -76,3 +76,15 @@ trl>=0.18.2,!=0.19.0,<=0.24.0 sentence-transformers cut_cross_entropy pillow + +# RAG: vector store, lexical index, document parsers. +# qdrant-client supports a pure-Python local mode (QdrantClient(path=...)) +# that we use to keep the studio install self-contained — no separate +# server. bm25s persists per-scope indices to disk. +qdrant-client>=1.12 +bm25s>=0.2 +pypdf>=4.0 +python-docx>=1.1 +beautifulsoup4>=4.12 +lxml>=5.0 +chardet>=5.2 diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index 6bb5d15e8e..433355f969 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -16,6 +16,7 @@ from routes.export import router as export_router from routes.training_history import router as training_history_router from routes.chat_history import router as chat_history_router from routes.providers import router as providers_router +from routes.rag import router as rag_router __all__ = [ "training_router", @@ -29,4 +30,5 @@ __all__ = [ "training_history_router", "chat_history_router", "providers_router", + "rag_router", ] diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index ed808040d2..648b57aab3 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -11,6 +11,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, ConfigDict, Field, ValidationError from auth.authentication import get_current_subject +from core.rag.ingestion import purge_all_thread_documents, purge_thread_documents from storage.studio_db import ( ChatMessageConflictError, CorruptSettingsError, @@ -229,6 +230,9 @@ async def delete_threads( payload: ChatDeleteRequest, current_subject: str = Depends(get_current_subject), ): + # rag_documents has no FK cascade to chat_threads, so purge their + # files + vectors + bm25 explicitly before deleting the threads. + purge_thread_documents(payload.ids) delete_chat_threads(payload.ids) return {"status": "deleted"} @@ -354,6 +358,7 @@ async def record_import_ledger( @router.delete("") async def clear_history(current_subject: str = Depends(get_current_subject)): + purge_all_thread_documents() clear_chat_history() return {"status": "deleted"} diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py new file mode 100644 index 0000000000..1fbda5b594 --- /dev/null +++ b/studio/backend/routes/rag.py @@ -0,0 +1,681 @@ +# 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). +""" + +from __future__ import annotations + +import asyncio +import json +import os +import queue as queue_module +import time +from pathlib import Path +from typing import Any, Literal, Optional +from uuid import uuid4 + +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, UploadFile +from fastapi.responses import StreamingResponse +from pydantic import BaseModel, Field + +from auth.authentication import get_current_subject, get_current_subject_sse + + +async def _sse_auth( + token: str | None = Query(None), + authorization: str | None = Header(None), +) -> str: + return await get_current_subject_sse(token, authorization) +from core.rag import embeddings, ingestion, reranker, retrieval, vector_store +from core.rag.vector_store import kb_scope, thread_scope +from loggers import get_logger +from storage.studio_db import get_connection +from utils.paths.storage_roots import ensure_dir, rag_uploads_root +from utils.rag.config import ( + RAG_MAX_UPLOAD_MB, + RAG_RERANK_CANDIDATE_K, + RAG_UPLOAD_EXTS, +) + +router = APIRouter() +logger = get_logger(__name__) + + +# ------------------------------------------------------------------ +# Pydantic schemas +# ------------------------------------------------------------------ + +class CreateKBRequest(BaseModel): + name: str = Field(min_length = 1, max_length = 200) + description: str | None = None + embedding_model: str | None = None + + +class KBResponse(BaseModel): + id: str + name: str + description: str | None + embedding_model: str + created_at: int + + +class KBListResponse(BaseModel): + knowledge_bases: list[KBResponse] + + +class DocumentResponse(BaseModel): + id: str + kb_id: str | None + thread_id: str | None + filename: str + content_type: str | None + status: str + num_chunks: int + byte_size: int + error: str | None + created_at: int + + +class DocumentListResponse(BaseModel): + documents: list[DocumentResponse] + + +class ThreadIndexSummary(BaseModel): + thread_id: str + title: str | None + num_documents: int + num_chunks: int + + +class ThreadIndexListResponse(BaseModel): + threads: list[ThreadIndexSummary] + + +class UploadResponse(BaseModel): + document_id: str + job_id: str + filename: str + + +class SearchRequest(BaseModel): + query: str = Field(min_length = 1, max_length = 4000) + kb_id: str | None = None + thread_id: str | None = None + top_k: int = Field(default = 10, ge = 1, le = 100) + mode: Literal["bm25", "dense", "hybrid"] = "hybrid" + document_ids: list[str] | None = None + enable_rerank: bool = False + reranker_model: str | None = None + + +class SearchHit(BaseModel): + chunk_id: str + document_id: str + chunk_index: int + text: str + score: float + page_number: int | None = None + filename: str | None = None + + +class SearchResponse(BaseModel): + hits: list[SearchHit] + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + +def _sanitize_filename(filename: str) -> str: + name = Path(filename).name.strip().replace("\x00", "") + return name or "document" + + +def _now_ms() -> int: + return int(time.time()) + + +def _row_to_kb(row: Any) -> KBResponse: + return KBResponse( + id = row["id"], + name = row["name"], + description = row["description"], + embedding_model = row["embedding_model"], + created_at = row["created_at"], + ) + + +def _row_to_document(row: Any) -> DocumentResponse: + return DocumentResponse( + id = row["id"], + kb_id = row["kb_id"], + thread_id = row["thread_id"], + filename = row["filename"], + content_type = row["content_type"], + status = row["status"], + num_chunks = row["num_chunks"], + byte_size = row["byte_size"], + error = row["error"], + created_at = row["created_at"], + ) + + +def _kb_or_404(kb_id: str) -> Any: + with get_connection() as conn: + row = conn.execute( + "SELECT * FROM rag_knowledge_bases WHERE id = ?", + (kb_id,), + ).fetchone() + if not row: + raise HTTPException(status_code = 404, detail = "Knowledge base not found") + return row + + +def _thread_or_404(thread_id: str) -> None: + with get_connection() as conn: + row = conn.execute( + "SELECT id FROM chat_threads WHERE id = ?", + (thread_id,), + ).fetchone() + if not row: + raise HTTPException(status_code = 404, detail = "Thread not found") + + +def _document_or_404(document_id: str) -> Any: + with get_connection() as conn: + row = conn.execute( + "SELECT * FROM rag_documents WHERE id = ?", + (document_id,), + ).fetchone() + if not row: + raise HTTPException(status_code = 404, detail = "Document not found") + return row + + +async def _save_upload(file: UploadFile) -> tuple[Path, str, int]: + filename = _sanitize_filename(file.filename or "document") + ext = Path(filename).suffix.lower() + if ext not in RAG_UPLOAD_EXTS: + allowed = ", ".join(sorted(RAG_UPLOAD_EXTS)) + raise HTTPException( + status_code = 400, + detail = f"Unsupported file type: {ext}. Allowed: {allowed}", + ) + upload_dir = ensure_dir(rag_uploads_root()) + stored_name = f"{uuid4().hex}_{Path(filename).stem}{ext}" + stored_path = upload_dir / stored_name + max_bytes = RAG_MAX_UPLOAD_MB * 1024 * 1024 + written = 0 + with open(stored_path, "wb") as f: + while True: + chunk = await file.read(1024 * 1024) + if not chunk: + break + written += len(chunk) + if written > max_bytes: + f.close() + stored_path.unlink(missing_ok = True) + raise HTTPException( + status_code = 413, + detail = f"File exceeds {RAG_MAX_UPLOAD_MB} MB limit", + ) + f.write(chunk) + if written == 0: + stored_path.unlink(missing_ok = True) + raise HTTPException(status_code = 400, detail = "Empty upload payload") + return stored_path, filename, written + + +def _start_ingestion( + *, + filename: str, + stored_path: Path, + byte_size: int, + content_type: str | None, + kb_id: str | None, + thread_id: str | None, + embedding_model: str, +) -> UploadResponse: + document_id = str(uuid4()) + with get_connection() as conn: + conn.execute( + """ + INSERT INTO rag_documents + (id, kb_id, thread_id, filename, content_type, stored_path, + status, num_chunks, byte_size, created_at) + VALUES (?, ?, ?, ?, ?, ?, 'pending', 0, ?, ?) + """, + ( + document_id, + kb_id, + thread_id, + filename, + content_type, + str(stored_path), + byte_size, + _now_ms(), + ), + ) + conn.commit() + job_id = ingestion.enqueue_ingestion( + document_id = document_id, + stored_path = stored_path, + kb_id = kb_id, + thread_id = thread_id, + embedding_model = embedding_model, + ) + return UploadResponse(document_id = document_id, job_id = job_id, filename = filename) + + +def _unlink_if_under_uploads(path: Path) -> None: + try: + real = Path(os.path.realpath(path)) + root = Path(os.path.realpath(rag_uploads_root())) + real.relative_to(root) + except (OSError, ValueError): + return + real.unlink(missing_ok = True) + + +# ------------------------------------------------------------------ +# Knowledge bases +# ------------------------------------------------------------------ + +@router.post("/knowledge-bases", response_model = KBResponse) +def create_knowledge_base( + payload: CreateKBRequest, + current_subject: str = Depends(get_current_subject), +) -> KBResponse: + from utils.rag.config import RAG_EMBEDDING_MODEL + + kb_id = str(uuid4()) + embedding_model = payload.embedding_model or RAG_EMBEDDING_MODEL + created_at = _now_ms() + with get_connection() as conn: + try: + conn.execute( + """ + INSERT INTO rag_knowledge_bases + (id, name, description, owner_user_id, embedding_model, created_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + kb_id, + payload.name, + payload.description, + current_subject, + embedding_model, + created_at, + ), + ) + conn.commit() + except Exception as exc: + raise HTTPException( + status_code = 409, + detail = f"Could not create KB: {exc}", + ) from exc + return KBResponse( + id = kb_id, + name = payload.name, + description = payload.description, + embedding_model = embedding_model, + created_at = created_at, + ) + + +@router.get("/knowledge-bases", response_model = KBListResponse) +def list_knowledge_bases( + current_subject: str = Depends(get_current_subject), +) -> KBListResponse: + with get_connection() as conn: + rows = conn.execute( + "SELECT * FROM rag_knowledge_bases ORDER BY created_at DESC" + ).fetchall() + return KBListResponse(knowledge_bases = [_row_to_kb(r) for r in rows]) + + +@router.delete("/knowledge-bases/{kb_id}") +def delete_knowledge_base( + kb_id: str, + current_subject: str = Depends(get_current_subject), +) -> dict: + _kb_or_404(kb_id) + with get_connection() as conn: + doc_rows = conn.execute( + "SELECT stored_path FROM rag_documents WHERE kb_id = ?", + (kb_id,), + ).fetchall() + conn.execute("DELETE FROM rag_knowledge_bases WHERE id = ?", (kb_id,)) + conn.commit() + for row in doc_rows: + _unlink_if_under_uploads(Path(row["stored_path"])) + ingestion.delete_scope_artifacts(kb_scope(kb_id)) + return {"ok": True} + + +# ------------------------------------------------------------------ +# Document upload (KB and per-thread) +# ------------------------------------------------------------------ + +@router.post("/knowledge-bases/{kb_id}/documents", response_model = UploadResponse) +async def upload_kb_document( + kb_id: str, + file: UploadFile, + current_subject: str = Depends(get_current_subject), +) -> UploadResponse: + kb_row = _kb_or_404(kb_id) + stored_path, filename, byte_size = await _save_upload(file) + return _start_ingestion( + filename = filename, + stored_path = stored_path, + byte_size = byte_size, + content_type = file.content_type, + kb_id = kb_id, + thread_id = None, + embedding_model = kb_row["embedding_model"], + ) + + +@router.post("/threads/{thread_id}/documents", response_model = UploadResponse) +async def upload_thread_document( + thread_id: str, + file: UploadFile, + current_subject: str = Depends(get_current_subject), +) -> UploadResponse: + from utils.rag.config import RAG_EMBEDDING_MODEL + + # 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. + stored_path, filename, byte_size = await _save_upload(file) + return _start_ingestion( + filename = filename, + stored_path = stored_path, + byte_size = byte_size, + content_type = file.content_type, + kb_id = None, + thread_id = thread_id, + embedding_model = RAG_EMBEDDING_MODEL, + ) + + +# ------------------------------------------------------------------ +# Document list / delete +# ------------------------------------------------------------------ + +@router.get("/knowledge-bases/{kb_id}/documents", response_model = DocumentListResponse) +def list_kb_documents( + kb_id: str, + current_subject: str = Depends(get_current_subject), +) -> DocumentListResponse: + _kb_or_404(kb_id) + with get_connection() as conn: + rows = conn.execute( + "SELECT * FROM rag_documents WHERE kb_id = ? ORDER BY created_at DESC", + (kb_id,), + ).fetchall() + return DocumentListResponse(documents = [_row_to_document(r) for r in rows]) + + +@router.get("/threads/{thread_id}/documents", response_model = DocumentListResponse) +def list_thread_documents( + thread_id: str, + current_subject: str = Depends(get_current_subject), +) -> DocumentListResponse: + with get_connection() as conn: + rows = conn.execute( + "SELECT * FROM rag_documents WHERE thread_id = ? ORDER BY created_at DESC", + (thread_id,), + ).fetchall() + return DocumentListResponse(documents = [_row_to_document(r) for r in rows]) + + +@router.delete("/documents/{document_id}") +def delete_document( + document_id: str, + current_subject: str = Depends(get_current_subject), +) -> dict: + row = _document_or_404(document_id) + scope = ( + kb_scope(row["kb_id"]) if row["kb_id"] else thread_scope(row["thread_id"]) + ) + with get_connection() as conn: + conn.execute("DELETE FROM rag_documents WHERE id = ?", (document_id,)) + conn.commit() + _unlink_if_under_uploads(Path(row["stored_path"])) + ingestion.delete_document_artifacts(document_id, scope) + return {"ok": True} + + +@router.get("/thread-indexes", response_model = ThreadIndexListResponse) +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. + """ + with get_connection() as conn: + rows = conn.execute( + """ + SELECT + d.thread_id AS thread_id, + t.title AS title, + COUNT(DISTINCT d.id) AS num_documents, + COALESCE(SUM(d.num_chunks), 0) AS num_chunks + FROM rag_documents d + LEFT JOIN chat_threads t ON t.id = d.thread_id + WHERE d.thread_id IS NOT NULL + GROUP BY d.thread_id, t.title + ORDER BY MAX(d.created_at) DESC + """ + ).fetchall() + return ThreadIndexListResponse( + threads = [ + ThreadIndexSummary( + thread_id = r["thread_id"], + title = r["title"], + num_documents = int(r["num_documents"]), + num_chunks = int(r["num_chunks"]), + ) + for r in rows + ] + ) + + +@router.delete("/threads/{thread_id}/documents") +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 Qdrant collection, the bm25 index, the + rag_documents/rag_chunks rows, and the uploaded files. The chat + thread itself is untouched. + """ + ingestion.purge_thread_documents([thread_id]) + return {"ok": True} + + +# ------------------------------------------------------------------ +# Ingestion job SSE +# ------------------------------------------------------------------ + +@router.get("/jobs/{job_id}/events") +async def job_events( + job_id: str, + request: Request, + current_subject: str = Depends(_sse_auth), +) -> StreamingResponse: + state = ingestion.get_job_state(job_id) + if state is None: + with get_connection() as conn: + row = conn.execute( + "SELECT * FROM rag_ingestion_jobs WHERE id = ?", + (job_id,), + ).fetchone() + if not row: + raise HTTPException(status_code = 404, detail = "Job not found") + return StreamingResponse( + _replay_terminal_state(row), + media_type = "text/event-stream", + ) + + consumer_queue = state.subscribe() + + async def stream(): + try: + initial = { + "type": "status", + "status": state.status, + "stage": state.stage, + "progress": state.progress, + } + yield f"data: {json.dumps(initial)}\n\n" + while True: + if await request.is_disconnected(): + break + try: + event = await asyncio.get_event_loop().run_in_executor( + None, + consumer_queue.get, + True, + 15.0, + ) + except queue_module.Empty: + yield ": keep-alive\n\n" + if state.status in ("completed", "failed"): + break + continue + yield f"data: {json.dumps(event)}\n\n" + if event.get("type") in ("complete", "error"): + break + finally: + state.unsubscribe(consumer_queue) + + return StreamingResponse(stream(), media_type = "text/event-stream") + + +async def _replay_terminal_state(row: Any): + payload = { + "type": "status", + "status": row["status"], + "stage": row["stage"], + "progress": row["progress"], + "error": row["error"], + } + yield f"data: {json.dumps(payload)}\n\n" + + +# ------------------------------------------------------------------ +# Search +# ------------------------------------------------------------------ + +@router.post("/search", response_model = SearchResponse) +def search( + payload: SearchRequest, + current_subject: str = Depends(get_current_subject), +) -> SearchResponse: + if bool(payload.kb_id) == bool(payload.thread_id): + raise HTTPException( + status_code = 400, + detail = "exactly one of kb_id or thread_id must be supplied", + ) + if payload.kb_id: + _kb_or_404(payload.kb_id) + scope = kb_scope(payload.kb_id) + else: + scope = thread_scope(payload.thread_id) + + # When reranking is opt-in, pull a wider candidate pool so the + # CrossEncoder has more to choose from before truncating to top_k. + candidate_k = ( + max(payload.top_k, RAG_RERANK_CANDIDATE_K) + if payload.enable_rerank + else payload.top_k + ) + + if payload.mode == "bm25": + hits = retrieval.retrieve_bm25(scope, payload.query, candidate_k) + elif payload.mode == "dense": + hits = retrieval.retrieve_dense( + scope, + payload.query, + candidate_k, + document_ids = payload.document_ids, + ) + else: + hits = retrieval.retrieve_hybrid( + scope, + payload.query, + k = candidate_k, + document_ids = payload.document_ids, + ) + + chunk_ids = [h.chunk_id for h in hits] + chunk_lookup: dict[str, dict] = {} + if chunk_ids: + placeholders = ",".join("?" for _ in chunk_ids) + with get_connection() as conn: + rows = conn.execute( + f""" + SELECT c.id AS chunk_id, c.document_id, c.chunk_index, c.text, + c.page_number, d.filename + FROM rag_chunks c + JOIN rag_documents d ON d.id = c.document_id + WHERE c.id IN ({placeholders}) + """, + chunk_ids, + ).fetchall() + for r in rows: + chunk_lookup[r["chunk_id"]] = dict(r) + + if payload.enable_rerank: + pairs = [ + (hit, chunk_lookup[hit.chunk_id]["text"]) + for hit in hits + if hit.chunk_id in chunk_lookup + ] + hits = reranker.rerank( + payload.query, + pairs, + model_name = payload.reranker_model, + top_k = payload.top_k, + ) + else: + hits = hits[: payload.top_k] + + out: list[SearchHit] = [] + for hit in hits: + meta = chunk_lookup.get(hit.chunk_id) + if not meta: + continue + out.append( + SearchHit( + chunk_id = hit.chunk_id, + document_id = meta["document_id"], + chunk_index = meta["chunk_index"], + text = meta["text"], + score = hit.score, + page_number = meta.get("page_number"), + filename = meta.get("filename"), + ) + ) + return SearchResponse(hits = out) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index de89b6cbd2..31bc67150c 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -205,6 +205,83 @@ 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. + conn.execute( + """ + CREATE TABLE IF NOT EXISTS rag_knowledge_bases ( + id TEXT NOT NULL PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + description TEXT, + owner_user_id TEXT, + embedding_model TEXT NOT NULL, + created_at INTEGER NOT NULL + ) + """ + ) + # 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. + conn.execute( + """ + CREATE TABLE IF NOT EXISTS rag_documents ( + id TEXT NOT NULL PRIMARY KEY, + kb_id TEXT REFERENCES rag_knowledge_bases(id) ON DELETE CASCADE, + thread_id TEXT, + filename TEXT NOT NULL, + content_type TEXT, + stored_path TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + num_chunks INTEGER NOT NULL DEFAULT 0, + byte_size INTEGER NOT NULL DEFAULT 0, + error TEXT, + created_at INTEGER NOT NULL, + CHECK ((kb_id IS NOT NULL) <> (thread_id IS NOT NULL)) + ) + """ + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_rag_documents_kb_id ON rag_documents(kb_id)" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_rag_documents_thread_id ON rag_documents(thread_id)" + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS rag_chunks ( + id TEXT NOT NULL PRIMARY KEY, + document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE, + chunk_index INTEGER NOT NULL, + text TEXT NOT NULL, + token_count INTEGER NOT NULL DEFAULT 0, + page_number INTEGER, + UNIQUE(document_id, chunk_index) + ) + """ + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_rag_chunks_document_id ON rag_chunks(document_id)" + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS rag_ingestion_jobs ( + id TEXT NOT NULL PRIMARY KEY, + document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE, + status TEXT NOT NULL DEFAULT 'pending', + progress REAL NOT NULL DEFAULT 0.0, + stage TEXT, + error TEXT, + started_at INTEGER, + finished_at INTEGER + ) + """ + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_rag_jobs_document_id ON rag_ingestion_jobs(document_id)" + ) def get_connection() -> sqlite3.Connection: diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 763d18bf3e..13e1c65ac4 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -120,6 +120,22 @@ def tensorboard_root() -> Path: return studio_root() / "runs" +def rag_root() -> Path: + return studio_root() / "rag" + + +def rag_uploads_root() -> Path: + return rag_root() / "uploads" + + +def rag_vectordb_root() -> Path: + return rag_root() / "qdrant" + + +def rag_bm25_root() -> Path: + return rag_root() / "bm25" + + def ensure_dir(path: Path) -> Path: path.mkdir(parents = True, exist_ok = True) return path @@ -261,6 +277,10 @@ def ensure_studio_directories() -> None: exports_root, auth_root, tensorboard_root, + rag_root, + rag_uploads_root, + rag_vectordb_root, + rag_bm25_root, ): ensure_dir(dir_fn()) _setup_cache_env() diff --git a/studio/backend/utils/rag/__init__.py b/studio/backend/utils/rag/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/utils/rag/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/utils/rag/config.py b/studio/backend/utils/rag/config.py new file mode 100644 index 0000000000..c2ae2a9af2 --- /dev/null +++ b/studio/backend/utils/rag/config.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import os + + +def _env_int(name: str, default: int) -> int: + raw = os.environ.get(name, "").strip() + if not raw: + return default + try: + return int(raw) + except ValueError: + return default + + +def _env_float(name: str, default: float) -> float: + raw = os.environ.get(name, "").strip() + if not raw: + return default + try: + return float(raw) + except ValueError: + return default + + +RAG_EMBEDDING_MODEL: str = ( + os.environ.get("UNSLOTH_RAG_EMBEDDING_MODEL", "").strip() + or "BAAI/bge-small-en-v1.5" +) + +RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512) +RAG_CHUNK_OVERLAP: int = _env_int("UNSLOTH_RAG_CHUNK_OVERLAP", 64) + +RAG_TOP_K_BM25: int = _env_int("UNSLOTH_RAG_TOP_K_BM25", 30) +RAG_TOP_K_DENSE: int = _env_int("UNSLOTH_RAG_TOP_K_DENSE", 30) +RAG_TOP_K_HYBRID: int = _env_int("UNSLOTH_RAG_TOP_K_HYBRID", 10) + +RAG_RRF_K: int = _env_int("UNSLOTH_RAG_RRF_K", 60) + +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" +) +RAG_RERANK_CANDIDATE_K: int = _env_int("UNSLOTH_RAG_RERANK_CANDIDATE_K", 50) +RAG_RERANK_BATCH_SIZE: int = _env_int("UNSLOTH_RAG_RERANK_BATCH_SIZE", 16) + +RAG_UPLOAD_EXTS: frozenset[str] = frozenset( + {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +) diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index c7bc0440bd..cbd7edf339 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -10,6 +10,7 @@ import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; +import { Route as knowledgeBasesRoute } from "./routes/knowledge-bases"; import { Route as loginRoute } from "./routes/login"; import { Route as onboardingRoute } from "./routes/onboarding"; import { Route as changePasswordRoute } from "./routes/change-password"; @@ -23,6 +24,7 @@ const routeTree = rootRoute.addChildren([ changePasswordRoute, gridTestRoute, settingsRoute, + knowledgeBasesRoute, studioRoute, chatRoute, exportRoute, diff --git a/studio/frontend/src/app/routes/knowledge-bases.tsx b/studio/frontend/src/app/routes/knowledge-bases.tsx new file mode 100644 index 0000000000..ee45828db8 --- /dev/null +++ b/studio/frontend/src/app/routes/knowledge-bases.tsx @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute, redirect } from "@tanstack/react-router"; +import { getPostAuthRoute } from "@/features/auth"; +import { useSettingsDialogStore } from "@/features/settings"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +// /knowledge-bases is a deep link to the settings modal's Knowledge +// Bases tab. Open it, then redirect home. Mirrors /settings. +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/knowledge-bases", + staticData: { title: "Knowledge Bases" }, + beforeLoad: async () => { + await requireAuth(); + useSettingsDialogStore.getState().openDialog("knowledge-bases"); + throw redirect({ to: getPostAuthRoute() }); + }, + component: () => null, +}); diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0c557f1b01..3928f79631 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -56,6 +56,12 @@ import { streamChatCompletions, validateModel, } from "./chat-api"; +import { + type SearchHit, + type SearchRequest, + search as ragSearch, +} from "@/features/rag/api/rag-api"; +import type { RagSource } from "./chat-settings-api"; import { createOpenAIContainer, listOpenAIContainers, @@ -65,6 +71,56 @@ import { isProviderKeyRotationError, } from "./providers-api"; +function extractMessageText(content: unknown): string { + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + const out: string[] = []; + for (const part of content) { + if ( + part && + typeof part === "object" && + (part as { type?: unknown }).type === "text" && + typeof (part as { text?: unknown }).text === "string" + ) { + out.push((part as { text: string }).text); + } + } + return out.join(" "); +} + +function buildRagRequest( + source: RagSource, + query: string, + resolvedThreadId: string | undefined, + enableRerank: boolean, + topK: number, +): SearchRequest | null { + const base: SearchRequest = { + query, + top_k: topK, + mode: "hybrid", + enable_rerank: enableRerank, + }; + if (source.kind === "thread") { + if (!resolvedThreadId) return null; + return { ...base, thread_id: resolvedThreadId }; + } + if (source.kind === "kb") { + return { ...base, kb_id: source.kbId }; + } + return null; +} + +function formatRagContext(hits: SearchHit[]): string { + const parts = hits.map((h) => { + const name = h.filename ?? `chunk ${h.chunk_index}`; + const pageAttr = + h.page_number != null ? ` page="${h.page_number}"` : ""; + return `\n${h.text}\n`; + }); + return `\nThe following documents may help answer the user's question:\n${parts.join("\n")}\n`; +} + /** Server-side usage data from llama-server (via stream_options.include_usage). */ interface ServerUsage { prompt_tokens: number; @@ -932,6 +988,49 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { content: safeSystemPrompt.trim(), }); } + + // 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. + const ragSource = runtime.ragSource; + if (ragSource.kind !== "off") { + const lastUser = [...outboundMessages] + .reverse() + .find((m) => m.role === "user"); + const queryText = lastUser + ? extractMessageText(lastUser.content) + : ""; + if (queryText.trim()) { + const ragReq = buildRagRequest( + ragSource, + queryText, + resolvedThreadId, + runtime.enableRerank, + runtime.ragTopK, + ); + if (ragReq) { + try { + const hits = await ragSearch(ragReq); + if (hits.length > 0) { + const block = formatRagContext(hits); + if ( + outboundMessages[0]?.role === "system" && + typeof outboundMessages[0].content === "string" + ) { + outboundMessages[0].content = + `${block}\n\n${outboundMessages[0].content}`; + } else { + outboundMessages.unshift({ role: "system", content: block }); + } + } + } catch (err) { + console.warn("RAG retrieval failed:", err); + } + } + } + } + let disabledToolGuard: string | null = null; const disabledToolGuardProviderType = externalProvider?.providerType; if ( diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts index 4208c90c76..0500e0fb6f 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -15,6 +15,11 @@ export interface PersistedChatPreset { params: PersistedInferenceParams; } +export type RagSource = + | { kind: "off" } + | { kind: "thread" } + | { kind: "kb"; kbId: string }; + export interface PersistedChatSettings { inferenceParams?: PersistedInferenceParams; customPresets?: PersistedChatPreset[]; @@ -26,6 +31,9 @@ export interface PersistedChatSettings { autoHealToolCalls?: boolean; maxToolCallsPerMessage?: number; toolCallTimeout?: number; + ragSource?: RagSource; + enableRerank?: boolean; + ragTopK?: number; } interface ChatSettingsResponse { diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 9cd4db2705..bef9fde24a 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -90,6 +90,23 @@ import { } from "./provider-capabilities"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import type { InferenceParams } from "./types/runtime"; +import type { RagSource } from "./api/chat-settings-api"; +import { DocumentRow } from "@/features/rag/components/document-row"; +import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog"; +import { useKnowledgeBases } from "@/features/rag/hooks/use-knowledge-bases"; +import { useThreadDocuments } from "@/features/rag/hooks/use-kb-documents"; +import { useRagStore } from "@/features/rag/stores/rag-store"; +import { Add01Icon, Delete02Icon } from "@hugeicons/core-free-icons"; + +function ragSourceLabel( + source: RagSource, + kbs: { id: string; name: string }[], +): string { + if (source.kind === "off") return "Off"; + if (source.kind === "thread") return "This thread's documents"; + const kb = kbs.find((k) => k.id === source.kbId); + return kb ? kb.name : "Knowledge base (missing)"; +} export { defaultInferenceParams, type Preset } from "./presets/preset-policy"; export type { InferenceParams } from "./types/runtime"; @@ -418,6 +435,21 @@ export function ChatSettingsPanel({ !isExternalModel || Boolean(providerCapabilities?.presencePenalty); const isMobile = useIsMobile(); const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; + const ragSource = useChatRuntimeStore((s) => s.ragSource); + const setRagSource = useChatRuntimeStore((s) => s.setRagSource); + const enableRerank = useChatRuntimeStore((s) => s.enableRerank); + const setEnableRerank = useChatRuntimeStore((s) => s.setEnableRerank); + const ragTopK = useChatRuntimeStore((s) => s.ragTopK); + const setRagTopK = useChatRuntimeStore((s) => s.setRagTopK); + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const { knowledgeBases, deleteKB } = useKnowledgeBases(); + const { documents: threadDocs, remove: removeThreadDoc } = useThreadDocuments( + ragSource.kind === "thread" ? activeThreadId : null, + ); + const clearThreadIndex = useRagStore((s) => s.clearThreadIndex); + const [kbCreateOpen, setKbCreateOpen] = useState(false); + const ragEnabled = ragSource.kind !== "off"; + const activeKbId = ragSource.kind === "kb" ? ragSource.kbId : null; const hasModelContent = !isExternalModel && (isGguf || Boolean(params.checkpoint)); const speculativeType = useChatRuntimeStore((s) => s.speculativeType); @@ -1166,6 +1198,181 @@ export function ChatSettingsPanel({ ) : null} + +
+
+ + + + + + + setRagSource({ kind: "off" })} + > + Off + + setRagSource({ kind: "thread" })} + > + This thread's documents + + {knowledgeBases.length > 0 ? ( + + ) : null} + {knowledgeBases.map((kb) => { + const isActive = kb.id === activeKbId; + return ( + + setRagSource({ kind: "kb", kbId: kb.id }) + } + > + {kb.name} + + + ); + })} + + setKbCreateOpen(true)} + className="text-muted-foreground" + > + + Create knowledge base… + + + +

+ Each message retrieves matching context from the selected + source before sending. +

+
+ setRagSource({ kind: "kb", kbId: kb.id })} + /> + {ragSource.kind === "thread" && activeThreadId ? ( +
+ + {threadDocs.length === 0 ? ( +

+ Attach a file using the + button in the composer to add + documents to this thread. +

+ ) : ( + <> +
+ {threadDocs.map((doc) => ( + { + void removeThreadDoc(doc.id); + }} + /> + ))} +
+ + + )} +
+ ) : null} +
+
+ + + {ragTopK} + +
+ v != null && setRagTopK(v)} + disabled={!ragEnabled} + /> +

+ Number of retrieved chunks passed to the model as context + (distinct from the sampling Top K below). Higher = more + grounding, more tokens. +

+
+
+
+ + Use reranker + + + Slower; uses GPU. Improves quality for fact-heavy questions. + +
+ +
+
+ + + + ); + })} {pendingAudio && (
@@ -854,7 +1023,7 @@ export function SharedComposer({ { diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index a00b53a44c..e7ce32e37d 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -19,6 +19,7 @@ import { loadChatSettingsWithLegacyImport, savePersistedChatSettingsPatch, } from "../utils/chat-settings-storage"; +import type { RagSource } from "../api/chat-settings-api"; const HF_TOKEN_KEY = "unsloth_hf_token"; export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; @@ -294,6 +295,9 @@ type ChatRuntimeStore = { } | null; modelLoading: boolean; activeNativePathToken: string | null; + ragSource: RagSource; + enableRerank: boolean; + ragTopK: number; hydratePersistedSettings: () => Promise; setModelLoading: (loading: boolean) => void; setModelRequiresTrustRemoteCode: (required: boolean) => void; @@ -337,6 +341,9 @@ type ChatRuntimeStore = { setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void; + setRagSource: (source: RagSource) => void; + setEnableRerank: (value: boolean) => void; + setRagTopK: (value: number) => void; }; type PersistedChatSettings = Awaited< @@ -352,7 +359,10 @@ type ScalarSettingKey = | "preserveThinking" | "autoHealToolCalls" | "maxToolCallsPerMessage" - | "toolCallTimeout"; + | "toolCallTimeout" + | "ragSource" + | "enableRerank" + | "ragTopK"; type PresetHydrationVersions = { customPresets: number; @@ -386,6 +396,9 @@ const SCALAR_SETTING_KEYS = [ "autoHealToolCalls", "maxToolCallsPerMessage", "toolCallTimeout", + "ragSource", + "enableRerank", + "ragTopK", ] as const satisfies readonly ScalarSettingKey[]; const inferenceParamMutationVersions = Object.fromEntries( @@ -590,6 +603,9 @@ export const useChatRuntimeStore = create((set, get) => ({ contextUsage: null, modelLoading: false, activeNativePathToken: null, + ragSource: { kind: "off" }, + enableRerank: false, + ragTopK: 5, hydratePersistedSettings: async () => { if (get().settingsHydrated) { return; @@ -789,6 +805,25 @@ export const useChatRuntimeStore = create((set, get) => ({ ); return { preserveThinking }; }), + setRagSource: (ragSource) => + set((state) => { + setScalarSettingVersion("ragSource", ragSource, state.ragSource); + return { ragSource }; + }), + setEnableRerank: (enableRerank) => + set((state) => { + setScalarSettingVersion( + "enableRerank", + enableRerank, + state.enableRerank, + ); + return { enableRerank }; + }), + setRagTopK: (ragTopK) => + set((state) => { + setScalarSettingVersion("ragTopK", ragTopK, state.ragTopK); + return { ragTopK }; + }), setToolsEnabled: (toolsEnabled, options) => set(() => { if (options?.persist !== false) { diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts new file mode 100644 index 0000000000..5368b5de26 --- /dev/null +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -0,0 +1,256 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch, getAuthToken } from "@/features/auth"; +import { apiUrl } from "@/lib/api-base"; +import { formatFastApiDetail } from "@/lib/format-fastapi-error"; + +export interface KnowledgeBase { + id: string; + name: string; + description: string | null; + embedding_model: string; + created_at: number; +} + +export interface RagDocument { + id: string; + kb_id: string | null; + thread_id: string | null; + filename: string; + content_type: string | null; + status: "pending" | "running" | "completed" | "failed"; + num_chunks: number; + byte_size: number; + error: string | null; + created_at: number; +} + +export interface UploadResponse { + document_id: string; + job_id: string; + filename: string; +} + +export interface SearchHit { + chunk_id: string; + document_id: string; + chunk_index: number; + text: string; + score: number; + page_number: number | null; + filename: string | null; +} + +export interface SearchRequest { + query: string; + kb_id?: string; + thread_id?: string; + top_k?: number; + mode?: "bm25" | "dense" | "hybrid"; + document_ids?: string[]; + enable_rerank?: boolean; + reranker_model?: string; +} + +export type JobEvent = + | { type: "status"; status: string; stage?: string | null; progress?: number; error?: string | null } + | { type: "progress"; stage: string; progress: number } + | { type: "complete"; num_chunks: number } + | { type: "error"; error: string }; + +function parseErrorText(status: number, body: unknown): string { + if (body && typeof body === "object") { + const detail = (body as { detail?: unknown }).detail; + const formatted = formatFastApiDetail(detail); + if (formatted) return formatted; + } + return `Request failed (${status})`; +} + +async function parseJsonOrThrow(response: Response): Promise { + const body = await response.json().catch(() => null); + if (!response.ok) { + throw new Error(parseErrorText(response.status, body)); + } + return body as T; +} + +async function throwOnError(response: Response): Promise { + if (response.ok) return; + const body = await response.json().catch(() => null); + throw new Error(parseErrorText(response.status, body)); +} + +// ------------------------------------------------------------------ +// Knowledge bases +// ------------------------------------------------------------------ + +export async function listKnowledgeBases(): Promise { + const response = await authFetch("/api/rag/knowledge-bases"); + const body = await parseJsonOrThrow<{ knowledge_bases: KnowledgeBase[] }>(response); + return body.knowledge_bases; +} + +export async function createKnowledgeBase(req: { + name: string; + description?: string; + embedding_model?: string; +}): Promise { + const response = await authFetch("/api/rag/knowledge-bases", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(req), + }); + return parseJsonOrThrow(response); +} + +export async function deleteKnowledgeBase(kbId: string): Promise { + const response = await authFetch( + `/api/rag/knowledge-bases/${encodeURIComponent(kbId)}`, + { method: "DELETE" }, + ); + await throwOnError(response); +} + +// ------------------------------------------------------------------ +// Documents +// ------------------------------------------------------------------ + +export async function listKBDocuments(kbId: string): Promise { + const response = await authFetch( + `/api/rag/knowledge-bases/${encodeURIComponent(kbId)}/documents`, + ); + const body = await parseJsonOrThrow<{ documents: RagDocument[] }>(response); + return body.documents; +} + +export async function listThreadDocuments(threadId: string): Promise { + const response = await authFetch( + `/api/rag/threads/${encodeURIComponent(threadId)}/documents`, + ); + const body = await parseJsonOrThrow<{ documents: RagDocument[] }>(response); + return body.documents; +} + +export async function uploadKBDocument( + kbId: string, + file: File, +): Promise { + const form = new FormData(); + form.append("file", file); + const response = await authFetch( + `/api/rag/knowledge-bases/${encodeURIComponent(kbId)}/documents`, + { method: "POST", body: form }, + ); + return parseJsonOrThrow(response); +} + +export async function uploadThreadDocument( + threadId: string, + file: File, +): Promise { + const form = new FormData(); + form.append("file", file); + const response = await authFetch( + `/api/rag/threads/${encodeURIComponent(threadId)}/documents`, + { method: "POST", body: form }, + ); + return parseJsonOrThrow(response); +} + +export async function deleteDocument(documentId: string): Promise { + const response = await authFetch( + `/api/rag/documents/${encodeURIComponent(documentId)}`, + { method: "DELETE" }, + ); + await throwOnError(response); +} + +export interface ThreadIndexSummary { + thread_id: string; + title: string | null; + num_documents: number; + num_chunks: number; +} + +export async function listThreadIndexes(): Promise { + const response = await authFetch("/api/rag/thread-indexes"); + const body = await parseJsonOrThrow<{ threads: ThreadIndexSummary[] }>(response); + return body.threads; +} + +export async function clearThreadDocuments(threadId: string): Promise { + const response = await authFetch( + `/api/rag/threads/${encodeURIComponent(threadId)}/documents`, + { method: "DELETE" }, + ); + await throwOnError(response); +} + +// ------------------------------------------------------------------ +// Search +// ------------------------------------------------------------------ + +export async function search(req: SearchRequest): Promise { + const response = await authFetch("/api/rag/search", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(req), + }); + const body = await parseJsonOrThrow<{ hits: SearchHit[] }>(response); + return body.hits; +} + +// ------------------------------------------------------------------ +// Ingestion SSE +// ------------------------------------------------------------------ + +/** + * Subscribe to an ingestion job's SSE event stream. + * + * EventSource cannot send custom headers, so we pass the bearer token + * as a `?token=…` query param. The backend's /jobs/{id}/events route + * accepts either the query param or the Authorization header. + * + * Returns an unsubscribe function that closes the EventSource. + */ +export function subscribeToJobEvents( + jobId: string, + handlers: { + onEvent?: (event: JobEvent) => void; + onError?: (error: Error) => void; + onClose?: () => void; + }, +): () => void { + const token = getAuthToken(); + const params = token ? `?token=${encodeURIComponent(token)}` : ""; + const url = apiUrl(`/api/rag/jobs/${encodeURIComponent(jobId)}/events${params}`); + const source = new EventSource(url); + + source.onmessage = (e) => { + try { + const parsed = JSON.parse(e.data) as JobEvent; + handlers.onEvent?.(parsed); + if (parsed.type === "complete" || parsed.type === "error") { + source.close(); + handlers.onClose?.(); + } + } catch (err) { + handlers.onError?.(err instanceof Error ? err : new Error(String(err))); + } + }; + + source.onerror = () => { + // Browser auto-reconnects unless we close explicitly. We close on + // any error so the consumer can decide whether to resubscribe. + source.close(); + handlers.onError?.(new Error("SSE connection lost")); + handlers.onClose?.(); + }; + + return () => { + source.close(); + handlers.onClose?.(); + }; +} diff --git a/studio/frontend/src/features/rag/components/document-row.tsx b/studio/frontend/src/features/rag/components/document-row.tsx new file mode 100644 index 0000000000..6e6fb226c4 --- /dev/null +++ b/studio/frontend/src/features/rag/components/document-row.tsx @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; +import { Delete02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import type { RagDocument } from "../api/rag-api"; + +const STATUS_VARIANT: Record< + RagDocument["status"], + "default" | "secondary" | "destructive" | "outline" +> = { + pending: "outline", + running: "secondary", + completed: "default", + failed: "destructive", +}; + +function humanBytes(bytes: number): string { + if (bytes < 1024) return `${bytes} B`; + if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`; + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; +} + +export function DocumentRow({ + doc, + onDelete, + rightSlot, + className, +}: { + doc: RagDocument; + onDelete?: () => void; + rightSlot?: React.ReactNode; + className?: string; +}) { + return ( +
+
+
+ + {doc.filename} + + + {doc.status} + +
+
+ {humanBytes(doc.byte_size)} + {doc.status === "completed" ? ( + {doc.num_chunks} chunks + ) : null} + {doc.error ? ( + {doc.error} + ) : null} +
+ {rightSlot} +
+ {onDelete ? ( + + ) : null} +
+ ); +} diff --git a/studio/frontend/src/features/rag/components/document-upload-dropzone.tsx b/studio/frontend/src/features/rag/components/document-upload-dropzone.tsx new file mode 100644 index 0000000000..6ada978f36 --- /dev/null +++ b/studio/frontend/src/features/rag/components/document-upload-dropzone.tsx @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; +import { Upload04Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useRef, useState } from "react"; + +const ACCEPTED = ".pdf,.txt,.md,.markdown,.docx,.html,.htm"; + +export function DocumentUploadDropzone({ + onFiles, + disabled, + className, +}: { + onFiles: (files: File[]) => void | Promise; + disabled?: boolean; + className?: string; +}) { + const inputRef = useRef(null); + const [isDragging, setIsDragging] = useState(false); + + const handleFiles = (files: FileList | null) => { + if (!files || files.length === 0 || disabled) return; + void onFiles(Array.from(files)); + }; + + return ( +
{ + e.preventDefault(); + if (!disabled) setIsDragging(true); + }} + onDragLeave={() => setIsDragging(false)} + onDrop={(e) => { + e.preventDefault(); + setIsDragging(false); + handleFiles(e.dataTransfer.files); + }} + > + +
+ Drop files here or + +
+
+ PDF, TXT, MD, DOCX, HTML +
+ { + handleFiles(e.target.files); + e.target.value = ""; + }} + /> +
+ ); +} diff --git a/studio/frontend/src/features/rag/components/ingestion-progress.tsx b/studio/frontend/src/features/rag/components/ingestion-progress.tsx new file mode 100644 index 0000000000..a5ab57937e --- /dev/null +++ b/studio/frontend/src/features/rag/components/ingestion-progress.tsx @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Progress } from "@/components/ui/progress"; +import { cn } from "@/lib/utils"; +import { useIngestionEvents } from "../hooks/use-ingestion-events"; + +const STAGE_LABELS: Record = { + queued: "Queued", + parse: "Parsing document", + load_model: "Loading embedder", + chunk: "Chunking text", + embed: "Embedding chunks", + done: "Indexing complete", +}; + +export function IngestionProgress({ + jobId, + className, +}: { + jobId: string; + className?: string; +}) { + const event = useIngestionEvents(jobId); + if (!event) { + return ( +
+ Starting… +
+ ); + } + + if (event.type === "error") { + return ( +
+ {event.error} +
+ ); + } + + if (event.type === "complete") { + return ( +
+ Indexed {event.num_chunks} chunks +
+ ); + } + + const stage = + "stage" in event && event.stage ? (event.stage as string) : "queued"; + const progress = + "progress" in event && typeof event.progress === "number" + ? event.progress + : 0; + const label = STAGE_LABELS[stage] ?? stage; + + return ( +
+
+ {label} + {Math.round(progress * 100)}% +
+ +
+ ); +} diff --git a/studio/frontend/src/features/rag/components/kb-create-dialog.tsx b/studio/frontend/src/features/rag/components/kb-create-dialog.tsx new file mode 100644 index 0000000000..b231c86549 --- /dev/null +++ b/studio/frontend/src/features/rag/components/kb-create-dialog.tsx @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { useState } from "react"; +import type { KnowledgeBase } from "../api/rag-api"; +import { useKnowledgeBases } from "../hooks/use-knowledge-bases"; + +export function KBCreateDialog({ + open, + onOpenChange, + onCreated, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + onCreated?: (kb: KnowledgeBase) => void; +}) { + const { createKB } = useKnowledgeBases(); + const [name, setName] = useState(""); + const [description, setDescription] = useState(""); + const [embeddingModel, setEmbeddingModel] = useState(""); + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(null); + + const reset = () => { + setName(""); + setDescription(""); + setEmbeddingModel(""); + setError(null); + setSubmitting(false); + }; + + const handleSubmit = async (e: React.FormEvent) => { + e.preventDefault(); + if (!name.trim() || submitting) return; + setSubmitting(true); + setError(null); + try { + const kb = await createKB({ + name: name.trim(), + description: description.trim() || undefined, + embedding_model: embeddingModel.trim() || undefined, + }); + onCreated?.(kb); + reset(); + onOpenChange(false); + } catch (err) { + setError(err instanceof Error ? err.message : String(err)); + setSubmitting(false); + } + }; + + return ( + { + if (!o) reset(); + onOpenChange(o); + }} + > + +
+ + Create knowledge base + + A knowledge base groups documents you can reuse across multiple + chat threads. + + +
+
+ + setName(e.target.value)} + placeholder="e.g. Internal docs" + autoFocus + required + /> +
+
+ + setDescription(e.target.value)} + placeholder="What's in this KB?" + /> +
+
+ + setEmbeddingModel(e.target.value)} + placeholder="Defaults to BAAI/bge-small-en-v1.5" + /> +
+ {error ? ( +
{error}
+ ) : null} +
+ + + + +
+ + + ); +} diff --git a/studio/frontend/src/features/rag/components/kb-detail-panel.tsx b/studio/frontend/src/features/rag/components/kb-detail-panel.tsx new file mode 100644 index 0000000000..033803eacc --- /dev/null +++ b/studio/frontend/src/features/rag/components/kb-detail-panel.tsx @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { ScrollArea } from "@/components/ui/scroll-area"; +import { Separator } from "@/components/ui/separator"; +import { useState } from "react"; +import type { KnowledgeBase } from "../api/rag-api"; +import { useKBDocuments } from "../hooks/use-kb-documents"; +import { DocumentRow } from "./document-row"; +import { DocumentUploadDropzone } from "./document-upload-dropzone"; +import { IngestionProgress } from "./ingestion-progress"; + +export function KBDetailPanel({ kb }: { kb: KnowledgeBase }) { + const { documents, loading, error, upload, remove } = useKBDocuments(kb.id); + const [activeJobsByDoc, setActiveJobsByDoc] = useState>( + {}, + ); + + const handleFiles = async (files: File[]) => { + for (const file of files) { + try { + const { documentId, jobId } = await upload(file); + setActiveJobsByDoc((prev) => ({ ...prev, [documentId]: jobId })); + } catch (err) { + console.error("upload failed", err); + } + } + }; + + return ( +
+
+

{kb.name}

+ {kb.description ? ( +

{kb.description}

+ ) : null} +

+ Embedding model: {kb.embedding_model} +

+
+ + + + + +
+

Documents

+ + {loading ? "Loading…" : `${documents.length} total`} + +
+ + {error ? ( +
{error}
+ ) : null} + + +
+ {documents.length === 0 && !loading ? ( +
+ No documents yet. Drop some files above to get started. +
+ ) : null} + {documents.map((doc) => { + const jobId = activeJobsByDoc[doc.id]; + const showProgress = + jobId && (doc.status === "pending" || doc.status === "running"); + return ( + { + void remove(doc.id); + }} + rightSlot={ + showProgress ? ( + + ) : null + } + /> + ); + })} +
+ +
+ ); +} diff --git a/studio/frontend/src/features/rag/components/kb-list.tsx b/studio/frontend/src/features/rag/components/kb-list.tsx new file mode 100644 index 0000000000..9b1649ce1e --- /dev/null +++ b/studio/frontend/src/features/rag/components/kb-list.tsx @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { ScrollArea } from "@/components/ui/scroll-area"; +import { cn } from "@/lib/utils"; +import { Add01Icon, Delete02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useState } from "react"; +import type { KnowledgeBase } from "../api/rag-api"; +import { useKnowledgeBases } from "../hooks/use-knowledge-bases"; +import { KBCreateDialog } from "./kb-create-dialog"; + +export function KBList({ + selectedId, + onSelect, +}: { + selectedId: string | null; + onSelect: (kb: KnowledgeBase | null) => void; +}) { + const { knowledgeBases, loading, error, deleteKB } = useKnowledgeBases(); + const [createOpen, setCreateOpen] = useState(false); + + return ( +
+
+

Knowledge bases

+ +
+ + {error ?
{error}
: null} + + +
+ {knowledgeBases.length === 0 && !loading ? ( +
+ No knowledge bases yet. +
+ ) : null} + {knowledgeBases.map((kb) => { + const isSelected = kb.id === selectedId; + return ( +
onSelect(kb)} + > +
+ {kb.name} + {kb.description ? ( + + {kb.description} + + ) : null} +
+ +
+ ); + })} +
+ + + onSelect(kb)} + /> +
+ ); +} diff --git a/studio/frontend/src/features/rag/components/thread-index-list.tsx b/studio/frontend/src/features/rag/components/thread-index-list.tsx new file mode 100644 index 0000000000..f4e986a8d1 --- /dev/null +++ b/studio/frontend/src/features/rag/components/thread-index-list.tsx @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { ScrollArea } from "@/components/ui/scroll-area"; +import { Delete02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useEffect } from "react"; +import { useRagStore } from "../stores/rag-store"; + +export function ThreadIndexList() { + const threadIndexes = useRagStore((s) => s.threadIndexes); + const loading = useRagStore((s) => s.threadIndexesLoading); + const loadThreadIndexes = useRagStore((s) => s.loadThreadIndexes); + const clearThreadIndex = useRagStore((s) => s.clearThreadIndex); + + useEffect(() => { + void loadThreadIndexes(); + }, [loadThreadIndexes]); + + return ( +
+
+

Thread documents

+ + {loading + ? "Loading…" + : `${threadIndexes.length} thread${threadIndexes.length === 1 ? "" : "s"}`} + +
+

+ Documents attached directly to a chat thread. Deleting a thread (from + the sidebar) also wipes its index. +

+ +
+ {threadIndexes.length === 0 && !loading ? ( +
+ No threads have attached documents. +
+ ) : null} + {threadIndexes.map((t) => ( +
+
+ + {t.title ?? Unsaved thread} + + + {t.num_documents} document + {t.num_documents === 1 ? "" : "s"} · {t.num_chunks} chunks + +
+ +
+ ))} +
+ +
+ ); +} diff --git a/studio/frontend/src/features/rag/hooks/use-ingestion-events.ts b/studio/frontend/src/features/rag/hooks/use-ingestion-events.ts new file mode 100644 index 0000000000..7727b253e9 --- /dev/null +++ b/studio/frontend/src/features/rag/hooks/use-ingestion-events.ts @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useEffect } from "react"; +import { useRagStore } from "../stores/rag-store"; + +/** + * Subscribe to an ingestion job's SSE events. Returns the latest event + * for that job from the global jobs map. Pass `null` to skip. + */ +export function useIngestionEvents(jobId: string | null) { + const event = useRagStore((s) => + jobId ? (s.jobs[jobId] ?? null) : null, + ); + const subscribeJob = useRagStore((s) => s.subscribeJob); + + useEffect(() => { + if (jobId) subscribeJob(jobId); + }, [jobId, subscribeJob]); + + return event; +} diff --git a/studio/frontend/src/features/rag/hooks/use-kb-documents.ts b/studio/frontend/src/features/rag/hooks/use-kb-documents.ts new file mode 100644 index 0000000000..01704a4555 --- /dev/null +++ b/studio/frontend/src/features/rag/hooks/use-kb-documents.ts @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useEffect } from "react"; +import { kbScopeKey, threadScopeKey, useRagStore } from "../stores/rag-store"; + +export function useKBDocuments(kbId: string | null) { + const scopeKey = kbId ? kbScopeKey(kbId) : ""; + const documents = useRagStore((s) => + scopeKey ? (s.documentsByScope[scopeKey] ?? []) : [], + ); + const loading = useRagStore((s) => (scopeKey ? !!s.docsLoading[scopeKey] : false)); + const error = useRagStore((s) => + scopeKey ? (s.docsError[scopeKey] ?? null) : null, + ); + const loadKBDocuments = useRagStore((s) => s.loadKBDocuments); + const uploadDocument = useRagStore((s) => s.uploadDocument); + const deleteDocument = useRagStore((s) => s.deleteDocument); + + useEffect(() => { + if (kbId) void loadKBDocuments(kbId); + }, [kbId, loadKBDocuments]); + + return { + documents, + loading, + error, + refresh: () => (kbId ? loadKBDocuments(kbId) : Promise.resolve()), + upload: (file: File) => + kbId + ? uploadDocument({ kind: "kb", kbId }, file) + : Promise.reject(new Error("no KB selected")), + remove: (documentId: string) => + scopeKey ? deleteDocument(documentId, scopeKey) : Promise.resolve(), + }; +} + +export function useThreadDocuments(threadId: string | null) { + const scopeKey = threadId ? threadScopeKey(threadId) : ""; + const documents = useRagStore((s) => + scopeKey ? (s.documentsByScope[scopeKey] ?? []) : [], + ); + const loading = useRagStore((s) => (scopeKey ? !!s.docsLoading[scopeKey] : false)); + const error = useRagStore((s) => + scopeKey ? (s.docsError[scopeKey] ?? null) : null, + ); + const loadThreadDocuments = useRagStore((s) => s.loadThreadDocuments); + const uploadDocument = useRagStore((s) => s.uploadDocument); + const deleteDocument = useRagStore((s) => s.deleteDocument); + + useEffect(() => { + if (threadId) void loadThreadDocuments(threadId); + }, [threadId, loadThreadDocuments]); + + return { + documents, + loading, + error, + refresh: () => + threadId ? loadThreadDocuments(threadId) : Promise.resolve(), + upload: (file: File) => + threadId + ? uploadDocument({ kind: "thread", threadId }, file) + : Promise.reject(new Error("no thread selected")), + remove: (documentId: string) => + scopeKey ? deleteDocument(documentId, scopeKey) : Promise.resolve(), + }; +} diff --git a/studio/frontend/src/features/rag/hooks/use-knowledge-bases.ts b/studio/frontend/src/features/rag/hooks/use-knowledge-bases.ts new file mode 100644 index 0000000000..98b4d9eee8 --- /dev/null +++ b/studio/frontend/src/features/rag/hooks/use-knowledge-bases.ts @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useEffect } from "react"; +import { useRagStore } from "../stores/rag-store"; + +export function useKnowledgeBases() { + const knowledgeBases = useRagStore((s) => s.knowledgeBases); + const loading = useRagStore((s) => s.kbsLoading); + const error = useRagStore((s) => s.kbsError); + const load = useRagStore((s) => s.loadKnowledgeBases); + const createKB = useRagStore((s) => s.createKB); + const deleteKB = useRagStore((s) => s.deleteKB); + + useEffect(() => { + if (knowledgeBases.length === 0 && !loading) { + void load(); + } + // Only run on mount — store-level cache prevents refetch loops. + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []); + + return { knowledgeBases, loading, error, refresh: load, createKB, deleteKB }; +} diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts new file mode 100644 index 0000000000..53c6e0a34d --- /dev/null +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -0,0 +1,245 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { create } from "zustand"; +import { + clearThreadDocuments as apiClearThreadDocuments, + createKnowledgeBase, + deleteDocument as apiDeleteDocument, + deleteKnowledgeBase as apiDeleteKB, + type JobEvent, + type KnowledgeBase, + listKBDocuments, + listKnowledgeBases, + listThreadDocuments, + listThreadIndexes, + type RagDocument, + subscribeToJobEvents, + type ThreadIndexSummary, + uploadKBDocument, + uploadThreadDocument, +} from "../api/rag-api"; + +interface RagStoreState { + knowledgeBases: KnowledgeBase[]; + kbsLoading: boolean; + kbsError: string | null; + + documentsByScope: Record; + docsLoading: Record; + docsError: Record; + + jobs: Record; + jobUnsubscribers: Record void>; + + threadIndexes: ThreadIndexSummary[]; + threadIndexesLoading: boolean; + + loadKnowledgeBases: () => Promise; + createKB: (req: { name: string; description?: string; embedding_model?: string }) => Promise; + deleteKB: (kbId: string) => Promise; + + loadKBDocuments: (kbId: string) => Promise; + loadThreadDocuments: (threadId: string) => Promise; + uploadDocument: ( + scope: { kind: "kb"; kbId: string } | { kind: "thread"; threadId: string }, + file: File, + ) => Promise<{ documentId: string; jobId: string }>; + deleteDocument: (documentId: string, scopeKey: string) => Promise; + + loadThreadIndexes: () => Promise; + clearThreadIndex: (threadId: string) => Promise; + + subscribeJob: (jobId: string, onComplete?: () => void) => void; +} + +function kbScopeKey(kbId: string): string { + return `kb:${kbId}`; +} +function threadScopeKey(threadId: string): string { + return `thread:${threadId}`; +} + +export const useRagStore = create((set, get) => ({ + knowledgeBases: [], + kbsLoading: false, + kbsError: null, + + documentsByScope: {}, + docsLoading: {}, + docsError: {}, + + jobs: {}, + jobUnsubscribers: {}, + + threadIndexes: [], + threadIndexesLoading: false, + + async loadKnowledgeBases() { + set({ kbsLoading: true, kbsError: null }); + try { + const kbs = await listKnowledgeBases(); + set({ knowledgeBases: kbs, kbsLoading: false }); + } catch (err) { + set({ + kbsLoading: false, + kbsError: err instanceof Error ? err.message : String(err), + }); + } + }, + + async createKB(req) { + const kb = await createKnowledgeBase(req); + set((state) => ({ knowledgeBases: [kb, ...state.knowledgeBases] })); + return kb; + }, + + async deleteKB(kbId) { + await apiDeleteKB(kbId); + set((state) => { + const scopeKey = kbScopeKey(kbId); + const { [scopeKey]: _docs, ...restDocs } = state.documentsByScope; + return { + knowledgeBases: state.knowledgeBases.filter((k) => k.id !== kbId), + documentsByScope: restDocs, + }; + }); + }, + + async loadKBDocuments(kbId) { + const key = kbScopeKey(kbId); + set((state) => ({ + docsLoading: { ...state.docsLoading, [key]: true }, + docsError: { ...state.docsError, [key]: null }, + })); + try { + const docs = await listKBDocuments(kbId); + set((state) => ({ + documentsByScope: { ...state.documentsByScope, [key]: docs }, + docsLoading: { ...state.docsLoading, [key]: false }, + })); + } catch (err) { + set((state) => ({ + docsLoading: { ...state.docsLoading, [key]: false }, + docsError: { + ...state.docsError, + [key]: err instanceof Error ? err.message : String(err), + }, + })); + } + }, + + async loadThreadDocuments(threadId) { + const key = threadScopeKey(threadId); + set((state) => ({ + docsLoading: { ...state.docsLoading, [key]: true }, + docsError: { ...state.docsError, [key]: null }, + })); + try { + const docs = await listThreadDocuments(threadId); + set((state) => ({ + documentsByScope: { ...state.documentsByScope, [key]: docs }, + docsLoading: { ...state.docsLoading, [key]: false }, + })); + } catch (err) { + set((state) => ({ + docsLoading: { ...state.docsLoading, [key]: false }, + docsError: { + ...state.docsError, + [key]: err instanceof Error ? err.message : String(err), + }, + })); + } + }, + + async uploadDocument(scope, file) { + const result = + scope.kind === "kb" + ? await uploadKBDocument(scope.kbId, file) + : await uploadThreadDocument(scope.threadId, file); + const scopeKey = + scope.kind === "kb" + ? kbScopeKey(scope.kbId) + : threadScopeKey(scope.threadId); + // Refresh the doc list so the new pending row appears. + if (scope.kind === "kb") { + void get().loadKBDocuments(scope.kbId); + } else { + void get().loadThreadDocuments(scope.threadId); + } + // Subscribe to the job; refresh docs on completion to pick up + // the final num_chunks / status. + get().subscribeJob(result.job_id, () => { + if (scope.kind === "kb") { + void get().loadKBDocuments(scope.kbId); + } else { + void get().loadThreadDocuments(scope.threadId); + } + }); + return { documentId: result.document_id, jobId: result.job_id, scopeKey } as { + documentId: string; + jobId: string; + }; + }, + + async deleteDocument(documentId, scopeKey) { + await apiDeleteDocument(documentId); + set((state) => { + const current = state.documentsByScope[scopeKey] ?? []; + return { + documentsByScope: { + ...state.documentsByScope, + [scopeKey]: current.filter((d) => d.id !== documentId), + }, + }; + }); + }, + + async loadThreadIndexes() { + set({ threadIndexesLoading: true }); + try { + const threads = await listThreadIndexes(); + set({ threadIndexes: threads, threadIndexesLoading: false }); + } catch { + set({ threadIndexesLoading: false }); + } + }, + + async clearThreadIndex(threadId) { + await apiClearThreadDocuments(threadId); + const scopeKey = threadScopeKey(threadId); + set((state) => { + const { [scopeKey]: _docs, ...restDocs } = state.documentsByScope; + return { + documentsByScope: restDocs, + threadIndexes: state.threadIndexes.filter( + (t) => t.thread_id !== threadId, + ), + }; + }); + }, + + subscribeJob(jobId, onComplete) { + const existing = get().jobUnsubscribers[jobId]; + if (existing) return; + const unsubscribe = subscribeToJobEvents(jobId, { + onEvent: (event) => { + set((state) => ({ jobs: { ...state.jobs, [jobId]: event } })); + if (event.type === "complete" || event.type === "error") { + onComplete?.(); + } + }, + onClose: () => { + set((state) => { + const { [jobId]: _gone, ...rest } = state.jobUnsubscribers; + return { jobUnsubscribers: rest }; + }); + }, + }); + set((state) => ({ + jobUnsubscribers: { ...state.jobUnsubscribers, [jobId]: unsubscribe }, + })); + }, +})); + +export { kbScopeKey, threadScopeKey }; diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 3b95cb2355..b8ae6bf8d8 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -11,6 +11,7 @@ import { cn } from "@/lib/utils"; import { Cancel01Icon, CloudIcon, + Database01Icon, Globe02Icon, HelpCircleIcon, Message01Icon, @@ -31,6 +32,7 @@ import { AppearanceTab } from "./tabs/appearance-tab"; import { ChatTab } from "./tabs/chat-tab"; import { ConnectionsTab } from "./tabs/connections-tab"; import { GeneralTab } from "./tabs/general-tab"; +import { KnowledgeBasesTab } from "./tabs/knowledge-bases-tab"; import { ProfileTab } from "./tabs/profile-tab"; interface TabDef { @@ -45,6 +47,12 @@ const TABS: TabDef[] = [ { id: "profile", label: "Profile", icon: UserIcon }, { id: "appearance", label: "Appearance", icon: PaintBrush02Icon }, { id: "chat", label: "Chat", icon: Message01Icon }, + { + id: "knowledge-bases", + label: "Knowledge Bases", + icon: Database01Icon, + badge: "New", + }, { id: "connections", label: "Connections", icon: CloudIcon, badge: "New" }, { id: "api-keys", label: "API", icon: Globe02Icon, badge: "New" }, { id: "about", label: "Help", icon: HelpCircleIcon }, @@ -60,6 +68,8 @@ function renderTab(tab: SettingsTab) { return ; case "chat": return ; + case "knowledge-bases": + return ; case "connections": return ; case "api-keys": @@ -81,6 +91,7 @@ export function SettingsDialog() { profile: null, appearance: null, chat: null, + "knowledge-bases": null, connections: null, "api-keys": null, about: null, diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts index 7ac5422ff1..8474f5f389 100644 --- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts +++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts @@ -8,6 +8,7 @@ export type SettingsTab = | "profile" | "appearance" | "chat" + | "knowledge-bases" | "connections" | "api-keys" | "about"; @@ -40,6 +41,7 @@ function loadInitialTab(): SettingsTab { "profile", "appearance", "chat", + "knowledge-bases", "connections", "api-keys", "about", diff --git a/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx b/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx new file mode 100644 index 0000000000..94cb90959e --- /dev/null +++ b/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Separator } from "@/components/ui/separator"; +import type { KnowledgeBase } from "@/features/rag/api/rag-api"; +import { KBDetailPanel } from "@/features/rag/components/kb-detail-panel"; +import { KBList } from "@/features/rag/components/kb-list"; +import { ThreadIndexList } from "@/features/rag/components/thread-index-list"; +import { useState } from "react"; + +export function KnowledgeBasesTab() { + const [selected, setSelected] = useState(null); + + return ( +
+
+

Knowledge bases

+

+ Create reusable document collections and pick one per chat thread to + ground answers in your own files. +

+
+ +
+
+ +
+ +
+ {selected ? ( + + ) : ( +
+ Select a knowledge base, or create a new one to get started. +
+ )} +
+
+ + +
+ ); +} diff --git a/tests/python/test_rag_bm25.py b/tests/python/test_rag_bm25.py new file mode 100644 index 0000000000..43c7504636 --- /dev/null +++ b/tests/python/test_rag_bm25.py @@ -0,0 +1,69 @@ +"""BM25 index lifecycle tests (skipped if bm25s is unavailable).""" + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + +pytest.importorskip("bm25s") + + +@pytest.fixture +def isolated_bm25_root(tmp_path, monkeypatch): + from utils.paths import storage_roots + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + # Reset module-level cache between tests. + from core.rag import bm25 + + bm25._cache.clear() + return tmp_path + + +def test_bm25_index_search_roundtrip(isolated_bm25_root): + from core.rag import bm25 + + scope = "kb_test" + chunks = [ + {"id": "c1", "text": "the quick brown fox jumps over the lazy dog"}, + {"id": "c2", "text": "machine learning models predict outputs from inputs"}, + {"id": "c3", "text": "fox terriers are small dogs"}, + ] + bm25.rebuild_index(scope, chunks) + results = bm25.search(scope, "fox", k = 3) + ids = [cid for cid, _ in results] + assert "c1" in ids + assert "c3" in ids + + +def test_bm25_empty_returns_empty(isolated_bm25_root): + from core.rag import bm25 + + assert bm25.search("kb_nonexistent", "anything", k = 5) == [] + + +def test_bm25_delete_scope(isolated_bm25_root): + from core.rag import bm25 + + scope = "kb_del" + chunks = [{"id": "a", "text": "alpha beta gamma"}] + bm25.rebuild_index(scope, chunks) + assert bm25.search(scope, "alpha", k = 1) + bm25.delete_scope(scope) + assert bm25.search(scope, "alpha", k = 1) == [] + + +def test_bm25_rebuild_replaces_old_corpus(isolated_bm25_root): + from core.rag import bm25 + + scope = "kb_replace" + bm25.rebuild_index(scope, [{"id": "old", "text": "alpha beta"}]) + bm25.rebuild_index(scope, [{"id": "new", "text": "gamma delta"}]) + results = bm25.search(scope, "alpha", k = 5) + ids = [cid for cid, _ in results] + assert "old" not in ids diff --git a/tests/python/test_rag_chunking.py b/tests/python/test_rag_chunking.py new file mode 100644 index 0000000000..f5a7b16380 --- /dev/null +++ b/tests/python/test_rag_chunking.py @@ -0,0 +1,84 @@ +"""Unit tests for RAG chunking — pure-python, no external deps.""" + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + +from core.rag.chunking import Chunk, chunk_pages +from core.rag.parsers import ParsedPage + + +def _wc_counter(text: str) -> int: + return max(1, len(text.split())) + + +def test_chunk_pages_splits_long_text(): + text = "Lorem ipsum dolor sit amet. " * 200 + chunks = chunk_pages( + [ParsedPage(text = text)], + max_tokens = 50, + overlap_tokens = 5, + token_counter = _wc_counter, + ) + assert len(chunks) > 1 + for chunk in chunks: + assert _wc_counter(chunk.text) <= 55 # max + small slack from atomic split granularity + + +def test_chunk_pages_short_text_is_one_chunk(): + text = "Just a short sentence." + chunks = chunk_pages( + [ParsedPage(text = text)], + max_tokens = 50, + overlap_tokens = 5, + token_counter = _wc_counter, + ) + assert len(chunks) == 1 + assert chunks[0].text == text + + +def test_chunk_pages_preserves_page_numbers(): + chunks = chunk_pages( + [ + ParsedPage(text = "Page one content here.", page_number = 1), + ParsedPage(text = "Page two content here.", page_number = 2), + ], + max_tokens = 50, + overlap_tokens = 0, + token_counter = _wc_counter, + ) + page_numbers = {c.page_number for c in chunks} + assert page_numbers == {1, 2} + + +def test_chunk_pages_no_empty_chunks(): + text = "\n\n\n\n\nReal content\n\n\n\n\n" + chunks = chunk_pages( + [ParsedPage(text = text)], + max_tokens = 50, + overlap_tokens = 0, + token_counter = _wc_counter, + ) + for chunk in chunks: + assert chunk.text.strip() + + +def test_chunk_pages_overlap_produces_repeated_tokens(): + # Build a list of unique numbered sentences so we can detect overlap. + sentences = [f"sentence-{i}" for i in range(40)] + text = " ".join(sentences) + chunks = chunk_pages( + [ParsedPage(text = text)], + max_tokens = 10, + overlap_tokens = 4, + token_counter = _wc_counter, + ) + if len(chunks) >= 2: + first_tail_words = set(chunks[0].text.split()[-4:]) + second_head_words = set(chunks[1].text.split()[:4]) + # At least one word should appear in both + assert first_tail_words & second_head_words diff --git a/tests/python/test_rag_parsers.py b/tests/python/test_rag_parsers.py new file mode 100644 index 0000000000..bd28ee1234 --- /dev/null +++ b/tests/python/test_rag_parsers.py @@ -0,0 +1,89 @@ +"""Document parser tests — each format skipped if its lib is unavailable.""" + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + + +def test_text_parser_utf8(tmp_path): + from core.rag.parsers import parse + + file = tmp_path / "sample.txt" + file.write_text("hello world\n\nsecond paragraph", encoding = "utf-8") + pages = parse(file) + assert len(pages) == 1 + assert "hello world" in pages[0].text + assert "second paragraph" in pages[0].text + + +def test_markdown_parser_treated_as_text(tmp_path): + from core.rag.parsers import parse + + file = tmp_path / "sample.md" + file.write_text("# Title\n\nBody text with **emphasis**.", encoding = "utf-8") + pages = parse(file) + assert pages and "Title" in pages[0].text + + +def test_unsupported_format_raises(tmp_path): + from core.rag.parsers import UnsupportedFormatError, parse + + file = tmp_path / "weird.xyz" + file.write_text("nope") + with pytest.raises(UnsupportedFormatError): + parse(file) + + +def test_html_parser_strips_scripts(tmp_path): + pytest.importorskip("bs4") + pytest.importorskip("lxml") + from core.rag.parsers import parse + + file = tmp_path / "sample.html" + file.write_text( + "

visible text

", + encoding = "utf-8", + ) + pages = parse(file) + assert pages + assert "visible text" in pages[0].text + assert "alert" not in pages[0].text + + +def test_pdf_parser_extracts_pages(tmp_path): + pypdf = pytest.importorskip("pypdf") + from pypdf import PdfWriter + + file = tmp_path / "tiny.pdf" + writer = PdfWriter() + writer.add_blank_page(width = 72, height = 72) + with open(file, "wb") as f: + writer.write(f) + from core.rag.parsers import parse + + # blank page yields no extractable text — should return [] without error + pages = parse(file) + assert isinstance(pages, list) + + +def test_docx_parser_extracts_paragraphs(tmp_path): + docx = pytest.importorskip("docx") + from docx import Document + + file = tmp_path / "sample.docx" + doc = Document() + doc.add_paragraph("First paragraph here.") + doc.add_paragraph("Second paragraph here.") + doc.save(str(file)) + from core.rag.parsers import parse + + pages = parse(file) + assert pages + assert "First paragraph" in pages[0].text + assert "Second paragraph" in pages[0].text diff --git a/tests/python/test_rag_reranker.py b/tests/python/test_rag_reranker.py new file mode 100644 index 0000000000..d63b94c089 --- /dev/null +++ b/tests/python/test_rag_reranker.py @@ -0,0 +1,57 @@ +"""Reranker tests — skipped if sentence_transformers is unavailable. + +These tests load a real CrossEncoder, so they're slow and gated under +the ``server`` marker so a default ``pytest`` run skips them. Force +with ``pytest -m server``. +""" + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + +pytest.importorskip("sentence_transformers") + + +def test_rerank_empty_returns_empty(): + from core.rag.reranker import rerank + + assert rerank("anything", []) == [] + + +@pytest.mark.server +def test_rerank_reorders_by_relevance(monkeypatch): + """Hide the relevant chunk at the back of the input and check it bubbles up.""" + monkeypatch.setenv("UNSLOTH_RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2") + from core.rag.reranker import rerank, unload + from core.rag.retrieval import Hit + + pairs = [ + (Hit("noise1", 0.0), "Cats are small carnivorous mammals."), + (Hit("noise2", 0.0), "The Eiffel Tower is in Paris, France."), + (Hit("noise3", 0.0), "Python is a programming language."), + (Hit("answer", 0.0), "The speed of light in vacuum is approximately 299792458 meters per second."), + ] + try: + ranked = rerank("How fast does light travel?", pairs, top_k = 2) + assert ranked + assert ranked[0].chunk_id == "answer" + finally: + unload() + + +@pytest.mark.server +def test_unload_clears_singleton(monkeypatch): + monkeypatch.setenv("UNSLOTH_RAG_RERANKER_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2") + from core.rag import reranker + from core.rag.retrieval import Hit + + reranker.rerank("q", [(Hit("a", 0.0), "some text")]) + assert reranker._model is not None + reranker.unload() + assert reranker._model is None diff --git a/tests/python/test_rag_retrieval.py b/tests/python/test_rag_retrieval.py new file mode 100644 index 0000000000..f589e5e731 --- /dev/null +++ b/tests/python/test_rag_retrieval.py @@ -0,0 +1,46 @@ +"""Unit tests for RAG RRF fusion — no external deps.""" + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + +from core.rag.retrieval import Hit, _rrf_fuse + + +def test_rrf_fuses_two_rankings(): + bm25 = [Hit("a", 10.0), Hit("b", 8.0), Hit("c", 5.0)] + dense = [Hit("c", 0.9), Hit("b", 0.8), Hit("d", 0.5)] + fused = _rrf_fuse([bm25, dense], rrf_k = 60, top_k = 3) + ids = [h.chunk_id for h in fused] + # b appears at rank 2 in both -> highest fused score + assert ids[0] == "b" + assert set(ids) == {"a", "b", "c"} or set(ids) == {"b", "c", "a"} + + +def test_rrf_top_k_limits_output(): + rankings = [ + [Hit(f"r1_{i}", 0.0) for i in range(20)], + [Hit(f"r2_{i}", 0.0) for i in range(20)], + ] + fused = _rrf_fuse(rankings, rrf_k = 60, top_k = 5) + assert len(fused) == 5 + + +def test_rrf_unique_ranking(): + # Single ranking — fused order matches input order. + ranking = [Hit("x", 0.0), Hit("y", 0.0), Hit("z", 0.0)] + fused = _rrf_fuse([ranking], rrf_k = 60, top_k = 3) + assert [h.chunk_id for h in fused] == ["x", "y", "z"] + + +def test_rrf_preserves_payload_from_first_ranking(): + a = Hit("a", 1.0, document_id = "doc1", chunk_index = 5) + b = Hit("a", 2.0, document_id = "doc2", chunk_index = 7) + fused = _rrf_fuse([[a], [b]], rrf_k = 60, top_k = 1) + # First sighting wins for payload (deterministic) + assert fused[0].document_id == "doc1" + assert fused[0].chunk_index == 5 diff --git a/tests/python/test_rag_vector_store.py b/tests/python/test_rag_vector_store.py new file mode 100644 index 0000000000..26bbdefd26 --- /dev/null +++ b/tests/python/test_rag_vector_store.py @@ -0,0 +1,84 @@ +"""Qdrant local-mode vector store tests (skipped if qdrant-client is unavailable).""" + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + +pytest.importorskip("qdrant_client") + + +@pytest.fixture +def isolated_qdrant(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + from core.rag import vector_store + + # Reset client cache so the fixture's tmp path is used. + vector_store._client = None + yield tmp_path + vector_store._client = None + + +def test_ensure_and_upsert_and_search(isolated_qdrant): + from core.rag import vector_store + + scope = "kb_test" + vector_store.ensure_collection(scope, dim = 4) + points = [ + { + "id": "p1", + "vector": [1.0, 0.0, 0.0, 0.0], + "payload": {"document_id": "doc1", "chunk_index": 0, "text": "first"}, + }, + { + "id": "p2", + "vector": [0.0, 1.0, 0.0, 0.0], + "payload": {"document_id": "doc1", "chunk_index": 1, "text": "second"}, + }, + ] + vector_store.upsert_chunks(scope, points) + results = vector_store.search(scope, [1.0, 0.0, 0.0, 0.0], top_k = 2) + assert results + assert results[0]["chunk_id"] == "p1" + + +def test_delete_scope_removes_collection(isolated_qdrant): + from core.rag import vector_store + + scope = "kb_to_delete" + vector_store.ensure_collection(scope, dim = 3) + assert vector_store.collection_exists(scope) + vector_store.delete_scope(scope) + assert not vector_store.collection_exists(scope) + + +def test_delete_document_removes_only_its_points(isolated_qdrant): + from core.rag import vector_store + + scope = "kb_doc_del" + vector_store.ensure_collection(scope, dim = 3) + vector_store.upsert_chunks( + scope, + [ + { + "id": "a", + "vector": [1.0, 0.0, 0.0], + "payload": {"document_id": "keep", "chunk_index": 0}, + }, + { + "id": "b", + "vector": [0.0, 1.0, 0.0], + "payload": {"document_id": "drop", "chunk_index": 0}, + }, + ], + ) + vector_store.delete_document(scope, "drop") + results = vector_store.search(scope, [0.0, 1.0, 0.0], top_k = 5) + doc_ids = {r["payload"]["document_id"] for r in results} + assert "drop" not in doc_ids + assert "keep" in doc_ids From c4b5889e535027f4804a91dd263276930d768ac9 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 11:10:13 +0400 Subject: [PATCH 002/122] Studio: layout-aware RAG parsers + heading-aware chunking (Phase 3A) Replace bare-pypdf/python-docx/BeautifulSoup extraction with Markdown- preserving parsers so the chunker can split on real heading boundaries instead of running paragraphs together. Parsers - pdf.py: pymupdf + pymupdf4llm.to_markdown() per page; pypdf kept as fallback when pymupdf can't open the file. - docx.py: mammoth.convert_to_html() + markdownify, with an explicit style_map so Title/Heading 1..6 become h1..h6 in the output. - html.py: BeautifulSoup pre-scrub (drop script/style) then markdownify so , ,
    convert faithfully. - text.py: signature update only; TXT/MD pass through unchanged. - parsers/__init__.py: new ParsedImage + ParseResult dataclass; parse() signature is now parse(path, *, want_images=False) -> ParseResult. ParseResult is iterable over .pages for backward compat. Chunker - chunking.py: prepend Markdown heading separators ("\n# " .. "\n#### ") to the priority list so heading-aware splits happen for free once the parsers emit Markdown. Ingestion - ingestion.py: single call site updated to consume ParseResult.pages. Deps (no-torch-runtime.txt) + pymupdf>=1.24, pymupdf4llm>=0.0.17, mammoth>=1.7, markdownify>=0.13 - pypdf kept as a fallback path. Tests - test_rag_parsers.py asserts Markdown headings survive PDF/DOCX/HTML extraction; also exercises ParseResult iteration backward-compat. - test_rag_chunking.py: new case verifying chunks start at Markdown heading boundaries when the input is Markdown. Foundation for Phase 3B-late (heading-aware spans for late chunking) and Phase 3B-multimodal (want_images=True enables image extraction in the same parser layer). No schema or opt-in flags in this commit. --- studio/backend/core/rag/chunking.py | 17 ++- studio/backend/core/rag/ingestion.py | 6 +- studio/backend/core/rag/parsers/__init__.py | 52 +++++++- studio/backend/core/rag/parsers/docx.py | 120 +++++++++++++++--- studio/backend/core/rag/parsers/html.py | 94 ++++++++++++-- studio/backend/core/rag/parsers/pdf.py | 106 +++++++++++++++- studio/backend/core/rag/parsers/text.py | 12 +- .../backend/requirements/no-torch-runtime.txt | 9 ++ tests/python/test_rag_chunking.py | 29 +++++ tests/python/test_rag_parsers.py | 94 ++++++++++---- 10 files changed, 475 insertions(+), 64 deletions(-) diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index e2a5c06493..4178b47c8a 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -106,7 +106,22 @@ def chunk_pages( max_tokens: int, overlap_tokens: int, token_counter: TokenCounter | None = None, - separators: tuple[str, ...] = ("\n\n", "\n", ". ", " ", ""), + 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### ", + "\n#### ", + "\n\n", + "\n", + ". ", + " ", + "", + ), ) -> list[Chunk]: """Split parsed pages into overlapping chunks. diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index a6aa7c720c..226fd53129 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -60,7 +60,11 @@ def _subprocess_worker( from core.rag.parsers import parse out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) - pages = parse(Path(stored_path)) + # want_images stays False for the text-only ingestion path; the + # multimodal path (Phase 3B-multimodal) will flip this based on + # the KB's mode. + parsed = parse(Path(stored_path), want_images = False) + pages = parsed.pages if not pages: out_queue.put({"type": "error", "error": "no extractable text in document"}) return diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py index 105e2a89d8..8b2bcff1ae 100644 --- a/studio/backend/core/rag/parsers/__init__.py +++ b/studio/backend/core/rag/parsers/__init__.py @@ -3,21 +3,67 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field 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. + """ text: str page_number: int | None = None +@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_bytes: bytes + mime_type: str + page_number: int | None = None + nearest_caption: str = "" + + +@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) + + def __iter__(self): + return iter(self.pages) + + def __len__(self): + return len(self.pages) + + def __bool__(self): + return bool(self.pages) or bool(self.images) + + class UnsupportedFormatError(ValueError): pass -def parse(path: Path) -> list[ParsedPage]: +def parse(path: Path, *, want_images: bool = False) -> ParseResult: suffix = path.suffix.lower() if suffix == ".pdf": from .pdf import extract @@ -29,4 +75,4 @@ def parse(path: Path) -> list[ParsedPage]: from .html import extract else: raise UnsupportedFormatError(f"Unsupported file type: {suffix}") - return extract(path) + return extract(path, want_images = want_images) diff --git a/studio/backend/core/rag/parsers/docx.py b/studio/backend/core/rag/parsers/docx.py index fedc401773..bf16a9052c 100644 --- a/studio/backend/core/rag/parsers/docx.py +++ b/studio/backend/core/rag/parsers/docx.py @@ -1,27 +1,113 @@ # 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. +""" + from __future__ import annotations +import logging +import re from pathlib import Path -from . import ParsedPage +from . import ParsedImage, ParsedPage, ParseResult + +logger = logging.getLogger(__name__) -def extract(path: Path) -> list[ParsedPage]: - from docx import Document +# 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. +_STYLE_MAP = """ +p[style-name='Title'] => h1.title:fresh +p[style-name='Subtitle'] => h2.subtitle:fresh +p[style-name='Heading 1'] => h1:fresh +p[style-name='Heading 2'] => h2:fresh +p[style-name='Heading 3'] => h3:fresh +p[style-name='Heading 4'] => h4:fresh +p[style-name='Heading 5'] => h5:fresh +p[style-name='Heading 6'] => h6:fresh +""" - document = Document(str(path)) - parts: list[str] = [] - for paragraph in document.paragraphs: - if paragraph.text and paragraph.text.strip(): - parts.append(paragraph.text) - for table in document.tables: - for row in table.rows: - cells = [cell.text.strip() for cell in row.cells if cell.text.strip()] - if cells: - parts.append(" | ".join(cells)) - text = "\n\n".join(parts).strip() - if not text: - return [] - return [ParsedPage(text = text, page_number = None)] + +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() + + +def extract(path: Path, *, want_images: bool = False) -> ParseResult: + import mammoth + + 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). + def _convert(image): + with image.open() as image_bytes: + blob = image_bytes.read() + mime = (image.content_type or "application/octet-stream").lower() + images.append( + ParsedImage( + image_bytes = blob, + mime_type = mime, + page_number = None, + nearest_caption = "", + ) + ) + return {"src": ""} + + 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: + result = mammoth.convert_to_html( + fp, + convert_image = convert_image, + style_map = _STYLE_MAP, + ) + for message in result.messages: + logger.debug("mammoth %s: %s", getattr(message, "type", "msg"), message.message) + + 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. + caption_pool = markdown[:1500] + images = [ + ParsedImage( + image_bytes = img.image_bytes, + mime_type = img.mime_type, + page_number = img.page_number, + nearest_caption = caption_pool, + ) + for img in images + ] + if not markdown: + return ParseResult(pages = [], images = images) + return ParseResult( + pages = [ParsedPage(text = markdown, page_number = None)], + images = images, + ) diff --git a/studio/backend/core/rag/parsers/html.py b/studio/backend/core/rag/parsers/html.py index 121cb6faf3..721e280736 100644 --- a/studio/backend/core/rag/parsers/html.py +++ b/studio/backend/core/rag/parsers/html.py @@ -1,26 +1,98 @@ # 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. +""" + from __future__ import annotations +import logging +import re from pathlib import Path +from urllib.parse import unquote, urlparse -from . import ParsedPage +from . import ParsedImage, ParsedPage, ParseResult + +logger = logging.getLogger(__name__) + +_SKIP_TAGS = ("script", "style", "noscript", "template") -_SKIP_TAGS = {"script", "style", "noscript", "template"} +def _collect_local_images(soup, html_path: Path) -> list[ParsedImage]: + images: list[ParsedImage] = [] + base_dir = html_path.parent + for tag in soup.find_all("img"): + 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. + continue + if not local_path.is_file(): + continue + try: + blob = local_path.read_bytes() + except OSError: + continue + suffix = local_path.suffix.lower().lstrip(".") + mime = { + "png": "image/png", + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "gif": "image/gif", + "webp": "image/webp", + "svg": "image/svg+xml", + }.get(suffix, f"image/{suffix or 'octet-stream'}") + caption = tag.get("alt") or tag.get("title") or "" + images.append( + ParsedImage( + image_bytes = blob, + mime_type = mime, + page_number = None, + nearest_caption = caption, + ) + ) + return images -def extract(path: Path) -> list[ParsedPage]: +def extract(path: Path, *, want_images: bool = False) -> ParseResult: from bs4 import BeautifulSoup + from markdownify import markdownify raw = path.read_bytes() soup = BeautifulSoup(raw, "lxml") - for tag in soup(_SKIP_TAGS): - tag.decompose() - text = soup.get_text(separator = "\n").strip() - lines = [line.strip() for line in text.splitlines() if line.strip()] - cleaned = "\n".join(lines) - if not cleaned: - return [] - return [ParsedPage(text = cleaned, page_number = None)] + for tag_name in _SKIP_TAGS: + for tag in soup.find_all(tag_name): + tag.decompose() + + images: list[ParsedImage] = [] + if want_images: + images = _collect_local_images(soup, path) + + md = markdownify( + str(soup), + heading_style = "ATX", + strip = list(_SKIP_TAGS), + ) + md = re.sub(r"\n{3,}", "\n\n", md).strip() + if not md: + return ParseResult(pages = [], images = images) + return ParseResult( + pages = [ParsedPage(text = md, page_number = None)], + images = images, + ) diff --git a/studio/backend/core/rag/parsers/pdf.py b/studio/backend/core/rag/parsers/pdf.py index 2a3f3fbda1..994b1aaed9 100644 --- a/studio/backend/core/rag/parsers/pdf.py +++ b/studio/backend/core/rag/parsers/pdf.py @@ -1,14 +1,101 @@ # 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. +""" + from __future__ import annotations +import logging from pathlib import Path -from . import ParsedPage +from . import ParsedImage, ParsedPage, ParseResult + +logger = logging.getLogger(__name__) -def extract(path: Path) -> list[ParsedPage]: +def _extract_with_pymupdf(path: Path, want_images: bool) -> ParseResult: + import pymupdf + import pymupdf4llm + + doc = pymupdf.open(str(path)) + try: + pages: list[ParsedPage] = [] + for page_index in range(len(doc)): + try: + md = pymupdf4llm.to_markdown( + doc, + pages = [page_index], + write_images = False, + ignore_images = True, + show_progress = False, + ) + except Exception: + # pymupdf4llm can choke on individual pages (rare). Fall + # back to plain text extraction for just that page. + md = doc[page_index].get_text("text") or "" + md = md.strip() + if md: + pages.append(ParsedPage(text = md, page_number = page_index + 1)) + + images: list[ParsedImage] = [] + if want_images: + images = _extract_images_pymupdf(doc, pages) + return ParseResult(pages = pages, images = images) + finally: + doc.close() + + +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} + out: list[ParsedImage] = [] + for page_index in range(len(doc)): + page_number = page_index + 1 + try: + image_list = doc[page_index].get_images(full = True) + except Exception: + continue + for img_info in image_list: + xref = img_info[0] + try: + extracted = doc.extract_image(xref) + except Exception: + continue + image_bytes = extracted.get("image") + ext = (extracted.get("ext") or "png").lower() + mime = { + "png": "image/png", + "jpg": "image/jpeg", + "jpeg": "image/jpeg", + "gif": "image/gif", + "webp": "image/webp", + "bmp": "image/bmp", + "tiff": "image/tiff", + }.get(ext, f"image/{ext}") + if not image_bytes: + continue + caption = (captions_by_page.get(page_number, "") or "")[:1500] + out.append( + ParsedImage( + image_bytes = image_bytes, + mime_type = mime, + page_number = page_number, + nearest_caption = caption, + ) + ) + return out + + +def _extract_with_pypdf_fallback(path: Path) -> ParseResult: from pypdf import PdfReader reader = PdfReader(str(path)) @@ -21,4 +108,17 @@ def extract(path: Path) -> list[ParsedPage]: text = text.strip() if text: pages.append(ParsedPage(text = text, page_number = index + 1)) - return pages + return ParseResult(pages = pages, images = []) + + +def extract(path: Path, *, want_images: bool = False) -> ParseResult: + try: + return _extract_with_pymupdf(path, want_images) + except Exception as exc: + logger.warning( + "pymupdf failed for %s (%s: %s); falling back to pypdf", + path, + type(exc).__name__, + exc, + ) + return _extract_with_pypdf_fallback(path) diff --git a/studio/backend/core/rag/parsers/text.py b/studio/backend/core/rag/parsers/text.py index 402b5f32f9..439a07440b 100644 --- a/studio/backend/core/rag/parsers/text.py +++ b/studio/backend/core/rag/parsers/text.py @@ -5,10 +5,11 @@ from __future__ import annotations from pathlib import Path -from . import ParsedPage +from . import ParsedPage, ParseResult -def extract(path: Path) -> list[ParsedPage]: +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") @@ -23,5 +24,8 @@ def extract(path: Path) -> list[ParsedPage]: text = raw.decode(encoding, errors = "replace") text = text.strip() if not text: - return [] - return [ParsedPage(text = text, page_number = None)] + return ParseResult(pages = [], images = []) + return ParseResult( + pages = [ParsedPage(text = text, page_number = None)], + images = [], + ) diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 36d1e0f5c4..ce879ccf94 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -83,6 +83,15 @@ pillow # server. bm25s persists per-scope indices to disk. qdrant-client>=1.12 bm25s>=0.2 +# RAG parsers (Phase 3A): layout-aware Markdown extraction so the chunker +# can split on real headings instead of running paragraphs together. +# pymupdf4llm preserves headings + pipe-tables; mammoth handles DOCX +# Heading styles; markdownify converts HTML /
/
    faithfully. +pymupdf>=1.24 +pymupdf4llm>=0.0.17 +mammoth>=1.7 +markdownify>=0.13 +# pypdf is kept as a fallback for malformed PDFs that defeat pymupdf. pypdf>=4.0 python-docx>=1.1 beautifulsoup4>=4.12 diff --git a/tests/python/test_rag_chunking.py b/tests/python/test_rag_chunking.py index f5a7b16380..835b96878d 100644 --- a/tests/python/test_rag_chunking.py +++ b/tests/python/test_rag_chunking.py @@ -82,3 +82,32 @@ def test_chunk_pages_overlap_produces_repeated_tokens(): second_head_words = set(chunks[1].text.split()[:4]) # At least one word should appear in both assert first_tail_words & second_head_words + + +def test_chunk_pages_splits_on_markdown_headings(): + # Phase 3A: heading separators take priority over paragraph breaks + # so chunks start at section boundaries when the parser emits + # Markdown. + md = ( + "# First Section\n\n" + + "alpha " * 30 + + "\n\n## Subsection A\n\n" + + "beta " * 30 + + "\n\n# Second Section\n\n" + + "gamma " * 30 + ) + chunks = chunk_pages( + [ParsedPage(text = md)], + max_tokens = 25, + overlap_tokens = 0, + token_counter = _wc_counter, + ) + # We expect multiple chunks and at least one to begin at a heading. + assert len(chunks) >= 2 + starts_at_heading = sum( + 1 for c in chunks if c.text.lstrip().startswith(("# ", "## ")) + ) + assert starts_at_heading >= 1, ( + f"expected at least one chunk to start at a Markdown heading; " + f"got starts: {[c.text[:20] for c in chunks]}" + ) diff --git a/tests/python/test_rag_parsers.py b/tests/python/test_rag_parsers.py index bd28ee1234..149a8b754a 100644 --- a/tests/python/test_rag_parsers.py +++ b/tests/python/test_rag_parsers.py @@ -1,4 +1,8 @@ -"""Document parser tests — each format skipped if its lib is unavailable.""" +"""Document parser tests — each format skipped if its lib is unavailable. + +Phase 3A: parsers now return ParseResult (iterable over .pages) and +emit Markdown so the chunker can split on heading boundaries. +""" import sys from pathlib import Path @@ -16,19 +20,22 @@ def test_text_parser_utf8(tmp_path): file = tmp_path / "sample.txt" file.write_text("hello world\n\nsecond paragraph", encoding = "utf-8") - pages = parse(file) - assert len(pages) == 1 - assert "hello world" in pages[0].text - assert "second paragraph" in pages[0].text + result = parse(file) + assert len(result) == 1 + assert "hello world" in result.pages[0].text + assert "second paragraph" in result.pages[0].text + assert result.images == [] -def test_markdown_parser_treated_as_text(tmp_path): +def test_markdown_parser_preserves_headings(tmp_path): from core.rag.parsers import parse file = tmp_path / "sample.md" file.write_text("# Title\n\nBody text with **emphasis**.", encoding = "utf-8") - pages = parse(file) - assert pages and "Title" in pages[0].text + result = parse(file) + assert result.pages + # Markdown should pass through unchanged — heading marker preserved. + assert "# Title" in result.pages[0].text def test_unsupported_format_raises(tmp_path): @@ -40,23 +47,41 @@ def test_unsupported_format_raises(tmp_path): parse(file) -def test_html_parser_strips_scripts(tmp_path): +def test_html_parser_emits_markdown_headings(tmp_path): pytest.importorskip("bs4") pytest.importorskip("lxml") + pytest.importorskip("markdownify") from core.rag.parsers import parse file = tmp_path / "sample.html" file.write_text( - "

    visible text

    ", + "" + "" + "

    Main Title

    " + "

    Sub Section

    " + "

    visible text

    " + "
    • one
    • two
    " + "", encoding = "utf-8", ) - pages = parse(file) - assert pages - assert "visible text" in pages[0].text - assert "alert" not in pages[0].text + result = parse(file) + assert result.pages + md = result.pages[0].text + # markdownify converts

    → '# ',

    → '## ' + assert "# Main Title" in md + assert "## Sub Section" in md + assert "visible text" in md + # script content scrubbed + assert "alert" not in md + # list items become Markdown bullets + assert "one" in md and "two" in md def test_pdf_parser_extracts_pages(tmp_path): + pytest.importorskip("pymupdf") + pytest.importorskip("pymupdf4llm") + from core.rag.parsers import parse + pypdf = pytest.importorskip("pypdf") from pypdf import PdfWriter @@ -65,25 +90,46 @@ def test_pdf_parser_extracts_pages(tmp_path): writer.add_blank_page(width = 72, height = 72) with open(file, "wb") as f: writer.write(f) - from core.rag.parsers import parse - # blank page yields no extractable text — should return [] without error - pages = parse(file) - assert isinstance(pages, list) + # Blank page yields no extractable text — should return empty pages + # without error. + result = parse(file) + assert isinstance(result.pages, list) + assert isinstance(result.images, list) -def test_docx_parser_extracts_paragraphs(tmp_path): - docx = pytest.importorskip("docx") +def test_docx_parser_emits_markdown_headings(tmp_path): + pytest.importorskip("docx") + pytest.importorskip("mammoth") + pytest.importorskip("markdownify") from docx import Document file = tmp_path / "sample.docx" doc = Document() + doc.add_heading("Top Level Heading", level = 1) doc.add_paragraph("First paragraph here.") + doc.add_heading("Sub Heading", level = 2) doc.add_paragraph("Second paragraph here.") doc.save(str(file)) from core.rag.parsers import parse - pages = parse(file) - assert pages - assert "First paragraph" in pages[0].text - assert "Second paragraph" in pages[0].text + result = parse(file) + assert result.pages + md = result.pages[0].text + # mammoth via _STYLE_MAP maps Heading 1/2 → h1/h2 → '# '/'## '. + assert "# Top Level Heading" in md + assert "## Sub Heading" in md + assert "First paragraph" in md + assert "Second paragraph" in md + + +def test_parse_result_is_iterable_for_backcompat(tmp_path): + """Code that does `for page in parse(path)` should keep working.""" + from core.rag.parsers import parse + + file = tmp_path / "sample.txt" + file.write_text("hello", encoding = "utf-8") + result = parse(file) + pages = list(result) + assert len(pages) == 1 + assert pages[0].text == "hello" From 4c1ab745d6ee0a17f938ef5fe8c345e535ebf194 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 11:54:38 +0400 Subject: [PATCH 003/122] Studio: schema + API plumbing for per-KB chunking strategy + mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 lays two orthogonal per-KB knobs in the data and API layers so follow-up commits (Phase 3B-late, Phase 3B-multimodal) only need to add their code path and UI selector, not schema or types. Schema (studio/backend/storage/studio_db.py) - rag_knowledge_bases gains chunking_strategy ('standard'|'late', default 'standard') and mode ('text'|'multimodal', default 'text'). Both are immutable after KB creation — changing either invalidates existing chunks because they were ingested through a specific pipeline. - rag_chunks gains kind ('text'|'image'|'caption', default 'text'), image_path (NULLABLE), linked_chunk_id (NULLABLE) — used by Phase 3B-multimodal to pair image chunks with their captions. - Idempotent ALTER TABLE additions for existing installs (mirrors the chat_threads display_name / *_code_exec_container_id pattern earlier in the file). API (studio/backend/routes/rag.py) - ChunkingStrategy + KBMode Literal aliases. - CreateKBRequest accepts both fields with backward-compat defaults. - KBResponse exposes both. - _validate_mode_combo rejects (multimodal, late) with 400 — no public open-weight embedder supports both at once. Surface the constraint early rather than failing silently during ingestion. Config (studio/backend/utils/rag/config.py) - RAG_EMBEDDER_MATRIX dict keyed by (mode, strategy) → embedder name. - resolve_embedder() helper falls back to RAG_EMBEDDING_MODEL for legacy KBs that pre-date the columns. - (multimodal, late) intentionally absent. Frontend (studio/frontend/src/features/rag/) - api/rag-api.ts: ChunkingStrategy + KBMode types; KnowledgeBase interface and createKnowledgeBase request type updated. - stores/rag-store.ts: createKB signature uses the shared request type. No user-visible UI changes yet — the only currently-usable combination is (text, standard), so adding one-option selectors would be UX noise. Phase 3B-late and Phase 3B-multimodal each add the relevant selector option as part of shipping the code path. --- studio/backend/routes/rag.py | 61 +++++++++++++++++-- studio/backend/storage/studio_db.py | 42 +++++++++++++ studio/backend/utils/rag/config.py | 26 ++++++++ .../frontend/src/features/rag/api/rag-api.ts | 17 +++++- .../src/features/rag/stores/rag-store.ts | 3 +- 5 files changed, 142 insertions(+), 7 deletions(-) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 1fbda5b594..d292db4ba6 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -59,10 +59,19 @@ logger = get_logger(__name__) # Pydantic schemas # ------------------------------------------------------------------ +ChunkingStrategy = Literal["standard", "late"] +KBMode = Literal["text", "multimodal"] + + 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" class KBResponse(BaseModel): @@ -70,6 +79,8 @@ class KBResponse(BaseModel): name: str description: str | None embedding_model: str + chunking_strategy: ChunkingStrategy + mode: KBMode created_at: int @@ -150,15 +161,44 @@ def _now_ms() -> int: 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" + ) + mode = row["mode"] if "mode" in keys else "text" return KBResponse( id = row["id"], name = row["name"], description = row["description"], embedding_model = row["embedding_model"], + chunking_strategy = chunking_strategy, + mode = mode, created_at = row["created_at"], ) +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. + """ + if mode == "multimodal" and chunking_strategy == "late": + raise HTTPException( + status_code = 400, + detail = ( + "Late chunking is not supported in multimodal mode — " + "the multimodal embedder does not expose per-token " + "embeddings. Pick 'standard' chunking or 'text' mode." + ), + ) + + def _row_to_document(row: Any) -> DocumentResponse: return DocumentResponse( id = row["id"], @@ -300,18 +340,27 @@ def create_knowledge_base( payload: CreateKBRequest, current_subject: str = Depends(get_current_subject), ) -> KBResponse: - from utils.rag.config import RAG_EMBEDDING_MODEL + from utils.rag.config import resolve_embedder + + _validate_mode_combo(payload.mode, payload.chunking_strategy) kb_id = str(uuid4()) - embedding_model = payload.embedding_model or RAG_EMBEDDING_MODEL + # 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. + embedding_model = ( + payload.embedding_model + or resolve_embedder(payload.mode, payload.chunking_strategy) + ) created_at = _now_ms() with get_connection() as conn: try: conn.execute( """ INSERT INTO rag_knowledge_bases - (id, name, description, owner_user_id, embedding_model, created_at) - VALUES (?, ?, ?, ?, ?, ?) + (id, name, description, owner_user_id, embedding_model, + chunking_strategy, mode, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( kb_id, @@ -319,6 +368,8 @@ def create_knowledge_base( payload.description, current_subject, embedding_model, + payload.chunking_strategy, + payload.mode, created_at, ), ) @@ -333,6 +384,8 @@ def create_knowledge_base( name = payload.name, description = payload.description, embedding_model = embedding_model, + chunking_strategy = payload.chunking_strategy, + mode = payload.mode, created_at = created_at, ) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 31bc67150c..c4541fbf9a 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -208,6 +208,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: # 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. conn.execute( """ CREATE TABLE IF NOT EXISTS rag_knowledge_bases ( @@ -216,10 +220,29 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: description TEXT, owner_user_id TEXT, embedding_model TEXT NOT NULL, + chunking_strategy TEXT NOT NULL DEFAULT 'standard', + mode TEXT NOT NULL DEFAULT 'text', created_at INTEGER NOT NULL ) """ ) + # 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. + kb_cols = { + row[1] + for row in conn.execute("PRAGMA table_info(rag_knowledge_bases)").fetchall() + } + if "chunking_strategy" not in kb_cols: + conn.execute( + "ALTER TABLE rag_knowledge_bases " + "ADD COLUMN chunking_strategy TEXT NOT NULL DEFAULT 'standard'" + ) + if "mode" not in kb_cols: + conn.execute( + "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 @@ -249,6 +272,11 @@ 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. conn.execute( """ CREATE TABLE IF NOT EXISTS rag_chunks ( @@ -258,10 +286,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: text TEXT NOT NULL, token_count INTEGER NOT NULL DEFAULT 0, page_number INTEGER, + kind TEXT NOT NULL DEFAULT 'text', + image_path TEXT, + linked_chunk_id TEXT, UNIQUE(document_id, chunk_index) ) """ ) + chunk_cols = { + row[1] for row in conn.execute("PRAGMA table_info(rag_chunks)").fetchall() + } + if "kind" not in chunk_cols: + conn.execute( + "ALTER TABLE rag_chunks ADD COLUMN kind TEXT NOT NULL DEFAULT 'text'" + ) + if "image_path" not in chunk_cols: + conn.execute("ALTER TABLE rag_chunks ADD COLUMN image_path TEXT") + if "linked_chunk_id" not in chunk_cols: + conn.execute("ALTER TABLE rag_chunks ADD COLUMN linked_chunk_id TEXT") conn.execute( "CREATE INDEX IF NOT EXISTS idx_rag_chunks_document_id ON rag_chunks(document_id)" ) diff --git a/studio/backend/utils/rag/config.py b/studio/backend/utils/rag/config.py index c2ae2a9af2..5f1a5f9bd5 100644 --- a/studio/backend/utils/rag/config.py +++ b/studio/backend/utils/rag/config.py @@ -31,6 +31,32 @@ 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. +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", + ("multimodal", "standard"): "BAAI/BGE-VL-base", +} + + +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. + """ + return RAG_EMBEDDER_MATRIX.get( + (mode, chunking_strategy), + RAG_EMBEDDING_MODEL, + ) + RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512) RAG_CHUNK_OVERLAP: int = _env_int("UNSLOTH_RAG_CHUNK_OVERLAP", 64) diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index 5368b5de26..14503db292 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -5,11 +5,16 @@ import { authFetch, getAuthToken } from "@/features/auth"; import { apiUrl } from "@/lib/api-base"; import { formatFastApiDetail } from "@/lib/format-fastapi-error"; +export type ChunkingStrategy = "standard" | "late"; +export type KBMode = "text" | "multimodal"; + export interface KnowledgeBase { id: string; name: string; description: string | null; embedding_model: string; + chunking_strategy: ChunkingStrategy; + mode: KBMode; created_at: number; } @@ -92,11 +97,19 @@ export async function listKnowledgeBases(): Promise { return body.knowledge_bases; } -export async function createKnowledgeBase(req: { +export interface CreateKnowledgeBaseRequest { name: string; description?: string; embedding_model?: string; -}): Promise { + // Phase 3: both default server-side to "standard" / "text" — clients + // that pre-date the field send the same payloads as before. + chunking_strategy?: ChunkingStrategy; + mode?: KBMode; +} + +export async function createKnowledgeBase( + req: CreateKnowledgeBaseRequest, +): Promise { const response = await authFetch("/api/rag/knowledge-bases", { method: "POST", headers: { "Content-Type": "application/json" }, diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index 53c6e0a34d..1d777ad582 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -5,6 +5,7 @@ import { create } from "zustand"; import { clearThreadDocuments as apiClearThreadDocuments, createKnowledgeBase, + type CreateKnowledgeBaseRequest, deleteDocument as apiDeleteDocument, deleteKnowledgeBase as apiDeleteKB, type JobEvent, @@ -36,7 +37,7 @@ interface RagStoreState { threadIndexesLoading: boolean; loadKnowledgeBases: () => Promise; - createKB: (req: { name: string; description?: string; embedding_model?: string }) => Promise; + createKB: (req: CreateKnowledgeBaseRequest) => Promise; deleteKB: (kbId: string) => Promise; loadKBDocuments: (kbId: string) => Promise; From 673b7f86bac4fdb99cdc971efc392346a3ca359e Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 12:18:09 +0400 Subject: [PATCH 004/122] Studio: late chunking opt-in per KB (Phase 3B-late) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a KB has chunking_strategy = 'late', ingestion takes a separate code path that embeds the full document in a single forward pass and mean-pools token embeddings per chunk span. Each chunk vector carries full-document context via the encoder's bidirectional attention — Jina's published technique, ~+6.5 nDCG@10 on long docs. Backend - chunking.py: new chunk_pages_with_spans() that joins all pages into a single full_doc, runs the existing recursive splitter, and returns per-chunk (char_start, char_end) offsets. Page-number metadata is recovered by overlap with the original page ranges so PDF citations still work. Existing chunk_pages() unchanged. - embeddings.py: new late_chunk_encode(doc_text, char_spans). Tokenizes the doc with return_offsets_mapping, runs the underlying transformer to get per-token last_hidden_state, then mean-pools per chunk span. When the doc exceeds the embedder's context, falls back to windowed late chunking with a 512-token overlap so cross-window context is partially preserved. - ingestion.py _subprocess_worker: branches on chunking_strategy. 'late' path: chunk_pages_with_spans -> late_chunk_encode -> one big chunks_batch message. 'standard' path unchanged. Both reuse the same parent-side pump. - ingestion.enqueue_ingestion: new chunking_strategy + mode kwargs; defaults to 'standard' / 'text' for legacy callers. embedder model resolved via resolve_embedder() from the (mode, strategy) matrix. - routes/rag.py: KB-doc upload reads chunking_strategy + mode from the KB row (defensive .get for pre-Phase-3 schemas) and threads them through _start_ingestion. Frontend - kb-create-dialog.tsx: new "Chunking strategy" select with Standard / Late options. Embedding-model placeholder switches to nomic when Late is picked. createKB request now carries chunking_strategy. - kb-list.tsx + chat-settings-sheet.tsx: small "⚡ Late" badge next to late-chunking KB names in the settings KB list and the chat sidebar dropdown so users see the mode at a glance. Tests - test_rag_late_chunking.py: pure-python tests for chunk_pages_with_spans (chunks index back into full_doc; page numbers inherited by overlap; pages joined with blank line). A server-marked test loads all-MiniLM-L6-v2 to exercise late_chunk_encode end-to-end. No multimodal yet; that's Phase 3B-multimodal (next PR). --- studio/backend/core/rag/chunking.py | 127 ++++++++-- studio/backend/core/rag/embeddings.py | 238 ++++++++++++++++++ studio/backend/core/rag/ingestion.py | 202 +++++++++++---- studio/backend/routes/rag.py | 16 ++ .../src/features/chat/chat-settings-sheet.tsx | 13 +- .../rag/components/kb-create-dialog.tsx | 45 +++- .../src/features/rag/components/kb-list.tsx | 12 +- tests/python/test_rag_late_chunking.py | 112 +++++++++ 8 files changed, 700 insertions(+), 65 deletions(-) create mode 100644 tests/python/test_rag_late_chunking.py diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py index 4178b47c8a..86488a9ef6 100644 --- a/studio/backend/core/rag/chunking.py +++ b/studio/backend/core/rag/chunking.py @@ -100,28 +100,31 @@ def _merge( return [c.strip() for c in chunks if c.strip()] +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### ", + "\n#### ", + "\n\n", + "\n", + ". ", + " ", + "", +) + + def chunk_pages( pages: list[ParsedPage], *, max_tokens: int, overlap_tokens: int, token_counter: TokenCounter | None = None, - 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### ", - "\n#### ", - "\n\n", - "\n", - ". ", - " ", - "", - ), + separators: tuple[str, ...] = DEFAULT_SEPARATORS, ) -> list[Chunk]: """Split parsed pages into overlapping chunks. @@ -142,3 +145,95 @@ def chunk_pages( ) ) return out + + +_PAGE_SEPARATOR = "\n\n" + + +def chunk_pages_with_spans( + pages: list[ParsedPage], + *, + max_tokens: int, + overlap_tokens: int, + 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`. + + 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. + """ + count = token_counter or _char_token_estimate + + parts: list[str] = [] + page_ranges: list[tuple[int, int, int | None]] = [] + cursor = 0 + for index, page in enumerate(pages): + parts.append(page.text) + start = cursor + end = cursor + len(page.text) + page_ranges.append((start, end, page.page_number)) + cursor = end + if index < len(pages) - 1: + cursor += len(_PAGE_SEPARATOR) + full_doc = _PAGE_SEPARATOR.join(parts) + + atomic = _atomic_split(full_doc, separators, max_tokens, count) + merged = _merge(atomic, max_tokens, overlap_tokens, count) + + chunks: list[Chunk] = [] + char_spans: list[tuple[int, int]] = [] + search_cursor = 0 + for piece in merged: + text = piece.strip() + if not text: + 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. + 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) + chunks.append( + Chunk( + text = text, + token_count = count(text), + page_number = page_number, + ) + ) + char_spans.append((idx, end_idx)) + # Advance past the *start* of this chunk so an overlapping + # next chunk can still be found. + search_cursor = idx + 1 + + return full_doc, chunks, char_spans + + +def _page_for_span( + start: int, + end: int, + page_ranges: list[tuple[int, int, int | None]], +) -> int | None: + for ps, pe, pn in page_ranges: + if start < pe and end > ps: + return pn + return None diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4953d48ffe..b3ea2de430 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -100,3 +100,241 @@ def token_counter(model_name: str | None = None): return max(1, len(text) // 4) return _count + + +# ------------------------------------------------------------------ +# Late chunking (Phase 3B-late) +# ------------------------------------------------------------------ + +_LATE_WINDOW_OVERLAP_TOKENS = 512 + + +def late_chunk_encode( + doc_text: str, + char_spans: list[tuple[int, int]], + *, + 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. + """ + import numpy as np + + if not char_spans: + return [] + model = get_embedder(model_name) + tokenizer = model.tokenizer + max_length = int(getattr(model, "max_seq_length", None) or 8192) + + encoded = tokenizer( + doc_text, + return_tensors = "pt", + return_offsets_mapping = True, + add_special_tokens = True, + truncation = False, + ) + offsets = encoded.pop("offset_mapping")[0].tolist() + n_tokens = int(encoded["input_ids"].shape[1]) + + if n_tokens <= max_length: + token_embeddings = _encode_tokens(model, encoded) + return _pool_spans( + token_embeddings, + offsets, + char_spans, + normalize = normalize, + np_module = np, + model = model, + doc_text = doc_text, + ) + + logger.info( + "Late chunking: doc has %d tokens > model max %d; using windowed pass", + n_tokens, + max_length, + ) + return _windowed_late_chunk_encode( + doc_text = doc_text, + char_spans = char_spans, + model = model, + max_length = max_length, + normalize = normalize, + np_module = np, + ) + + +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 + device = next(transformer.parameters()).device + inputs_on_device = {k: v.to(device) for k, v in encoded.items()} + with torch.no_grad(): + outputs = transformer(**inputs_on_device) + return outputs.last_hidden_state[0].detach().cpu().numpy() + + +def _pool_spans( + token_embeddings, + offsets, + char_spans, + *, + normalize: bool, + np_module, + model, + 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). + """ + 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. + indices = [ + i - token_index_offset + for i, (ts, te) in enumerate(offsets) + if te > ts and te > char_start and ts < char_end + ] + 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, + convert_to_numpy = True, + show_progress_bar = False, + ) + vectors.append(vec) + continue + pooled = token_embeddings[indices].mean(axis = 0) + if normalize: + denom = float(np_module.linalg.norm(pooled)) + if denom > 0: + pooled = pooled / denom + vectors.append(pooled) + return vectors + + +def _windowed_late_chunk_encode( + *, + doc_text: str, + char_spans: list[tuple[int, int]], + model, + max_length: int, + 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. + """ + import torch + + tokenizer = model.tokenizer + transformer = model[0].auto_model + device = next(transformer.parameters()).device + + full = tokenizer( + doc_text, + return_tensors = "pt", + return_offsets_mapping = True, + add_special_tokens = False, + truncation = False, + ) + all_input_ids = full["input_ids"][0] + all_offsets = full["offset_mapping"][0].tolist() + 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: + end = min(pos + max_length, n_tokens) + windows.append((pos, end)) + if end >= n_tokens: + break + pos += stride + + # Cache window → token embeddings (only encode when needed). + window_embeddings: dict[int, "np_module.ndarray"] = {} + + def _window_embeddings(window_index: int): + if window_index in window_embeddings: + return window_embeddings[window_index] + ws, we = windows[window_index] + win_ids = all_input_ids[ws:we].unsqueeze(0).to(device) + win_attn = torch.ones_like(win_ids) + with torch.no_grad(): + outputs = transformer(input_ids = win_ids, attention_mask = win_attn) + emb = outputs.last_hidden_state[0].detach().cpu().numpy() + window_embeddings[window_index] = emb + return emb + + 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) + if te > ts and te > char_start and ts < char_end + ] + if not chunk_token_indices: + vec = model.encode( + doc_text[char_start:char_end], + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + 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): + overlap = sum(1 for ti in chunk_token_indices if ws <= ti < we) + if overlap > best_overlap: + best_overlap = overlap + best_window = wi + ws, _we = windows[best_window] + emb = _window_embeddings(best_window) + local_indices = [ + ti - ws + for ti in chunk_token_indices + if ws <= ti < ws + emb.shape[0] + ] + if not local_indices: + vec = model.encode( + doc_text[char_start:char_end], + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + vectors.append(vec) + continue + pooled = emb[local_indices].mean(axis = 0) + if normalize: + denom = float(np_module.linalg.norm(pooled)) + if denom > 0: + pooled = pooled / denom + vectors.append(pooled) + return vectors diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 226fd53129..2d99c7e4d7 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -54,74 +54,172 @@ def _subprocess_worker( overlap: int, batch_size: int, out_queue: Any, + chunking_strategy: str = "standard", + mode: str = "text", ) -> None: try: - from core.rag.chunking import chunk_pages + from core.rag.chunking import chunk_pages, chunk_pages_with_spans from core.rag.parsers import parse out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) - # want_images stays False for the text-only ingestion path; the - # multimodal path (Phase 3B-multimodal) will flip this based on - # the KB's mode. - parsed = parse(Path(stored_path), want_images = False) + # want_images is True only for multimodal KBs. The image side of + # the pipeline lands in Phase 3B-multimodal; for now the parser + # collects the bytes anyway in case we want them later, but only + # the text pages are consumed. + parsed = parse(Path(stored_path), want_images = (mode == "multimodal")) pages = parsed.pages if not pages: out_queue.put({"type": "error", "error": "no extractable text in document"}) return out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1}) - from core.rag.embeddings import get_embedder, token_counter + from core.rag.embeddings import ( + get_embedder, + late_chunk_encode, + token_counter, + ) model = get_embedder(model_name) counter = token_counter(model_name) dim = int(model.get_sentence_embedding_dimension()) out_queue.put({"type": "dim", "dim": dim}) - out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) - chunks = chunk_pages( - pages, - max_tokens = chunk_size, - overlap_tokens = overlap, - token_counter = counter, - ) - if not chunks: - out_queue.put({"type": "error", "error": "chunker produced no chunks"}) - return - - total = len(chunks) - for i in range(0, total, batch_size): - batch = chunks[i : i + batch_size] - vectors = model.encode( - [c.text for c in batch], + if chunking_strategy == "late": + _run_late_chunking( + pages = pages, + chunk_size = chunk_size, + overlap = overlap, + counter = counter, + model_name = model_name, + late_chunk_encode = late_chunk_encode, + out_queue = out_queue, + ) + else: + _run_standard_chunking( + pages = pages, + chunk_size = chunk_size, + overlap = overlap, + counter = counter, batch_size = batch_size, - normalize_embeddings = True, - convert_to_numpy = True, - show_progress_bar = False, + model = model, + chunk_pages = chunk_pages, + out_queue = out_queue, ) - out_queue.put( - { - "type": "chunks_batch", - "first_index": i, - "chunks": [ - { - "text": c.text, - "token_count": c.token_count, - "page_number": c.page_number, - } - for c in batch - ], - "vectors": vectors.tolist(), - } - ) - progress = 0.3 + 0.65 * min(1.0, (i + len(batch)) / total) - out_queue.put({"type": "progress", "stage": "embed", "progress": progress}) - - out_queue.put({"type": "complete", "num_chunks": total}) except Exception as exc: # noqa: BLE001 logger.exception("ingestion subprocess failed") out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"}) +def _run_standard_chunking( + *, + pages, + chunk_size, + overlap, + counter, + batch_size, + model, + chunk_pages, + out_queue, +) -> None: + out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) + chunks = chunk_pages( + pages, + max_tokens = chunk_size, + overlap_tokens = overlap, + token_counter = counter, + ) + if not chunks: + out_queue.put({"type": "error", "error": "chunker produced no chunks"}) + return + + total = len(chunks) + for i in range(0, total, batch_size): + batch = chunks[i : i + batch_size] + vectors = model.encode( + [c.text for c in batch], + batch_size = batch_size, + normalize_embeddings = True, + convert_to_numpy = True, + show_progress_bar = False, + ) + out_queue.put( + { + "type": "chunks_batch", + "first_index": i, + "chunks": [ + { + "text": c.text, + "token_count": c.token_count, + "page_number": c.page_number, + } + for c in batch + ], + "vectors": vectors.tolist(), + } + ) + progress = 0.3 + 0.65 * min(1.0, (i + len(batch)) / total) + out_queue.put({"type": "progress", "stage": "embed", "progress": progress}) + + out_queue.put({"type": "complete", "num_chunks": total}) + + +def _run_late_chunking( + *, + pages, + chunk_size, + overlap, + counter, + model_name, + 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. + """ + from core.rag.chunking import chunk_pages_with_spans + + out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) + full_doc, chunks, char_spans = chunk_pages_with_spans( + pages, + max_tokens = chunk_size, + overlap_tokens = overlap, + token_counter = counter, + ) + if not chunks: + out_queue.put({"type": "error", "error": "chunker produced no chunks"}) + return + + out_queue.put({"type": "progress", "stage": "embed", "progress": 0.4}) + vectors = late_chunk_encode( + full_doc, + char_spans, + model_name = model_name, + normalize = True, + ) + + out_queue.put({"type": "progress", "stage": "embed", "progress": 0.9}) + out_queue.put( + { + "type": "chunks_batch", + "first_index": 0, + "chunks": [ + { + "text": c.text, + "token_count": c.token_count, + "page_number": c.page_number, + } + for c in chunks + ], + "vectors": [v.tolist() for v in vectors], + } + ) + out_queue.put({"type": "complete", "num_chunks": len(chunks)}) + + # ------------------------------------------------------------------ # Job manager (parent side) # ------------------------------------------------------------------ @@ -390,14 +488,26 @@ def enqueue_ingestion( kb_id: str | None = None, thread_id: str | None = None, embedding_model: str | None = None, + 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. """ + from utils.rag.config import resolve_embedder + scope = _scope_for(kb_id, thread_id) - model_name = embedding_model or RAG_EMBEDDING_MODEL + model_name = ( + embedding_model + or resolve_embedder(mode, chunking_strategy) + or RAG_EMBEDDING_MODEL + ) job_id = str(uuid4()) with get_connection() as conn: conn.execute( @@ -424,6 +534,8 @@ def enqueue_ingestion( RAG_CHUNK_OVERLAP, RAG_EMBED_BATCH_SIZE, out_queue, + chunking_strategy, + mode, ), daemon = True, ) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index d292db4ba6..bbf0b22834 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -289,6 +289,8 @@ def _start_ingestion( kb_id: str | None, thread_id: str | None, embedding_model: str, + chunking_strategy: str = "standard", + mode: str = "text", ) -> UploadResponse: document_id = str(uuid4()) with get_connection() as conn: @@ -317,6 +319,8 @@ def _start_ingestion( kb_id = kb_id, thread_id = thread_id, embedding_model = embedding_model, + chunking_strategy = chunking_strategy, + mode = mode, ) return UploadResponse(document_id = document_id, job_id = job_id, filename = filename) @@ -432,6 +436,16 @@ 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. + 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" + ) + mode = kb_row["mode"] if "mode" in kb_keys else "text" return _start_ingestion( filename = filename, stored_path = stored_path, @@ -440,6 +454,8 @@ async def upload_kb_document( kb_id = kb_id, thread_id = None, embedding_model = kb_row["embedding_model"], + chunking_strategy = chunking_strategy, + mode = mode, ) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index bef9fde24a..5803832120 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1232,6 +1232,7 @@ export function ChatSettingsPanel({ ) : null} {knowledgeBases.map((kb) => { const isActive = kb.id === activeKbId; + const isLate = kb.chunking_strategy === "late"; return ( - {kb.name} + + {kb.name} + {isLate ? ( + + ⚡ Late + + ) : null} + + + + + + + ); +} diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index 1d777ad582..3cf9d98d66 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -15,6 +15,9 @@ import { listThreadDocuments, listThreadIndexes, type RagDocument, + type ReingestKBOptions, + reingestKnowledgeBase as apiReingestKB, + reingestThreadDocuments as apiReingestThread, subscribeToJobEvents, type ThreadIndexSummary, uploadKBDocument, @@ -51,6 +54,9 @@ interface RagStoreState { loadThreadIndexes: () => Promise; clearThreadIndex: (threadId: string) => Promise; + reingestKB: (kbId: string, opts?: ReingestKBOptions) => Promise; + reingestThread: (threadId: string) => Promise; + subscribeJob: (jobId: string, onComplete?: () => void) => void; } @@ -220,6 +226,35 @@ export const useRagStore = create((set, get) => ({ }); }, + async reingestKB(kbId, opts) { + const response = await apiReingestKB(kbId, opts ?? {}); + // Refresh the KB list so updated chunking_strategy / mode / embedder + // values flow back into the UI, and the doc list so old chunk + // counts reset to 0 until each job completes. + void get().loadKnowledgeBases(); + void get().loadKBDocuments(kbId); + // Subscribe to every new job so progress chips render and the doc + // list refreshes on completion (mirrors uploadDocument's pattern). + for (const jobId of response.job_ids) { + get().subscribeJob(jobId, () => { + void get().loadKBDocuments(kbId); + }); + } + return response.job_ids; + }, + + async reingestThread(threadId) { + const response = await apiReingestThread(threadId); + void get().loadThreadDocuments(threadId); + void get().loadThreadIndexes(); + for (const jobId of response.job_ids) { + get().subscribeJob(jobId, () => { + void get().loadThreadDocuments(threadId); + }); + } + return response.job_ids; + }, + subscribeJob(jobId, onComplete) { const existing = get().jobUnsubscribers[jobId]; if (existing) return; diff --git a/tests/python/test_rag_reingest.py b/tests/python/test_rag_reingest.py new file mode 100644 index 0000000000..70508a2397 --- /dev/null +++ b/tests/python/test_rag_reingest.py @@ -0,0 +1,57 @@ +"""Reingest endpoint tests (Backfill UX). + +Full end-to-end reingest needs a running studio + a real embedder; that's +covered manually via the curl smoke flow in the plan. Here we cover the +parts that are testable without external models: payload validation and +the (multimodal, late) constraint propagation. +""" + +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + + +def test_reingest_request_accepts_all_optional_fields(): + from routes.rag import ReingestKBRequest + + empty = ReingestKBRequest() + assert empty.chunking_strategy is None + assert empty.mode is None + assert empty.embedding_model is None + + partial = ReingestKBRequest(chunking_strategy = "late") + assert partial.chunking_strategy == "late" + assert partial.mode is None + + +def test_reingest_request_rejects_unknown_strategy(): + from routes.rag import ReingestKBRequest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + ReingestKBRequest(chunking_strategy = "telekinetic") + + +def test_reingest_request_rejects_unknown_mode(): + from routes.rag import ReingestKBRequest + from pydantic import ValidationError + + with pytest.raises(ValidationError): + ReingestKBRequest(mode = "augmented") + + +def test_constraint_still_enforced_for_reingest_combos(): + """The combination guard is shared with create — verify it still bites.""" + from fastapi import HTTPException + + from routes.rag import _validate_mode_combo + + with pytest.raises(HTTPException) as excinfo: + _validate_mode_combo("multimodal", "late") + assert excinfo.value.status_code == 400 From 2b85a165e690794377708d6408afcc9611cceba6 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 12:45:00 +0400 Subject: [PATCH 007/122] Studio: global RAG ingestion toast stack (Phase 2C) Adds a floating progress-card stack mounted at the app root that watches useRagStore.jobs and renders one card per in-flight ingestion job, regardless of which page the user is on. Wraps the existing IngestionProgress component so the progress UI stays consistent with the per-doc chips in the KB detail panel and chat sidebar. - Terminal cards (complete/error) linger for 4s then auto-dismiss. - A manual dismiss button is always available. - Reduced-motion preference is respected (no slide animation). - Positioned bottom-right, z-50, pointer-events-none container so clicks pass through to the page underneath. Mounted in app/routes/__root.tsx next to the existing SettingsDialog so it's visible across every authenticated route. Closes the last deferred item from Phase 2. --- studio/frontend/src/app/routes/__root.tsx | 2 + .../rag/components/ingestion-toast-stack.tsx | 114 ++++++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 47bff815e6..69f9c26fb0 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -5,6 +5,7 @@ import { AppSidebar } from "@/components/app-sidebar"; import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; +import { IngestionToastStack } from "@/features/rag/components/ingestion-toast-stack"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; import { useTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; @@ -114,6 +115,7 @@ function RootLayout() { return ( + {hideNavbar ? (
    diff --git a/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx b/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx new file mode 100644 index 0000000000..ad824c7111 --- /dev/null +++ b/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Cancel01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { AnimatePresence, motion, useReducedMotion } from "motion/react"; +import { useEffect, useState } from "react"; +import { useRagStore } from "../stores/rag-store"; +import { IngestionProgress } from "./ingestion-progress"; + +/** + * Floating progress stack for in-flight RAG ingestion jobs. + * + * Watches the rag-store `jobs` map (populated by `subscribeJob`) and + * renders one card per job that's neither completed nor errored. On + * completion the card stays for a few seconds with a success state + * before fading out, so users notice the run finished even if they + * weren't watching the per-doc progress chips. + * + * Mounted once at the app root (`__root.tsx`) so it follows users + * across pages. + */ + +const DISMISS_DELAY_MS = 4000; + +export function IngestionToastStack() { + const jobs = useRagStore((s) => s.jobs); + const reduced = useReducedMotion(); + // Per-job timeout handles so terminal toasts auto-dismiss. + const [dismissedJobs, setDismissedJobs] = useState>( + () => new Set(), + ); + + // Schedule auto-dismiss for jobs that have reached a terminal state. + useEffect(() => { + const timers: ReturnType[] = []; + for (const [jobId, event] of Object.entries(jobs)) { + if ( + (event.type === "complete" || event.type === "error") && + !dismissedJobs.has(jobId) + ) { + timers.push( + setTimeout(() => { + setDismissedJobs((prev) => { + const next = new Set(prev); + next.add(jobId); + return next; + }); + }, DISMISS_DELAY_MS), + ); + } + } + return () => timers.forEach(clearTimeout); + }, [jobs, dismissedJobs]); + + const visible = Object.entries(jobs).filter( + ([jobId]) => !dismissedJobs.has(jobId), + ); + if (visible.length === 0) return null; + + return ( +
    + + {visible.map(([jobId, event]) => { + const isTerminal = + event.type === "complete" || event.type === "error"; + const isError = event.type === "error"; + return ( + +
    +
    + + {isError + ? "Ingestion failed" + : event.type === "complete" + ? "Indexed" + : "Indexing document"} + + +
    + +
    +     + ); + })} +     +
    + ); +} From ca83bea538b53e069454644fcc984969d71e71b0 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 12:47:16 +0400 Subject: [PATCH 008/122] Studio: app-level RAG defaults for new KBs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Power users can now set their preferred chunking strategy / mode / embedder once in Settings → Knowledge Bases and have new KBs use those values by default, instead of toggling on every create. Backend - routes/rag.py: - GET /api/rag/defaults returns the stored RagDefaults (or sensible fallbacks when nothing is set: standard / text / null embedder). - PUT /api/rag/defaults is PATCH-style — only fields present in the body overwrite. The (multimodal, late) constraint is enforced here too, so users can't poison the defaults with a combination the create path would reject. - Persistence reuses the existing chat_settings store via upsert_chat_settings_merge; the values live under a single rag.defaults key as a nested JSON dict. Frontend - rag-api.ts: getRagDefaults / setRagDefaults wrappers + RagDefaults + UpdateRagDefaultsRequest types. - rag-store.ts: defaults state, loadDefaults / updateDefaults actions. loadDefaults swallows errors so a missing endpoint just leaves defaults null. - rag-defaults-section.tsx (new): self-contained mode + strategy + embedding-model controls, persists on change. Used in the Settings KB tab below the ThreadIndexList section. - knowledge-bases-tab.tsx: mounts RagDefaultsSection below thread indexes with a separator. - kb-create-dialog.tsx: loads defaults on open and prefills the form with them (falls back to hard-coded standard / text when defaults haven't loaded yet). reset() returns to the latest defaults rather than the hard-coded ones. --- studio/backend/routes/rag.py | 75 ++++++++- .../frontend/src/features/rag/api/rag-api.ts | 28 ++++ .../rag/components/kb-create-dialog.tsx | 42 ++++- .../rag/components/rag-defaults-section.tsx | 148 ++++++++++++++++++ .../src/features/rag/stores/rag-store.ts | 26 +++ .../settings/tabs/knowledge-bases-tab.tsx | 3 + 6 files changed, 314 insertions(+), 8 deletions(-) create mode 100644 studio/frontend/src/features/rag/components/rag-defaults-section.tsx diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index dddb6c66b5..dc0122f44a 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -43,7 +43,11 @@ async def _sse_auth( from core.rag import embeddings, ingestion, reranker, retrieval, vector_store from core.rag.vector_store import kb_scope, thread_scope from loggers import get_logger -from storage.studio_db import get_connection +from storage.studio_db import ( + get_connection, + list_chat_settings, + upsert_chat_settings_merge, +) from utils.paths.storage_roots import ensure_dir, rag_uploads_root from utils.rag.config import ( RAG_MAX_UPLOAD_MB, @@ -407,6 +411,75 @@ def list_knowledge_bases( return KBListResponse(knowledge_bases = [_row_to_kb(r) for r in rows]) +class RagDefaults(BaseModel): + chunking_strategy: ChunkingStrategy = "standard" + mode: KBMode = "text" + embedding_model: str | None = None + + +class UpdateRagDefaultsRequest(BaseModel): + """Patch shape — only fields present overwrite stored values.""" + chunking_strategy: ChunkingStrategy | None = None + mode: KBMode | None = None + embedding_model: str | None = None + + +_DEFAULTS_KEY = "rag.defaults" + + +def _load_rag_defaults() -> RagDefaults: + settings = list_chat_settings() + raw = settings.get(_DEFAULTS_KEY) or {} + if not isinstance(raw, dict): + raw = {} + return RagDefaults( + chunking_strategy = raw.get("chunking_strategy") or "standard", + mode = raw.get("mode") or "text", + embedding_model = raw.get("embedding_model"), + ) + + +@router.get("/defaults", response_model = RagDefaults) +def get_rag_defaults( + current_subject: str = Depends(get_current_subject), +) -> RagDefaults: + return _load_rag_defaults() + + +@router.put("/defaults", response_model = RagDefaults) +def set_rag_defaults( + payload: UpdateRagDefaultsRequest, + current_subject: str = Depends(get_current_subject), +) -> RagDefaults: + 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. + if payload.embedding_model is None: + new_embedder = current.embedding_model + elif payload.embedding_model.strip() == "": + new_embedder = None + else: + new_embedder = payload.embedding_model.strip() + _validate_mode_combo(new_mode, new_strategy) + + upsert_chat_settings_merge( + { + _DEFAULTS_KEY: { + "chunking_strategy": new_strategy, + "mode": new_mode, + "embedding_model": new_embedder, + } + } + ) + return RagDefaults( + chunking_strategy = new_strategy, + mode = new_mode, + embedding_model = new_embedder, + ) + + class ReingestKBRequest(BaseModel): """All fields optional — omitting one keeps the KB's current value.""" chunking_strategy: ChunkingStrategy | None = None diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index d6c18faf3b..121d7494bb 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -246,6 +246,34 @@ export async function reingestThreadDocuments( return parseJsonOrThrow(response); } +export interface RagDefaults { + chunking_strategy: ChunkingStrategy; + mode: KBMode; + embedding_model: string | null; +} + +export async function getRagDefaults(): Promise { + const response = await authFetch("/api/rag/defaults"); + return parseJsonOrThrow(response); +} + +export interface UpdateRagDefaultsRequest { + chunking_strategy?: ChunkingStrategy; + mode?: KBMode; + embedding_model?: string | null; +} + +export async function setRagDefaults( + payload: UpdateRagDefaultsRequest, +): Promise { + const response = await authFetch("/api/rag/defaults", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + return parseJsonOrThrow(response); +} + // ------------------------------------------------------------------ // Search // ------------------------------------------------------------------ diff --git a/studio/frontend/src/features/rag/components/kb-create-dialog.tsx b/studio/frontend/src/features/rag/components/kb-create-dialog.tsx index c56b8b38f3..22e73afc77 100644 --- a/studio/frontend/src/features/rag/components/kb-create-dialog.tsx +++ b/studio/frontend/src/features/rag/components/kb-create-dialog.tsx @@ -19,13 +19,14 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; -import { useState } from "react"; +import { useEffect, useState } from "react"; import type { ChunkingStrategy, KBMode, KnowledgeBase, } from "../api/rag-api"; import { useKnowledgeBases } from "../hooks/use-knowledge-bases"; +import { useRagStore } from "../stores/rag-store"; export function KBCreateDialog({ open, @@ -37,21 +38,48 @@ export function KBCreateDialog({ onCreated?: (kb: KnowledgeBase) => void; }) { const { createKB } = useKnowledgeBases(); + const defaults = useRagStore((s) => s.defaults); + const loadDefaults = useRagStore((s) => s.loadDefaults); + + // Cache the initial values so reset() puts us back to the latest + // saved defaults rather than the hard-coded ones. + const initialStrategy: ChunkingStrategy = + defaults?.chunking_strategy ?? "standard"; + const initialMode: KBMode = defaults?.mode ?? "text"; + const initialEmbedder = defaults?.embedding_model ?? ""; + const [name, setName] = useState(""); const [description, setDescription] = useState(""); - const [embeddingModel, setEmbeddingModel] = useState(""); + const [embeddingModel, setEmbeddingModel] = useState(initialEmbedder); const [chunkingStrategy, setChunkingStrategy] = - useState("standard"); - const [mode, setMode] = useState("text"); + useState(initialStrategy); + const [mode, setMode] = useState(initialMode); const [submitting, setSubmitting] = useState(false); const [error, setError] = useState(null); + // Fetch defaults on first open and re-sync the form when they arrive. + useEffect(() => { + if (open && !defaults) { + void loadDefaults(); + } + }, [open, defaults, loadDefaults]); + + useEffect(() => { + if (open && defaults) { + setChunkingStrategy(defaults.chunking_strategy); + setMode(defaults.mode); + setEmbeddingModel(defaults.embedding_model ?? ""); + } + // Intentional: only when `open` flips, not on every defaults change. + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [open]); + const reset = () => { setName(""); setDescription(""); - setEmbeddingModel(""); - setChunkingStrategy("standard"); - setMode("text"); + setEmbeddingModel(defaults?.embedding_model ?? ""); + setChunkingStrategy(defaults?.chunking_strategy ?? "standard"); + setMode(defaults?.mode ?? "text"); setError(null); setSubmitting(false); }; diff --git a/studio/frontend/src/features/rag/components/rag-defaults-section.tsx b/studio/frontend/src/features/rag/components/rag-defaults-section.tsx new file mode 100644 index 0000000000..fd3be615e9 --- /dev/null +++ b/studio/frontend/src/features/rag/components/rag-defaults-section.tsx @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { useEffect, useState } from "react"; +import type { ChunkingStrategy, KBMode } from "../api/rag-api"; +import { useRagStore } from "../stores/rag-store"; + +/** + * Defaults for newly-created KBs, surfaced in the Settings → Knowledge + * Bases tab. Values pre-fill the KB create dialog so power users can + * pick their preferred strategy/mode once instead of toggling per KB. + * + * Backed by chat_settings via PUT /api/rag/defaults. The selects + * enforce the same (multimodal, late) constraint as the create + * dialog. + */ +export function RagDefaultsSection() { + const defaults = useRagStore((s) => s.defaults); + const loadDefaults = useRagStore((s) => s.loadDefaults); + const updateDefaults = useRagStore((s) => s.updateDefaults); + + const [chunkingStrategy, setChunkingStrategy] = + useState("standard"); + const [mode, setMode] = useState("text"); + const [embeddingModel, setEmbeddingModel] = useState(""); + const [error, setError] = useState(null); + + // Load defaults on mount; reflect them in local state. + useEffect(() => { + void loadDefaults(); + }, [loadDefaults]); + + useEffect(() => { + if (defaults) { + setChunkingStrategy(defaults.chunking_strategy); + setMode(defaults.mode); + setEmbeddingModel(defaults.embedding_model ?? ""); + } + }, [defaults]); + + const lateDisabled = mode === "multimodal"; + const multimodalDisabled = chunkingStrategy === "late"; + + const persist = (patch: { + chunking_strategy?: ChunkingStrategy; + mode?: KBMode; + embedding_model?: string | null; + }) => { + setError(null); + void updateDefaults(patch).catch((err) => { + setError(err instanceof Error ? err.message : String(err)); + }); + }; + + return ( +
    +
    +

    Defaults for new knowledge bases

    +

    + Pre-fills the KB create dialog. Existing KBs keep their own + settings — use the Reconfigure button to change those. +

    +
    +
    +
    + + +
    +
    + + +
    +
    +
    + + setEmbeddingModel(e.target.value)} + onBlur={() => persist({ embedding_model: embeddingModel })} + placeholder="Leave blank to use the matrix default" + /> +
    + {error ?
    {error}
    : null} +
    + ); +} diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index 3cf9d98d66..09a740fbe7 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -8,18 +8,22 @@ import { type CreateKnowledgeBaseRequest, deleteDocument as apiDeleteDocument, deleteKnowledgeBase as apiDeleteKB, + getRagDefaults as apiGetRagDefaults, type JobEvent, type KnowledgeBase, listKBDocuments, listKnowledgeBases, listThreadDocuments, listThreadIndexes, + type RagDefaults, type RagDocument, type ReingestKBOptions, reingestKnowledgeBase as apiReingestKB, reingestThreadDocuments as apiReingestThread, + setRagDefaults as apiSetRagDefaults, subscribeToJobEvents, type ThreadIndexSummary, + type UpdateRagDefaultsRequest, uploadKBDocument, uploadThreadDocument, } from "../api/rag-api"; @@ -57,6 +61,10 @@ interface RagStoreState { reingestKB: (kbId: string, opts?: ReingestKBOptions) => Promise; reingestThread: (threadId: string) => Promise; + defaults: RagDefaults | null; + loadDefaults: () => Promise; + updateDefaults: (patch: UpdateRagDefaultsRequest) => Promise; + subscribeJob: (jobId: string, onComplete?: () => void) => void; } @@ -82,6 +90,8 @@ export const useRagStore = create((set, get) => ({ threadIndexes: [], threadIndexesLoading: false, + defaults: null, + async loadKnowledgeBases() { set({ kbsLoading: true, kbsError: null }); try { @@ -243,6 +253,22 @@ export const useRagStore = create((set, get) => ({ return response.job_ids; }, + async loadDefaults() { + try { + const defaults = await apiGetRagDefaults(); + set({ defaults }); + } catch { + // Defaults endpoint is best-effort; a 401 / network blip just + // leaves defaults null and dialogs fall back to hard-coded + // ('standard', 'text', no embedder override). + } + }, + + async updateDefaults(patch) { + const defaults = await apiSetRagDefaults(patch); + set({ defaults }); + }, + async reingestThread(threadId) { const response = await apiReingestThread(threadId); void get().loadThreadDocuments(threadId); diff --git a/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx b/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx index 94cb90959e..0f5a3c290c 100644 --- a/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx @@ -5,6 +5,7 @@ import { Separator } from "@/components/ui/separator"; import type { KnowledgeBase } from "@/features/rag/api/rag-api"; import { KBDetailPanel } from "@/features/rag/components/kb-detail-panel"; import { KBList } from "@/features/rag/components/kb-list"; +import { RagDefaultsSection } from "@/features/rag/components/rag-defaults-section"; import { ThreadIndexList } from "@/features/rag/components/thread-index-list"; import { useState } from "react"; @@ -38,6 +39,8 @@ export function KnowledgeBasesTab() { + + ); } From ee1ff2bb506e03074270eb97130f9049a10a7a68 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 12:53:28 +0400 Subject: [PATCH 009/122] Studio: per-thread RAG chunking/mode overrides MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Threads can now opt into late chunking or multimodal mode independently of the KBs they reference. Per-thread settings persist in chat_settings under thread::rag and fall back to the app-level defaults when the thread hasn't set anything explicitly. Backend (routes/rag.py) - ThreadRagSettings / UpdateThreadRagSettingsRequest Pydantic models. - GET/PUT /api/rag/threads/{thread_id}/settings backed by chat_settings (upsert_chat_settings_merge). Same (multimodal, late) constraint enforcement as the create + defaults endpoints. - POST /api/rag/threads/{thread_id}/reingest now accepts the same body shape — if any field is set, the new settings are persisted via set_thread_rag_settings BEFORE the reingest, so subsequent uploads pick up the change too. - upload_thread_document reads the per-thread settings and passes them through to _start_ingestion, replacing the previous hard-coded ('standard', 'text', RAG_EMBEDDING_MODEL) defaults. Frontend - rag-api.ts: ThreadRagSettings type + getThreadRagSettings / setThreadRagSettings wrappers. reingestThreadDocuments now accepts optional UpdateThreadRagSettingsRequest opts. - rag-store.ts: threadSettings map keyed by threadId, plus loadThreadSettings / updateThreadSettings actions. reingestThread refreshes the local settings copy when opts were supplied. - chat-settings-sheet.tsx Retrieval section: when source = thread, shows side-by-side Mode + Chunking selects above the documents list. Selecting a different value: - persists immediately if the thread has no docs - prompts "Re-index N documents?" if docs exist; on Yes calls reingestThread with the new opts, on No reverts the select The (multimodal, late) constraint is enforced via per-option disabled + tooltip, matching the KB create dialog. --- studio/backend/routes/rag.py | 133 ++++++++++++++++-- .../src/features/chat/chat-settings-sheet.tsx | 121 +++++++++++++++- .../frontend/src/features/rag/api/rag-api.ts | 39 ++++- .../src/features/rag/stores/rag-store.ts | 46 +++++- 4 files changed, 325 insertions(+), 14 deletions(-) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index dc0122f44a..88ce2d84b2 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -480,6 +480,89 @@ def set_rag_defaults( ) +class ThreadRagSettings(BaseModel): + chunking_strategy: ChunkingStrategy = "standard" + mode: KBMode = "text" + embedding_model: str | None = None + + +class UpdateThreadRagSettingsRequest(BaseModel): + chunking_strategy: ChunkingStrategy | None = None + mode: KBMode | None = None + embedding_model: str | None = None + + +def _thread_settings_key(thread_id: str) -> str: + return f"thread:{thread_id}:rag" + + +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. + """ + settings = list_chat_settings() + raw = settings.get(_thread_settings_key(thread_id)) or {} + if not isinstance(raw, dict): + raw = {} + fallback = _load_rag_defaults() + return ThreadRagSettings( + chunking_strategy = ( + raw.get("chunking_strategy") or fallback.chunking_strategy + ), + mode = raw.get("mode") or fallback.mode, + embedding_model = raw.get("embedding_model") or fallback.embedding_model, + ) + + +@router.get( + "/threads/{thread_id}/settings", + response_model = ThreadRagSettings, +) +def get_thread_rag_settings( + thread_id: str, + current_subject: str = Depends(get_current_subject), +) -> ThreadRagSettings: + return _load_thread_settings(thread_id) + + +@router.put( + "/threads/{thread_id}/settings", + response_model = ThreadRagSettings, +) +def set_thread_rag_settings( + thread_id: str, + payload: UpdateThreadRagSettingsRequest, + current_subject: str = Depends(get_current_subject), +) -> ThreadRagSettings: + current = _load_thread_settings(thread_id) + new_strategy = payload.chunking_strategy or current.chunking_strategy + new_mode = payload.mode or current.mode + if payload.embedding_model is None: + new_embedder = current.embedding_model + elif payload.embedding_model.strip() == "": + new_embedder = None + else: + new_embedder = payload.embedding_model.strip() + _validate_mode_combo(new_mode, new_strategy) + + upsert_chat_settings_merge( + { + _thread_settings_key(thread_id): { + "chunking_strategy": new_strategy, + "mode": new_mode, + "embedding_model": new_embedder, + } + } + ) + return ThreadRagSettings( + chunking_strategy = new_strategy, + mode = new_mode, + embedding_model = new_embedder, + ) + + class ReingestKBRequest(BaseModel): """All fields optional — omitting one keeps the KB's current value.""" chunking_strategy: ChunkingStrategy | None = None @@ -619,20 +702,45 @@ def reingest_knowledge_base( ) def reingest_thread_documents( thread_id: str, + payload: UpdateThreadRagSettingsRequest | None = None, current_subject: str = Depends(get_current_subject), ) -> ReingestResponse: - """Rebuild a thread's RAG index using the current defaults. + """Rebuild a thread's RAG index. - No body — per-thread strategy/mode overrides aren't exposed in v1. + 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. """ - from utils.rag.config import RAG_EMBEDDING_MODEL + from utils.rag.config import resolve_embedder + if payload is None: + payload = UpdateThreadRagSettingsRequest() + if ( + payload.chunking_strategy is not None + 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, + current_subject = current_subject, + ) + else: + settings = _load_thread_settings(thread_id) + + embedder = settings.embedding_model or resolve_embedder( + settings.mode, + settings.chunking_strategy, + ) return _reingest_scope( kb_id = None, thread_id = thread_id, - chunking_strategy = "standard", - mode = "text", - embedding_model = RAG_EMBEDDING_MODEL, + chunking_strategy = settings.chunking_strategy, + mode = settings.mode, + embedding_model = embedder, ) @@ -696,12 +804,19 @@ async def upload_thread_document( file: UploadFile, current_subject: str = Depends(get_current_subject), ) -> UploadResponse: - from utils.rag.config import RAG_EMBEDDING_MODEL + 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. 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, + settings.chunking_strategy, + ) return _start_ingestion( filename = filename, stored_path = stored_path, @@ -709,7 +824,9 @@ async def upload_thread_document( content_type = file.content_type, kb_id = None, thread_id = thread_id, - embedding_model = RAG_EMBEDDING_MODEL, + embedding_model = embedder, + chunking_strategy = settings.chunking_strategy, + mode = settings.mode, ) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index f6f0100d4d..cf5fb78677 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -96,6 +96,10 @@ import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog"; import { useKnowledgeBases } from "@/features/rag/hooks/use-knowledge-bases"; import { useThreadDocuments } from "@/features/rag/hooks/use-kb-documents"; import { useRagStore } from "@/features/rag/stores/rag-store"; +import type { + ChunkingStrategy as RagChunkingStrategy, + KBMode, +} from "@/features/rag/api/rag-api"; import { Add01Icon, Delete02Icon } from "@hugeicons/core-free-icons"; function ragSourceLabel( @@ -448,6 +452,49 @@ export function ChatSettingsPanel({ ); const clearThreadIndex = useRagStore((s) => s.clearThreadIndex); const reingestThread = useRagStore((s) => s.reingestThread); + const threadSettings = useRagStore((s) => + activeThreadId ? s.threadSettings[activeThreadId] : undefined, + ); + const loadThreadSettings = useRagStore((s) => s.loadThreadSettings); + const updateThreadSettings = useRagStore((s) => s.updateThreadSettings); + const ragDefaults = useRagStore((s) => s.defaults); + + // Load this thread's RAG settings once when the sheet sees a thread + // for the first time. Updates re-render automatically via the store. + useEffect(() => { + if (ragSource.kind === "thread" && activeThreadId && !threadSettings) { + void loadThreadSettings(activeThreadId); + } + }, [ragSource.kind, activeThreadId, threadSettings, loadThreadSettings]); + + const effectiveThreadChunking: RagChunkingStrategy = + threadSettings?.chunking_strategy ?? + ragDefaults?.chunking_strategy ?? + "standard"; + const effectiveThreadMode: KBMode = + threadSettings?.mode ?? ragDefaults?.mode ?? "text"; + + const applyThreadSettingChange = ( + patch: { chunking_strategy?: RagChunkingStrategy; mode?: KBMode }, + ) => { + if (!activeThreadId) return; + if (threadDocs.length === 0) { + // No existing chunks to invalidate — just persist. + void updateThreadSettings(activeThreadId, patch); + return; + } + const ok = window.confirm( + `Re-index ${threadDocs.length} document${threadDocs.length === 1 ? "" : "s"} ` + + `with the new settings? Existing chunks will be deleted and rebuilt.`, + ); + if (ok) { + void reingestThread(activeThreadId, patch); + } else { + // User declined — refresh the store so the select snaps back + // to the unchanged settings. + void loadThreadSettings(activeThreadId); + } + }; const [kbCreateOpen, setKbCreateOpen] = useState(false); const ragEnabled = ragSource.kind !== "off"; const activeKbId = ragSource.kind === "kb" ? ragSource.kbId : null; @@ -1311,7 +1358,76 @@ export function ChatSettingsPanel({ onCreated={(kb) => setRagSource({ kind: "kb", kbId: kb.id })} /> {ragSource.kind === "thread" && activeThreadId ? ( -
    + <> +
    +
    + + +
    +
    + + +
    +
    +

    + Changing either setting will re-index this thread's existing + documents. +

    +
    @@ -1368,7 +1484,8 @@ export function ChatSettingsPanel({
    )} -
    + + ) : null}
    diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index 121d7494bb..e0ae25607d 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -232,15 +232,52 @@ export async function reingestKnowledgeBase( return parseJsonOrThrow(response); } +export interface ThreadRagSettings { + chunking_strategy: ChunkingStrategy; + mode: KBMode; + embedding_model: string | null; +} + +export interface UpdateThreadRagSettingsRequest { + chunking_strategy?: ChunkingStrategy; + mode?: KBMode; + embedding_model?: string | null; +} + +export async function getThreadRagSettings( + threadId: string, +): Promise { + const response = await authFetch( + `/api/rag/threads/${encodeURIComponent(threadId)}/settings`, + ); + return parseJsonOrThrow(response); +} + +export async function setThreadRagSettings( + threadId: string, + payload: UpdateThreadRagSettingsRequest, +): Promise { + const response = await authFetch( + `/api/rag/threads/${encodeURIComponent(threadId)}/settings`, + { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }, + ); + return parseJsonOrThrow(response); +} + export async function reingestThreadDocuments( threadId: string, + opts: UpdateThreadRagSettingsRequest = {}, ): Promise { const response = await authFetch( `/api/rag/threads/${encodeURIComponent(threadId)}/reingest`, { method: "POST", headers: { "Content-Type": "application/json" }, - body: "{}", + body: JSON.stringify(opts), }, ); return parseJsonOrThrow(response); diff --git a/studio/frontend/src/features/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index 09a740fbe7..02d6422748 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -9,6 +9,7 @@ import { deleteDocument as apiDeleteDocument, deleteKnowledgeBase as apiDeleteKB, getRagDefaults as apiGetRagDefaults, + getThreadRagSettings as apiGetThreadSettings, type JobEvent, type KnowledgeBase, listKBDocuments, @@ -21,9 +22,12 @@ import { reingestKnowledgeBase as apiReingestKB, reingestThreadDocuments as apiReingestThread, setRagDefaults as apiSetRagDefaults, + setThreadRagSettings as apiSetThreadSettings, subscribeToJobEvents, type ThreadIndexSummary, + type ThreadRagSettings, type UpdateRagDefaultsRequest, + type UpdateThreadRagSettingsRequest, uploadKBDocument, uploadThreadDocument, } from "../api/rag-api"; @@ -59,12 +63,22 @@ interface RagStoreState { clearThreadIndex: (threadId: string) => Promise; reingestKB: (kbId: string, opts?: ReingestKBOptions) => Promise; - reingestThread: (threadId: string) => Promise; + reingestThread: ( + threadId: string, + opts?: UpdateThreadRagSettingsRequest, + ) => Promise; defaults: RagDefaults | null; loadDefaults: () => Promise; updateDefaults: (patch: UpdateRagDefaultsRequest) => Promise; + threadSettings: Record; + loadThreadSettings: (threadId: string) => Promise; + updateThreadSettings: ( + threadId: string, + patch: UpdateThreadRagSettingsRequest, + ) => Promise; + subscribeJob: (jobId: string, onComplete?: () => void) => void; } @@ -92,6 +106,8 @@ export const useRagStore = create((set, get) => ({ defaults: null, + threadSettings: {}, + async loadKnowledgeBases() { set({ kbsLoading: true, kbsError: null }); try { @@ -269,10 +285,15 @@ export const useRagStore = create((set, get) => ({ set({ defaults }); }, - async reingestThread(threadId) { - const response = await apiReingestThread(threadId); + async reingestThread(threadId, opts) { + const response = await apiReingestThread(threadId, opts ?? {}); void get().loadThreadDocuments(threadId); void get().loadThreadIndexes(); + if (opts) { + // The reingest endpoint persists the new settings as a side + // effect; refresh the local copy so the UI reflects them. + void get().loadThreadSettings(threadId); + } for (const jobId of response.job_ids) { get().subscribeJob(jobId, () => { void get().loadThreadDocuments(threadId); @@ -281,6 +302,25 @@ export const useRagStore = create((set, get) => ({ return response.job_ids; }, + async loadThreadSettings(threadId) { + try { + const settings = await apiGetThreadSettings(threadId); + set((state) => ({ + threadSettings: { ...state.threadSettings, [threadId]: settings }, + })); + } catch { + // Best-effort — falls back to defaults UI-side when missing. + } + }, + + async updateThreadSettings(threadId, patch) { + const settings = await apiSetThreadSettings(threadId, patch); + set((state) => ({ + threadSettings: { ...state.threadSettings, [threadId]: settings }, + })); + return settings; + }, + subscribeJob(jobId, onComplete) { const existing = get().jobUnsubscribers[jobId]; if (existing) return; From 7e069816a165883ccb86dc5796c4a0dd059b3296 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 12:53:41 +0400 Subject: [PATCH 010/122] Studio: end-to-end multimodal RAG integration test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit @pytest.mark.server integration coverage that drives the ingestion subprocess in-process and asserts the full multimodal pipeline produces image + caption chunks with proper kind / image_path / pair_group metadata. Default pytest runs skip; explicitly: pytest -m server tests/python/test_rag_multimodal_integration.py Generates a small PDF (text + PNG figure + caption paragraph) via pymupdf, points UNSLOTH_STUDIO_HOME at tmp_path, resets the embedder singleton, and calls _subprocess_worker with mode='multimodal'. Asserts: - at least one chunk of each kind (text / image / caption) is emitted - image chunks carry a real on-disk image_path under tmp_path - image + caption chunks share a pair_group A second test confirms BGE-VL produces text and image vectors of the same dimension — sanity-check for the shared-space assumption that the multimodal retrieval path relies on. Downloads BGE-VL-base (~600 MB) on first run, so the test is gated behind the existing server marker rather than the default suite. --- .../python/test_rag_multimodal_integration.py | 173 ++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 tests/python/test_rag_multimodal_integration.py diff --git a/tests/python/test_rag_multimodal_integration.py b/tests/python/test_rag_multimodal_integration.py new file mode 100644 index 0000000000..36654b8ad8 --- /dev/null +++ b/tests/python/test_rag_multimodal_integration.py @@ -0,0 +1,173 @@ +"""End-to-end multimodal RAG integration test. + +Marked `server` so default pytest runs skip it — downloads BGE-VL-base +(~600 MB) on first run and exercises the real embedding stack. Run +explicitly with: + + ~/.unsloth/studio/unsloth_studio/bin/python -m pytest \ + tests/python/test_rag_multimodal_integration.py -v -m server + +Exercises the ingestion subprocess worker in-process (with a regular +queue rather than mp.Queue) so we cover the parse → chunk → load +embedder → encode_images → emit chunks_batch path without spinning up +a child process. The parent-side chunk insertion is covered separately +by test_rag_multimodal.py. +""" + +import os +import queue as queue_module +import sys +from io import BytesIO +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + + +@pytest.mark.server +def test_multimodal_subprocess_emits_image_and_caption_chunks( + tmp_path, + monkeypatch, +): + pymupdf = pytest.importorskip("pymupdf") + pytest.importorskip("pymupdf4llm") + pytest.importorskip("sentence_transformers") + pytest.importorskip("PIL") + pytest.importorskip("torch") + + # Use a tmp studio root so the ingest subprocess writes images + # somewhere isolated. + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base") + monkeypatch.setenv("UNSLOTH_RAG_CHUNK_SIZE", "200") + monkeypatch.setenv("UNSLOTH_RAG_CHUNK_OVERLAP", "20") + + # Reset module-level caches so the new env vars take effect. + import importlib + + import utils.rag.config as rag_config + + importlib.reload(rag_config) + from core.rag import embeddings as embeddings_module + + embeddings_module._model = None + embeddings_module._model_name = None + + # Generate a small PDF with text + one embedded image. + from PIL import Image + + img = Image.new("RGB", (96, 64), (200, 100, 50)) + img_buf = BytesIO() + img.save(img_buf, format = "PNG") + img_bytes = img_buf.getvalue() + + doc = pymupdf.open() + page = doc.new_page(width = 612, height = 792) + page.insert_text( + (72, 100), + "Architecture overview\n\nThe following diagram shows our system.", + fontsize = 11, + ) + image_rect = pymupdf.Rect(72, 200, 168, 264) + page.insert_image(image_rect, stream = img_bytes) + page.insert_text( + (72, 290), + "Figure 1: the architecture diagram described above.", + fontsize = 11, + ) + pdf_path = tmp_path / "sample.pdf" + doc.save(str(pdf_path)) + doc.close() + + # Drive the subprocess worker in-process with a regular queue. + from core.rag.ingestion import _subprocess_worker + + out_queue: "queue_module.Queue[dict]" = queue_module.Queue() + _subprocess_worker( + stored_path = str(pdf_path), + model_name = "BAAI/BGE-VL-base", + chunk_size = 200, + overlap = 20, + batch_size = 4, + out_queue = out_queue, + chunking_strategy = "standard", + mode = "multimodal", + document_id = "test-doc-1", + ) + + # Drain everything (the queue is in-process so order is stable). + events: list[dict] = [] + while not out_queue.empty(): + events.append(out_queue.get_nowait()) + + # The worker must emit at least one chunks_batch and exactly one + # terminal complete/error event. + assert any(e["type"] == "chunks_batch" for e in events) + terminals = [e for e in events if e["type"] in ("complete", "error")] + assert len(terminals) == 1, terminals + assert terminals[0]["type"] == "complete" + + # Collect all chunks across batches. + all_chunks: list[dict] = [] + for e in events: + if e["type"] == "chunks_batch": + all_chunks.extend(e["chunks"]) + + kinds = [c.get("kind") for c in all_chunks] + assert "text" in kinds, "expected at least one text chunk" + assert "image" in kinds, "expected at least one image chunk" + # The PDF has a paragraph immediately after the image, so caption + # pairing should fire. + assert "caption" in kinds, "expected at least one caption chunk" + + # Image chunks must carry a file path that exists on disk under + # the tmp studio root. + image_chunks = [c for c in all_chunks if c.get("kind") == "image"] + for chunk in image_chunks: + assert chunk.get("image_path"), chunk + path_on_disk = Path(chunk["image_path"]) + assert path_on_disk.is_file() + assert str(path_on_disk).startswith(str(tmp_path)) + + # Paired image + caption chunks share a pair_group. + pair_groups: dict[str, list[str]] = {} + for chunk in all_chunks: + group = chunk.get("pair_group") + if group: + pair_groups.setdefault(group, []).append(chunk.get("kind", "")) + paired = [ + kinds + for kinds in pair_groups.values() + if "image" in kinds and "caption" in kinds + ] + assert paired, f"expected an image/caption pair, got groups={pair_groups}" + + +@pytest.mark.server +def test_text_and_image_vectors_share_dimension(monkeypatch): + """BGE-VL is a shared-space embedder — sanity-check before relying on it.""" + pytest.importorskip("sentence_transformers") + pytest.importorskip("PIL") + pytest.importorskip("torch") + monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base") + from core.rag import embeddings as embeddings_module + + embeddings_module._model = None + embeddings_module._model_name = None + + from PIL import Image + + img = Image.new("RGB", (32, 32), (50, 150, 200)) + buf = BytesIO() + img.save(buf, format = "PNG") + + image_vectors = embeddings_module.encode_images([buf.getvalue()]) + text_vectors = embeddings_module.encode(["a blue square"]) + + assert image_vectors[0].shape == text_vectors[0].shape, ( + f"text dim {text_vectors[0].shape} != image dim {image_vectors[0].shape}" + ) From c74fc13ebcf4d21e1ef63f473d145c728505c55c Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 13:41:45 +0400 Subject: [PATCH 011/122] Studio: RAG-as-tool composer button (Phase 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Promotes RAG to a first-class composer toggle alongside Think / Web Search / Code, with tool-use semantics on local models that support tools and a pre-fetch fallback on external providers. The model decides when to call `search_knowledge_base` on local inference; on external providers retrieval still fires before each message (the existing pre-fetch path), gated on the same button. Backend - core/rag/tool.py (new): search_knowledge_base handler + JSON-schema tool spec. Resolves scope (kb_id wins over thread_id) from the request's rag_scope, runs retrieve_hybrid + optional rerank, then hydrates filename / page_number / text from sqlite and formats as numbered Markdown citations ('[1] file.pdf (page 5): ...') for the LLM to cite. Empty scope returns a user-facing hint; empty results return a clear no-match message instead of an empty string. - core/inference/tools.py: SEARCH_KNOWLEDGE_BASE_TOOL added to ALL_TOOLS (lazy import keeps tools.py importable on inference paths that never touch RAG). execute_tool() gains a tool_context parameter that carries per-request extras the LLM doesn't see (currently just rag_scope). The new 'search_knowledge_base' branch dispatches to the handler with scope unpacked from tool_context. - core/inference/llama_cpp.py + safetensors_agentic.py + orchestrator.py: thread tool_context through generate_chat_completion_ with_tools / run_safetensors_tool_loop / execute_tool. Both local backends (GGUF llama-server and safetensors agentic) carry the same context object. - models/inference.py: ChatCompletionRequest gains optional rag_scope: dict ({kb_id?, thread_id?, enable_rerank?, default_top_k?, reranker_model?}). Ignored unless 'search_knowledge_base' is in enabled_tools. - routes/inference.py: both the GGUF and safetensors call sites for generate_chat_completion_with_tools forward payload.rag_scope into tool_context. Frontend - chat-runtime-store.ts: global ragToolEnabled boolean + setter + CHAT_RAG_TOOL_ENABLED_KEY localStorage, mirroring toolsEnabled / codeToolsEnabled. Settings-hydration migration auto-flips ragToolEnabled=true for pre-Phase-4 users who already had ragSource set, so existing RAG users don't silently lose retrieval on upgrade. - shared-composer.tsx: new 'RAG' pill button after Images (uses lucide BookOpenIcon, composer-pill-btn style, data-active toggle). Disabled when no model is loaded. Toggling on from ragSource='off' auto-flips source to 'thread' so the sidebar lands ready-to-go. - chat-adapter.ts: * The existing pre-fetch block is now gated on ragToolEnabled AND only fires when the tool path isn't viable (external provider OR local model without tool-use support). Tool-capable local models skip pre-fetch and let the LLM decide. * The local-model body assembly adds 'search_knowledge_base' to enabled_tools and packs ragSource + enableRerank + ragTopK into a rag_scope object the backend tool handler consumes. - chat-settings-sheet.tsx: entire Retrieval CollapsibleSection is wrapped in {ragToolEnabled && ...} so it hides when the button is off — the button is now the single on/off control. The 'Off' option is removed from the Source dropdown (the button handles that). Default open when shown so settings are one click away. Tests - test_rag_tool_handler.py: handler covers empty query, missing scope, kb_id > thread_id precedence, thread-only path, citation formatting (numbered + page numbers + unknown source); tool spec shape (function/name/required); execute_tool dispatch with and without tool_context; ALL_TOOLS includes the new spec without dropping the existing ones. Verification scope - Local GGUF with tools: toggle button on, upload doc, ask about doc content → assistant emits a search_knowledge_base tool call card (rendered by the existing ToolFallback component since no custom UI exists yet — that's a v2 nice-to-have). - External provider (Anthropic / OpenAI / etc.): same button, same UX, but uses the pre-fetch path under the hood. - Migration: pre-existing ragSource != off → button initializes ON so retrieval keeps working. --- studio/backend/core/inference/llama_cpp.py | 2 + studio/backend/core/inference/orchestrator.py | 2 + .../core/inference/safetensors_agentic.py | 2 + studio/backend/core/inference/tools.py | 33 ++- studio/backend/core/rag/tool.py | 179 ++++++++++++++++ studio/backend/models/inference.py | 10 + studio/backend/routes/inference.py | 10 + .../src/features/chat/api/chat-adapter.ts | 41 +++- .../src/features/chat/chat-settings-sheet.tsx | 10 +- .../src/features/chat/shared-composer.tsx | 32 ++- .../chat/stores/chat-runtime-store.ts | 27 +++ tests/python/test_rag_tool_handler.py | 193 ++++++++++++++++++ 12 files changed, 531 insertions(+), 10 deletions(-) create mode 100644 studio/backend/core/rag/tool.py create mode 100644 tests/python/test_rag_tool_handler.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bf8a3c04df..ba438d9ea6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4428,6 +4428,7 @@ class LlamaCppBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + tool_context: Optional[dict] = None, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -5080,6 +5081,7 @@ class LlamaCppBackend: cancel_event = cancel_event, timeout = _effective_timeout, session_id = session_id, + tool_context = tool_context, ) yield { diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 7e7d7026f6..4cfd87f1a1 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -838,6 +838,7 @@ class InferenceOrchestrator: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + tool_context: Optional[dict] = None, use_adapter: Optional[Union[bool, str]] = None, **_unused, ): @@ -895,6 +896,7 @@ class InferenceOrchestrator: max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, + tool_context = tool_context, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 73bb3d090a..edfae5ce16 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -105,6 +105,7 @@ def run_safetensors_tool_loop( max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, + tool_context: Optional[dict] = None, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -340,6 +341,7 @@ def run_safetensors_tool_loop( cancel_event = cancel_event, timeout = eff_timeout, session_id = session_id, + tool_context = tool_context, ) except Exception as exc: logger.exception("Tool %s raised: %s", tool_name, exc) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0e9cce7c3e..d378e861d1 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -502,7 +502,20 @@ TERMINAL_TOOL = { }, } -ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, 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. +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()] _TIMEOUT_UNSET = object() @@ -514,12 +527,17 @@ def execute_tool( cancel_event = None, timeout: int | None = _TIMEOUT_UNSET, session_id: str | None = None, + tool_context: dict | None = None, ) -> str: """Execute a tool by name with the given arguments. Returns result as a string. ``timeout``: int sets per-call limit in seconds, ``None`` means no limit, unset (default) uses ``_EXEC_TIMEOUT`` (300 s). ``session_id``: optional thread/session ID for per-conversation sandbox isolation. + ``tool_context``: optional per-request extras the LLM does not see (RAG scope, + future per-tool overrides). Keys consumed: + - ``rag_scope``: ``{kb_id?, thread_id?, enable_rerank?, default_top_k?, + reranker_model?}`` — consumed by ``search_knowledge_base``. """ logger.info( f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}" @@ -539,6 +557,19 @@ def execute_tool( return _bash_exec( arguments.get("command", ""), cancel_event, effective_timeout, session_id ) + if name == "search_knowledge_base": + from core.rag.tool import search_knowledge_base + + scope = (tool_context or {}).get("rag_scope") or {} + return search_knowledge_base( + query = arguments.get("query", ""), + top_k = arguments.get("top_k"), + scope_kb_id = scope.get("kb_id"), + scope_thread_id = scope.get("thread_id"), + enable_rerank = bool(scope.get("enable_rerank")), + reranker_model = scope.get("reranker_model"), + default_top_k = int(scope.get("default_top_k") or 5), + ) return f"Unknown tool: {name}" diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py new file mode 100644 index 0000000000..83a2599857 --- /dev/null +++ b/studio/backend/core/rag/tool.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`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. +""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +SEARCH_KNOWLEDGE_BASE_TOOL = { + "type": "function", + "function": { + "name": "search_knowledge_base", + "description": ( + "Search the user's attached documents for information relevant to " + "the user's question. Call this when the user references content " + "from their docs, asks fact-heavy questions, or needs grounded " + "citations. Returns numbered chunks with source filenames; cite " + "them in your reply as [1], [2], etc." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": ( + "A focused search query — phrase it as the question " + "you want answered, not as a keyword list." + ), + }, + "top_k": { + "type": "integer", + "minimum": 1, + "maximum": 20, + "description": ( + "How many chunks to retrieve (default 5). Higher = " + "more grounding, more tokens." + ), + }, + }, + "required": ["query"], + }, + }, +} + + +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. + """ + if not hits: + return ( + "No matching chunks were found in the attached documents. " + "Either nothing in this scope is relevant, or no documents " + "have been ingested yet." + ) + lines: list[str] = [] + for index, hit in enumerate(hits, start = 1): + name = hit.get("filename") or "unknown source" + page = hit.get("page_number") + suffix = f" (page {page})" if page is not None else "" + text = (hit.get("text") or "").strip() + lines.append(f"[{index}] {name}{suffix}: {text}") + return "\n\n".join(lines) + + +def search_knowledge_base( + *, + query: str, + top_k: int | None = None, + scope_kb_id: str | None = None, + scope_thread_id: str | None = None, + enable_rerank: bool = False, + reranker_model: str | None = None, + default_top_k: int = 5, +) -> 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. + """ + if not query or not query.strip(): + return "Error: empty query." + + if not scope_kb_id and not scope_thread_id: + return ( + "No knowledge base or thread documents are configured for " + "retrieval. Ask the user to upload a document or select a " + "knowledge base in the chat settings." + ) + + from core.rag import retrieval + from core.rag.vector_store import kb_scope, thread_scope + from storage.studio_db import get_connection + + scope = ( + kb_scope(scope_kb_id) if scope_kb_id else thread_scope(scope_thread_id) + ) + k = top_k if top_k is not None else default_top_k + + if enable_rerank: + from utils.rag.config import RAG_RERANK_CANDIDATE_K + + candidate_k = max(k, RAG_RERANK_CANDIDATE_K) + else: + candidate_k = k + + try: + hits = retrieval.retrieve_hybrid(scope, query.strip(), k = candidate_k) + except Exception as exc: # noqa: BLE001 + logger.exception("search_knowledge_base retrieval failed") + return f"Error: retrieval failed ({type(exc).__name__})." + + chunk_ids = [h.chunk_id for h in hits] + lookup: dict[str, dict] = {} + if chunk_ids: + placeholders = ",".join("?" for _ in chunk_ids) + with get_connection() as conn: + rows = conn.execute( + f""" + SELECT c.id AS chunk_id, c.text, c.page_number, + c.kind, d.filename + FROM rag_chunks c + JOIN rag_documents d ON d.id = c.document_id + WHERE c.id IN ({placeholders}) + """, + chunk_ids, + ).fetchall() + for row in rows: + lookup[row["chunk_id"]] = dict(row) + + if enable_rerank and hits: + from core.rag import reranker + + pairs = [ + (hit, lookup[hit.chunk_id]["text"]) + for hit in hits + if hit.chunk_id in lookup + ] + try: + hits = reranker.rerank( + query.strip(), + pairs, + model_name = reranker_model, + top_k = k, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("rerank failed in search_knowledge_base: %s", exc) + hits = hits[:k] + else: + hits = hits[:k] + + # Image-kind hits don't carry LLM-friendly text — skip them. The + # paired caption (linked_chunk_id) usually surfaces separately. + formatted = [ + lookup[hit.chunk_id] + for hit in hits + if hit.chunk_id in lookup and lookup[hit.chunk_id].get("kind") != "image" + ] + return _format_hits_for_llm(formatted) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b5626951c4..091dba4f8d 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -686,6 +686,16 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.", ) + rag_scope: Optional[dict] = Field( + None, + description = ( + "[x-unsloth] Per-request context the `search_knowledge_base` tool " + "consumes when the LLM invokes it. Shape: " + "{kb_id?: str, thread_id?: str, enable_rerank?: bool, " + "default_top_k?: int, reranker_model?: str}. Ignored unless " + "'search_knowledge_base' is in enabled_tools." + ), + ) cancel_id: Optional[str] = Field( None, description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 02270ab405..7c2a4e06bb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2483,6 +2483,11 @@ async def openai_chat_completions( if payload.tool_call_timeout is not None else 300, session_id = payload.session_id, + tool_context = ( + {"rag_scope": payload.rag_scope} + if payload.rag_scope + else None + ), ) _tool_sentinel = object() @@ -2950,6 +2955,11 @@ async def openai_chat_completions( def sf_generate_with_tools(): return backend.generate_chat_completion_with_tools( + tool_context = ( + {"rag_scope": payload.rag_scope} + if payload.rag_scope + else None + ), messages = _sf_chat_messages, tools = _sf_tools_to_use, system_prompt = _sf_system_prompt or "", diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 3928f79631..aa134d8b76 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -993,8 +993,22 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // 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. + // + // Phase 4: pre-fetch only fires when the new RAG button is on AND + // we can't register `search_knowledge_base` as a real tool — + // i.e., external providers or local models that don't expose + // tool-use. Local models with tool support take the tool-call + // path further down (see `enabled_tools` assembly), and the LLM + // decides per turn whether to invoke retrieval. const ragSource = runtime.ragSource; - if (ragSource.kind !== "off") { + const ragToolEnabled = runtime.ragToolEnabled; + const ragToolPathTaken = + ragToolEnabled && supportsTools && !isExternalRequest; + if ( + ragToolEnabled + && ragSource.kind !== "off" + && !ragToolPathTaken + ) { const lastUser = [...outboundMessages] .reverse() .find((m) => m.role === "user"); @@ -1617,13 +1631,36 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(supportsPreserveThinking ? { preserve_thinking: preserveThinking } : {}), - ...(supportsTools && (toolsEnabled || codeToolsEnabled) + ...(supportsTools + && (toolsEnabled || codeToolsEnabled || ragToolPathTaken) ? { enable_tools: true, enabled_tools: [ ...(toolsEnabled ? ["web_search"] : []), ...(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. + ...(ragToolPathTaken + ? { + rag_scope: { + kb_id: + ragSource.kind === "kb" + ? ragSource.kbId + : null, + thread_id: + ragSource.kind === "thread" + ? (resolvedThreadId ?? null) + : null, + enable_rerank: runtime.enableRerank, + default_top_k: runtime.ragTopK, + }, + } + : {}), auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, max_tool_calls_per_message: diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index cf5fb78677..dc57de469f 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -441,6 +441,7 @@ export function ChatSettingsPanel({ const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; const ragSource = useChatRuntimeStore((s) => s.ragSource); const setRagSource = useChatRuntimeStore((s) => s.setRagSource); + const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled); const enableRerank = useChatRuntimeStore((s) => s.enableRerank); const setEnableRerank = useChatRuntimeStore((s) => s.setEnableRerank); const ragTopK = useChatRuntimeStore((s) => s.ragTopK); @@ -1246,7 +1247,8 @@ export function ChatSettingsPanel({ ) : null} - + {ragToolEnabled ? ( +
    + ) : null} )} + {/* 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. */} +
    {dictationSupported && ( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index e7ce32e37d..d1f84c42ae 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -26,6 +26,7 @@ export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled"; +export const CHAT_RAG_TOOL_ENABLED_KEY = "unsloth_chat_rag_tool_enabled"; // External provider selection is encoded into `params.checkpoint` as // `external::::`. PersistedChatSettings deliberately @@ -264,6 +265,7 @@ type ChatRuntimeStore = { */ supportsBuiltinImageGeneration: boolean; toolsEnabled: boolean; + ragToolEnabled: boolean; codeToolsEnabled: boolean; imageToolsEnabled: boolean; toolStatus: string | null; @@ -344,6 +346,7 @@ type ChatRuntimeStore = { setRagSource: (source: RagSource) => void; setEnableRerank: (value: boolean) => void; setRagTopK: (value: number) => void; + setRagToolEnabled: (value: boolean) => void; }; type PersistedChatSettings = Awaited< @@ -578,6 +581,10 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsBuiltinCodeExecution: false, supportsBuiltinImageGeneration: false, toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false), + // Phase 4: RAG button defaults off. Migration nudge happens after + // settings hydration, when the persisted ragSource becomes visible — + // see hydratePersistedSettings. + ragToolEnabled: loadBool(CHAT_RAG_TOOL_ENABLED_KEY, false), codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false), imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), toolStatus: null, @@ -630,6 +637,21 @@ export const useChatRuntimeStore = create((set, get) => ({ ), ...getHydratedSettingsState(settings, state, hydrationVersions), }; + // Phase 4 migration: pre-existing users with ragSource set + // before the RAG button shipped should keep getting RAG — + // auto-flip ragToolEnabled so the button starts ON for them. + // The CHAT_RAG_TOOL_ENABLED_KEY localStorage write makes the + // migration stick across reloads. + const hydratedRagSource = + (nextState.ragSource as RagSource | undefined) ?? state.ragSource; + if ( + hydratedRagSource && + hydratedRagSource.kind !== "off" && + !state.ragToolEnabled + ) { + nextState.ragToolEnabled = true; + saveBool(CHAT_RAG_TOOL_ENABLED_KEY, true); + } return nextState; }); } catch { @@ -831,6 +853,11 @@ export const useChatRuntimeStore = create((set, get) => ({ } return { toolsEnabled }; }), + setRagToolEnabled: (ragToolEnabled) => + set(() => { + saveBool(CHAT_RAG_TOOL_ENABLED_KEY, ragToolEnabled); + return { ragToolEnabled }; + }), setCodeToolsEnabled: (codeToolsEnabled) => set(() => { saveBool(CHAT_CODE_TOOLS_ENABLED_KEY, codeToolsEnabled); diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py new file mode 100644 index 0000000000..552719ada8 --- /dev/null +++ b/tests/python/test_rag_tool_handler.py @@ -0,0 +1,193 @@ +"""Unit tests for the `search_knowledge_base` tool handler (Phase 4).""" + +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" +if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + + +def _make_hit(chunk_id: str): + """Minimal stand-in for retrieval.Hit — just needs .chunk_id.""" + class _Hit: + pass + h = _Hit() + h.chunk_id = chunk_id + h.score = 1.0 + h.kind = "text" + h.document_id = None + h.chunk_index = 0 + return h + + +def test_empty_query_returns_error(): + from core.rag.tool import search_knowledge_base + + result = search_knowledge_base(query = "", scope_thread_id = "t-1") + assert result.startswith("Error:") + assert "empty" in result.lower() + + +def test_missing_scope_returns_user_facing_hint(): + from core.rag.tool import search_knowledge_base + + result = search_knowledge_base( + query = "anything", + scope_kb_id = None, + scope_thread_id = None, + ) + assert "No knowledge base" in result + assert "thread documents" in result + + +def test_kb_takes_precedence_over_thread(): + """When both kb_id and thread_id are passed, kb_id wins.""" + from core.rag import tool + + captured = {} + + def _stub_retrieve(scope, query, k): + captured["scope"] = scope + return [] + + with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]), + "retrieve_hybrid", + _stub_retrieve): + result = tool.search_knowledge_base( + query = "x", + scope_kb_id = "kb-abc", + scope_thread_id = "thread-xyz", + ) + + assert captured["scope"].startswith("kb_") + assert "kb-abc" in captured["scope"] + assert "thread" not in captured["scope"].split("kb_")[1] + + +def test_thread_scope_when_only_thread_set(): + from core.rag import tool + + captured = {} + + def _stub_retrieve(scope, query, k): + captured["scope"] = scope + return [] + + with patch.object(tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]), + "retrieve_hybrid", + _stub_retrieve): + tool.search_knowledge_base( + query = "x", + scope_thread_id = "thread-xyz", + ) + + assert captured["scope"].startswith("thread_") + + +def test_empty_results_message_is_user_facing(): + from core.rag.tool import _format_hits_for_llm + + result = _format_hits_for_llm([]) + assert "No matching chunks" in result + + +def test_format_hits_produces_numbered_citations(): + from core.rag.tool import _format_hits_for_llm + + hits = [ + {"filename": "alpha.pdf", "page_number": 3, "text": "first body"}, + {"filename": "beta.md", "page_number": None, "text": "second body"}, + ] + result = _format_hits_for_llm(hits) + assert "[1] alpha.pdf (page 3): first body" in result + assert "[2] beta.md: second body" in result + # Each hit on its own paragraph so the LLM can cite cleanly. + assert "\n\n" in result + + +def test_format_hits_handles_unknown_source(): + from core.rag.tool import _format_hits_for_llm + + hits = [{"filename": None, "page_number": None, "text": "orphan"}] + result = _format_hits_for_llm(hits) + assert "[1] unknown source: orphan" in result + + +def test_tool_spec_shape_is_openai_compatible(): + from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL + + assert SEARCH_KNOWLEDGE_BASE_TOOL["type"] == "function" + fn = SEARCH_KNOWLEDGE_BASE_TOOL["function"] + assert fn["name"] == "search_knowledge_base" + assert "query" in fn["parameters"]["required"] + assert "top_k" in fn["parameters"]["properties"] + # Description should hint at when to call so the LLM picks it up + # appropriately. Don't lock the exact wording. + assert "documents" in fn["description"].lower() + + +def test_execute_tool_dispatches_to_search_knowledge_base(): + """tools.execute_tool should route 'search_knowledge_base' correctly.""" + from core.inference import tools + + called = {} + + def _stub(*, query, top_k = None, scope_kb_id = None, scope_thread_id = None, + enable_rerank = False, reranker_model = None, default_top_k = 5): + called["query"] = query + called["top_k"] = top_k + called["scope_kb_id"] = scope_kb_id + called["scope_thread_id"] = scope_thread_id + called["enable_rerank"] = enable_rerank + called["default_top_k"] = default_top_k + return "stub-result" + + with patch("core.rag.tool.search_knowledge_base", _stub): + result = tools.execute_tool( + "search_knowledge_base", + {"query": "hello", "top_k": 7}, + tool_context = { + "rag_scope": { + "kb_id": "kb-1", + "enable_rerank": True, + "default_top_k": 3, + } + }, + ) + assert result == "stub-result" + assert called["query"] == "hello" + assert called["top_k"] == 7 + assert called["scope_kb_id"] == "kb-1" + assert called["scope_thread_id"] is None + assert called["enable_rerank"] is True + assert called["default_top_k"] == 3 + + +def test_execute_tool_handles_missing_tool_context(): + """tool_context=None should still dispatch without crashing.""" + from core.inference import tools + + def _stub(*, query, **_kwargs): + return f"got: {query}" + + with patch("core.rag.tool.search_knowledge_base", _stub): + result = tools.execute_tool( + "search_knowledge_base", + {"query": "ping"}, + tool_context = None, + ) + assert result == "got: ping" + + +def test_all_tools_includes_rag(): + from core.inference.tools import ALL_TOOLS + + names = [t["function"]["name"] for t in ALL_TOOLS] + assert "search_knowledge_base" in names + assert "web_search" in names # regression — we shouldn't have removed the others + assert "python" in names From a3ad6015bc3c6fd8a840da4789b6ca6a706a1afc Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 13:47:34 +0400 Subject: [PATCH 012/122] Studio: fix tsc errors in Phase 4 frontend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two errors surfaced by the frontend build (`tsc -b`) after the Phase 4 commit c74fc13eb landed: src/features/chat/chat-settings-sheet.tsx(449,9): TS2451 — cannot redeclare block-scoped 'activeThreadId'. src/features/rag/stores/rag-store.ts(326,9): TS2774 — condition will always return true since this function is always defined. Fixes - chat-settings-sheet.tsx: an `activeThreadId` declaration already existed near the code-exec section (line 636). Phase 2B's RAG retrieval-section block introduced a second declaration at line 449. The earlier one is needed for the Retrieval block; drop the later redeclaration — downstream code still resolves it via lexical scope. - rag-store.ts subscribeJob: `get().jobUnsubscribers[jobId]` indexes a `Record void>`. Without `noUncheckedIndexedAccess` TS infers the result as the function type (never undefined), so `if (existing)` is always-truthy. Replace with `if (jobId in get().jobUnsubscribers) return;` — same semantics, satisfies TS. --- studio/frontend/src/features/chat/chat-settings-sheet.tsx | 3 ++- studio/frontend/src/features/rag/stores/rag-store.ts | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index dc57de469f..f954670c42 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -633,7 +633,8 @@ export function ChatSettingsPanel({ activeExternalProvider.baseUrl, ) && activeExternalProvider.providerType === "openai"; - const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + // (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/rag/stores/rag-store.ts b/studio/frontend/src/features/rag/stores/rag-store.ts index 02d6422748..a0434c448f 100644 --- a/studio/frontend/src/features/rag/stores/rag-store.ts +++ b/studio/frontend/src/features/rag/stores/rag-store.ts @@ -322,8 +322,10 @@ export const useRagStore = create((set, get) => ({ }, subscribeJob(jobId, onComplete) { - const existing = get().jobUnsubscribers[jobId]; - if (existing) return; + // Object indexing in TS returns V (not V|undefined) without + // noUncheckedIndexedAccess, so we test membership explicitly to + // avoid the "always-truthy function reference" lint. + if (jobId in get().jobUnsubscribers) return; const unsubscribe = subscribeToJobEvents(jobId, { onEvent: (event) => { set((state) => ({ jobs: { ...state.jobs, [jobId]: event } })); From 810b2e27f3f5d81087e27bd9ebd0e42c75def3b9 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 14:02:10 +0400 Subject: [PATCH 013/122] Studio: fix React #185 update-depth loop in RAG additions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two useEffects added in Phase 2C and Phase 4 followed the anti-pattern of calling a state setter inside the effect and listing the setter's output in the dep array. Both could re-fire indefinitely when the selected slice changed shape on each render — which the chat page hit on first load. chat-settings-sheet.tsx (thread settings loader) - Before: `useEffect(load, [..., threadSettings, ...])` with `if (!threadSettings) load()` inside. After load, the Zustand selector returned a freshly-constructed slice, dep changed, effect re-ran. If anything in between caused threadSettings to briefly flicker undefined (e.g. a race during initial hydration or a fast subsequent thread switch), the load fired again and the cycle repeated. - After: ref-guarded by activeThreadId — `threadSettingsLoadedRef` tracks which threadId has been loaded; the effect deps shrink to `[ragSource.kind, activeThreadId, loadThreadSettings]`, all stable per-thread, removing the feedback loop. ingestion-toast-stack.tsx (terminal-job auto-dismiss) - Before: `useEffect(..., [jobs, dismissedJobs])` with `setDismissedJobs(prev => new Set(prev).add(jobId))` inside the scheduled setTimeout. Each setter creates a new Set reference; the dep change re-triggers the effect, which clears and reschedules timers. Under fast SSE event arrival or a strict-mode double-mount, the scheduler runs faster than its cleanup and React caps the depth. - After: dismissedJobs is read via a ref (kept in sync at the top of the component); the effect only depends on `[jobs]`. A `scheduledJobsRef` prevents duplicate timer scheduling for the same job across multiple effect runs, and the setDismissedJobs updater no-ops when the job is already dismissed. No behavior change for the happy path — toasts still auto-dismiss after DISMISS_DELAY_MS; thread settings still load on first sight of a thread. --- .../src/features/chat/chat-settings-sheet.tsx | 16 ++++++++++++---- .../rag/components/ingestion-toast-stack.tsx | 19 +++++++++++++++---- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index f954670c42..439b6088c7 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -460,13 +460,21 @@ export function ChatSettingsPanel({ const updateThreadSettings = useRagStore((s) => s.updateThreadSettings); const ragDefaults = useRagStore((s) => s.defaults); - // Load this thread's RAG settings once when the sheet sees a thread - // for the first time. Updates re-render automatically via the store. + // 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. + const threadSettingsLoadedRef = useRef(null); useEffect(() => { - if (ragSource.kind === "thread" && activeThreadId && !threadSettings) { + if ( + ragSource.kind === "thread" + && activeThreadId + && threadSettingsLoadedRef.current !== activeThreadId + ) { + threadSettingsLoadedRef.current = activeThreadId; void loadThreadSettings(activeThreadId); } - }, [ragSource.kind, activeThreadId, threadSettings, loadThreadSettings]); + }, [ragSource.kind, activeThreadId, loadThreadSettings]); const effectiveThreadChunking: RagChunkingStrategy = threadSettings?.chunking_strategy ?? diff --git a/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx b/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx index ad824c7111..7e35180280 100644 --- a/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx +++ b/studio/frontend/src/features/rag/components/ingestion-toast-stack.tsx @@ -5,7 +5,7 @@ import { Button } from "@/components/ui/button"; import { Cancel01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion, useReducedMotion } from "motion/react"; -import { useEffect, useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { useRagStore } from "../stores/rag-store"; import { IngestionProgress } from "./ingestion-progress"; @@ -33,16 +33,27 @@ export function IngestionToastStack() { ); // Schedule auto-dismiss for jobs that have reached a terminal state. + // Ref-tracked so `dismissedJobs` isn't a useEffect dep — the setter + // fires *inside* the effect, and depending on its output here is a + // recipe for update-depth loops if the scheduler ever runs faster + // than the cleanup. We snapshot the latest dismissed set into a ref + // and read from it inside the scheduling loop instead. + const scheduledJobsRef = useRef>(new Set()); + const dismissedJobsRef = useRef>(dismissedJobs); + dismissedJobsRef.current = dismissedJobs; useEffect(() => { const timers: ReturnType[] = []; for (const [jobId, event] of Object.entries(jobs)) { if ( - (event.type === "complete" || event.type === "error") && - !dismissedJobs.has(jobId) + (event.type === "complete" || event.type === "error") + && !dismissedJobsRef.current.has(jobId) + && !scheduledJobsRef.current.has(jobId) ) { + scheduledJobsRef.current.add(jobId); timers.push( setTimeout(() => { setDismissedJobs((prev) => { + if (prev.has(jobId)) return prev; const next = new Set(prev); next.add(jobId); return next; @@ -52,7 +63,7 @@ export function IngestionToastStack() { } } return () => timers.forEach(clearTimeout); - }, [jobs, dismissedJobs]); + }, [jobs]); const visible = Object.entries(jobs).filter( ([jobId]) => !dismissedJobs.has(jobId), From 5edbc99916eb8da85e2b24019fc69a34ab2774c5 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 15:24:48 +0400 Subject: [PATCH 014/122] Studio: stable empty-array sentinel in RAG document selectors Fresh [] literals from useRagStore selectors triggered Zustand's Object.is snapshot check on every render, causing React #185 on the chat page when documentsByScope[key] was unpopulated. --- .../src/features/rag/hooks/use-kb-documents.ts | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/features/rag/hooks/use-kb-documents.ts b/studio/frontend/src/features/rag/hooks/use-kb-documents.ts index 01704a4555..84919921fb 100644 --- a/studio/frontend/src/features/rag/hooks/use-kb-documents.ts +++ b/studio/frontend/src/features/rag/hooks/use-kb-documents.ts @@ -2,12 +2,20 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useEffect } from "react"; +import type { RagDocument } from "../api/rag-api"; import { kbScopeKey, threadScopeKey, useRagStore } from "../stores/rag-store"; +// Module-scope sentinel so the selector returns a stable reference +// when the scope key isn't populated yet. A `[]` literal in the +// selector returns a new array on every call → Zustand's Object.is +// snapshot check flags it as changed → re-render → selector reruns +// → new `[]` → infinite loop → React error #185. +const EMPTY_DOCS: RagDocument[] = []; + export function useKBDocuments(kbId: string | null) { const scopeKey = kbId ? kbScopeKey(kbId) : ""; const documents = useRagStore((s) => - scopeKey ? (s.documentsByScope[scopeKey] ?? []) : [], + scopeKey ? (s.documentsByScope[scopeKey] ?? EMPTY_DOCS) : EMPTY_DOCS, ); const loading = useRagStore((s) => (scopeKey ? !!s.docsLoading[scopeKey] : false)); const error = useRagStore((s) => @@ -38,7 +46,7 @@ export function useKBDocuments(kbId: string | null) { export function useThreadDocuments(threadId: string | null) { const scopeKey = threadId ? threadScopeKey(threadId) : ""; const documents = useRagStore((s) => - scopeKey ? (s.documentsByScope[scopeKey] ?? []) : [], + scopeKey ? (s.documentsByScope[scopeKey] ?? EMPTY_DOCS) : EMPTY_DOCS, ); const loading = useRagStore((s) => (scopeKey ? !!s.docsLoading[scopeKey] : false)); const error = useRagStore((s) => From 510509318ae8320dcd79cbc288e777a7fb8ee69b Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 15:59:57 +0400 Subject: [PATCH 015/122] Studio: mirror RAG pill in in-thread composer The Phase 4 RAG button only existed in shared-composer; the assistant-ui in-thread composer renders its own pill set, so RAG was missing once a thread had messages. --- .../src/components/assistant-ui/thread.tsx | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 431e568205..597ce7717c 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -64,6 +64,7 @@ import { flushResourcesSync } from "@assistant-ui/tap"; import { ArrowDownIcon, ArrowUpIcon, + BookOpenIcon, ChevronLeftIcon, ChevronRightIcon, DownloadIcon, @@ -908,6 +909,44 @@ 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. +const RagToggle: FC = () => { + const modelLoaded = useChatRuntimeStore( + (s) => !!s.params.checkpoint && !s.modelLoading, + ); + const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled); + const setRagToolEnabled = useChatRuntimeStore((s) => s.setRagToolEnabled); + const ragSource = useChatRuntimeStore((s) => s.ragSource); + const setRagSource = useChatRuntimeStore((s) => s.setRagSource); + const disabled = !modelLoaded; + return ( + + ); +}; + const ToolStatusDisplay: FC = () => { const toolStatus = useChatRuntimeStore((s) => s.toolStatus); const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); @@ -980,6 +1019,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({ +
    From 30856de7390ba9421eebeba0c834c927d0026994 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 16:13:05 +0400 Subject: [PATCH 016/122] Studio: route in-thread doc uploads through RAG ingest When ragToolEnabled is on, the in-thread composer's + button now opens a doc-only picker and uploads selected files to the per-thread RAG pipeline (matching SharedComposer). Pending chips show Uploading -> Indexing -> Ready status; Send is blocked while any doc is in-flight. Previously these files were base64'd inline as native attachments, which bypassed RAG entirely. --- .../src/components/assistant-ui/thread.tsx | 99 ++++++++++-- .../chat/components/pending-doc-chips.tsx | 62 ++++++++ .../chat/hooks/use-thread-doc-uploads.ts | 146 ++++++++++++++++++ 3 files changed, 297 insertions(+), 10 deletions(-) create mode 100644 studio/frontend/src/features/chat/components/pending-doc-chips.tsx create mode 100644 studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 597ce7717c..440a6817c7 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -23,6 +23,12 @@ import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python"; import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal"; import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { PendingDocChips } from "@/features/chat/components/pending-doc-chips"; +import { + DOCUMENT_ACCEPT, + isDocumentFile, + useThreadDocUploads, +} from "@/features/chat/hooks/use-thread-doc-uploads"; import { IntentAwareScrollProvider, useIntentAwareAutoScroll, @@ -71,6 +77,7 @@ import { GlobeIcon, HeadphonesIcon, ImageIcon, + PaperclipIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, @@ -306,8 +313,17 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { ), ); const hasPendingAudio = useChatRuntimeStore((s) => Boolean(s.pendingAudioName)); + const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled); + const { pendingDocs, addDoc, removeDoc, clearDocs, isIndexing } = + useThreadDocUploads(); const hasSendableContent = composerText.trim().length > 0 || hasAttachments || hasPendingAudio; + const sendBlocked = + disabled || + !hasSendableContent || + isComposing || + hasPendingAttachments || + isIndexing; const handleSubmit = useCallback( (event: FormEvent) => { @@ -315,18 +331,31 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { disabled || !hasSendableContent || isComposingRef.current || - hasPendingAttachments + hasPendingAttachments || + isIndexing ) { event.preventDefault(); + return; } + // Send is going through — drop the chips. The docs themselves + // remain in the backend thread KB and stay searchable. + clearDocs(); }, - [disabled, hasPendingAttachments, hasSendableContent, isComposingRef], + [ + disabled, + hasPendingAttachments, + hasSendableContent, + isComposingRef, + isIndexing, + clearDocs, + ], ); const composerContent = ( <> + = ({ disabled }) => { {...inputProps} /> - !hasSendableContent || isComposingRef.current || hasPendingAttachments + !hasSendableContent || + isComposingRef.current || + hasPendingAttachments || + isIndexing } + ragModeOn={ragToolEnabled} + onAddDoc={addDoc} /> ); @@ -1005,14 +1037,61 @@ const ToolStatusDisplay: FC = () => { ); }; -const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({ - disabled, - blockSend, +// 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. +const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({ + onSelect, }) => { + const inputRef = useRef(null); + return ( + <> + { + const files = e.target.files; + if (!files) return; + for (let i = 0; i < files.length; i++) { + const f = files[i]; + if (f && isDocumentFile(f)) onSelect(f); + } + // Reset so re-selecting the same file fires onChange. + e.target.value = ""; + }} + /> + inputRef.current?.click()} + > + + + + ); +}; + +const ComposerAction: FC<{ + disabled?: boolean; + blockSend?: () => boolean; + ragModeOn?: boolean; + onAddDoc?: (file: File) => void; +}> = ({ disabled, blockSend, ragModeOn, onAddDoc }) => { return (
    - + {ragModeOn && onAddDoc ? ( + + ) : ( + + )} diff --git a/studio/frontend/src/features/chat/components/pending-doc-chips.tsx b/studio/frontend/src/features/chat/components/pending-doc-chips.tsx new file mode 100644 index 0000000000..d65d02c129 --- /dev/null +++ b/studio/frontend/src/features/chat/components/pending-doc-chips.tsx @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { FileTextIcon, XIcon } from "lucide-react"; +import type { FC } from "react"; + +import { cn } from "@/lib/utils"; +import type { PendingDoc } from "../hooks/use-thread-doc-uploads"; + +interface PendingDocChipsProps { + docs: PendingDoc[]; + 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 ( +
    + {docs.map((doc) => { + const statusLabel = + doc.status === "uploading" + ? "Uploading…" + : doc.status === "ingesting" + ? "Indexing…" + : doc.status === "error" + ? (doc.errorMessage ?? "Failed") + : "Ready"; + const statusClass = + doc.status === "error" + ? "text-destructive" + : doc.status === "ready" + ? "text-muted-foreground" + : "text-muted-foreground italic"; + return ( +
    + +
    + {doc.file.name} + + {statusLabel} + +
    + +
    + ); + })} +
    + ); +}; 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 new file mode 100644 index 0000000000..5e77cbbc0d --- /dev/null +++ b/studio/frontend/src/features/chat/hooks/use-thread-doc-uploads.ts @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useCallback, useState } from "react"; +import { toast } from "sonner"; + +import { subscribeToJobEvents } from "@/features/rag/api/rag-api"; +import { useRagStore } from "@/features/rag/stores/rag-store"; +import { useChatRuntimeStore } from "../stores/chat-runtime-store"; + +export type PendingDoc = { + id: string; + file: File; + status: "uploading" | "ingesting" | "ready" | "error"; + jobId?: string; + documentId?: string; + errorMessage?: string; +}; + +const DOCUMENT_EXTENSIONS = new Set([ + ".pdf", + ".txt", + ".md", + ".markdown", + ".docx", + ".html", + ".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"; + +export function isDocumentFile(file: File): boolean { + const lower = file.name.toLowerCase(); + const dot = lower.lastIndexOf("."); + if (dot < 0) return false; + return DOCUMENT_EXTENSIONS.has(lower.slice(dot)); +} + +export interface UseThreadDocUploadsResult { + pendingDocs: PendingDoc[]; + addDoc: (file: File) => void; + removeDoc: (id: string) => void; + clearDocs: () => void; + 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. +export function useThreadDocUploads(): UseThreadDocUploadsResult { + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const [pendingDocs, setPendingDocs] = useState([]); + + const addDoc = useCallback( + (file: File) => { + if (!activeThreadId) { + toast.error("Attach to a thread first"); + return; + } + const id = crypto.randomUUID(); + setPendingDocs((prev) => [...prev, { id, file, status: "uploading" }]); + const uploadDocument = useRagStore.getState().uploadDocument; + uploadDocument({ kind: "thread", threadId: activeThreadId }, file) + .then(({ documentId, jobId }) => { + setPendingDocs((prev) => + prev.map((d) => + d.id === id + ? { ...d, status: "ingesting", jobId, documentId } + : d, + ), + ); + subscribeToJobEvents(jobId, { + onEvent: (event) => { + if (event.type === "complete") { + setPendingDocs((prev) => + prev.map((d) => + 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. + if (useChatRuntimeStore.getState().ragSource.kind === "off") { + useChatRuntimeStore + .getState() + .setRagSource({ kind: "thread" }); + } + } else if (event.type === "error") { + setPendingDocs((prev) => + prev.map((d) => + d.id === id + ? { ...d, status: "error", errorMessage: event.error } + : d, + ), + ); + } + }, + }); + }) + .catch((err: unknown) => { + const message = err instanceof Error ? err.message : "Upload failed"; + setPendingDocs((prev) => + prev.map((d) => + d.id === id + ? { ...d, status: "error", errorMessage: message } + : d, + ), + ); + toast.error(`Document upload failed: ${message}`); + }); + }, + [activeThreadId], + ); + + const removeDoc = useCallback( + (id: string) => { + setPendingDocs((prev) => { + const doc = prev.find((d) => d.id === id); + if (doc?.documentId) { + void useRagStore + .getState() + .deleteDocument(doc.documentId, `thread:${activeThreadId ?? ""}`) + .catch(() => { + // Best effort — the chip is going away regardless. + }); + } + return prev.filter((d) => d.id !== id); + }); + }, + [activeThreadId], + ); + + const clearDocs = useCallback(() => setPendingDocs([]), []); + + const isIndexing = pendingDocs.some( + (d) => d.status === "uploading" || d.status === "ingesting", + ); + + return { pendingDocs, addDoc, removeDoc, clearDocs, isIndexing }; +} From 6c694d2ef8fb5d6adc62f2c61d4dd48d4103ffef Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 24 May 2026 16:23:26 +0400 Subject: [PATCH 017/122] Studio: move RAG deps to dedicated rag.txt and install in normal path no-torch-runtime.txt is only consumed in NO_TORCH (Intel Mac GGUF-only) mode, so qdrant-client / bm25s / pymupdf etc. were never installed in the normal install path. Split them into rag.txt and add a step to install_python_stack.py that installs it after studio deps, skipped only when NO_TORCH is set. --- .../backend/requirements/no-torch-runtime.txt | 21 --------------- studio/backend/requirements/rag.txt | 27 +++++++++++++++++++ studio/install_python_stack.py | 13 +++++++++ 3 files changed, 40 insertions(+), 21 deletions(-) create mode 100644 studio/backend/requirements/rag.txt diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index ce879ccf94..c33ebf4d94 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -76,24 +76,3 @@ trl>=0.18.2,!=0.19.0,<=0.24.0 sentence-transformers cut_cross_entropy pillow - -# RAG: vector store, lexical index, document parsers. -# qdrant-client supports a pure-Python local mode (QdrantClient(path=...)) -# that we use to keep the studio install self-contained — no separate -# server. bm25s persists per-scope indices to disk. -qdrant-client>=1.12 -bm25s>=0.2 -# RAG parsers (Phase 3A): layout-aware Markdown extraction so the chunker -# can split on real headings instead of running paragraphs together. -# pymupdf4llm preserves headings + pipe-tables; mammoth handles DOCX -# Heading styles; markdownify converts HTML /

/
    faithfully. -pymupdf>=1.24 -pymupdf4llm>=0.0.17 -mammoth>=1.7 -markdownify>=0.13 -# pypdf is kept as a fallback for malformed PDFs that defeat pymupdf. -pypdf>=4.0 -python-docx>=1.1 -beautifulsoup4>=4.12 -lxml>=5.0 -chardet>=5.2 diff --git a/studio/backend/requirements/rag.txt b/studio/backend/requirements/rag.txt new file mode 100644 index 0000000000..11bbf906e5 --- /dev/null +++ b/studio/backend/requirements/rag.txt @@ -0,0 +1,27 @@ +# Studio RAG dependencies. +# Installed by studio/install_python_stack.py in the normal (with-torch) +# path. Skipped in NO_TORCH (Intel Mac GGUF-only) mode because RAG +# embedding relies on sentence-transformers, which requires torch. + +# Vector store + lexical index. qdrant-client supports a pure-Python +# local mode (QdrantClient(path=...)) that we use to keep the studio +# install self-contained — no separate server. bm25s persists per-scope +# indices to disk. +qdrant-client>=1.12 +bm25s>=0.2 + +# Layout-aware Markdown extraction (Phase 3A) so the chunker can split +# on real headings instead of running paragraphs together. pymupdf4llm +# preserves headings + pipe-tables; mammoth handles DOCX Heading styles; +# markdownify converts HTML /
/