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