4f24b12cc9* Add customizable RAG embedding model setting and reorganize settings tabs Chat with files, project sources, and knowledge bases previously always embedded with unsloth/bge-small-en-v1.5. This adds a Settings option to pick any Hugging Face embedding model (or local path), with HF search autocomplete, server-side verification that the repo is actually an embedding model, and a save anyway escape hatch for offline or local models. The setting persists in app_settings and applies at runtime to both the sentence-transformers and llama-server GGUF embedder backends without a restart. Also reorganizes the General settings tab: Documents & RAG sits above Uploads, Helper LLM moved above the danger zone, and Model auto-switch (OpenAI API) moved to the bottom of the API tab. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support local model paths on the GGUF embedder and normalize default saves Found by simulation testing of the embedding model setting: Local paths saved as the embedding model now work on the llama-server GGUF backend (the default backend on macOS and CPU). A path to a .gguf file is used directly and a directory is scanned for a variant-matching non-mmproj .gguf, with a clear error when none exists. Previously a local path was sent to the HF hub API and failed with a repo lookup error. Saving the default model explicitly no longer stores an override, so is_custom stays false and the UI does not show a reset button for the default value. * Address review: stale-vector handling, GGUF derivation, save-time guards Review follow-ups, each verified by new tests: Re-uploading a document after an embedding model change now re-indexes instead of deduping by content hash. Documents record the embedder that produced their vectors (lazy embedding_model column, NULL legacy rows keep deduping) and a mismatch replaces the old document. A vector width change no longer bricks the dense index. ensure_vec drops and recreates chunks_vec when the dim changes (old vectors are in a foreign space and only block inserts) and search_dense returns empty on a width mismatch instead of surfacing a vec0 error, so lexical search keeps working until documents are re-uploaded. Saving a local sentence-transformers folder with no .gguf now returns 409 with a clear message when the install embeds via llama-server, instead of failing at first index. force still saves. A custom RAG_EMBEDDING_MODEL env without RAG_EMBED_GGUF_REPO now derives the -GGUF companion repo instead of silently keeping the bge GGUF on CPU and macOS installs. The resolved GGUF path is tagged with the repo captured at entry, so a setting change during a download cannot mark the old model as current. GGUF repo detection matches gguf as a whole name segment rather than a substring, hf_token is trimmed before verification, and the settings combobox drops a redundant state mirror of its controlled value. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shrink embedding model font to 11px in the input and dropdown The combobox wrapper applies className to the outer input group, so the size utility must target the inner input element; the previous text-xs never reached it and the field rendered at the browser default. * Show curated unsloth embedding models when the search field is empty The empty-query listing was the global top-downloads page, which holds no unsloth mirrors for the unsloth-first float to reorder, so the dropdown opened on third-party models. Match the model picker: curated unsloth listing when empty, whole-Hub search once a query is typed. * Address review: settings resilience and index consistency Keep the last known embedding model on settings store errors, remove the re-entrant dim lock in the llama-server backend, accept local GGUF saves and verify GGUF availability for HF repos on that backend, match local path embedders exactly in model list filters, drop same-width stale vectors from dense search, pin the embedder per ingestion job, and only replace completed documents after the re-index succeeds. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Consolidate the GGUF repo derivation tests * Trim to a single core embedding-model test * Address review: GGUF repo saves and cache race Accept a GGUF-named HF repo on the llama-server backend by verifying GGUF availability instead of the sentence-transformers metadata gate, and guard the settings cache with a generation counter so a read overlapping a save cannot repopulate it with the pre-save value. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
117 lines
6.2 KiB
Python
117 lines
6.2 KiB
Python
# 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
|
|
import re
|
|
|
|
DEFAULT_EMBEDDING_MODEL = "unsloth/bge-small-en-v1.5"
|
|
EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL)
|
|
# 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"))
|
|
|
|
# Whole-document context: a thread-attached file under the token budget is injected
|
|
# in full (every chunk, in order) instead of top-K retrieval; above it, use retrieval.
|
|
THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1"
|
|
WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000"))
|
|
|
|
UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
|
|
# Reject uploads larger than this, so one pathological file can't drive unbounded parse
|
|
# + vision work at ingest. 0 disables the cap. Default 200 MB.
|
|
MAX_UPLOAD_BYTES = int(os.environ.get("RAG_MAX_UPLOAD_BYTES", str(200 * 1024 * 1024)))
|
|
|
|
# Extract PDF text as layout-aware Markdown (pymupdf4llm) instead of flat text, so
|
|
# tables, headings and lists survive into chunks and retrieval. Falls back to plain
|
|
# PyMuPDF text when off, when pymupdf4llm is missing, or when extraction fails.
|
|
PDF_MARKDOWN = os.environ.get("RAG_PDF_MARKDOWN", "1") == "1"
|
|
|
|
# Figure captioning via the loaded vision model: detected figures are transcribed +
|
|
# described so they become searchable. On by default, a no-op without a vision model;
|
|
# the chat's "Describe figures & charts" toggle overrides it per upload.
|
|
CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "1") == "1"
|
|
# Total per-document tile budget (figure-bearing pages are tiled, see below).
|
|
CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "24"))
|
|
CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "60"))
|
|
# Larger than a one-line caption since captions transcribe every label. FIGURE_DPI is
|
|
# high enough to keep small box/axis labels legible when tiles are rendered.
|
|
CAPTION_MAX_TOKENS = int(os.environ.get("RAG_CAPTION_MAX_TOKENS", "768"))
|
|
FIGURE_DPI = int(os.environ.get("RAG_FIGURE_DPI", "200"))
|
|
# Figure pages are tiled into an overlapping ROWS x COLS grid of high-DPI tiles (plus
|
|
# an optional full page), so small labels and every sub-figure are covered without
|
|
# exact region detection. MAX_PAGES bounds figure pages; MAX_IMAGES bounds total tiles.
|
|
FIGURE_TILE_ROWS = int(os.environ.get("RAG_FIGURE_TILE_ROWS", "2"))
|
|
FIGURE_TILE_COLS = int(os.environ.get("RAG_FIGURE_TILE_COLS", "2"))
|
|
FIGURE_TILE_OVERLAP = float(os.environ.get("RAG_FIGURE_TILE_OVERLAP", "0.12"))
|
|
FIGURE_FULLPAGE = os.environ.get("RAG_FIGURE_FULLPAGE", "1") == "1"
|
|
CAPTION_MAX_PAGES = int(os.environ.get("RAG_CAPTION_MAX_PAGES", "4"))
|
|
|
|
# Scanned-PDF OCR: a page with little extractable text is rendered and transcribed by
|
|
# the vision model so it becomes searchable. Needs a vision model, else skipped (page
|
|
# stays empty). MIN_CHARS is the text length below which a page is treated as scanned.
|
|
OCR_SCANNED = os.environ.get("RAG_OCR_SCANNED", "1") == "1"
|
|
OCR_MIN_CHARS = int(os.environ.get("RAG_OCR_MIN_CHARS", "16"))
|
|
OCR_MAX_PAGES = int(os.environ.get("RAG_OCR_MAX_PAGES", "20"))
|
|
OCR_DPI = int(os.environ.get("RAG_OCR_DPI", "150"))
|
|
OCR_TIMEOUT_S = float(os.environ.get("RAG_OCR_TIMEOUT_S", "60"))
|
|
OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048"))
|
|
|
|
# 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")
|
|
|
|
|
|
def effective_embedding_model() -> str:
|
|
"""The embedding model actually in use: the persisted Settings override when
|
|
one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a
|
|
Settings change applies without a restart."""
|
|
try:
|
|
from utils.embedding_model_settings import get_rag_embedding_model
|
|
return get_rag_embedding_model()
|
|
except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot)
|
|
return EMBEDDING_MODEL
|
|
|
|
|
|
def _names_gguf(model: str) -> bool:
|
|
"""True when "gguf" appears as a whole name segment, so plain substrings
|
|
like "bigguf" don't count."""
|
|
return "gguf" in re.split(r"[^a-z0-9]+", model.lower())
|
|
|
|
|
|
def effective_gguf_repo() -> str:
|
|
"""GGUF repo for the llama-server backend, tracking the effective model.
|
|
|
|
An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom
|
|
model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its
|
|
``-GGUF`` companion repo (the unsloth convention the default pair follows),
|
|
or is used as-is when it already names a GGUF repo.
|
|
"""
|
|
if "RAG_EMBED_GGUF_REPO" in os.environ:
|
|
return EMBED_GGUF_REPO
|
|
model = effective_embedding_model()
|
|
if model == DEFAULT_EMBEDDING_MODEL:
|
|
return EMBED_GGUF_REPO
|
|
if _names_gguf(model):
|
|
return model
|
|
return f"{model}-GGUF"
|
|
|
|
|
|
# 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"))
|