diff --git a/studio/backend/core/rag/bm25.py b/studio/backend/core/rag/bm25.py index 28b991d168..508d75129d 100644 --- a/studio/backend/core/rag/bm25.py +++ b/studio/backend/core/rag/bm25.py @@ -64,6 +64,10 @@ def rebuild_index(scope: str, chunks: list[dict]) -> None: tokens = bm25s.tokenize(texts, show_progress = False) retriever = bm25s.BM25() retriever.index(tokens, show_progress = False) + # Drop stale files from any previous build/library version before + # writing the new index so we never mix old + new artifacts in the + # scope dir (bm25s.BM25.save does not unlink files it does not write). + delete_scope(scope) ensure_dir(base) retriever.save(str(base)) _ids_path(scope).write_text(json.dumps(ids)) @@ -79,8 +83,21 @@ def _load(scope: str) -> tuple[Any, list[str]] | None: return _cache[scope] import bm25s - retriever = bm25s.BM25.load(str(_scope_dir(scope)), load_corpus = False) - ids = json.loads(_ids_path(scope).read_text()) + try: + retriever = bm25s.BM25.load(str(_scope_dir(scope)), load_corpus = False) + ids = json.loads(_ids_path(scope).read_text()) + except (FileNotFoundError, OSError, json.JSONDecodeError, ValueError) as exc: + # Corrupt or partially-written index files: treat as a + # missing index so search returns empty and a future + # re-ingest can rebuild cleanly. Log so the failure + # is visible without crashing the request. + logger.warning( + "bm25 index unreadable for scope %s (%s: %s); treating as missing", + scope, + type(exc).__name__, + exc, + ) + return None _cache[scope] = (retriever, ids) return _cache[scope] diff --git a/studio/backend/core/rag/scope.py b/studio/backend/core/rag/scope.py new file mode 100644 index 0000000000..270693faba --- /dev/null +++ b/studio/backend/core/rag/scope.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Scope identifiers + per-scope embedder resolution. + +A "scope" is the namespace key carried by vector rows and BM25 +indexes — ``kb_`` for stand-alone Knowledge Bases or +``thread_`` for per-thread document sets. + +`resolve_scope_embedder` looks up which embedder populated a scope's +vectors so the query side can re-use the same model. Lives in +``core/rag`` rather than ``routes`` so the inference-side tool +handler can call it without a route-→-core import cycle. +""" + +from __future__ import annotations + +from storage.studio_db import get_connection, list_chat_settings +from utils.rag.config import resolve_embedder + +# Persisted chat-settings keys. Defined here so the resolver (core) +# and the route handlers (routes/rag.py) share one source of truth. +RAG_DEFAULTS_KEY = "rag.defaults" + + +def thread_settings_key(thread_id: str) -> str: + return f"thread:{thread_id}:rag" + + +def resolve_scope_embedder(scope: str) -> str | None: + """Return the embedder used to populate ``scope``'s vector rows. + + Resolution order: + - ``kb_`` → ``rag_knowledge_bases.embedding_model`` column. + - ``thread_`` → per-thread override → app-level defaults + override → ``RAG_EMBEDDER_MATRIX[(mode, chunking)]``. + + Returns ``None`` for unrecognised scope strings; callers treat + ``None`` as "fall back to the configured default embedder". + """ + if scope.startswith("kb_"): + kb_id = scope[len("kb_") :] + with get_connection() as conn: + row = conn.execute( + "SELECT embedding_model FROM rag_knowledge_bases WHERE id = ?", + (kb_id,), + ).fetchone() + return row["embedding_model"] if row else None + + if scope.startswith("thread_"): + thread_id = scope[len("thread_") :] + all_settings = list_chat_settings() + defaults = all_settings.get(RAG_DEFAULTS_KEY) or {} + if not isinstance(defaults, dict): + defaults = {} + per_thread = all_settings.get(thread_settings_key(thread_id)) or {} + if not isinstance(per_thread, dict): + per_thread = {} + explicit = per_thread.get("embedding_model") or defaults.get("embedding_model") + if explicit: + return explicit + mode = per_thread.get("mode") or defaults.get("mode") or "text" + chunking_strategy = ( + per_thread.get("chunking_strategy") + or defaults.get("chunking_strategy") + or "standard" + ) + return resolve_embedder(mode, chunking_strategy) + + return None diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index 5ec1f039ef..9156ca8324 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -125,12 +125,9 @@ def search_knowledge_base( else: candidate_k = k - # Resolve the scope's embedder (the one that populated its - # vectors). Inline import to avoid a backend-route → core - # dependency cycle at module load. - from routes.rag import _resolve_scope_embedder + from core.rag.scope import resolve_scope_embedder - scope_embedder = _resolve_scope_embedder(scope) + scope_embedder = resolve_scope_embedder(scope) logger.info( "search_knowledge_base: scope=%s embedder=%s top_k=%d min_score=%.3f rerank=%s query=%r", diff --git a/studio/backend/core/rag/vector_store.py b/studio/backend/core/rag/vector_store.py index 64b1725c45..40536ff688 100644 --- a/studio/backend/core/rag/vector_store.py +++ b/studio/backend/core/rag/vector_store.py @@ -13,8 +13,8 @@ A "scope" is `kb_` for standalone knowledge bases or indexed; queries filter by scope before computing distances so different scopes can hold vectors of different dimensions without breaking the cosine math. The per-scope embedder resolver -(routes/rag.py:_resolve_scope_embedder) guarantees one embedder per -scope, so dims within a scope are always consistent. +(core/rag/scope.py:resolve_scope_embedder) guarantees one embedder +per scope, so dims within a scope are always consistent. """ from __future__ import annotations diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 504eb50b9e..f80a3b6fc4 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -181,33 +181,10 @@ def _now_ms() -> int: return int(time.time()) -def _resolve_scope_embedder(scope: str) -> str | None: - """Look up the embedder that populated `scope`'s vector rows. - - Returns the model name to use for query-side embedding so the - similarity math doesn't mix vector spaces (Qwen3-VL 2048-d - documents vs. bge-small 384-d query would crash with a shape - mismatch). Returns None for unknown scopes; callers should treat - None as "fall back to the default embedder". - """ - from utils.rag.config import resolve_embedder - - if scope.startswith("kb_"): - kb_id = scope[len("kb_") :] - with get_connection() as conn: - row = conn.execute( - "SELECT embedding_model FROM rag_knowledge_bases WHERE id = ?", - (kb_id,), - ).fetchone() - return row["embedding_model"] if row else None - if scope.startswith("thread_"): - thread_id = scope[len("thread_") :] - settings = _load_thread_settings(thread_id) - return settings.embedding_model or resolve_embedder( - settings.mode, - settings.chunking_strategy, - ) - return None +# Per-scope embedder lookup lives in core/rag/scope.py so the +# inference-side tool handler can use it without a routes-→-core +# import cycle. +from core.rag.scope import resolve_scope_embedder as _resolve_scope_embedder # noqa: E402 def _row_to_kb(row: Any) -> KBResponse: @@ -295,6 +272,8 @@ def _document_or_404(document_id: str) -> Any: async def _save_upload(file: UploadFile) -> tuple[Path, str, int]: + import anyio + filename = _sanitize_filename(file.filename or "document") ext = Path(filename).suffix.lower() if ext not in RAG_UPLOAD_EXTS: @@ -308,20 +287,27 @@ async def _save_upload(file: UploadFile) -> tuple[Path, str, int]: stored_path = upload_dir / stored_name max_bytes = RAG_MAX_UPLOAD_MB * 1024 * 1024 written = 0 - 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) + # anyio.open_file routes each write through a worker thread so a + # multi-MB upload doesn't stall concurrent requests on the event + # loop. The async-with handles close on both happy and error + # paths; the outer try/except cleans up the partial file after + # the file handle is closed (Windows refuses unlink on an open fd). + try: + async with await anyio.open_file(stored_path, "wb") as f: + while True: + chunk = await file.read(1024 * 1024) + if not chunk: + break + written += len(chunk) + if written > max_bytes: + raise HTTPException( + status_code = 413, + detail = f"File exceeds {RAG_MAX_UPLOAD_MB} MB limit", + ) + await f.write(chunk) + except HTTPException: + stored_path.unlink(missing_ok = True) + raise if written == 0: stored_path.unlink(missing_ok = True) raise HTTPException(status_code = 400, detail = "Empty upload payload")