diff --git a/.gitignore b/.gitignore index a839633790..50a442fbc8 100644 --- a/.gitignore +++ b/.gitignore @@ -235,3 +235,6 @@ package-lock.json !studio/backend/core/data_recipe/oxc-validator/package-lock.json !studio/package-lock.json llama.cpp/ +/.Codex +/.gemini +/.antigravitycli diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 6ddcbc8e0b..04eefbd1b0 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/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2d95112d6d..f1331219cc 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -616,7 +616,19 @@ class LlamaCppBackend: 3. unload_model() — terminates llama-server subprocess """ - def __init__(self): + def __init__(self, kill_orphans: bool = True): + """Construct a backend wrapper around llama-server. + + ``kill_orphans`` (default True): at construction time, reap any + llama-server processes lingering from a prior studio crash. Safe + for the global singleton (only one LlamaCppBackend exists at + startup). Pass ``False`` for short-lived secondary instances + spawned alongside an already-running chat-model server (e.g. + the RAG captioner helper, `_run_with_helper`) — otherwise the + constructor will kill the parent's healthy chat model because + it can't distinguish "another instance's healthy server" from + "a stale process". + """ self._process: Optional[subprocess.Popen] = None self._port: Optional[int] = None self._model_identifier: Optional[str] = None @@ -699,7 +711,8 @@ class LlamaCppBackend: # to decide whether to wait for the VRAM reclaim to finish. self._last_kill_monotonic: float = 0.0 - self._kill_orphaned_servers() + if kill_orphans: + self._kill_orphaned_servers() atexit.register(self._cleanup) # ── Properties ──────────────────────────────────────────────── @@ -4227,9 +4240,12 @@ class LlamaCppBackend: # without triggering a retry storm. Cancel during both # prefill and streaming is handled by the watcher thread # which closes the response, unblocking any httpx read. + # 300 s headroom for large models (30B+) re-prefilling after + # a tool call that returned a long result (e.g. RAG chunks + # with images) — prior 120 s was tripping on Gemma-4-31B. prefill_timeout = httpx.Timeout( connect = 30, - read = 120.0, + read = 300.0, write = 10, pool = 10, ) @@ -4452,6 +4468,7 @@ class LlamaCppBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + tool_context: Optional[dict] = None, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -5121,6 +5138,7 @@ class LlamaCppBackend: cancel_event = cancel_event, timeout = _effective_timeout, session_id = session_id, + tool_context = tool_context, ) yield { diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 7e7d7026f6..4cfd87f1a1 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -838,6 +838,7 @@ class InferenceOrchestrator: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + tool_context: Optional[dict] = None, use_adapter: Optional[Union[bool, str]] = None, **_unused, ): @@ -895,6 +896,7 @@ class InferenceOrchestrator: max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, + tool_context = tool_context, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 73bb3d090a..edfae5ce16 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -105,6 +105,7 @@ def run_safetensors_tool_loop( max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, + tool_context: Optional[dict] = None, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -340,6 +341,7 @@ def run_safetensors_tool_loop( cancel_event = cancel_event, timeout = eff_timeout, session_id = session_id, + tool_context = tool_context, ) except Exception as exc: logger.exception("Tool %s raised: %s", tool_name, exc) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 9572a2169a..48d70aa67f 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -511,7 +511,15 @@ TERMINAL_TOOL = { }, } -ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL] + +# Lazy import: don't pull rag stack on inference paths that never see RAG. +def _get_rag_tool_spec(): + from core.rag.tool import SEARCH_KNOWLEDGE_BASE_TOOL + + return SEARCH_KNOWLEDGE_BASE_TOOL + + +ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()] # OpenAI's function.name regex: ^[a-zA-Z0-9_-]{1,64}$ -- enforced before @@ -609,12 +617,17 @@ def execute_tool( cancel_event = None, timeout: int | None = _TIMEOUT_UNSET, session_id: str | None = None, + tool_context: dict | None = None, ) -> str: """Execute a tool by name with the given arguments. Returns result as a string. ``timeout``: int sets per-call limit in seconds, ``None`` means no limit, unset (default) uses ``_EXEC_TIMEOUT`` (300 s). ``session_id``: optional thread/session ID for per-conversation sandbox isolation. + ``tool_context``: optional per-request extras the LLM does not see (RAG scope, + future per-tool overrides). Keys consumed: + - ``rag_scope``: ``{kb_id?, thread_id?, enable_rerank?, default_top_k?, + reranker_model?, min_score?, mode?}`` — consumed by ``search_knowledge_base``. """ logger.info( f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}" @@ -653,6 +666,23 @@ def execute_tool( return _bash_exec( arguments.get("command", ""), cancel_event, effective_timeout, session_id ) + if name == "search_knowledge_base": + from core.rag.tool import search_knowledge_base + + scope = (tool_context or {}).get("rag_scope") or {} + raw_mode = scope.get("mode") + mode = raw_mode if raw_mode in ("bm25", "dense", "hybrid") else "hybrid" + return search_knowledge_base( + query = arguments.get("query", ""), + top_k = arguments.get("top_k"), + scope_kb_id = scope.get("kb_id"), + scope_thread_id = scope.get("thread_id"), + enable_rerank = bool(scope.get("enable_rerank")), + reranker_model = scope.get("reranker_model"), + default_top_k = int(scope.get("default_top_k") or 5), + min_score = float(scope.get("min_score") or 0.0), + mode = mode, + ) return f"Unknown tool: {name}" 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/authorization.py b/studio/backend/core/rag/authorization.py new file mode 100644 index 0000000000..dbd1787785 --- /dev/null +++ b/studio/backend/core/rag/authorization.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Subject-scoped authorization for RAG document preview routes. + +Used by `/api/rag/documents/{document_id}/file` and +`/api/rag/documents/{document_id}/preview-target` to enforce that the +current authenticated subject is allowed to see a given document and +chunk. Existence and authorization failures collapse to a single 404 +so the API does not leak document IDs to a non-owner. +""" + +from __future__ import annotations + +import sqlite3 + +from fastapi import HTTPException + +from storage.studio_db import get_connection + +_NOT_FOUND_DETAIL = "Document not found" + + +def document_for_subject_or_404( + document_id: str, + current_subject: str, +) -> sqlite3.Row: + """Return the `rag_documents` row if `current_subject` may access it. + + Authorization rules: + + - KB documents: the document's KB must have + `rag_knowledge_bases.owner_user_id == current_subject`. A KB with a + NULL owner is not accessible through this helper (legacy pre-auth + rows must be migrated or accessed via admin tooling). + + - Thread documents: thread-scoped RAG documents are gated by an + explicit single-user invariant for Studio's current release. The + `chat_threads` table does not yet carry an `owner_user_id` column, + so we cannot bind a thread to a specific subject in the schema. + The helper still requires (a) an authenticated subject (enforced + by the route's `Depends(get_current_subject)`) and (b) that the + referenced thread actually exists in `chat_threads`. A missing + thread row collapses to 404 so a non-existent thread cannot + silently grant access through a dangling `thread_id`. + # TODO(thread-owner): once `chat_threads.owner_user_id` exists, + # join through it the same way KB documents do and drop the + # single-user invariant. Update the test + # `tests/test_rag_authorization.py::test_thread_doc_other_user_404` + # to assert per-user isolation rather than thread existence. + + Both not-found and not-authorized raise `HTTPException(404)` with the + same detail string. Callers must NOT distinguish the two cases in + their response, to avoid leaking document existence to a non-owner. + + Returns the document row so the caller can read `stored_path`, + `filename`, `content_type`, etc. without re-querying. + """ + if not document_id or not current_subject: + raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL) + + with get_connection() as conn: + row = conn.execute( + "SELECT * FROM rag_documents WHERE id = ?", + (document_id,), + ).fetchone() + if row is None: + raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL) + + kb_id = row["kb_id"] + thread_id = row["thread_id"] + + if kb_id is not None: + owner_row = conn.execute( + "SELECT owner_user_id FROM rag_knowledge_bases WHERE id = ?", + (kb_id,), + ).fetchone() + if owner_row is None: + raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL) + owner = owner_row["owner_user_id"] + if owner is None or owner != current_subject: + raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL) + return row + + if thread_id is not None: + # Single-user invariant (see TODO above). We require the + # thread row to exist; an unknown thread_id is treated as + # not-found, not as silent grant. + thread_row = conn.execute( + "SELECT id FROM chat_threads WHERE id = ?", + (thread_id,), + ).fetchone() + if thread_row is None: + raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL) + return row + + # Documents must belong to either a KB or a thread (DB CHECK + # constraint enforces XOR on insert); a row that satisfies + # neither is corrupt — treat as 404. + raise HTTPException(status_code = 404, detail = _NOT_FOUND_DETAIL) + + +def chunk_belongs_to_document(chunk_id: str, document_id: str) -> bool: + """True iff `chunk_id` exists in `rag_chunks` for `document_id`. + + Used by `/preview-target?chunk_id=...` after the caller has + already established subject authorization for `document_id`. Does + NOT perform authorization itself: callers MUST call + `document_for_subject_or_404(document_id, ...)` first, otherwise a + valid `chunk_id` from another subject's document would leak via a + `True` return. + """ + if not chunk_id or not document_id: + return False + + with get_connection() as conn: + row = conn.execute( + "SELECT 1 FROM rag_chunks WHERE id = ? AND document_id = ?", + (chunk_id, document_id), + ).fetchone() + return row is not None diff --git a/studio/backend/core/rag/bm25.py b/studio/backend/core/rag/bm25.py new file mode 100644 index 0000000000..0d45ba2420 --- /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 index (rebuild on change; bm25s has no cheap incremental insert). + +Each scope dir holds the bm25s files + ids.json mapping row index → chunk_id. +""" + +from __future__ import annotations + +import json +import shutil +import threading +from pathlib import Path +from typing import Any + +from loggers import get_logger +from utils.paths.storage_roots import ensure_dir, rag_bm25_root + +logger = get_logger(__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 scope's BM25 from full chunk list. Empty list deletes the index.""" + 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) + # bm25s.BM25.save does not unlink stale files; clear the dir first. + delete_scope(scope) + 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 + + 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/partial index: treat as missing so re-ingest rebuilds cleanly. + 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] + + +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/captioner.py b/studio/backend/core/rag/captioner.py new file mode 100644 index 0000000000..b488bd7dcd --- /dev/null +++ b/studio/backend/core/rag/captioner.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Figure captioning for RAG ingestion. + +Two captioning sources, tried in order: + +1. The user's currently-loaded chat VLM — when ``vlm_url`` / ``vlm_model`` + are provided by the parent process (it probes ``llama_cpp.is_vision`` + at enqueue time). Captions go through that model's OpenAI-compatible + ``/v1/chat/completions`` endpoint as base64 ``image_url``. + +2. A helper llama-server fallback that loads the pre-cached + ``unsloth/gemma-4-E2B-it-GGUF`` (gemma-3n family, multimodal) with + its mmproj for vision. Spawned for the lifetime of a + ``caption_images`` call, unloaded before return so no llama-server + process leaks past ingestion. + +Defensive: any per-image failure returns an empty string; total +captioner unavailability (no chat VLM + helper load failure) returns +empty strings for every image. The caller (``_stream_image_chunks``) +falls back to the parser's page-text caption in that case. +""" + +from __future__ import annotations + +import base64 +from io import BytesIO +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +_PROMPT = ( + "This image is a region cropped from a PDF page that contains a " + "single figure (schematic, chart, diagram, table, photo, or " + "their combination). Describe the figure's structure and content " + "in <=80 words. Focus on factual visible content: axes, labels, " + "arrow labels, box labels, legends, visible text in the figure, " + "and what entities are connected to what. Do not speculate beyond " + "what is visible and do not describe the page header/footer or " + "body paragraphs." +) +_MAX_NEW_TOKENS = 200 +# Downscale large images so the base64 payload stays manageable; the chat +# model's prefill cost scales with image-tile count, not pixel count, but +# very large inputs still bloat the JSON body. 1600 px on the long side +# matches PR #5351's chat-composer extractor. +_MAX_IMAGE_SIZE = 1600 +_REQUEST_TIMEOUT_SECONDS = 120.0 + +# Helper VLM (used when no vision-capable chat model is loaded). +# Matches the model pre-cached by precache_helper_gguf() at studio +# startup so the captioner doesn't have to wait on a fresh download. +_HELPER_REPO = "unsloth/gemma-4-E2B-it-GGUF" +_HELPER_VARIANT = "UD-Q4_K_XL" +_HELPER_MODEL_NAME = "helper" + + +def _image_to_data_url(blob: bytes) -> str: + from PIL import Image + + img = Image.open(BytesIO(blob)).convert("RGB") + if max(img.size) > _MAX_IMAGE_SIZE: + img.thumbnail((_MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE)) + buf = BytesIO() + img.save(buf, format = "JPEG", quality = 88) + encoded = base64.b64encode(buf.getvalue()).decode("ascii") + return f"data:image/jpeg;base64,{encoded}" + + +def _load_helper_vlm() -> Optional[tuple[Any, str, str]]: + """Spawn a private LlamaCppBackend with the helper VLM + mmproj. + + Returns ``(backend, base_url, model_name)`` on success, ``None`` on + failure. The caller is responsible for unloading the backend when + done (so the helper doesn't outlive the ingestion subprocess). + """ + try: + from core.inference.llama_cpp import LlamaCppBackend + + # kill_orphans=False is critical: the global singleton is + # already running the user's chat-model llama-server. Killing + # "orphans" here would reap that healthy chat process because + # the orphan-killer can't tell two LlamaCppBackend instances + # apart by PID ownership. + backend = LlamaCppBackend(kill_orphans = False) + logger.info( + "RAG captioner: loading helper VLM as fallback", + repo = _HELPER_REPO, + variant = _HELPER_VARIANT, + ) + ok = backend.load_model( + hf_repo = _HELPER_REPO, + hf_variant = _HELPER_VARIANT, + model_identifier = f"rag-captioner:{_HELPER_REPO}:{_HELPER_VARIANT}", + is_vision = True, + n_ctx = 4096, + n_gpu_layers = -1, + ) + if not ok: + logger.warning("RAG captioner: helper VLM failed to start") + return None + return backend, backend.base_url, _HELPER_MODEL_NAME + except Exception as exc: # noqa: BLE001 + logger.warning("RAG captioner: helper VLM load raised", error = str(exc)) + return None + + +def _post_one(client: Any, endpoint: str, model: str, blob: bytes) -> str: + """POST one image to the OpenAI-compatible endpoint, return caption.""" + data_url = _image_to_data_url(blob) + payload = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": _PROMPT}, + { + "type": "image_url", + "image_url": {"url": data_url}, + }, + ], + } + ], + "max_tokens": _MAX_NEW_TOKENS, + "temperature": 0.0, + # Reasoning models (gemma-4, qwen3-thinking, etc.) burn the whole + # token budget on output and emit empty visible content + # — useless for a short image caption. Disable thinking for this + # request only; the user's chat sessions stay unaffected. + "chat_template_kwargs": {"enable_thinking": False}, + } + response = client.post(endpoint, json = payload) + response.raise_for_status() + data = response.json() + content = data.get("choices", [{}])[0].get("message", {}).get("content", "") + return content.strip() if isinstance(content, str) else "" + + +def caption_images( + image_bytes_list: list[bytes], + *, + vlm_url: Optional[str] = None, + vlm_model: Optional[str] = None, +) -> list[str]: + """Generate one short caption per image; same-length output. + + Tries the loaded chat VLM first (``vlm_url`` + ``vlm_model``). If + those are missing, spawns the helper VLM, captions, and unloads it + before returning. On any failure returns ``""`` for the affected + image. Never raises. + """ + logger.info( + "caption_images: invoked", + n_images = len(image_bytes_list), + vlm_url = vlm_url, + vlm_model = vlm_model, + ) + if not image_bytes_list: + return [] + + import httpx + + helper_backend: Optional[Any] = None + try: + # Resolve endpoint + model: chat VLM if available, else helper. + if vlm_url and vlm_model: + endpoint = f"{vlm_url.rstrip('/')}/v1/chat/completions" + model_name = vlm_model + else: + loaded = _load_helper_vlm() + if loaded is None: + logger.warning( + "caption_images: helper load failed, returning empty captions" + ) + return ["" for _ in image_bytes_list] + helper_backend, helper_base_url, helper_model_name = loaded + endpoint = f"{helper_base_url.rstrip('/')}/v1/chat/completions" + model_name = helper_model_name + + out: list[str] = [] + with httpx.Client(timeout = _REQUEST_TIMEOUT_SECONDS) as client: + for idx, blob in enumerate(image_bytes_list): + try: + out.append(_post_one(client, endpoint, model_name, blob)) + except Exception as exc: # noqa: BLE001 + logger.warning( + "caption_images: per-image request failed", + idx = idx, + endpoint = endpoint, + error = str(exc), + ) + out.append("") + non_empty = sum(1 for c in out if c.strip()) + logger.info( + "caption_images: complete", + total = len(out), + non_empty = non_empty, + ) + return out + finally: + # Always tear down the helper if we spawned one. Chat VLM (when + # provided by the parent) is left alone — it's not ours to manage. + if helper_backend is not None: + try: + helper_backend.unload_model() + logger.info("RAG captioner: helper VLM unloaded") + except Exception as exc: # noqa: BLE001 + logger.warning("RAG captioner: helper unload failed", error = str(exc)) diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py new file mode 100644 index 0000000000..8ea5ca57fd --- /dev/null +++ b/studio/backend/core/rag/chunking.py @@ -0,0 +1,309 @@ +# 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 re +from dataclasses import dataclass +from typing import Callable + +from .parsers import ParsedPage + +# Match "Figure 1:", "Figure 1.2:", "Fig. 3.", "Table 4:" etc. at line-start, +# tolerating leading bold markers. Used to break chunks BEFORE such captions +# so the caption ends up at the start of its own chunk — dense embeddings +# pool over the whole chunk, so figure references buried at the end get +# diluted by surrounding body text. +_FIGURE_BOUNDARY_RE = re.compile( + # Number forms covered: "1", "12", "1.2", "B.1" (appendix-style), + # tolerating bold wrappers around either the label or the number. + r"^\**(?:Figure|Fig\.|Table|Tab\.)\s+[A-Z]?\.?\d+(?:\.\d+)?\**[\.:]", + re.MULTILINE | re.IGNORECASE, +) + + +def _split_at_figure_boundaries(text: str) -> list[str]: + """Split markdown at the start of each figure / table caption. + + Each segment starts with either the original text head or a "Figure N:" / + "Table N:" line, so the caption anchors the embedding of the chunk it + lands in. Returns the original text as a single-element list when no + captions are found. + """ + matches = list(_FIGURE_BOUNDARY_RE.finditer(text)) + if not matches: + return [text] + segments: list[str] = [] + last = 0 + for m in matches: + if m.start() > last: + segments.append(text[last : m.start()]) + last = m.start() + segments.append(text[last:]) + return [s for s in segments if s.strip()] + + +@dataclass(frozen = True) +class Chunk: + text: str + token_count: int + page_number: int | None = None + source_page_index: int | None = None + page_char_start: int | None = None + page_char_end: int | None = None + line_start: int | None = None + line_end: int | None = None + + +TokenCounter = Callable[[str], int] + + +def _char_token_estimate(text: str) -> int: + 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 + 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 into <= max_tokens chunks with overlap.""" + 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 _line_bounds(text: str, start: int, end: int) -> tuple[int, int]: + """Return 1-based inclusive line numbers for a page-local span.""" + line_start = text.count("\n", 0, start) + 1 + line_end = text.count("\n", 0, max(start, end - 1)) + 1 + return line_start, line_end + + +def _locate_piece( + page_text: str, + piece: str, + search_cursor: int, +) -> tuple[int | None, int | None, int | None, int | None, int]: + idx = page_text.find(piece, search_cursor) + if idx < 0: + idx = page_text.find(piece) + if idx < 0: + return None, None, None, None, search_cursor + end = idx + len(piece) + line_start, line_end = _line_bounds(page_text, idx, end) + return idx, end, line_start, line_end, idx + 1 + + +# Markdown headings first so layout-aware parser output splits at sections. +DEFAULT_SEPARATORS: tuple[str, ...] = ( + "\n# ", + "\n## ", + "\n### ", + "\n#### ", + "\n\n", + "\n", + ". ", + " ", + "", +) + + +def chunk_pages( + pages: list[ParsedPage], + *, + max_tokens: int, + overlap_tokens: int, + token_counter: TokenCounter | None = None, + separators: tuple[str, ...] = DEFAULT_SEPARATORS, +) -> list[Chunk]: + """Split pages independently so page_number stays attached to chunks.""" + count = token_counter or _char_token_estimate + out: list[Chunk] = [] + for page_index, page in enumerate(pages): + search_cursor = 0 + for segment in _split_at_figure_boundaries(page.text): + atomic = _atomic_split(segment, separators, max_tokens, count) + merged = _merge(atomic, max_tokens, overlap_tokens, count) + for piece in merged: + start, end, line_start, line_end, search_cursor = _locate_piece( + page.text, + piece, + search_cursor, + ) + out.append( + Chunk( + text = piece, + token_count = count(piece), + page_number = page.page_number, + source_page_index = page_index, + page_char_start = start, + page_char_end = end, + line_start = line_start, + line_end = line_end, + ) + ) + return out + + +_PAGE_SEPARATOR = "\n\n" + + +def chunk_pages_with_spans( + pages: list[ParsedPage], + *, + max_tokens: int, + overlap_tokens: int, + token_counter: TokenCounter | None = None, + separators: tuple[str, ...] = DEFAULT_SEPARATORS, +) -> tuple[str, list[Chunk], list[tuple[int, int]]]: + """Late-chunking variant: joins pages so the embedder sees the whole doc. + + Returns ``(full_doc, chunks, char_spans)``; ``char_spans[i]`` is the + (start, end) char offset of ``chunks[i].text`` inside ``full_doc``. + Page numbers are recovered by overlap with the original page ranges. + """ + count = token_counter or _char_token_estimate + + parts: list[str] = [] + page_ranges: list[tuple[int, int, int, int | None]] = [] + cursor = 0 + for index, page in enumerate(pages): + parts.append(page.text) + start = cursor + end = cursor + len(page.text) + page_ranges.append((start, end, index, page.page_number)) + cursor = end + if index < len(pages) - 1: + cursor += len(_PAGE_SEPARATOR) + full_doc = _PAGE_SEPARATOR.join(parts) + + atomic: list[str] = [] + for segment in _split_at_figure_boundaries(full_doc): + atomic.extend(_atomic_split(segment, separators, max_tokens, count)) + merged = _merge(atomic, max_tokens, overlap_tokens, count) + + chunks: list[Chunk] = [] + char_spans: list[tuple[int, int]] = [] + search_cursor = 0 + for piece in merged: + text = piece.strip() + if not text: + continue + idx = full_doc.find(text, search_cursor) + if idx < 0: + # Overlap can push past a chunk's true start; restart from head. + idx = full_doc.find(text) + if idx < 0: + continue + end_idx = idx + len(text) + page_locator = _page_for_span(idx, end_idx, page_ranges) + source_page_index: int | None = None + page_number: int | None = None + page_char_start: int | None = None + page_char_end: int | None = None + line_start: int | None = None + line_end: int | None = None + if page_locator is not None: + page_start, page_end, page_idx, page_no = page_locator + source_page_index = page_idx + page_number = page_no + page_char_start = max(0, idx - page_start) + page_char_end = min(page_end, end_idx) - page_start + line_start, line_end = _line_bounds( + pages[page_idx].text, + page_char_start, + page_char_end, + ) + chunks.append( + Chunk( + text = text, + token_count = count(text), + page_number = page_number, + source_page_index = source_page_index, + page_char_start = page_char_start, + page_char_end = page_char_end, + line_start = line_start, + line_end = line_end, + ) + ) + char_spans.append((idx, end_idx)) + # Advance past start (not end) so overlapping next chunk is findable. + search_cursor = idx + 1 + + return full_doc, chunks, char_spans + + +def _page_for_span( + start: int, + end: int, + page_ranges: list[tuple[int, int, int, int | None]], +) -> tuple[int, int, int, int | None] | None: + for ps, pe, page_index, page_number in page_ranges: + if start < pe and end > ps: + return ps, pe, page_index, page_number + return None diff --git a/studio/backend/core/rag/db.py b/studio/backend/core/rag/db.py new file mode 100644 index 0000000000..954bffb0ad --- /dev/null +++ b/studio/backend/core/rag/db.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Lazy process-wide rag.db connection with sqlite-vec loaded.""" + +from __future__ import annotations + +import sqlite3 +import threading +from pathlib import Path + +from loggers import get_logger +from utils.paths.storage_roots import ensure_dir, rag_root + +logger = get_logger(__name__) + +_conn: sqlite3.Connection | None = None +_conn_lock = threading.Lock() + + +def rag_db_path() -> Path: + return rag_root() / "rag.db" + + +def _load_sqlite_vec(conn: sqlite3.Connection) -> None: + try: + conn.enable_load_extension(True) + except AttributeError as exc: + raise RuntimeError( + "This Python build cannot load SQLite extensions " + "(connection.enable_load_extension is unavailable). RAG " + "requires sqlite-vec, which loads as a SQLite extension. " + "Re-install studio via install.sh so the venv uses uv's " + "managed Python (python-build-standalone), compiled with " + "--enable-loadable-sqlite-extensions." + ) from exc + import sqlite_vec + + sqlite_vec.load(conn) + conn.enable_load_extension(False) + + +def _ensure_schema(conn: sqlite3.Connection) -> None: + conn.executescript( + """ + CREATE TABLE IF NOT EXISTS rag_vectors ( + chunk_id TEXT PRIMARY KEY, + scope TEXT NOT NULL, + document_id TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + kind TEXT NOT NULL DEFAULT 'text', + dim INTEGER NOT NULL, + vector BLOB NOT NULL, + payload_json TEXT NOT NULL DEFAULT '{}' + ); + CREATE INDEX IF NOT EXISTS idx_rag_vectors_scope + ON rag_vectors(scope); + CREATE INDEX IF NOT EXISTS idx_rag_vectors_scope_doc + ON rag_vectors(scope, document_id); + """ + ) + conn.commit() + + +def get_rag_connection() -> sqlite3.Connection: + global _conn + with _conn_lock: + if _conn is None: + ensure_dir(rag_root()) + conn = sqlite3.connect( + str(rag_db_path()), + check_same_thread = False, + ) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode = WAL") + _load_sqlite_vec(conn) + _ensure_schema(conn) + _conn = conn + logger.info("RAG vector store opened", path = str(rag_db_path())) + return _conn + + +def _reset_for_tests() -> None: + global _conn + with _conn_lock: + if _conn is not None: + try: + _conn.close() + except Exception: + pass + _conn = None diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py new file mode 100644 index 0000000000..aa0052c5cc --- /dev/null +++ b/studio/backend/core/rag/embeddings.py @@ -0,0 +1,456 @@ +# 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 embedder singleton. Independent of the chat InferenceBackend.""" + +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: + logger.info("Loading RAG embedder: %s", model_name) + + # BGE-VL's ST shim breaks across ST versions; load via AutoModel. + if model_name.startswith("BAAI/BGE-VL"): + return _BGEVLAdapter(model_name) + + from unsloth import FastSentenceTransformer + + # trust_remote_code: nomic-embed-text-v1.5 needs custom modeling for 8K ctx. + return FastSentenceTransformer.from_pretrained( + model_name, + for_inference = True, + trust_remote_code = True, + ) + + +class _BGEVLAdapter: + """SentenceTransformer-shaped adapter over BGE-VL's AutoModel.""" + + def __init__(self, hf_model_name: str): + from transformers import AutoModel + import torch + + self._model = AutoModel.from_pretrained( + hf_model_name, + trust_remote_code = True, + ) + # Required: BGE-VL's encode() raises without an installed processor. + self._model.set_processor(hf_model_name) + device = "cuda" if torch.cuda.is_available() else "cpu" + self._model.to(device).eval() + self._device = device + self._dim: int | None = None + + def _normalize(self, tensor): + import torch.nn.functional as F + + return F.normalize(tensor, p = 2.0, dim = -1) + + # CLIP positional embedding cap; longer text triggers shape mismatch. + _CLIP_TEXT_MAX_TOKENS = 77 + + def encode( + self, + inputs, + *, + batch_size: int = 32, + normalize_embeddings: bool = True, + convert_to_numpy: bool = True, + show_progress_bar: bool = False, + **_ignored, + ): + import io + + import numpy as np + import torch + from PIL import Image + + if inputs is None or len(inputs) == 0: + return np.zeros( + (0, self.get_sentence_embedding_dimension()), dtype = np.float32 + ) + + sample = inputs[0] + is_image = isinstance(sample, Image.Image) or isinstance( + sample, (bytes, bytearray) + ) + + chunks_out = [] + for start in range(0, len(inputs), batch_size): + batch = list(inputs[start : start + batch_size]) + if is_image: + # BGE-VL's internal data_process re-opens each item with + # Image.open(...), which needs a file-like (has .read()) + # or a path — NOT a pre-opened PIL Image. Pass BytesIO so + # the model's own opener works. PIL Images get rebuffered + # via an in-memory PNG round-trip. + file_likes: list[Any] = [] + for b in batch: + if isinstance(b, (bytes, bytearray)): + file_likes.append(io.BytesIO(b)) + elif isinstance(b, Image.Image): + buf = io.BytesIO() + b.save(buf, format = "PNG") + buf.seek(0) + file_likes.append(buf) + else: + file_likes.append(b) + with torch.no_grad(): + vecs = self._model.encode(images = file_likes) + else: + vecs = self._encode_text_truncated([str(t) for t in batch]) + if normalize_embeddings: + vecs = self._normalize(vecs) + chunks_out.append(vecs.detach().cpu()) + + out = torch.cat(chunks_out, dim = 0) + return out.numpy() if convert_to_numpy else out + + def _encode_text_truncated(self, texts: list[str]): + """Truncate to CLIP's 77-token limit; long text in multimodal mode is lossy.""" + import torch + + tokenizer = self._get_text_tokenizer() + inputs = tokenizer( + texts, + return_tensors = "pt", + padding = True, + truncation = True, + max_length = self._CLIP_TEXT_MAX_TOKENS, + ) + inputs = {k: v.to(self._device) for k, v in inputs.items()} + if any(len(t.split()) > 30 for t in texts): + logger.info( + "BGE-VL text encode: truncating chunks to %d tokens (CLIP cap)", + self._CLIP_TEXT_MAX_TOKENS, + ) + with torch.no_grad(): + return self._model.get_text_features(**inputs) + + def _get_text_tokenizer(self): + processor = getattr(self._model, "processor", None) + if processor is not None: + tok = getattr(processor, "tokenizer", None) + if tok is not None: + return tok + tok = getattr(self._model, "tokenizer", None) + if tok is not None: + return tok + raise AttributeError("BGE-VL adapter could not locate a text tokenizer") + + def get_sentence_embedding_dimension(self) -> int: + if self._dim is None: + v = self.encode(["dim-probe"], batch_size = 1) + self._dim = int(v.shape[-1]) + return self._dim + + def tokenize(self, texts): + return self._get_text_tokenizer()( + texts, + return_tensors = "pt", + padding = True, + ) + + +def get_embedder(model_name: str | None = None) -> Any: + 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 encode_images( + image_bytes_list: list[bytes], + *, + model_name: str | None = None, + batch_size: int | None = None, + normalize: bool = True, +): + """Embed image bytes via a CLIP-family multimodal encoder.""" + from io import BytesIO + + from PIL import Image + + if not image_bytes_list: + return [] + model = get_embedder(model_name) + images = [Image.open(BytesIO(b)).convert("RGB") for b in image_bytes_list] + return model.encode( + images, + 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 token-count callable backed by the embedder's tokenizer.""" + 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 + + +# --- Late chunking (Jina technique) --- + +_LATE_WINDOW_OVERLAP_TOKENS = 512 + + +def late_chunk_encode( + doc_text: str, + char_spans: list[tuple[int, int]], + *, + model_name: str | None = None, + normalize: bool = True, +): + """Single forward pass over the doc, mean-pool token embeddings per chunk span.""" + import numpy as np + + if not char_spans: + return [] + model = get_embedder(model_name) + tokenizer = model.tokenizer + max_length = int(getattr(model, "max_seq_length", None) or 8192) + + encoded = tokenizer( + doc_text, + return_tensors = "pt", + return_offsets_mapping = True, + add_special_tokens = True, + truncation = False, + ) + offsets = encoded.pop("offset_mapping")[0].tolist() + n_tokens = int(encoded["input_ids"].shape[1]) + + if n_tokens <= max_length: + token_embeddings = _encode_tokens(model, encoded) + return _pool_spans( + token_embeddings, + offsets, + char_spans, + normalize = normalize, + np_module = np, + model = model, + doc_text = doc_text, + ) + + logger.info( + "Late chunking: doc has %d tokens > model max %d; using windowed pass", + n_tokens, + max_length, + ) + return _windowed_late_chunk_encode( + doc_text = doc_text, + char_spans = char_spans, + model = model, + max_length = max_length, + normalize = normalize, + np_module = np, + ) + + +def _encode_tokens(model, encoded): + import torch + + transformer = model[0].auto_model + device = next(transformer.parameters()).device + inputs_on_device = {k: v.to(device) for k, v in encoded.items()} + with torch.no_grad(): + outputs = transformer(**inputs_on_device) + return outputs.last_hidden_state[0].detach().cpu().numpy() + + +def _pool_spans( + token_embeddings, + offsets, + char_spans, + *, + normalize: bool, + np_module, + model, + doc_text: str, + token_index_offset: int = 0, +): + """Mean-pool token embeddings per (char_start, char_end) span.""" + vectors = [] + n_rows = token_embeddings.shape[0] + for char_start, char_end in char_spans: + # Skip special tokens whose offsets are (0, 0). + indices = [ + i - token_index_offset + for i, (ts, te) in enumerate(offsets) + if te > ts and te > char_start and ts < char_end + ] + indices = [i for i in indices if 0 <= i < n_rows] + if not indices: + vec = model.encode( + doc_text[char_start:char_end], + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + vectors.append(vec) + continue + pooled = token_embeddings[indices].mean(axis = 0) + if normalize: + denom = float(np_module.linalg.norm(pooled)) + if denom > 0: + pooled = pooled / denom + vectors.append(pooled) + return vectors + + +def _windowed_late_chunk_encode( + *, + doc_text: str, + char_spans: list[tuple[int, int]], + model, + max_length: int, + normalize: bool, + np_module, +): + """Doc > ctx window: pool each chunk against the window containing most of its tokens.""" + import torch + + tokenizer = model.tokenizer + transformer = model[0].auto_model + device = next(transformer.parameters()).device + + full = tokenizer( + doc_text, + return_tensors = "pt", + return_offsets_mapping = True, + add_special_tokens = False, + truncation = False, + ) + all_input_ids = full["input_ids"][0] + all_offsets = full["offset_mapping"][0].tolist() + n_tokens = int(all_input_ids.shape[0]) + stride = max(1, max_length - _LATE_WINDOW_OVERLAP_TOKENS) + + windows: list[tuple[int, int]] = [] + pos = 0 + while pos < n_tokens: + end = min(pos + max_length, n_tokens) + windows.append((pos, end)) + if end >= n_tokens: + break + pos += stride + + window_embeddings: dict[int, "np_module.ndarray"] = {} + + def _window_embeddings(window_index: int): + if window_index in window_embeddings: + return window_embeddings[window_index] + ws, we = windows[window_index] + win_ids = all_input_ids[ws:we].unsqueeze(0).to(device) + win_attn = torch.ones_like(win_ids) + with torch.no_grad(): + outputs = transformer(input_ids = win_ids, attention_mask = win_attn) + emb = outputs.last_hidden_state[0].detach().cpu().numpy() + window_embeddings[window_index] = emb + return emb + + vectors = [] + for char_start, char_end in char_spans: + chunk_token_indices = [ + i + for i, (ts, te) in enumerate(all_offsets) + if te > ts and te > char_start and ts < char_end + ] + if not chunk_token_indices: + vec = model.encode( + doc_text[char_start:char_end], + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + vectors.append(vec) + continue + best_window = 0 + best_overlap = 0 + for wi, (ws, we) in enumerate(windows): + overlap = sum(1 for ti in chunk_token_indices if ws <= ti < we) + if overlap > best_overlap: + best_overlap = overlap + best_window = wi + ws, _we = windows[best_window] + emb = _window_embeddings(best_window) + local_indices = [ + ti - ws for ti in chunk_token_indices if ws <= ti < ws + emb.shape[0] + ] + if not local_indices: + vec = model.encode( + doc_text[char_start:char_end], + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + vectors.append(vec) + continue + pooled = emb[local_indices].mean(axis = 0) + if normalize: + denom = float(np_module.linalg.norm(pooled)) + if denom > 0: + pooled = pooled / denom + vectors.append(pooled) + return vectors diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py new file mode 100644 index 0000000000..f3030c22ba --- /dev/null +++ b/studio/backend/core/rag/ingestion.py @@ -0,0 +1,1000 @@ +# 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 ingestion pipeline. + +Spawn-subprocess per job (parse/chunk/embed); parent persists chunks, +vectors, and rebuilds BM25 on completion. Only the parent opens rag.db. +""" + +from __future__ import annotations + +import json +import multiprocessing as mp +import queue as queue_module +import sqlite3 +import threading +import time +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from loggers import get_logger +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 = get_logger(__name__) + +_CTX = mp.get_context("spawn") +_QUEUE_TIMEOUT_SECONDS = 300 + + +# --- Subprocess worker --- + +_MIME_TO_EXT = { + "image/png": ".png", + "image/jpeg": ".jpg", + "image/jpg": ".jpg", + "image/gif": ".gif", + "image/webp": ".webp", + "image/bmp": ".bmp", + "image/tiff": ".tiff", + "image/svg+xml": ".svg", +} + + +def _subprocess_worker( + stored_path: str, + model_name: str, + chunk_size: int, + overlap: int, + batch_size: int, + out_queue: Any, + chunking_strategy: str = "standard", + mode: str = "text", + document_id: str = "", + vlm_url: str | None = None, + vlm_model: str | None = None, + enable_captions: bool = True, +) -> None: + # Spawned subprocess: structlog isn't configured here (the parent's + # setup runs in the FastAPI process only), so configure it the same + # way so captioner / parser logs render as JSON like the rest, not + # structlog's default dev ConsoleRenderer. + try: + import os as _os + + from loggers.config import LogConfig + + LogConfig.setup_logging( + env = _os.getenv("ENVIRONMENT_TYPE", "production"), + ) + except Exception: # noqa: BLE001 + pass + try: + from core.rag.captioner import caption_images + from core.rag.chunking import chunk_pages + from core.rag.parsers import inline_image_captions, parse + + out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05}) + # Always extract images so we can caption + splice for both + # modes. Text mode uses the captions inline in markdown; multimodal + # additionally embeds the raw images as image-kind chunks. + parsed = parse(Path(stored_path), want_images = True) + pages = parsed.pages + if not pages and not parsed.images: + out_queue.put( + {"type": "error", "error": "no extractable content in document"} + ) + return + + # Caption figures once (chat VLM if available, else helper VLM + # fallback), then splice captions into the page markdown so the + # chunker indexes them like any other text. Multimodal mode also + # passes these same captions through to _stream_image_chunks + # below — no duplicate VLM calls per image. + captions: list[str] = [] + if parsed.images and enable_captions: + out_queue.put( + {"type": "progress", "stage": "caption_images", "progress": 0.08} + ) + captions = caption_images( + [img.image_bytes for img in parsed.images], + vlm_url = vlm_url, + vlm_model = vlm_model, + ) + pages = inline_image_captions(pages, parsed.images, captions) + + out_queue.put( + { + "type": "document_pages", + "pages": [ + { + "page_index": index, + "page_number": page.page_number, + "text": page.text, + "char_count": len(page.text), + "line_count": len(page.text.splitlines()), + } + for index, page in enumerate(pages) + ], + } + ) + + out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1}) + from core.rag.embeddings import ( + get_embedder, + late_chunk_encode, + token_counter, + ) + + model = get_embedder(model_name) + counter = token_counter(model_name) + dim = int(model.get_sentence_embedding_dimension()) + out_queue.put({"type": "dim", "dim": dim}) + + if chunking_strategy == "late": + _run_late_chunking( + pages = pages, + stored_path = Path(stored_path), + chunk_size = chunk_size, + overlap = overlap, + counter = counter, + model_name = model_name, + late_chunk_encode = late_chunk_encode, + out_queue = out_queue, + ) + return + + text_count = _run_standard_chunking( + pages = pages, + stored_path = Path(stored_path), + chunk_size = chunk_size, + overlap = overlap, + counter = counter, + batch_size = batch_size, + model = model, + chunk_pages = chunk_pages, + out_queue = out_queue, + send_complete = False, + ) + image_count = 0 + if mode == "multimodal" and parsed.images and document_id: + image_count = _stream_image_chunks( + images = parsed.images, + document_id = document_id, + model_name = model_name, + out_queue = out_queue, + first_index = text_count, + precomputed_captions = captions, + ) + out_queue.put({"type": "complete", "num_chunks": text_count + image_count}) + except Exception as exc: # noqa: BLE001 + logger.exception("ingestion subprocess failed") + out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"}) + + +def _run_standard_chunking( + *, + pages, + stored_path, + chunk_size, + overlap, + counter, + batch_size, + model, + chunk_pages, + out_queue, + send_complete: bool = True, +) -> int: + """Stream text chunks; returns count. send_complete=False when images follow.""" + out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) + chunks = chunk_pages( + pages, + max_tokens = chunk_size, + overlap_tokens = overlap, + token_counter = counter, + ) + if not chunks: + if send_complete: + out_queue.put({"type": "error", "error": "chunker produced no chunks"}) + return 0 + from core.rag.locators import pdf_regions_for_chunks + + pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks) + + 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, + "source_page_index": c.source_page_index, + "page_char_start": c.page_char_start, + "page_char_end": c.page_char_end, + "line_start": c.line_start, + "line_end": c.line_end, + "pdf_regions": pdf_regions[i + offset], + "kind": "text", + } + for offset, c in enumerate(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}) + + if send_complete: + out_queue.put({"type": "complete", "num_chunks": total}) + return total + + +def _stream_image_chunks( + *, + images, + document_id: str, + model_name: str, + out_queue, + first_index: int, + precomputed_captions: list[str] | None = None, +) -> int: + """Persist images, emit image+caption chunks; pairs share pair_group. + + ``precomputed_captions`` come from the parent's earlier + caption_images call (used so we don't VLM-caption the same images + twice — once for markdown splicing, once for the caption-kind chunk). + If absent we fall back to each image's nearest_caption (page text). + """ + from core.rag.embeddings import encode, encode_images + from utils.paths.storage_roots import ensure_dir, rag_uploads_root + + if not images: + return 0 + + out_queue.put({"type": "progress", "stage": "extract_images", "progress": 0.85}) + + img_dir = ensure_dir(rag_uploads_root() / "images" / document_id) + + paths: list[str] = [] + bytes_for_encoding: list[bytes] = [] + captions: list[str] = [] + pages: list[int | None] = [] + pre_caps = precomputed_captions or [] + for idx, img in enumerate(images): + ext = _MIME_TO_EXT.get(img.mime_type, ".bin") + path = img_dir / f"img-{idx:04d}{ext}" + try: + path.write_bytes(img.image_bytes) + except OSError: + logger.warning("failed to save image; skipping", path = str(path)) + continue + paths.append(str(path)) + bytes_for_encoding.append(img.image_bytes) + vlm_cap = pre_caps[idx].strip() if idx < len(pre_caps) and pre_caps[idx] else "" + captions.append(vlm_cap or (img.nearest_caption or "")) + pages.append(img.page_number) + + if not paths: + return 0 + + image_vectors = encode_images(bytes_for_encoding, model_name = model_name) + + caption_to_image: list[int] = [i for i, cap in enumerate(captions) if cap.strip()] + if caption_to_image: + caption_vectors_arr = encode( + [captions[i] for i in caption_to_image], + model_name = model_name, + ) + caption_vectors = caption_vectors_arr.tolist() + else: + caption_vectors = [] + + out_chunks: list[dict] = [] + out_vectors: list[list[float]] = [] + cap_iter = iter(zip(caption_to_image, caption_vectors)) + next_cap = next(cap_iter, None) + + for idx, (path, page, caption) in enumerate(zip(paths, pages, captions)): + group_id = f"img-{idx:04d}" + out_chunks.append( + { + "text": caption[:1000] if caption else "", + "token_count": 0, + "page_number": page, + "kind": "image", + "image_path": path, + "pair_group": group_id, + } + ) + out_vectors.append(image_vectors[idx].tolist()) + if next_cap is not None and next_cap[0] == idx: + _cap_index, cap_vec = next_cap + out_chunks.append( + { + "text": caption, + "token_count": max(1, len(caption.split())), + "page_number": page, + "kind": "caption", + "image_path": None, + "pair_group": group_id, + } + ) + out_vectors.append(cap_vec) + next_cap = next(cap_iter, None) + + out_queue.put( + { + "type": "chunks_batch", + "first_index": first_index, + "chunks": out_chunks, + "vectors": out_vectors, + } + ) + out_queue.put({"type": "progress", "stage": "extract_images", "progress": 0.95}) + return len(out_chunks) + + +def _run_late_chunking( + *, + pages, + stored_path, + chunk_size, + overlap, + counter, + model_name, + late_chunk_encode, + out_queue, +) -> None: + """Chunk once, embed in one pass, ship all chunks in one chunks_batch.""" + from core.rag.chunking import chunk_pages_with_spans + from core.rag.locators import pdf_regions_for_chunks + + out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2}) + full_doc, chunks, char_spans = chunk_pages_with_spans( + pages, + max_tokens = chunk_size, + overlap_tokens = overlap, + token_counter = counter, + ) + if not chunks: + out_queue.put({"type": "error", "error": "chunker produced no chunks"}) + return + + out_queue.put({"type": "progress", "stage": "embed", "progress": 0.4}) + vectors = late_chunk_encode( + full_doc, + char_spans, + model_name = model_name, + normalize = True, + ) + pdf_regions = pdf_regions_for_chunks(stored_path, pages, chunks) + + out_queue.put({"type": "progress", "stage": "embed", "progress": 0.9}) + out_queue.put( + { + "type": "chunks_batch", + "first_index": 0, + "chunks": [ + { + "text": c.text, + "token_count": c.token_count, + "page_number": c.page_number, + "source_page_index": c.source_page_index, + "page_char_start": c.page_char_start, + "page_char_end": c.page_char_end, + "line_start": c.line_start, + "line_end": c.line_end, + "pdf_regions": pdf_regions[index], + "kind": "text", + } + for index, c in enumerate(chunks) + ], + "vectors": [v.tolist() for v in vectors], + } + ) + out_queue.put({"type": "complete", "num_chunks": len(chunks)}) + + +# --- 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.cancelled = False + self.proc: Any = None + self.out_queue: Any = 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 + vector_store; return [{id, text}] for BM25.""" + rows: list[tuple] = [] + points: list[dict] = [] + bm25_rows: list[dict] = [] + pair_groups: dict[str, list[str]] = {} + + for offset, (meta, vec) in enumerate(zip(chunks_meta, vectors)): + chunk_index = first_index + offset + chunk_id = str(uuid4()) + kind = meta.get("kind", "text") + image_path = meta.get("image_path") + pair_group = meta.get("pair_group") + if pair_group: + pair_groups.setdefault(pair_group, []).append(chunk_id) + rows.append( + ( + chunk_id, + document_id, + chunk_index, + meta["text"], + meta["token_count"], + meta["page_number"], + kind, + image_path, + meta.get("source_page_index"), + meta.get("page_char_start"), + meta.get("page_char_end"), + meta.get("line_start"), + meta.get("line_end"), + json.dumps(meta.get("pdf_regions") or [], separators = (",", ":")) + if meta.get("pdf_regions") + else None, + ) + ) + points.append( + { + "id": chunk_id, + "vector": vec, + "payload": { + "document_id": document_id, + "chunk_index": chunk_index, + "text": meta["text"], + "page_number": meta["page_number"], + "kind": kind, + "image_path": image_path, + "source_page_index": meta.get("source_page_index"), + "page_char_start": meta.get("page_char_start"), + "page_char_end": meta.get("page_char_end"), + "line_start": meta.get("line_start"), + "line_end": meta.get("line_end"), + "pdf_regions": meta.get("pdf_regions") or [], + }, + } + ) + if kind in ("text", "caption") and meta["text"]: + 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, + kind, image_path, source_page_index, page_char_start, + page_char_end, line_start, line_end, pdf_regions_json) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + rows, + ) + # Link only when exactly two members in a pair_group. + for ids in pair_groups.values(): + if len(ids) != 2: + continue + id_a, id_b = ids + conn.execute( + "UPDATE rag_chunks SET linked_chunk_id = ? WHERE id = ?", + (id_b, id_a), + ) + conn.execute( + "UPDATE rag_chunks SET linked_chunk_id = ? WHERE id = ?", + (id_a, id_b), + ) + conn.commit() + vector_store.upsert_chunks(scope, points) + return bm25_rows + + +def _replace_document_pages(document_id: str, pages: list[dict]) -> None: + now = int(time.time()) + rows = [ + ( + document_id, + int(page["page_index"]), + page.get("page_number"), + page.get("text") or "", + int(page.get("char_count", len(page.get("text") or ""))), + int(page.get("line_count", len((page.get("text") or "").splitlines()))), + now, + ) + for page in pages + ] + with get_connection() as conn: + doc_row = conn.execute( + "SELECT 1 FROM rag_documents WHERE id = ?", + (document_id,), + ).fetchone() + if doc_row is None: + raise sqlite3.IntegrityError("FOREIGN KEY constraint failed") + conn.execute( + "DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,) + ) + if rows: + conn.executemany( + """ + INSERT INTO rag_document_pages + (document_id, page_index, page_number, text, char_count, + line_count, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, + rows, + ) + conn.commit() + + +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 until subprocess completes/errors/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: + if state.cancelled: + break + 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 == "__cancel__": + state.cancelled = True + break + 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 == "document_pages": + try: + _replace_document_pages( + state.document_id, + list(msg.get("pages") or []), + ) + except sqlite3.IntegrityError as exc: + final_error = ( + f"document was removed before ingestion finished ({exc})" + ) + break + elif mtype == "chunks_batch": + if embedding_dim is None: + 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) + try: + bm25_rows = _insert_chunks_and_collect_for_bm25( + state.document_id, + state.scope, + int(msg["first_index"]), + msg["chunks"], + msg["vectors"], + ) + except sqlite3.IntegrityError as exc: + # rag_documents row was deleted mid-ingest (user removed + # the chip / cleared the index). Fail the job cleanly + # rather than crashing the pump thread. + final_error = ( + f"document was removed before ingestion finished ({exc})" + ) + break + 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", mtype = repr(mtype)) + finally: + proc.join(timeout = 30) + if proc.is_alive(): + proc.terminate() + proc.join(timeout = 5) + + finished_at = int(time.time()) + if state.cancelled: + # User cancelled mid-flight. The route-side deleteDocument removes the + # row, file, and chunk artifacts; here we just mark terminal and notify + # subscribers so the SSE stream closes cleanly. + _update_document_row(state.document_id, status = "cancelled") + _update_job_row( + state.job_id, + status = "cancelled", + stage = "cancelled", + finished_at = finished_at, + ) + state.status = "cancelled" + state.push_event({"type": "cancelled"}) + return + if final_status == "completed": + full_scope_chunks = _all_scope_chunks(state.scope) + bm25.rebuild_index(state.scope, full_scope_chunks) + _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 _probe_loaded_vlm() -> tuple[str | None, str | None]: + """Best-effort: return (base_url, model_name) when a vision-capable + chat model is currently loaded via llama-server; (None, None) otherwise. + + Used to caption figures with the user's chat VLM instead of loading + a dedicated captioning model — no extra VRAM, no extra download. + Only the llama-server backend is supported today; transformers / + unsloth in-process VLMs would need a different bridge. + """ + try: + # The singleton getter lives in routes.inference, not the + # llama_cpp module. Importing from the wrong place silently + # returned None for every probe — captioner always fell back + # to the helper VLM even when the chat model was vision-capable. + from routes.inference import get_llama_cpp_backend + except Exception as exc: + logger.warning("RAG probe: get_llama_cpp_backend import failed", error = str(exc)) + return None, None + try: + backend = get_llama_cpp_backend() + except Exception as exc: + logger.warning("RAG probe: get_llama_cpp_backend() raised", error = str(exc)) + return None, None + if not getattr(backend, "is_loaded", False): + return None, None + if not getattr(backend, "is_vision", False): + return None, None + base_url = getattr(backend, "base_url", None) + model_id = getattr(backend, "model_identifier", None) + if not base_url or not model_id: + return None, None + return base_url, model_id + + +def enqueue_ingestion( + document_id: str, + stored_path: Path, + *, + kb_id: str | None = None, + thread_id: str | None = None, + embedding_model: str | None = None, + chunking_strategy: str = "standard", + mode: str = "text", + enable_captions: bool = True, +) -> str: + """Create the job row, spawn the subprocess, start the pump; return job_id.""" + from utils.rag.config import resolve_embedder + + scope = _scope_for(kb_id, thread_id) + model_name = ( + embedding_model + or resolve_embedder(mode, chunking_strategy) + or RAG_EMBEDDING_MODEL + ) + # Probe the loaded chat backend so the subprocess can route figure + # captioning to the user's own vision model (no extra VRAM). Runs + # for both modes — text mode splices captions into markdown, and + # multimodal mode additionally feeds them to the image-vector + # encoder. If no vision chat model is loaded, the subprocess falls + # back to the helper VLM (pre-cached at studio startup). Skipped + # entirely when captioning is disabled for this upload. + vlm_url: str | None = None + vlm_model: str | None = None + if enable_captions: + vlm_url, vlm_model = _probe_loaded_vlm() + if vlm_url: + logger.info( + "RAG ingest: will caption figures via loaded chat VLM", + vlm_model = vlm_model, + vlm_url = vlm_url, + ) + else: + logger.info( + "RAG ingest: no vision-capable chat model loaded; " + "subprocess will use the helper gemma-3n VLM fallback." + ) + else: + logger.info("RAG ingest: figure captioning disabled for this upload") + 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() + state.out_queue = out_queue + proc = _CTX.Process( + target = _subprocess_worker, + args = ( + str(stored_path), + model_name, + RAG_CHUNK_SIZE, + RAG_CHUNK_OVERLAP, + RAG_EMBED_BATCH_SIZE, + out_queue, + chunking_strategy, + mode, + document_id, + vlm_url, + vlm_model, + enable_captions, + ), + daemon = True, + ) + proc.start() + state.proc = proc + 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 cancel_ingestion(job_id: str) -> bool: + """Stop an in-flight ingestion: wake the pump via a sentinel and kill the + worker subprocess so it stops consuming GPU/CPU. Returns False if the job + is unknown or already terminal. Artifact/row cleanup is the caller's job + (the route deletes the document).""" + state = get_job_state(job_id) + if state is None: + return False + if state.status in ("completed", "failed", "cancelled"): + return False + state.cancelled = True + if state.out_queue is not None: + try: + state.out_queue.put_nowait({"type": "__cancel__"}) + except Exception: + pass + proc = state.proc + if proc is not None and proc.is_alive(): + proc.terminate() + return True + + +def delete_document_artifacts(document_id: str, scope: str) -> None: + """Drop the doc's vectors, rebuild BM25. Caller deletes the rag_documents row.""" + vector_store.delete_document(scope, document_id) + remaining = _all_scope_chunks(scope) + if remaining: + 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: + """Drop RAG artifacts for the given thread ids (no FK cascade to chat_threads).""" + 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.""" + 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/locators.py b/studio/backend/core/rag/locators.py new file mode 100644 index 0000000000..db0dae319a --- /dev/null +++ b/studio/backend/core/rag/locators.py @@ -0,0 +1,500 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backfill and PDF-region helpers for durable RAG chunk locators.""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from loggers import get_logger +from storage.studio_db import get_connection + +from . import vector_store +from .parsers import ParsedPage, parse +from .vector_store import kb_scope, thread_scope + +logger = get_logger(__name__) + + +@dataclass(frozen = True) +class LocatorMatch: + page_index: int + page_number: int | None + start: int + end: int + line_start: int + line_end: int + + +@dataclass(frozen = True) +class BackfillResult: + document_id: str + total_chunks: int + matched: int + already_located: int + ambiguous: int + missing: int + skipped: int + regions_matched: int + pages_refreshed: int + + +def _line_bounds(text: str, start: int, end: int) -> tuple[int, int]: + line_start = text.count("\n", 0, start) + 1 + line_end = text.count("\n", 0, max(start, end - 1)) + 1 + return line_start, line_end + + +def _find_exact(page_text: str, needle: str) -> list[tuple[int, int]]: + if not needle: + return [] + out: list[tuple[int, int]] = [] + cursor = 0 + while True: + idx = page_text.find(needle, cursor) + if idx < 0: + break + out.append((idx, idx + len(needle))) + cursor = idx + 1 + return out + + +def _normalize_with_map(text: str) -> tuple[str, list[int], list[int]]: + chars: list[str] = [] + starts: list[int] = [] + ends: list[int] = [] + last_space = False + for idx, ch in enumerate(text): + if ch.isspace(): + if chars and not last_space: + chars.append(" ") + starts.append(idx) + ends.append(idx + 1) + elif chars and last_space: + ends[-1] = idx + 1 + last_space = True + continue + chars.append(ch.casefold()) + starts.append(idx) + ends.append(idx + 1) + last_space = False + + first = 0 + while first < len(chars) and chars[first] == " ": + first += 1 + last = len(chars) + while last > first and chars[last - 1] == " ": + last -= 1 + return "".join(chars[first:last]), starts[first:last], ends[first:last] + + +def _find_normalized(page_text: str, needle: str) -> list[tuple[int, int]]: + norm_page, starts, ends = _normalize_with_map(page_text) + norm_needle, _needle_starts, _needle_ends = _normalize_with_map(needle) + if not norm_page or not norm_needle: + return [] + out: list[tuple[int, int]] = [] + cursor = 0 + while True: + idx = norm_page.find(norm_needle, cursor) + if idx < 0: + break + end_idx = idx + len(norm_needle) - 1 + if 0 <= idx < len(starts) and 0 <= end_idx < len(ends): + out.append((starts[idx], ends[end_idx])) + cursor = idx + 1 + return out + + +def _locate_unique( + text: str, pages: list[ParsedPage] +) -> tuple[LocatorMatch | None, str]: + text = (text or "").strip() + if not text: + return None, "missing" + + matches: list[LocatorMatch] = [] + for page_index, page in enumerate(pages): + for start, end in _find_exact(page.text, text): + line_start, line_end = _line_bounds(page.text, start, end) + matches.append( + LocatorMatch( + page_index = page_index, + page_number = page.page_number, + start = start, + end = end, + line_start = line_start, + line_end = line_end, + ) + ) + if len(matches) == 1: + return matches[0], "matched" + if len(matches) > 1: + return None, "ambiguous" + + for page_index, page in enumerate(pages): + for start, end in _find_normalized(page.text, text): + line_start, line_end = _line_bounds(page.text, start, end) + matches.append( + LocatorMatch( + page_index = page_index, + page_number = page.page_number, + start = start, + end = end, + line_start = line_start, + line_end = line_end, + ) + ) + if len(matches) == 1: + return matches[0], "matched" + if len(matches) > 1: + return None, "ambiguous" + return None, "missing" + + +def _replace_document_pages(document_id: str, pages: list[ParsedPage]) -> None: + now = int(time.time()) + rows = [ + ( + document_id, + index, + page.page_number, + page.text, + len(page.text), + len(page.text.splitlines()), + now, + ) + for index, page in enumerate(pages) + ] + with get_connection() as conn: + conn.execute( + "DELETE FROM rag_document_pages WHERE document_id = ?", (document_id,) + ) + if rows: + conn.executemany( + """ + INSERT INTO rag_document_pages + (document_id, page_index, page_number, text, char_count, + line_count, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, + rows, + ) + conn.commit() + + +def _region_anchor(page_text: str, match: LocatorMatch) -> str | None: + segment = page_text[match.start : match.end] + words = [w.strip(" \t\r\n*#`[]()") for w in segment.split()] + words = [w for w in words if len(w) >= 2] + if len(words) < 3: + return None + anchor = " ".join(words[: min(16, len(words))]) + return anchor if len(anchor) >= 12 else None + + +def _normalized_occurrences(haystack: str, needle: str) -> int: + norm_haystack, _starts, _ends = _normalize_with_map(haystack) + norm_needle, _needle_starts, _needle_ends = _normalize_with_map(needle) + if not norm_haystack or not norm_needle: + return 0 + count = 0 + cursor = 0 + while True: + idx = norm_haystack.find(norm_needle, cursor) + if idx < 0: + return count + count += 1 + cursor = idx + 1 + + +def pdf_regions_for_match( + pdf_path: Path, + pages: list[ParsedPage], + match: LocatorMatch, +) -> list[dict[str, Any]]: + """Return normalized PDF rectangles for a unique chunk match. + + Regions are intentionally conservative: no PyMuPDF, no page, no + unique anchor, or no positive-area rectangles all produce an empty + list rather than guessed highlights. + """ + if pdf_path.suffix.lower() != ".pdf": + return [] + if match.page_index < 0 or match.page_index >= len(pages): + return [] + anchor = _region_anchor(pages[match.page_index].text, match) + if not anchor: + return [] + + try: + import pymupdf + except Exception: + return [] + + try: + doc = pymupdf.open(str(pdf_path)) + except Exception: + return [] + + try: + return _pdf_regions_for_match_doc(doc, pages, match, anchor) + finally: + doc.close() + + +def _pdf_regions_for_match_doc( + doc: Any, + pages: list[ParsedPage], + match: LocatorMatch, + anchor: str, +) -> list[dict[str, Any]]: + try: + if match.page_index >= len(doc): + return [] + page = doc[match.page_index] + raw_text = page.get_text("text") or "" + if _normalized_occurrences(raw_text, anchor) != 1: + return [] + rects = page.search_for(anchor) or [] + page_rect = page.rect + page_width = float(page_rect.width) + page_height = float(page_rect.height) + if page_width <= 0 or page_height <= 0: + return [] + + out: list[dict[str, Any]] = [] + for rect in rects: + width = max(0.0, float(rect.x1 - rect.x0)) + height = max(0.0, float(rect.y1 - rect.y0)) + if width <= 0 or height <= 0: + continue + out.append( + { + "pageIndex": match.page_index, + "pageNumber": match.page_number, + "x": max(0.0, min(1.0, float(rect.x0) / page_width)), + "y": max(0.0, min(1.0, float(rect.y0) / page_height)), + "width": max(0.0, min(1.0, width / page_width)), + "height": max(0.0, min(1.0, height / page_height)), + "confidence": "exact", + "source": "pymupdf-search", + } + ) + return out + except Exception: + return [] + + +def pdf_regions_for_chunks( + pdf_path: Path, + pages: list[ParsedPage], + chunks: list[Any], +) -> list[list[dict[str, Any]]]: + if pdf_path.suffix.lower() != ".pdf": + return [[] for _ in chunks] + try: + import pymupdf + + doc = pymupdf.open(str(pdf_path)) + except Exception: + return [[] for _ in chunks] + + regions: list[list[dict[str, Any]]] = [] + try: + for chunk in chunks: + page_index = getattr(chunk, "source_page_index", None) + start = getattr(chunk, "page_char_start", None) + end = getattr(chunk, "page_char_end", None) + if page_index is None or start is None or end is None: + regions.append([]) + continue + if page_index < 0 or page_index >= len(pages): + regions.append([]) + continue + line_start, line_end = _line_bounds(pages[page_index].text, start, end) + match = LocatorMatch( + page_index = int(page_index), + page_number = getattr(chunk, "page_number", None), + start = int(start), + end = int(end), + line_start = line_start, + line_end = line_end, + ) + anchor = _region_anchor(pages[match.page_index].text, match) + if not anchor: + regions.append([]) + continue + regions.append(_pdf_regions_for_match_doc(doc, pages, match, anchor)) + return regions + finally: + doc.close() + + +def _scope_for_document(kb_id: str | None, thread_id: str | None) -> str | None: + if kb_id: + return kb_scope(kb_id) + if thread_id: + return thread_scope(thread_id) + return None + + +def _update_vector_payloads( + scope: str | None, updates: dict[str, dict[str, Any]] +) -> None: + if not scope or not updates: + return + try: + vector_store.update_chunk_payload_fields(scope, updates) + except Exception as exc: + logger.warning( + "RAG locator backfill: vector payload update failed", + error = str(exc), + ) + + +def backfill_document_locators(document_id: str, stored_path: Path) -> BackfillResult: + parsed = parse(stored_path, want_images = False) + pages = parsed.pages + _replace_document_pages(document_id, pages) + + with get_connection() as conn: + doc_row = conn.execute( + "SELECT kb_id, thread_id FROM rag_documents WHERE id = ?", + (document_id,), + ).fetchone() + if doc_row is None: + return BackfillResult(document_id, 0, 0, 0, 0, 0, 0, 0, len(pages)) + + rows = conn.execute( + """ + SELECT id, text, kind, page_number, source_page_index, + page_char_start, page_char_end, line_start, line_end, + pdf_regions_json + FROM rag_chunks + WHERE document_id = ? + ORDER BY chunk_index ASC + """, + (document_id,), + ).fetchall() + + scope = _scope_for_document(doc_row["kb_id"], doc_row["thread_id"]) + total = len(rows) + matched = 0 + already_located = 0 + ambiguous = 0 + missing = 0 + skipped = 0 + regions_matched = 0 + sql_updates: list[tuple[Any, ...]] = [] + vector_updates: dict[str, dict[str, Any]] = {} + + for row in rows: + kind = row["kind"] or "text" + text = row["text"] or "" + if kind not in ("text", "caption") or not text.strip(): + skipped += 1 + continue + + existing_complete = ( + row["source_page_index"] is not None + and row["page_char_start"] is not None + and row["page_char_end"] is not None + and row["line_start"] is not None + and row["line_end"] is not None + ) + + match: LocatorMatch | None + status: str + if existing_complete: + already_located += 1 + page_index = int(row["source_page_index"]) + if 0 <= page_index < len(pages): + match = LocatorMatch( + page_index = page_index, + page_number = row["page_number"], + start = int(row["page_char_start"]), + end = int(row["page_char_end"]), + line_start = int(row["line_start"]), + line_end = int(row["line_end"]), + ) + else: + match = None + status = "already_located" + else: + match, status = _locate_unique(text, pages) + if status == "matched" and match is not None: + matched += 1 + elif status == "ambiguous": + ambiguous += 1 + continue + else: + missing += 1 + continue + + if match is None: + continue + + regions = pdf_regions_for_match(stored_path, pages, match) + regions_json = json.dumps(regions, separators = (",", ":")) if regions else None + if regions: + regions_matched += 1 + + if status == "matched" or (regions and not row["pdf_regions_json"]): + sql_updates.append( + ( + match.page_number, + match.page_index, + match.start, + match.end, + match.line_start, + match.line_end, + regions_json, + row["id"], + ) + ) + vector_updates[row["id"]] = { + "page_number": match.page_number, + "source_page_index": match.page_index, + "page_char_start": match.start, + "page_char_end": match.end, + "line_start": match.line_start, + "line_end": match.line_end, + "pdf_regions": regions, + } + + if sql_updates: + with get_connection() as conn: + conn.executemany( + """ + UPDATE rag_chunks + SET page_number = COALESCE(page_number, ?), + source_page_index = ?, + page_char_start = ?, + page_char_end = ?, + line_start = ?, + line_end = ?, + pdf_regions_json = COALESCE(?, pdf_regions_json) + WHERE id = ? + """, + sql_updates, + ) + conn.commit() + _update_vector_payloads(scope, vector_updates) + + return BackfillResult( + document_id = document_id, + total_chunks = total, + matched = matched, + already_located = already_located, + ambiguous = ambiguous, + missing = missing, + skipped = skipped, + regions_matched = regions_matched, + pages_refreshed = len(pages), + ) diff --git a/studio/backend/core/rag/parsers/__init__.py b/studio/backend/core/rag/parsers/__init__.py new file mode 100644 index 0000000000..99e565a123 --- /dev/null +++ b/studio/backend/core/rag/parsers/__init__.py @@ -0,0 +1,181 @@ +# 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 re +from dataclasses import dataclass, field +from pathlib import Path + +# Same shape as the chunker's figure-boundary regex but captures the +# figure label (Figure / Fig. / Table / Tab.) AND the number so we can +# attribute a VLM caption back to a specific figure on a multi-figure +# page. +_FIGURE_LINE_RE = re.compile( + r"^(?P\**)(?P