diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 7e8d52525c..c6c2e1fc37 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -269,7 +269,8 @@ jobs: tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ tests/utils/test_attention_masks.py \ - tests/utils/test_trunc_normal_patch.py + tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py python -m pytest --collect-only -q "$RUNNER_TEMP/unsloth-zoo/tests/" - name: import_fixes drift detectors (18 tests, HARD GATE) @@ -333,11 +334,9 @@ jobs: python -m pytest -v --tb=short tests/test_callback_signature_drift.py - name: unsloth Bucket-A — CPU tests not in Repo tests (CPU) - # 16 tests across 5 files. They live inside tests/saving/ and - # tests/utils/, both of which Repo tests (CPU) excludes via --ignore - # because their sibling files need real GPUs / real HF weights. - # The five files below are pure-Python + AST/protobuf/regex tests - # that run cleanly on CPU. Env inherited from the job block. + # CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/ + # that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model + # loads; run cleanly here (transformers/torch installed). run: | python -m pytest -q --tb=short \ tests/saving/test_save_shell_injection.py \ @@ -345,11 +344,12 @@ jobs: tests/saving/test_fix_sentencepiece_gguf_robustness.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ + tests/python/test_fast_language_model_text_only.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn # requires CUDA + dev toolchain, which the CPU-only ubuntu-latest - # runner does not have. The other 15 Bucket-A tests pass cleanly. + # runner does not have. The other Bucket-A tests pass cleanly. - name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU) # 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 8e202d161a..55c2c551a3 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -758,6 +758,7 @@ class InferenceBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + rag_scope: Optional[dict] = None, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -807,6 +808,7 @@ class InferenceBackend: max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, + rag_scope = rag_scope, ) def generate_chat_response( diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index ec3a216258..953c1d3141 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -50,6 +50,8 @@ from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) from core.inference.tool_call_parser import ( + RAG_MAX_SEARCHES_PER_TURN, + RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, ) @@ -4303,6 +4305,7 @@ class LlamaCppBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + rag_scope: Optional[dict] = None, seed: Optional[int] = None, disable_parallel_tool_use: bool = False, ) -> Generator[dict, None, None]: @@ -4314,12 +4317,21 @@ class LlamaCppBackend: {"type": "content", "text": "token"} -- streamed content tokens (cumulative) {"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative) """ - from core.inference.tools import execute_tool + from core.inference.tools import build_rag_autoinject, execute_tool if not self.is_loaded: raise RuntimeError("llama-server is not loaded") conversation = list(messages) + + # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits + # the same tool card + citations a real call would. + _auto = build_rag_autoinject(conversation, rag_scope) + if _auto: + for _ev in _auto["events"]: + yield _ev + conversation.extend(_auto["messages"]) + url = f"{self.base_url}/v1/chat/completions" _accumulated_completion_tokens = 0 _accumulated_predicted_ms = 0.0 @@ -4356,6 +4368,9 @@ class LlamaCppBackend: _MAX_BUFFER_CHARS = 32 _append_budget_exhausted_nudge = True + # RAG: cap knowledge-base searches per assistant turn. The controller is + # tool-agnostic, so this gate stays in the loop. + _kb_search_count = 0 # ── Re-prompt on plan-without-action ───────────────── # When the model describes what it intends to do (forward-looking @@ -4996,13 +5011,23 @@ class LlamaCppBackend: yield decision.tool_start_event() _effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout - result = execute_tool( - decision.tool_name, - decision.arguments, - cancel_event = cancel_event, - timeout = _effective_timeout, - session_id = session_id, - ) + # RAG: cap paraphrased KB re-searches that slip past the dup guard. + if ( + decision.tool_name == "search_knowledge_base" + and _kb_search_count >= RAG_MAX_SEARCHES_PER_TURN + ): + result = RAG_SEARCH_CAP_NUDGE + else: + result = execute_tool( + decision.tool_name, + decision.arguments, + cancel_event = cancel_event, + timeout = _effective_timeout, + session_id = session_id, + rag_scope = rag_scope, + ) + if decision.tool_name == "search_knowledge_base": + _kb_search_count += 1 completion = tool_controller.record_result(decision, result) yield completion.tool_end_event() conversation.append(completion.tool_message()) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 74a011093a..e394b342f0 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -821,6 +821,7 @@ class InferenceOrchestrator: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + rag_scope: Optional[dict] = None, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, **_unused, @@ -881,6 +882,7 @@ class InferenceOrchestrator: max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, + rag_scope = rag_scope, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 7d498a170d..7942edb6d7 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -23,6 +23,8 @@ from loggers import get_logger from core.inference.tool_call_parser import ( _TOOL_ALL_PATS, BUDGET_EXHAUSTED_NUDGE, + RAG_MAX_SEARCHES_PER_TURN, + RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, parse_tool_calls_from_text, strip_tool_markup, @@ -143,6 +145,7 @@ def run_safetensors_tool_loop( max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, + rag_scope: Optional[dict] = None, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -167,11 +170,23 @@ def run_safetensors_tool_loop( * ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}`` """ conversation = list(messages) + + # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search. + from core.inference.tools import build_rag_autoinject + + _auto = build_rag_autoinject(conversation, rag_scope) + if _auto: + for _ev in _auto["events"]: + yield _ev + conversation.extend(_auto["messages"]) + unrestricted_tools = not tools tool_controller = ToolLoopController( tools = None if unrestricted_tools else tools, auto_heal_tool_calls = auto_heal_tool_calls, ) + # RAG: cap knowledge-base searches per assistant turn (controller-agnostic). + kb_search_count = 0 final_attempt_done = False next_call_id = 0 @@ -498,17 +513,27 @@ def run_safetensors_tool_loop( yield decision.tool_start_event() eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout - try: - result = execute_tool( - decision.tool_name, - decision.arguments, - cancel_event = cancel_event, - timeout = eff_timeout, - session_id = session_id, - ) - except Exception as exc: - logger.exception("Tool %s raised: %s", decision.tool_name, exc) - result = f"Error: tool raised an exception: {exc}" + # RAG: cap paraphrased KB re-searches that slip past the dup guard. + if ( + decision.tool_name == "search_knowledge_base" + and kb_search_count >= RAG_MAX_SEARCHES_PER_TURN + ): + result = RAG_SEARCH_CAP_NUDGE + else: + try: + result = execute_tool( + decision.tool_name, + decision.arguments, + cancel_event = cancel_event, + timeout = eff_timeout, + session_id = session_id, + rag_scope = rag_scope, + ) + except Exception as exc: + logger.exception("Tool %s raised: %s", decision.tool_name, exc) + result = f"Error: tool raised an exception: {exc}" + if decision.tool_name == "search_knowledge_base": + kb_search_count += 1 completion = tool_controller.record_result(decision, result) yield completion.tool_end_event() diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index ccb777db1c..8d5d45269e 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -64,6 +64,15 @@ BUDGET_EXHAUSTED_NUDGE = ( "any more tools." ) +# The exact-args dup guard misses paraphrased re-searches, so also cap executed +# KB searches per turn, then nudge. +RAG_MAX_SEARCHES_PER_TURN = 3 +RAG_SEARCH_CAP_NUDGE = ( + "You have already searched the knowledge base several times this turn. " + "Do not search again. Answer the question using the passages already " + "retrieved above; if they do not contain the answer, say so plainly." +) + # Pre-compiled patterns reused by ``parse_tool_calls_from_text``. _TC_JSON_START_RE = re.compile(r"\s*\{") diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index 667ec7a33e..cb751ede3d 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -234,9 +234,11 @@ def is_tool_error(result: str) -> bool: def strip_result_for_model(result: str) -> str: - """Remove frontend-only image sentinels before feeding the model.""" - if "__IMAGES__:" in result: - return result.split("__IMAGES__:", 1)[0].rstrip() + """Remove frontend-only sentinels (image paths, RAG source map) before + feeding the result back to the model.""" + for sentinel in ("__IMAGES__:", "__RAG_SOURCES__:"): + if sentinel in result: + result = result.split(sentinel, 1)[0].rstrip() return result diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 023cab3e5e..b27fa6ff73 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -39,6 +39,9 @@ logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes +# Splits the UI source-map from the result; loops strip it (like __IMAGES__). +RAG_SOURCES_SENTINEL = "\n__RAG_SOURCES__:" + # Import these at module level so the preexec_fn closure triggers no imports in # the forked child (which can deadlock multi-threaded servers). _libc = None @@ -539,7 +542,41 @@ RENDER_HTML_TOOL = { }, } -ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, RENDER_HTML_TOOL] +# Duplicated (not imported from core.rag.tool) so the registry never pulls in +# the RAG stack; dispatch imports it lazily. +SEARCH_KNOWLEDGE_BASE_TOOL = { + "type": "function", + "function": { + "name": "search_knowledge_base", + "description": ( + "Search the user's uploaded documents and knowledge bases for " + "relevant passages. Use this whenever the question may be answered " + "by the attached documents, then cite the returned chunks." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Natural-language search query.", + }, + "top_k": { + "type": "integer", + "description": "Max chunks to return.", + }, + }, + "required": ["query"], + }, + }, +} + +ALL_TOOLS = [ + WEB_SEARCH_TOOL, + PYTHON_TOOL, + TERMINAL_TOOL, + RENDER_HTML_TOOL, + SEARCH_KNOWLEDGE_BASE_TOOL, +] # OpenAI's function.name regex ^[a-zA-Z0-9_-]{1,64}$, enforced before streaming. @@ -652,14 +689,19 @@ def execute_tool( cancel_event = None, timeout: int | None = _TIMEOUT_UNSET, session_id: str | None = None, + rag_scope: dict | None = None, ) -> str: """Execute a tool by name with the given arguments; returns a string. ``timeout``: int seconds, ``None`` = no limit, unset = ``_EXEC_TIMEOUT``. ``session_id``: optional ID for per-conversation sandbox isolation. + ``rag_scope``: hidden per-request RAG context the model never sees; consumed + by ``search_knowledge_base``. """ logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}") effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout + if name == "search_knowledge_base": + return _search_knowledge_base(arguments, rag_scope) if name == "render_html": return _render_html_result(arguments) if name.startswith(MCP_TOOL_PREFIX): @@ -696,6 +738,208 @@ def execute_tool( return f"Unknown tool: {name}" +def _opt_int(v) -> int | None: + try: + return int(v) if v is not None else None + except (TypeError, ValueError): + return None + + +def _scope_retrieval_kwargs(scope: dict) -> dict: + """Retrieval mode from rag_scope; candidate pools and RRF come from config.""" + mode = scope.get("mode") + return {"mode": mode if mode in ("hybrid", "dense", "lexical") else "hybrid"} + + +def _search_knowledge_base(arguments: dict, rag_scope: dict | None) -> str: + """Run the RAG search bound to the hidden per-request ``rag_scope`` (the model + supplies only ``query``/``top_k``). Lazy import; missing sqlite-vec degrades + to a friendly message.""" + scope = rag_scope or {} + query = (arguments or {}).get("query", "") + if not query or not str(query).strip(): + return "Error: query is empty." + try: + from storage import rag_db + if not rag_db.RAG_AVAILABLE: + return "Knowledge base search is unavailable on this server." + from core.rag.tool import search_knowledge_base_with_sources + except Exception as exc: # noqa: BLE001 + logger.warning("RAG tool unavailable: %s", exc) + return "Knowledge base search is unavailable on this server." + + top_k = _opt_int((arguments or {}).get("top_k") or scope.get("default_top_k")) + text, sources = search_knowledge_base_with_sources( + query = str(query), + scope_kb_id = scope.get("kb_id"), + scope_thread_id = scope.get("thread_id"), + top_k = top_k, + **_scope_retrieval_kwargs(scope), + ) + # Append the UI source-map after the sentinel; loops strip it before the model. + if sources: + import json as _json + return text + RAG_SOURCES_SENTINEL + _json.dumps(sources, ensure_ascii = False) + return text + + +# Forced first-pass RAG retrieval: a high cosine floor keeps it precise (fires on +# on-topic queries, skips weak ones) and helps small models that under-call the tool. +# Tunable via RAG_AUTOINJECT_MIN_SCORE. +_AUTOINJECT_DEFAULT_FLOOR = 0.70 + + +def _autoinject_enabled() -> bool: + return os.environ.get("RAG_AUTOINJECT", "1").strip().lower() not in ( + "0", + "false", + "no", + "off", + ) + + +def _autoinject_floor() -> float: + raw = os.environ.get("RAG_AUTOINJECT_MIN_SCORE") + if raw is not None: + try: + return float(raw) + except ValueError: + pass + return _AUTOINJECT_DEFAULT_FLOOR + + +# Lean: injecting the full top_k every turn prefills thousands of tokens. +_AUTOINJECT_DEFAULT_TOP_K = 4 + + +def _autoinject_top_k() -> int: + raw = os.environ.get("RAG_AUTOINJECT_TOP_K") + if raw is not None: + try: + return max(1, int(raw)) + except ValueError: + pass + return _AUTOINJECT_DEFAULT_TOP_K + + +def _last_user_text(conversation: list[dict]) -> str: + """Plain text of the most recent user turn (text parts only).""" + for msg in reversed(conversation): + if msg.get("role") != "user": + continue + content = msg.get("content") + if isinstance(content, str): + return content.strip() + if isinstance(content, list): + parts = [ + p.get("text", "") + for p in content + if isinstance(p, dict) and p.get("type") in ("text", "input_text") + ] + return " ".join(t for t in parts if t).strip() + return "" + return "" + + +def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> dict | None: + """Pre-retrieve the latest user turn; if a hit clears the cosine floor return + ``{"events": [...], "messages": [...]}`` to splice into the loop, else ``None``. + Toggle via ``rag_scope.autoinject`` (else env ``RAG_AUTOINJECT``); floor via + ``rag_scope.autoinject_min_score`` (else env ``RAG_AUTOINJECT_MIN_SCORE``). + + Also the small-model fallback: models below ~4B often answer from memory + instead of calling ``search_knowledge_base``, so forcing retrieval here keeps + attachments consulted regardless of model size.""" + if not rag_scope: + return None + enabled = rag_scope.get("autoinject") + if enabled is None: + enabled = _autoinject_enabled() + if not enabled: + return None + query = _last_user_text(conversation) + if not query: + return None + try: + from storage import rag_db + if not rag_db.RAG_AVAILABLE: + return None + from core.rag.tool import search_for_autoinject + except Exception as exc: # noqa: BLE001 + logger.warning("RAG auto-inject unavailable: %s", exc) + return None + + floor_override = rag_scope.get("autoinject_min_score") + floor = float(floor_override) if floor_override is not None else _autoinject_floor() + # Cap at the lean top_k, but honor a lower user setting. + lean_k = _autoinject_top_k() + sidebar_k = _opt_int(rag_scope.get("default_top_k")) + top_k = min(sidebar_k, lean_k) if sidebar_k is not None else lean_k + try: + found = search_for_autoinject( + query = query, + scope_kb_id = rag_scope.get("kb_id"), + scope_thread_id = rag_scope.get("thread_id"), + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG auto-inject retrieval failed: %s", exc) + return None + if not found: + logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + return None + + text, sources = found + import json as _json + import uuid as _uuid + + call_id = "rag_auto_" + _uuid.uuid4().hex[:12] + args = {"query": query} + full_result = text + RAG_SOURCES_SENTINEL + _json.dumps(sources, ensure_ascii = False) + events = [ + {"type": "status", "text": f"Searching documents: {query[:60]}"}, + { + "type": "tool_start", + "tool_name": "search_knowledge_base", + "tool_call_id": call_id, + "arguments": args, + }, + { + "type": "tool_end", + "tool_name": "search_knowledge_base", + "tool_call_id": call_id, + "result": full_result, + }, + {"type": "status", "text": ""}, + ] + messages = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": "search_knowledge_base", + "arguments": _json.dumps(args, ensure_ascii = False), + }, + } + ], + }, + { + "role": "tool", + "name": "search_knowledge_base", + "tool_call_id": call_id, + "content": text, + }, + ] + logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80]) + return {"events": events, "messages": messages} + + _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion) # Raw download cap > _MAX_PAGE_CHARS because SSR pages embed large # sections stripped during conversion; 512 KB reaches article content even diff --git a/studio/backend/core/rag/__init__.py b/studio/backend/core/rag/__init__.py new file mode 100644 index 0000000000..5dda48afc9 --- /dev/null +++ b/studio/backend/core/rag/__init__.py @@ -0,0 +1,16 @@ +# 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 core package. Import submodules lazily; keep this free of top-level +submodule imports to avoid pulling in heavy deps.""" + +__all__ = [ + "config", + "parsers", + "chunking", + "embeddings", + "store", + "retrieval", + "tool", + "ingestion", +] diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py new file mode 100644 index 0000000000..be8e341064 --- /dev/null +++ b/studio/backend/core/rag/captioner.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Caption figures with the loaded vision model and splice the text into the page +so images are searchable via the normal FTS5 + dense path. No-op (never raises) +without a vision model or on failure; gated by ``config.CAPTION_IMAGES``.""" + +from __future__ import annotations + +import base64 +import logging + +from . import config + +logger = logging.getLogger(__name__) + +_CAPTION_PROMPT = ( + "Describe this figure or image from a document in one or two concise " + "sentences, for search indexing. State what it depicts (e.g. a diagram, " + "chart, table or photo) and its key content. Do not add commentary." +) + + +def vision_endpoint() -> tuple[str, str] | None: + """``(base_url, model)`` for a loaded vision GGUF model, else None.""" + try: + from routes.inference import get_llama_cpp_backend + backend = get_llama_cpp_backend() + if getattr(backend, "is_loaded", False) and getattr(backend, "is_vision", False): + return backend.base_url, "local" + except Exception: # noqa: BLE001 - never let discovery break ingestion + return None + return None + + +def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + import httpx + + data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii") + payload = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": _CAPTION_PROMPT}, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + "max_tokens": 200, + "temperature": 0.2, + "stream": False, + # Off: thinking models would spend the budget reasoning, returning "". + "chat_template_kwargs": {"enable_thinking": False}, + } + try: + r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout) + r.raise_for_status() + text = r.json()["choices"][0]["message"]["content"] + return text.strip() or None + except Exception: # noqa: BLE001 - a failed caption is non-fatal + logger.debug("caption request failed", exc_info = True) + return None + + +def caption_images( + images: list, *, endpoint: tuple[str, str] | None = None +) -> dict[int, list[str]]: + """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when + disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``.""" + if not config.CAPTION_IMAGES or not images: + return {} + ep = endpoint or vision_endpoint() + if ep is None: + return {} + base_url, model = ep + + out: dict[int, list[str]] = {} + for img in images[: config.CAPTION_MAX_IMAGES]: + image_bytes = getattr(img, "image_bytes", None) + if not image_bytes: + continue + caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S) + if caption: + page = getattr(img, "page_number", None) or 0 + out.setdefault(int(page), []).append(caption) + return out + + +def splice_captions(pages: list, captions: dict[int, list[str]]) -> list: + """Append captions to their page's text so the chunker indexes them, keeping + figures attributable in retrieved chunks. Returns new ``Page`` objects.""" + if not captions: + return pages + from .parsers import Page + + out: list = [] + for page in pages: + caps = captions.get(page.page_number or 0) + if not caps: + out.append(page) + continue + extra = "".join(f"\n\n[Figure on page {page.page_number}: {c}]" for c in caps) + text = page.text + extra + out.append(Page(text = text, page_number = page.page_number, char_count = len(text))) + return out diff --git a/studio/backend/core/rag/chunking.py b/studio/backend/core/rag/chunking.py new file mode 100644 index 0000000000..c64acb4c60 --- /dev/null +++ b/studio/backend/core/rag/chunking.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 + +"""Page-aware recursive-separator chunking with token overlap. Each chunk records +its ``[page_char_start, page_char_end)`` span and ``source_page_index``, used by +the locator pass to highlight it on the PDF page.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +from .parsers import Page + +TokenCounter = Callable[[str], int] +SEPARATORS = ("\n# ", "\n## ", "\n### ", "\n\n", "\n", ". ", " ", "") + + +@dataclass(frozen = True) +class Chunk: + text: str + token_count: int + page_number: int | None + source_page_index: int + chunk_index: int + page_char_start: int + page_char_end: int + + +def _split(text: str, seps: tuple[str, ...], max_tokens: int, count: TokenCounter) -> list[str]: + """Recursively split into pieces each <= max_tokens (best effort). Pieces + rejoin to ``text`` exactly, so offsets are a running length.""" + if count(text) <= max_tokens: + return [text] + for i, sep in enumerate(seps): + parts = list(text) if sep == "" else text.split(sep) + if len(parts) <= 1: + continue + if sep: # re-attach the separator + parts = [p + sep for p in parts[:-1]] + parts[-1:] + out: list[str] = [] + for p in parts: + out.extend( + [p] if count(p) <= max_tokens else _split(p, seps[i + 1 :], max_tokens, count) + ) + return [p for p in out if p] + n = max(1, max_tokens * 4) + return [text[j : j + n] for j in range(0, len(text), n)] + + +def _merge( + pieces: list[str], starts: list[int], max_tokens: int, overlap: int, count: TokenCounter +) -> list[tuple[str, int, int]]: + """Greedy-merge pieces into <= max_tokens chunks with token overlap. + ``starts[i]`` is ``pieces[i]``'s page char offset; returns + ``(chunk_text, char_start, char_end)`` spans.""" + chunks: list[tuple[str, int, int]] = [] + buf: list[str] = [] + buf_starts: list[int] = [] + buf_tok = 0 + + def _flush() -> None: + raw = "".join(buf) + stripped = raw.strip() + if not stripped: + return + lead = len(raw) - len(raw.lstrip()) + trail = len(raw) - len(raw.rstrip()) + start = buf_starts[0] + lead + end = buf_starts[0] + len(raw) - trail + chunks.append((stripped, start, end)) + + for piece, start in zip(pieces, starts): + pt = count(piece) + if buf and buf_tok + pt > max_tokens: + _flush() + # Bound the carry so carry + this piece fits max_tokens; else a full + # overlap before a near-max piece overflows the embedder. + carry_budget = min(overlap, max(0, max_tokens - pt)) + carry, carry_starts, run = [], [], 0 + for prev, prev_start in zip(reversed(buf), reversed(buf_starts)): + if run + count(prev) > carry_budget: + break + carry.insert(0, prev) + carry_starts.insert(0, prev_start) + run += count(prev) + buf, buf_starts, buf_tok = carry, carry_starts, run + buf.append(piece) + buf_starts.append(start) + buf_tok += pt + if buf: + _flush() + return chunks + + +def chunk_pages( + pages: list[Page], *, max_tokens: int, overlap: int, count: TokenCounter +) -> list[Chunk]: + """Split each page into overlapping chunks, tracking per-page char offsets.""" + out: list[Chunk] = [] + for page_index, page in enumerate(pages): + pieces = _split(page.text, SEPARATORS, max_tokens, count) + # _split preserves offsets, so a running cursor gives exact ones. + starts: list[int] = [] + cursor = 0 + for piece in pieces: + starts.append(cursor) + cursor += len(piece) + for text, char_start, char_end in _merge(pieces, starts, max_tokens, overlap, count): + out.append( + Chunk( + text = text, + token_count = count(text), + page_number = page.page_number, + source_page_index = page_index, + chunk_index = len(out), + page_char_start = char_start, + page_char_end = char_end, + ) + ) + return out diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py new file mode 100644 index 0000000000..993423683c --- /dev/null +++ b/studio/backend/core/rag/config.py @@ -0,0 +1,41 @@ +# 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 config; every value is env-overridable.""" + +from __future__ import annotations + +import os + +EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "unsloth/bge-small-en-v1.5") +# Under bge's 512 limit, leaving headroom for the 2 special tokens (else overflow: +# llama-server 500s, ST truncates). Keep <= embedder_max - ~12. +CHUNK_TOKENS = int(os.environ.get("RAG_CHUNK_TOKENS", "500")) +CHUNK_OVERLAP = int(os.environ.get("RAG_CHUNK_OVERLAP", "64")) +TOP_K_LEXICAL = int(os.environ.get("RAG_TOP_K_LEXICAL", "30")) +TOP_K_DENSE = int(os.environ.get("RAG_TOP_K_DENSE", "30")) +TOP_K_HYBRID = int(os.environ.get("RAG_TOP_K_HYBRID", "10")) +RRF_K = int(os.environ.get("RAG_RRF_K", "60")) + +UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} + +# Figure captioning via the loaded vision model; off by default since each caption +# is a model call. MAX_IMAGES bounds per-doc cost. +CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "0") == "1" +CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "8")) +CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30")) + +# Embedder backend. "auto": sentence-transformers on a CUDA/ROCm GPU (torch fp16 +# wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes +# the vectors, so the index must be rebuilt. +EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto") +# llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this +# tiny model) and exact vs fp32, for ~30MB more on disk. +EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF") +EMBED_GGUF_VARIANT = os.environ.get("RAG_EMBED_GGUF_VARIANT", "F16") +EMBED_DEVICE = os.environ.get("RAG_EMBED_DEVICE", "auto") # "auto" | "gpu" | "cpu" +EMBED_HOST = os.environ.get("RAG_EMBED_HOST", "127.0.0.1") +EMBED_PORT = int(os.environ.get("RAG_EMBED_PORT", "0")) # 0 = auto-pick a free port +EMBED_BATCH = int(os.environ.get("RAG_EMBED_BATCH", "64")) +EMBED_STARTUP_TIMEOUT_S = float(os.environ.get("RAG_EMBED_STARTUP_TIMEOUT_S", "120")) +EMBED_REQUEST_TIMEOUT_S = float(os.environ.get("RAG_EMBED_REQUEST_TIMEOUT_S", "60")) diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py new file mode 100644 index 0000000000..2fadbcc6c6 --- /dev/null +++ b/studio/backend/core/rag/embed_llama_server.py @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGUF embedder over the bundled llama.cpp, served via HTTP (no torch). + +Opt-in (``RAG_EMBED_BACKEND=llama-server``). Runs a dedicated +``llama-server --embedding`` subprocess on its own port and calls its OpenAI-style +``/v1/embeddings`` + ``/tokenize``, fully isolated from the chat backend. + +Device is ``auto`` (GPU when present, else CPU, falling back to CPU if a GPU start +fails); ``RAG_EMBED_DEVICE`` forces it. We call only llama_cpp's *static* helpers +(no torch), copying the instance-coupled bits locally, since constructing a +``LlamaCppBackend`` runs an ``__init__`` reaper that kills any Studio llama-server +-- so each request re-spawns ours if it died (self-heal). +""" + +from __future__ import annotations + +import atexit +import logging +import os +import subprocess +import threading +import time +from functools import lru_cache +from pathlib import Path + +import httpx +import numpy as np + +from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs + +from . import config + +logger = logging.getLogger(__name__) + +# httpx transport errors meaning "the server is gone" -> trigger a respawn. +_TRANSPORT_ERRORS = ( + httpx.ConnectError, + httpx.ReadError, + httpx.RemoteProtocolError, + httpx.WriteError, +) + + +class LlamaServerBackend: + """Manages a llama.cpp embedding subprocess and talks to it over HTTP.""" + + def __init__(self) -> None: + # Lifecycle (spawn/restart/kill) is serialized; HTTP requests are not. + self._lifecycle_lock = threading.Lock() + self._process: subprocess.Popen | None = None + self._port: int | None = None + self._stdout_lines: list[str] = [] + self._stdout_thread: threading.Thread | None = None + self._dim: int | None = None + self._dim_lock = threading.Lock() + self._model_path: str | None = None + self._binary: str | None = None + # Sticky after an auto GPU start fails: later spawns stay on CPU. + self._force_cpu = False + # Pooled client; requests pass full URLs, so a respawn's new port needs + # no rebuild. + self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S) + atexit.register(self._shutdown) + + @property + def _base_url(self) -> str: + return f"http://{config.EMBED_HOST}:{self._port}" + + def _resolve_binary(self) -> str: + """Find llama-server, verify embeddings support, cache it. Raises if + missing/unsupported.""" + if self._binary is not None: + return self._binary + from core.inference.llama_cpp import LlamaCppBackend + + binary = LlamaCppBackend._find_llama_server_binary() + if not binary: + raise RuntimeError( + "llama-server binary not found; cannot use RAG_EMBED_BACKEND=" + "llama-server. Install llama.cpp or set LLAMA_SERVER_PATH / " + "UNSLOTH_LLAMA_CPP_PATH." + ) + self._assert_embedding_support(binary) + self._binary = binary + return binary + + @staticmethod + @lru_cache(maxsize = 8) + def _help_text(binary: str) -> str: + """`llama-server --help`, cached. Ignore exit code (some builds exit + non-zero on --help).""" + try: + proc = subprocess.run( + [binary, "--help"], + capture_output = True, + text = True, + timeout = 30, + **windows_hidden_subprocess_kwargs(), + ) + return (proc.stdout or "") + (proc.stderr or "") + except Exception as e: # noqa: BLE001 + logger.warning("could not run `llama-server --help`: %s", e) + return "" + + def _assert_embedding_support(self, binary: str) -> None: + help_text = self._help_text(binary) + # Empty help -> assume capable (a missing flag still fails at spawn). + if help_text and "--embedding" not in help_text: + raise RuntimeError( + "the bundled llama-server build lacks --embedding support; " + "RAG_EMBED_BACKEND=llama-server requires an embeddings-capable build" + ) + + def _resolve_model_path(self) -> str: + """Download (or cache-hit) the variant-matching, non-mmproj GGUF embedder, + returning its local path.""" + if self._model_path is not None: + return self._model_path + from huggingface_hub import hf_hub_download, list_repo_files + + repo = config.EMBED_GGUF_REPO + token = os.environ.get("HF_TOKEN") or None + files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")] + files = [f for f in files if "mmproj" not in f.lower()] + if not files: + raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}") + variant = config.EMBED_GGUF_VARIANT.lower() + match = [f for f in files if variant in f.lower()] or files + filename = sorted(match, key = len)[0] + logger.info("resolving GGUF embedder %s/%s", repo, filename) + self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token) + return self._model_path + + # Min free VRAM (MiB) for the embedder; below this, auto stays on CPU. + _MIN_GPU_FREE_MIB = 1024 + + def _use_gpu(self) -> bool: + """``RAG_EMBED_DEVICE``: ``gpu``/``cpu`` force it; ``auto`` uses a GPU when + present. A sticky CPU fallback (after an auto GPU start fails) wins.""" + dev = config.EMBED_DEVICE.lower() + if dev == "gpu": + return True + if dev == "cpu" or self._force_cpu: + return False + return self._gpu_available() # auto + + @staticmethod + def _gpu_available() -> bool: + """Apple Metal, or an NVIDIA/ROCm GPU with enough free VRAM. Reuses + llama_cpp's static probe (nvidia-smi first, so the common path needs no + torch).""" + from utils.hardware import is_apple_silicon + + if is_apple_silicon(): + return True # bundled mac build offloads to Metal + from core.inference.llama_cpp import LlamaCppBackend + + gpus = LlamaCppBackend._get_gpu_free_memory() # [(idx, free_mib)], honors CVD + return any(free >= LlamaServerBackend._MIN_GPU_FREE_MIB for _, free in gpus) + + def _build_cmd(self, binary: str, model_path: str, port: int, *, use_gpu: bool) -> list[str]: + # No --embd-normalize (not in every build; we normalize in Python to match + # the ST path). --fit off: don't auto-resize ctx/offload to device memory. + cmd = [ + binary, + "-m", + model_path, + "--host", + config.EMBED_HOST, + "--port", + str(port), + "--embedding", + "--pooling", + "cls", + "--fit", + "off", + ] + # -1 offloads every layer (matches the chat server); 0 keeps it on CPU. + cmd += ["-ngl", "-1" if use_gpu else "0"] + return cmd + + def _build_env(self, binary: str, *, use_gpu: bool) -> dict[str, str]: + env = child_env_without_native_path_secret() + env["LLAMA_SET_ROWS"] = "1" # ggml set_rows fast path + if use_gpu: + self._add_linux_cuda_libs(env, str(Path(binary).parent)) + else: + # Blank devices so a CUDA build stays on CPU and reserves no VRAM. + env["CUDA_VISIBLE_DEVICES"] = "" + return env + + @staticmethod + def _add_linux_cuda_libs(env: dict[str, str], binary_dir: str) -> None: + """Best-effort LD_LIBRARY_PATH so the prebuilt binary finds CUDA libs.""" + import glob + import platform + import sys + + if sys.platform == "win32": + return # Windows resolves CUDA via PATH in the inherited env. + arch = platform.machine() + lib_dirs = [binary_dir] + for pattern in ( + os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", "cu*", "lib"), + os.path.join(sys.prefix, "lib", "python*", "site-packages", "nvidia", "cudnn", "lib"), + ): + lib_dirs.extend(d for d in glob.glob(pattern) if os.path.isdir(d)) + for cuda_lib in ( + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{arch}-linux/lib", + "/usr/local/cuda-12/lib64", + ): + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) + existing = env.get("LD_LIBRARY_PATH", "") + joined = ":".join(lib_dirs) + env["LD_LIBRARY_PATH"] = f"{joined}:{existing}" if existing else joined + + def _drain_stdout(self, proc: subprocess.Popen) -> None: + """Drain the child's stdout so its pipe buffer never deadlocks; keep the + tail for crash diagnostics.""" + try: + for line in proc.stdout: # type: ignore[union-attr] + line = line.rstrip() + if line: + self._stdout_lines.append(line) + if len(self._stdout_lines) > 200: + del self._stdout_lines[:-200] + logger.debug("[llama-embed] %s", line) + except Exception: # noqa: BLE001 - drain thread must never raise + pass + + def _spawn(self) -> None: + """Start the embed server (caller holds the lock). On ``auto``, a failed + GPU start falls back to CPU once; explicit ``gpu``/``cpu`` do not.""" + use_gpu = self._use_gpu() + try: + self._spawn_once(use_gpu) + except RuntimeError: + auto = config.EMBED_DEVICE.lower() not in ("gpu", "cpu") + if use_gpu and auto: + logger.warning("embed server GPU start failed; falling back to CPU") + self._force_cpu = True + self._spawn_once(False) + else: + raise + + def _spawn_once(self, use_gpu: bool) -> None: + binary = self._resolve_binary() + model_path = self._resolve_model_path() + port = config.EMBED_PORT or self._find_free_port() + env = self._build_env(binary, use_gpu = use_gpu) + cmd = self._build_cmd(binary, model_path, port, use_gpu = use_gpu) + logger.info( + "starting llama-server embedder (%s): %s", + "gpu" if use_gpu else "cpu", + " ".join(cmd), + ) + self._stdout_lines = [] + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **windows_hidden_subprocess_kwargs(), + ) + self._process = proc + self._port = port + self._stdout_thread = threading.Thread( + target = self._drain_stdout, + args = (proc,), + daemon = True, + name = "llama-embed-stdout", + ) + self._stdout_thread.start() + if not self._wait_for_health(config.EMBED_STARTUP_TIMEOUT_S): + tail = "\n".join(self._stdout_lines[-30:]) + self._kill_process() + raise RuntimeError( + f"llama-server embedder failed to become healthy. Last output:\n{tail[:2000]}" + ) + + @staticmethod + def _find_free_port() -> int: + from core.inference.llama_cpp import LlamaCppBackend + return LlamaCppBackend._find_free_port() + + def _wait_for_health( + self, + timeout: float, + interval: float = 0.5, + ) -> bool: + """Poll /health until 200; bail early if the process exits.""" + deadline = time.monotonic() + timeout + url = f"{self._base_url}/health" + while time.monotonic() < deadline: + if not self._process_alive(): + code = None if self._process is None else self._process.returncode + logger.error("llama-server embedder exited early (code %s)", code) + return False + try: + if httpx.get(url, timeout = 2.0).status_code == 200: + return True + except (*_TRANSPORT_ERRORS, httpx.TimeoutException): + pass + time.sleep(interval) + logger.error("llama-server embedder health check timed out after %ss", timeout) + return False + + def _process_alive(self) -> bool: + return self._process is not None and self._process.poll() is None + + def _ensure_ready(self) -> None: + """Guarantee a live server, (re)spawning if needed. Double-checked so the + alive path takes no lock; self-heals after the chat reaper kills us.""" + if self._process_alive(): + return + with self._lifecycle_lock: + if self._process_alive(): + return + self._kill_process() + self._spawn() + + def _restart(self) -> None: + with self._lifecycle_lock: + self._kill_process() + self._spawn() + + def _kill_process(self) -> None: + proc = self._process + if proc is None: + return + try: + proc.terminate() + proc.wait(timeout = 5) + except subprocess.TimeoutExpired: + logger.warning("llama-server embedder did not exit on SIGTERM; killing") + proc.kill() + try: + proc.wait(timeout = 5) + except Exception: # noqa: BLE001 + pass + except Exception as e: # noqa: BLE001 + logger.warning("error killing llama-server embedder: %s", e) + finally: + self._process = None + if self._stdout_thread is not None: + self._stdout_thread.join(timeout = 2) + self._stdout_thread = None + + def _shutdown(self) -> None: + try: + self._kill_process() + finally: + try: + self._client.close() + except Exception: # noqa: BLE001 + pass + + def _post(self, path: str, payload: dict) -> dict: + """POST to the server, restarting once and retrying on a dropped connection + (the reaper may have killed us) or a timeout (the bundled build sometimes + wedges a request); a fresh server unsticks both.""" + last_exc: Exception | None = None + for attempt in range(2): + self._ensure_ready() + try: + resp = self._client.post(f"{self._base_url}{path}", json = payload) + resp.raise_for_status() + return resp.json() + except (*_TRANSPORT_ERRORS, httpx.TimeoutException) as e: + last_exc = e + if attempt == 0: + self._restart() + continue + except httpx.HTTPStatusError as e: + body = e.response.text[:500] if e.response is not None else "" + raise RuntimeError( + f"llama-server embedder POST {path} -> {e.response.status_code}: {body}" + ) from e + raise RuntimeError(f"llama-server embedder POST {path} failed after retry") from last_exc + + def encode( + self, + texts, + *, + model_name = None, + normalize = True, + ): + """Embed texts -> (N, dim) float32. ``model_name`` is ignored (the GGUF is + fixed by config). Normalizes in Python to match the ST backend.""" + n = len(texts) + if n == 0: + return np.zeros((0, self.dim()), dtype = np.float32) + rows: list[list[float]] = [] + batch = max(1, config.EMBED_BATCH) + for start in range(0, n, batch): + chunk = list(texts[start : start + batch]) + data = self._post( + "/v1/embeddings", + {"input": chunk, "model": "embedding", "encoding_format": "float"}, + ) + items = data.get("data", []) + if len(items) != len(chunk): + raise RuntimeError( + f"embedder returned {len(items)} vectors for {len(chunk)} inputs" + ) + # OpenAI spec lets the server reorder; sort back by index. + items = sorted(items, key = lambda d: d.get("index", 0)) + rows.extend(d["embedding"] for d in items) + arr = np.asarray(rows, dtype = np.float32) + if arr.ndim != 2: + raise RuntimeError(f"embedder returned ragged vectors: shape {arr.shape}") + if normalize: + norms = np.linalg.norm(arr, axis = 1, keepdims = True) + norms[norms == 0] = 1.0 + arr = arr / norms + return arr + + def dim(self, *, model_name = None) -> int: + """Embedding width, probed once via a 1-text encode and cached.""" + if self._dim is not None: + return self._dim + with self._dim_lock: + if self._dim is None: + vec = self.encode(["x"], normalize = False) + self._dim = int(vec.shape[1]) + return self._dim + + def warm(self, *, model_name = None) -> None: + """Start the server and probe dim off the request path.""" + self._ensure_ready() + self.dim() + + def token_counter(self, *, model_name = None): + """Count tokens via the GGUF's /tokenize so chunk sizing matches the + embedder. Cached per text.""" + + @lru_cache(maxsize = 4096) + def _count(text: str) -> int: + data = self._post("/tokenize", {"content": text, "add_special": False}) + return len(data.get("tokens", [])) + + return _count diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py new file mode 100644 index 0000000000..4e76e4fcaa --- /dev/null +++ b/studio/backend/core/rag/embeddings.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Dense embedder facade dispatching to a process-wide backend from +``config.EMBED_BACKEND`` (``auto`` picks by hardware): ``sentence-transformers`` +(torch) or ``llama-server`` (GGUF, no torch). + +Backends produce different vectors, so switching requires rebuilding the index. We +degrade to llama.cpp rather than crash when ST breaks on a machine: an init-time +probe falls back before any vector is produced (so spaces can't mix), and a +runtime ``encode`` failure swaps the process to llama-server for the rest of its +life (KBs already embedded with ST should then be reindexed). +""" + +from __future__ import annotations + +import logging +import os +import threading +from functools import lru_cache +from typing import Callable + +from utils.hardware.hardware import DeviceType, get_device + +from . import config + +logger = logging.getLogger(__name__) + +# "false" silences the fast tokenizer's fork warning; encode() flips it to "true" +# only during a batch tokenize (rayon speedup), then restores it. +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") + +_lock = threading.Lock() +# Serializes encode/tokenize (HF fast tokenizer isn't thread-safe). Separate from +# _lock so a long encode never blocks a reload. +_compute_lock = threading.Lock() +_model = None +_name: str | None = None + + +# Studio device -> torch device string. Apple has no torch device -> CPU. +_TORCH_DEVICE = {DeviceType.CUDA: "cuda", DeviceType.XPU: "xpu"} + + +def _device() -> str: + return _TORCH_DEVICE.get(get_device(), "cpu") + + +def _get(model_name: str | None = None): + """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 + for a ~1.5x speedup at negligible accuracy loss.""" + global _model, _name + name = model_name or config.EMBEDDING_MODEL + with _lock: + if _model is None or _name != name: + from sentence_transformers import SentenceTransformer + + device = _device() + logger.info("loading embedding model %s on %s", name, device) + _model = SentenceTransformer( + name, device = device, model_kwargs = {"torch_dtype": "float16"} + ) + _name = name + return _model + + +@lru_cache(maxsize = 1) +def _inference_ctx_factory(): + """``torch.inference_mode`` if torch imports, else ``nullcontext``. Returns the + factory so each call gets a fresh single-use guard.""" + try: + import torch + return torch.inference_mode + except Exception: # noqa: BLE001 - torch may be missing or broken + from contextlib import nullcontext + return nullcontext + + +def _inference_ctx(): + return _inference_ctx_factory()() + + +def _st_encode( + texts: list[str], + *, + model_name: str | None = None, + normalize: bool = True, +): + """ST encode -> (N, dim) float32. Serialized (fast-tokenizer borrow check), + under inference_mode when torch is present, with rayon enabled for the call.""" + model = _get(model_name) + with _compute_lock: + os.environ["TOKENIZERS_PARALLELISM"] = "true" + try: + with _inference_ctx(): + out = model.encode( + texts, + normalize_embeddings = normalize, + convert_to_numpy = True, + show_progress_bar = False, + ) + finally: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + # fp16 weights yield fp16 output; store float32 for sqlite-vec + stable cosine. + if hasattr(out, "astype"): + out = out.astype("float32", copy = False) + return out + + +def _st_dim(model_name: str | None = None) -> int: + return _get(model_name).get_sentence_embedding_dimension() + + +def _st_token_counter(model_name: str | None = None) -> Callable[[str], int]: + """Token counter using the model's tokenizer, under the compute lock (the same + fast tokenizer backs encode and isn't thread-safe), with rayon enabled for the + call. Mirrors ``_st_encode``.""" + tok = _get(model_name).tokenizer + + def _count(t: str) -> int: + with _compute_lock: + os.environ["TOKENIZERS_PARALLELISM"] = "true" + try: + return len(tok.encode(t, add_special_tokens = False)) + finally: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + return _count + + +class _SentenceTransformersBackend: + """Default backend; delegates to the module-level ST helpers so the ``_get`` + monkeypatch in tests keeps working.""" + + def encode( + self, + texts, + *, + model_name = None, + normalize = True, + ): + try: + return _st_encode(texts, model_name = model_name, normalize = normalize) + except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure + # ST loaded but this encode blew up; swap the process to the llama-server + # embedder (so later encodes stay in one space) and retry. + fallback = _switch_to_llama_fallback(st_err) + if fallback is None: + raise + return fallback.encode(texts, model_name = model_name, normalize = normalize) + + def token_counter(self, *, model_name = None): + return _st_token_counter(model_name) + + def dim(self, *, model_name = None): + return _st_dim(model_name) + + def warm(self, *, model_name = None): + _get(model_name) + + +_backend_lock = threading.Lock() +_backend = None +_backend_key: str | None = None + +_ST_ALIASES = frozenset({"sentence-transformers", "sentence_transformers", "st"}) +_LLAMA_ALIASES = frozenset( + {"llama-server", "llama_server", "llama", "llama.cpp", "llamacpp", "gguf"} +) +_AUTO_ALIASES = frozenset({"auto", ""}) + + +def _resolve_auto() -> str: + """Pick a backend for ``auto``: sentence-transformers when a CUDA/ROCm GPU is + present (torch fp16 wins bulk indexing), else the torch-free GGUF llama-server + -- or ST if its binary is missing. GPU check is torch-free (nvidia-smi).""" + from core.inference.llama_cpp import LlamaCppBackend + + if LlamaCppBackend._get_gpu_free_memory(): + return "sentence-transformers" + if LlamaCppBackend._find_llama_server_binary(): + return "llama-server" + return "sentence-transformers" + + +def _try_make_llama_backend(): + """A llama-server GGUF embedding backend if its binary is present, else None. + Construction is lazy -- no server starts until warm.""" + from core.inference.llama_cpp import LlamaCppBackend + + if not LlamaCppBackend._find_llama_server_binary(): + return None + from .embed_llama_server import LlamaServerBackend + + return LlamaServerBackend() + + +def _build_st_backend_or_fallback(): + """Build the ST backend, probing it by loading the model now. If the probe + raises (no torch, CUDA mismatch, bad wheel) and the GGUF llama-server embedder + is available, fall back to it. The probe runs before any vector is produced, so + this never mixes spaces. Re-raises if no embedder can start.""" + backend = _SentenceTransformersBackend() + try: + backend.warm(model_name = None) + return backend + except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure + fallback = _try_make_llama_backend() + if fallback is None: + raise + logger.warning( + "sentence-transformers embedder unavailable (%s); falling back to the " + "llama-server GGUF embedder", + st_err, + ) + return fallback + + +def _switch_to_llama_fallback(err): + """An ST encode failed at runtime even though the model had loaded. Swap the + process embedder to llama-server so every later encode stays in one space, and + return it (None if no binary). Vectors written before the swap were ST, so any + KB already embedded with ST should be reindexed.""" + global _backend, _backend_key + with _backend_lock: + if not isinstance(_backend, _SentenceTransformersBackend): + return _backend # another thread already swapped (or was never ST) + fallback = _try_make_llama_backend() + if fallback is None: + return None + logger.warning( + "sentence-transformers encode failed (%s); switching to the llama-server " + "embedder for the rest of this process. Reindex any knowledge base that " + "was already embedded with sentence-transformers.", + err, + ) + _backend = fallback + _backend_key = (config.EMBED_BACKEND or "auto").strip().lower() + return fallback + + +def _get_backend(): + """The process-wide embedding backend for ``config.EMBED_BACKEND``, built once. + Cached by the raw config value, so ``auto`` detection runs only on a miss and a + config change rebuilds it.""" + global _backend, _backend_key + raw = (config.EMBED_BACKEND or "auto").strip().lower() + with _backend_lock: + if _backend is not None and _backend_key == raw: + return _backend + key = _resolve_auto() if raw in _AUTO_ALIASES else raw + if key in _ST_ALIASES: + _backend = _build_st_backend_or_fallback() + elif key in _LLAMA_ALIASES: + # Imported lazily so the ST path never imports llama plumbing. + from .embed_llama_server import LlamaServerBackend + _backend = LlamaServerBackend() + else: + raise ValueError( + f"Unknown RAG_EMBED_BACKEND={config.EMBED_BACKEND!r}; expected " + "'auto', 'sentence-transformers' or 'llama-server'" + ) + _backend_key = raw + return _backend + + +def _reset_backend() -> None: + """Drop the cached backend (test teardown / re-init).""" + global _backend, _backend_key + with _backend_lock: + _backend = None + _backend_key = None + + +def warm(model_name: str | None = None) -> None: + """Eagerly load the embedder so the first real request isn't slow.""" + _get_backend().warm(model_name = model_name) + + +def encode( + texts: list[str], + *, + model_name: str | None = None, + normalize: bool = True, +): + """Embed texts into an (N, dim) float32 numpy array.""" + return _get_backend().encode(texts, model_name = model_name, normalize = normalize) + + +def dim(model_name: str | None = None) -> int: + """Embedding dimension for the (loaded) model.""" + return _get_backend().dim(model_name = model_name) + + +def token_counter(model_name: str | None = None) -> Callable[[str], int]: + """Callable counting tokens with the embedder's own tokenizer.""" + return _get_backend().token_counter(model_name = model_name) diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py new file mode 100644 index 0000000000..5ac0639f11 --- /dev/null +++ b/studio/backend/core/rag/ingestion.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""In-process threaded ingestion: parse -> chunk -> embed -> store. +``start_ingestion`` returns ``(document_id, job_id)`` immediately and runs on a +daemon thread, pushing progress onto a per-job queue (streamed as SSE by +``job_events``). Documents are deduped by content hash per scope.""" + +from __future__ import annotations + +import hashlib +import logging +import os +import queue +import threading + +from storage import rag_db + +from . import captioner, chunking, config, embeddings, parsers, store + +logger = logging.getLogger(__name__) + +# Per-job event queues, drained by job_events; ``None`` ends the stream. +_jobs: dict[str, "queue.Queue"] = {} +_jobs_lock = threading.Lock() + +_EMBED_BATCH = 64 # bounds peak memory + + +def _sha256_file(path: str) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + h.update(block) + return h.hexdigest() + + +def _emit(job_id: str, event: dict) -> None: + with _jobs_lock: + q = _jobs.get(job_id) + if q is not None: + q.put(event) + + +def _set_job( + conn, + job_id: str, + *, + status: str | None = None, + stage: str | None = None, + progress: float | None = None, + error: str | None = None, +) -> None: + conn.execute( + "UPDATE ingestion_jobs SET " + "status=COALESCE(?, status), " + "stage=COALESCE(?, stage), " + "progress=COALESCE(?, progress), " + "error=COALESCE(?, error) " + "WHERE id=?", + (status, stage, progress, error, job_id), + ) + conn.commit() + + +def _progress(conn, job_id: str, stage: str, progress: float) -> None: + _set_job(conn, job_id, status = "running", stage = stage, progress = progress) + _emit(job_id, {"type": "progress", "stage": stage, "progress": progress}) + + +def _embed_all(texts: list[str], model_name: str | None): + """Embed texts in batches into a flat vector list.""" + vectors: list = [] + for i in range(0, len(texts), _EMBED_BATCH): + batch = texts[i : i + _EMBED_BATCH] + out = embeddings.encode(batch, model_name = model_name, normalize = True) + vectors.extend(out) + return vectors + + +def _run( + job_id: str, document_id: str, scope: str, stored_path: str, model_name: str | None +) -> None: + conn = rag_db.get_connection() + try: + _progress(conn, job_id, "parsing", 0.1) + pages = parsers.parse(stored_path) + if config.CAPTION_IMAGES and stored_path.lower().endswith(".pdf"): + # Caption figures, splice into page text (no-op without a vision model). + try: + figures = parsers.render_pdf_figures( + stored_path, max_figures = config.CAPTION_MAX_IMAGES + ) + except Exception: + logger.warning("figure rendering failed for job %s", job_id, exc_info = True) + figures = [] + if figures: + _progress(conn, job_id, "captioning", 0.2) + captions = captioner.caption_images(figures) + pages = captioner.splice_captions(pages, captions) + + _progress(conn, job_id, "chunking", 0.3) + count = embeddings.token_counter(model_name) + chunks = chunking.chunk_pages( + pages, + max_tokens = config.CHUNK_TOKENS, + overlap = config.CHUNK_OVERLAP, + count = count, + ) + if not chunks: + store.set_document_status(conn, document_id, "completed", num_chunks = 0) + _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) + _emit(job_id, {"type": "complete", "num_chunks": 0}) + return + + _progress(conn, job_id, "embedding", 0.5) + vectors = _embed_all([c.text for c in chunks], model_name) + + # Locate each chunk's highlight regions (non-PDFs/failures yield none). + regions = None + if stored_path.lower().endswith(".pdf"): + try: + from . import locators + regions = locators.pdf_regions_for_chunks(stored_path, pages, chunks) + except Exception: + logger.warning("pdf region location failed for job %s", job_id, exc_info = True) + regions = None + + _progress(conn, job_id, "storing", 0.9) + store.add_chunks(conn, scope, document_id, chunks, vectors, regions) + store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks)) + + _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) + _emit(job_id, {"type": "complete", "num_chunks": len(chunks)}) + except Exception as exc: # noqa: BLE001 - report any failure to the client + logger.exception("ingestion job %s failed", job_id) + try: + store.set_document_status(conn, document_id, "failed", error = str(exc)) + _set_job(conn, job_id, status = "failed", stage = "error", error = str(exc)) + except Exception: # noqa: BLE001 + logger.exception("failed to record ingestion failure for job %s", job_id) + _emit(job_id, {"type": "error", "stage": "error", "error": str(exc)}) + finally: + conn.close() + _emit(job_id, None) + + +def start_ingestion( + scope: str, + kb_id: str | None, + thread_id: str | None, + filename: str, + stored_path: str, + *, + model_name: str | None = None, +) -> tuple[str, str]: + """Create the document + job rows and spawn the worker, returning + ``(document_id, job_id)``. A duplicate content hash in this scope returns the + existing id with an already-completed job (no re-ingest).""" + ext = os.path.splitext(stored_path)[1].lower() + if ext not in config.UPLOAD_EXTS: + raise ValueError(f"unsupported file type: {ext}") + + sha = _sha256_file(stored_path) + conn = rag_db.get_connection() + try: + existing = store.document_by_hash(conn, scope, sha) + if existing is not None: + job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) + with _jobs_lock: + _jobs[job_id] = queue.Queue() + _emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True}) + _emit(job_id, None) + return existing, job_id + + document_id = store.create_document( + conn, + scope = scope, + filename = filename, + sha256 = sha, + kb_id = kb_id, + thread_id = thread_id, + status = "pending", + stored_path = stored_path, + ) + job_id = _new_job(conn, document_id, scope) + finally: + conn.close() + + with _jobs_lock: + _jobs[job_id] = queue.Queue() + threading.Thread( + target = _run, + args = (job_id, document_id, scope, stored_path, model_name), + daemon = True, + ).start() + return document_id, job_id + + +def _new_job( + conn, + document_id: str, + scope: str, + *, + status: str = "pending", + progress: float = 0.0, +) -> str: + import uuid + from datetime import datetime, timezone + + job_id = str(uuid.uuid4()) + conn.execute( + "INSERT INTO ingestion_jobs(id, document_id, scope, status, stage, progress, created_at) " + "VALUES(?,?,?,?,?,?,?)", + ( + job_id, + document_id, + scope, + status, + None, + progress, + datetime.now(timezone.utc).isoformat(), + ), + ) + conn.commit() + return job_id + + +def job_events(job_id: str): + """Yield job events for SSE; ends when the worker signals completion.""" + with _jobs_lock: + q = _jobs.get(job_id) + if q is None: + return + while True: + event = q.get() + if event is None: + break + yield event + with _jobs_lock: + _jobs.pop(job_id, None) + + +def get_job_status(job_id: str) -> dict | None: + """Read the persisted ingestion job row (status / stage / progress / error).""" + conn = rag_db.get_connection() + try: + row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone() + return dict(row) if row else None + finally: + conn.close() diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py new file mode 100644 index 0000000000..57c0487486 --- /dev/null +++ b/studio/backend/core/rag/locators.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Map a chunk to highlight rectangles on its page (computed at ingest). + +The chunk's leading phrase is anchored in the page word list (``get_text("words")``), +so matching survives ligatures and dehyphenation that glyph-exact ``search_for`` +misses. Matched words union per line into rects normalized to 0..1. Missing +PyMuPDF, a too-short anchor, or no unique match yields no regions (never a guess). +""" + +from __future__ import annotations + +import unicodedata +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +# Anchor: up to MAX interior words from the chunk's start, shrunk toward MIN +# to recover a unique match. +MAX_ANCHOR_WORDS = 12 +MIN_ANCHOR_WORDS = 4 + + +@dataclass(frozen = True) +class LocatorMatch: + page_index: int + page_number: int | None + start: int + end: int + + +def _norm_token(token: str) -> str: + """Canonical match form: NFKC (decomposes ligatures), casefold, strip + surrounding punctuation/markdown. "" if punctuation-only.""" + token = unicodedata.normalize("NFKC", token).casefold() + return token.strip(" \t\r\n*#`[]()_.,;:!?\"'“”‘’-–—…|/\\") + + +def _anchor_tokens(page_text: str, match: LocatorMatch) -> list[str]: + """Normalized anchor tokens from the chunk's leading span. Drops first and last + token (boundaries often slice mid-word) when long enough.""" + segment = page_text[match.start : match.end] + raw = segment.split() + if len(raw) >= MIN_ANCHOR_WORDS + 2: + raw = raw[1:-1] + tokens = [t for t in (_norm_token(w) for w in raw) if t] + return tokens[:MAX_ANCHOR_WORDS] + + +def _find_subsequences(haystack: list[str], needle: list[str]) -> list[int]: + """Start indices where ``needle`` occurs consecutively in ``haystack``.""" + n, m = len(haystack), len(needle) + if m == 0 or m > n: + return [] + first = needle[0] + out: list[int] = [] + for i in range(n - m + 1): + if haystack[i] == first and haystack[i : i + m] == needle: + out.append(i) + return out + + +def _locate(page_words: list, needle: list[str]) -> list[int] | None: + """Matched word indices for the best anchor, or None. Tries the full anchor + then shorter prefixes, taking the first that matches exactly once; else the + first hit if still ambiguous.""" + # Skip punctuation-only words so they never break a phrase. + tokens: list[str] = [] + idx_map: list[int] = [] + for j, w in enumerate(page_words): + t = _norm_token(w[4]) + if t: + tokens.append(t) + idx_map.append(j) + + ambiguous_first: list[int] | None = None + for size in range(len(needle), MIN_ANCHOR_WORDS - 1, -1): + sub = needle[:size] + hits = _find_subsequences(tokens, sub) + if len(hits) == 1: + p = hits[0] + return [idx_map[p + k] for k in range(size)] + if hits and ambiguous_first is None: + p = hits[0] + ambiguous_first = [idx_map[p + k] for k in range(size)] + return ambiguous_first + + +def _rects_from_words(page_words: list, indices: list[int], pw: float, ph: float): + """Union matched words per (block, line) into normalized page rectangles.""" + lines: dict[tuple, list[float]] = {} + for j in indices: + w = page_words[j] + x0, y0, x1, y1 = float(w[0]), float(w[1]), float(w[2]), float(w[3]) + key = (w[5], w[6]) # block, line + box = lines.get(key) + if box is None: + lines[key] = [x0, y0, x1, y1] + else: + box[0], box[1] = min(box[0], x0), min(box[1], y0) + box[2], box[3] = max(box[2], x1), max(box[3], y1) + + out: list[dict[str, Any]] = [] + for x0, y0, x1, y1 in lines.values(): + w = x1 - x0 + h = y1 - y0 + if w <= 0 or h <= 0: + continue + out.append( + { + "x": max(0.0, min(1.0, x0 / pw)), + "y": max(0.0, min(1.0, y0 / ph)), + "width": max(0.0, min(1.0, w / pw)), + "height": max(0.0, min(1.0, h / ph)), + } + ) + return out + + +def _regions_for_match(doc: Any, page_text: str, match: LocatorMatch) -> list[dict[str, Any]]: + try: + if match.page_index < 0 or match.page_index >= len(doc): + return [] + needle = _anchor_tokens(page_text, match) + if len(needle) < MIN_ANCHOR_WORDS: + return [] + page = doc[match.page_index] + page_words = page.get_text("words") or [] + if not page_words: + return [] + indices = _locate(page_words, needle) + if not indices: + return [] + pw = float(page.rect.width) + ph = float(page.rect.height) + if pw <= 0 or ph <= 0: + return [] + rects = _rects_from_words(page_words, indices, pw, ph) + for r in rects: + r["pageIndex"] = match.page_index + r["pageNumber"] = match.page_number + return rects + except Exception: + return [] + + +def pdf_regions_for_chunks(pdf_path: Path, pages: list, chunks: list) -> list[list[dict[str, Any]]]: + """Region rects per chunk (parallel to ``chunks``), keyed off each chunk's + ``source_page_index`` / ``page_char_start`` / ``page_char_end``. Non-PDFs and + failures yield [], never an exception.""" + pdf_path = Path(pdf_path) + 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 + match = LocatorMatch( + page_index = int(page_index), + page_number = getattr(chunk, "page_number", None), + start = int(start), + end = int(end), + ) + regions.append(_regions_for_match(doc, pages[page_index].text, match)) + return regions + finally: + doc.close() diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py new file mode 100644 index 0000000000..84da941762 --- /dev/null +++ b/studio/backend/core/rag/parsers.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Document parsing -> list[Page], one dispatch with lazy optional deps. + +PDFs keep per-page boundaries (``page_number``); txt/md/docx/html return a single +page. ``parse(path, want_images=True)`` also returns embedded images. Heavy imports +are lazy, so importing this module never fails on a missing dep. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from html.parser import HTMLParser + +logger = logging.getLogger(__name__) + + +@dataclass(frozen = True) +class Page: + """A unit of extracted text. ``page_number`` is 1-based (None if N/A).""" + + text: str + page_number: int | None = None + char_count: int = 0 + + +@dataclass(frozen = True) +class ParsedImage: + """A raster image embedded in a document (PDF only).""" + + image_bytes: bytes + page_number: int | None + xref: int + + +def _page(text: str, page_number: int | None) -> Page: + return Page(text = text, page_number = page_number, char_count = len(text)) + + +class _Stripper(HTMLParser): + """Collect visible text, skipping