ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio RAG: remove multimodal image embedding and the mode field/selector

Browse source 
Matches #5910's text-only footprint. Removes image-vector embedding
(encode_images, _stream_image_chunks, the _BGEVLAdapter CLIP shim), the
multimodal `mode`/KBMode concept + VL embedders (single text embedder
now), the mode selector UI across the KB dialogs + thread settings, the
MM badges, the /images serving route, and the dead image rendering in
the search tool card. Captioning (figure text spliced into markdown)
stays — #5910 keeps it too. DB mode/image columns left dormant (no
migration). RagDefaultsSection dropped (no controls left).
This commit is contained in:
Roland Tannous 2026-06-02 13:37:29 +04:00
commit 2366104c4f
27 changed files with 36 additions and 1019 deletions
Download patch file Download diff file Expand all files Collapse all files

156
studio/backend/core/rag/embeddings.py
View file

@ -22,10 +22,6 @@ _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.
@ -36,133 +32,6 @@ def _load(model_name: str) -> Any:
)
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 data_process re-opens each item via Image.open(...),
# which needs a file-like (.read()) or path — NOT a pre-opened PIL
# Image. Pass BytesIO; PIL Images get rebuffered via an in-memory PNG.
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
@ -206,31 +75,6 @@ def encode(
)
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)

132
studio/backend/core/rag/ingestion.py
View file

@ -58,8 +58,6 @@ def _subprocess_worker(
overlap: int,
batch_size: int,
out_queue: Any,
mode: str = "text",
document_id: str = "",
vlm_url: str | None = None,
vlm_model: str | None = None,
enable_captions: bool = True,
@ -83,8 +81,8 @@ def _subprocess_worker(
from core.rag.parsers import inline_image_captions, parse
out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05})
# Always extract images to caption + splice for both modes. Text mode uses
# captions inline in markdown; multimodal also embeds raw images as image-kind chunks.
# Extract images so figures can be captioned and spliced into the page
# markdown, where the chunker indexes them like any other text.
parsed = parse(Path(stored_path), want_images = True)
pages = parsed.pages
if not pages and not parsed.images:
@ -95,7 +93,6 @@ def _subprocess_worker(
# Caption figures once (chat VLM if available, else helper VLM), then splice
# captions into the page markdown so the chunker indexes them like any text.
# Multimodal reuses these captions in _stream_image_chunks below — no duplicate VLM calls.
captions: list[str] = []
if parsed.images and enable_captions:
out_queue.put(
@ -147,17 +144,7 @@ def _subprocess_worker(
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})
out_queue.put({"type": "complete", "num_chunks": text_count})
except Exception as exc: # noqa: BLE001
logger.exception("ingestion subprocess failed")
out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"})
@ -232,111 +219,6 @@ def _run_standard_chunking(
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)
# --- Job manager (parent side) ---
@ -754,17 +636,15 @@ def enqueue_ingestion(
kb_id: str | None = None,
thread_id: str | None = None,
embedding_model: str | None = None,
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) or RAG_EMBEDDING_MODEL
model_name = embedding_model or resolve_embedder() or RAG_EMBEDDING_MODEL
# Probe the loaded chat backend so the subprocess can caption figures with the
# user's own vision model (no extra VRAM). Runs for both modes — text splices
# captions into markdown, multimodal also feeds them to the image-vector encoder.
# user's own vision model (no extra VRAM). Text splices captions into markdown.
# No vision chat model loaded → falls back to the helper VLM (pre-cached at startup).
# Skipped when captioning is disabled for this upload.
vlm_url: str | None = None
@ -811,8 +691,6 @@ def enqueue_ingestion(
RAG_CHUNK_OVERLAP,
RAG_EMBED_BATCH_SIZE,
out_queue,
mode,
document_id,
vlm_url,
vlm_model,
enable_captions,

3
studio/backend/core/rag/scope.py
View file

@ -38,7 +38,6 @@ def resolve_scope_embedder(scope: str) -> str | None:
explicit = per_thread.get("embedding_model") or defaults.get("embedding_model")
if explicit:
return explicit
mode = per_thread.get("mode") or defaults.get("mode") or "text"
return resolve_embedder(mode)
return resolve_embedder()
return None

9
studio/backend/core/rag/tool.py
View file

@ -120,15 +120,6 @@ def _format_hits_for_llm(hits: list[dict], start_id: int = 0) -> str:
tokens = hit.get("token_count")
if tokens:
attrs.append(f'tokens="{tokens}"')
kind = hit.get("kind")
if kind and kind != "text":
attrs.append(f'kind="{_xml_attr(kind)}"')
image_path = hit.get("image_path")
if kind == "image" and image_path and document_id:
# Mirror routes/rag.py search-response shape so the frontend tool card
# can render the image inline via the same route.
image_url = f"/api/rag/images/{document_id}/{Path(image_path).name}"
attrs.append(f'image_url="{_xml_attr(image_url)}"')
text = (hit.get("text") or "").strip()
blocks.append(f"<chunk {' '.join(attrs)}>\n{text}\n</chunk>")
return "\n\n".join(blocks)

93
studio/backend/routes/rag.py
View file

@ -63,14 +63,11 @@ logger = get_logger(__name__)
# --- Pydantic schemas ---
KBMode = Literal["text", "multimodal"]
class CreateKBRequest(BaseModel):
name: str = Field(min_length = 1, max_length = 200)
description: str | None = None
embedding_model: str | None = None
mode: KBMode = "text"
class KBResponse(BaseModel):
@ -78,7 +75,6 @@ class KBResponse(BaseModel):
name: str
description: str | None
embedding_model: str
mode: KBMode
created_at: int
@ -169,14 +165,11 @@ from core.rag.scope import resolve_scope_embedder as _resolve_scope_embedder #
def _row_to_kb(row: Any) -> KBResponse:
keys = row.keys() if hasattr(row, "keys") else ()
mode = row["mode"] if "mode" in keys else "text"
return KBResponse(
id = row["id"],
name = row["name"],
description = row["description"],
embedding_model = row["embedding_model"],
mode = mode,
created_at = row["created_at"],
)
@ -280,7 +273,6 @@ def _start_ingestion(
kb_id: str | None,
thread_id: str | None,
embedding_model: str,
mode: str = "text",
caption_images: bool = True,
content_hash: str | None = None,
) -> UploadResponse:
@ -339,7 +331,6 @@ def _start_ingestion(
kb_id = kb_id,
thread_id = thread_id,
embedding_model = embedding_model,
mode = mode,
enable_captions = caption_images,
)
return UploadResponse(document_id = document_id, job_id = job_id, filename = filename)
@ -366,8 +357,7 @@ def create_knowledge_base(
from utils.rag.config import resolve_embedder
kb_id = str(uuid4())
# No override: resolve the embedder from the KB mode.
embedding_model = payload.embedding_model or resolve_embedder(payload.mode)
embedding_model = payload.embedding_model or resolve_embedder()
created_at = _now_ms()
with closing_connection() as conn:
try:
@ -375,8 +365,8 @@ def create_knowledge_base(
"""
INSERT INTO rag_knowledge_bases
(id, name, description, owner_user_id, embedding_model,
mode, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
created_at)
VALUES (?, ?, ?, ?, ?, ?)
""",
(
kb_id,
@ -384,7 +374,6 @@ def create_knowledge_base(
payload.description,
current_subject,
embedding_model,
payload.mode,
created_at,
),
)
@ -399,7 +388,6 @@ def create_knowledge_base(
name = payload.name,
description = payload.description,
embedding_model = embedding_model,
mode = payload.mode,
created_at = created_at,
)
@ -416,14 +404,12 @@ def list_knowledge_bases(
class RagDefaults(BaseModel):
mode: KBMode = "text"
embedding_model: str | None = None
class UpdateRagDefaultsRequest(BaseModel):
"""Patch shape — only fields present overwrite stored values."""
mode: KBMode | None = None
embedding_model: str | None = None
@ -436,7 +422,6 @@ def _load_rag_defaults() -> RagDefaults:
if not isinstance(raw, dict):
raw = {}
return RagDefaults(
mode = raw.get("mode") or "text",
embedding_model = raw.get("embedding_model"),
)
@ -461,7 +446,7 @@ def warmup_rag_embedder(
from utils.rag.config import resolve_embedder
defaults = _load_rag_defaults()
model_name = defaults.embedding_model or resolve_embedder(defaults.mode)
model_name = defaults.embedding_model or resolve_embedder()
try:
embeddings.get_embedder(model_name)
except Exception as exc: # noqa: BLE001
@ -477,7 +462,6 @@ def set_rag_defaults(
current_subject: str = Depends(get_current_subject),
) -> RagDefaults:
current = _load_rag_defaults()
new_mode = payload.mode or current.mode
# PATCH-style: empty string clears, null/missing keeps current.
if payload.embedding_model is None:
new_embedder = current.embedding_model
@ -489,24 +473,20 @@ def set_rag_defaults(
upsert_chat_settings_merge(
{
_DEFAULTS_KEY: {
"mode": new_mode,
"embedding_model": new_embedder,
}
}
)
return RagDefaults(
mode = new_mode,
embedding_model = new_embedder,
)
class ThreadRagSettings(BaseModel):
mode: KBMode = "text"
embedding_model: str | None = None
class UpdateThreadRagSettingsRequest(BaseModel):
mode: KBMode | None = None
embedding_model: str | None = None
# Reingest-only (not persisted); omit or None keeps captioning on.
caption_images: bool | None = None
@ -524,7 +504,6 @@ def _load_thread_settings(thread_id: str) -> ThreadRagSettings:
raw = {}
fallback = _load_rag_defaults()
return ThreadRagSettings(
mode = raw.get("mode") or fallback.mode,
embedding_model = raw.get("embedding_model") or fallback.embedding_model,
)
@ -550,7 +529,6 @@ def set_thread_rag_settings(
current_subject: str = Depends(get_current_subject),
) -> ThreadRagSettings:
current = _load_thread_settings(thread_id)
new_mode = payload.mode or current.mode
if payload.embedding_model is None:
new_embedder = current.embedding_model
elif payload.embedding_model.strip() == "":
@ -561,13 +539,11 @@ def set_thread_rag_settings(
upsert_chat_settings_merge(
{
_thread_settings_key(thread_id): {
"mode": new_mode,
"embedding_model": new_embedder,
}
}
)
return ThreadRagSettings(
mode = new_mode,
embedding_model = new_embedder,
)
@ -575,7 +551,6 @@ def set_thread_rag_settings(
class ReingestKBRequest(BaseModel):
"""All fields optional — omitting one keeps the KB's current value."""
mode: KBMode | None = None
embedding_model: str | None = None
# Not persisted on the KB; omit or None keeps captioning on for the rebuild.
caption_images: bool | None = None
@ -590,7 +565,6 @@ def _reingest_scope(
*,
kb_id: str | None,
thread_id: str | None,
mode: str,
embedding_model: str,
caption_images: bool = True,
) -> ReingestResponse:
@ -638,7 +612,6 @@ def _reingest_scope(
kb_id = kb_id,
thread_id = thread_id,
embedding_model = embedding_model,
mode = mode,
caption_images = caption_images,
)
job_ids.append(upload.job_id)
@ -655,34 +628,24 @@ def reingest_knowledge_base(
payload: ReingestKBRequest,
current_subject: str = Depends(get_current_subject),
) -> ReingestResponse:
from utils.rag.config import resolve_embedder
kb_row = _kb_or_404(kb_id)
keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
current_mode = kb_row["mode"] if "mode" in keys else "text"
current_embedder = kb_row["embedding_model"]
new_mode = payload.mode or current_mode
new_embedder = payload.embedding_model or (
current_embedder if new_mode == current_mode else resolve_embedder(new_mode)
)
new_embedder = payload.embedding_model or current_embedder
with closing_connection() as conn:
conn.execute(
"""
UPDATE rag_knowledge_bases
SET mode = ?, embedding_model = ?
SET embedding_model = ?
WHERE id = ?
""",
(new_mode, new_embedder, kb_id),
(new_embedder, kb_id),
)
conn.commit()
return _reingest_scope(
kb_id = kb_id,
thread_id = None,
mode = new_mode,
embedding_model = new_embedder,
caption_images = payload.caption_images is not False,
)
@ -702,7 +665,7 @@ def reingest_thread_documents(
if payload is None:
payload = UpdateThreadRagSettingsRequest()
if payload.mode is not None or payload.embedding_model is not None:
if payload.embedding_model is not None:
settings = set_thread_rag_settings(
thread_id,
payload,
@ -711,11 +674,10 @@ def reingest_thread_documents(
else:
settings = _load_thread_settings(thread_id)
embedder = settings.embedding_model or resolve_embedder(settings.mode)
embedder = settings.embedding_model or resolve_embedder()
return _reingest_scope(
kb_id = None,
thread_id = thread_id,
mode = settings.mode,
embedding_model = embedder,
caption_images = payload.caption_images is not False,
)
@ -752,9 +714,6 @@ async def upload_kb_document(
) -> UploadResponse:
kb_row = _kb_or_404(kb_id)
stored_path, filename, byte_size, content_hash = await _save_upload(file)
# Tolerate pre-Phase-3 rows missing mode.
kb_keys = kb_row.keys() if hasattr(kb_row, "keys") else ()
mode = kb_row["mode"] if "mode" in kb_keys else "text"
return _start_ingestion(
filename = filename,
stored_path = stored_path,
@ -763,7 +722,6 @@ async def upload_kb_document(
kb_id = kb_id,
thread_id = None,
embedding_model = kb_row["embedding_model"],
mode = mode,
caption_images = caption_images,
content_hash = content_hash,
)
@ -781,7 +739,7 @@ async def upload_thread_document(
# No chat_threads check — fresh threads aren't persisted until first run.
stored_path, filename, byte_size, content_hash = await _save_upload(file)
settings = _load_thread_settings(thread_id)
embedder = settings.embedding_model or resolve_embedder(settings.mode)
embedder = settings.embedding_model or resolve_embedder()
return _start_ingestion(
filename = filename,
stored_path = stored_path,
@ -790,7 +748,6 @@ async def upload_thread_document(
kb_id = None,
thread_id = thread_id,
embedding_model = embedder,
mode = settings.mode,
caption_images = caption_images,
content_hash = content_hash,
)
@ -826,28 +783,6 @@ def list_thread_documents(
return DocumentListResponse(documents = [_row_to_document(r) for r in rows])
@router.get("/images/{document_id}/{filename}")
def get_rag_image(
document_id: str,
filename: str,
current_subject: str = Depends(get_current_subject),
) -> FileResponse:
"""Serve an extracted image; realpath-check against the uploads root."""
document_for_subject_or_404(document_id, current_subject)
if "/" in filename or "\\" in filename or filename.startswith("."):
raise HTTPException(status_code = 400, detail = "Invalid filename")
root = Path(os.path.realpath(rag_uploads_root() / "images"))
candidate = rag_uploads_root() / "images" / document_id / filename
try:
real = Path(os.path.realpath(candidate))
real.relative_to(root)
except (OSError, ValueError) as exc:
raise HTTPException(status_code = 404, detail = "Image not found") from exc
if not real.is_file():
raise HTTPException(status_code = 404, detail = "Image not found")
return FileResponse(str(real))
@router.delete("/documents/{document_id}")
def delete_document(
document_id: str,
@ -1375,10 +1310,6 @@ def get_document_preview_target(
chunk_kind: PreviewChunkKind = chunk_row["kind"] or "text" # type: ignore[assignment]
image_url: str | None = None
if chunk_kind == "image" and chunk_row["image_path"]:
image_url = (
f"/api/rag/images/{doc_row['id']}/" f"{Path(chunk_row['image_path']).name}"
)
return PreviewTargetResponse(
**base,
@ -1600,10 +1531,6 @@ def search(
continue
kind = meta.get("kind", "text") or "text"
image_url: str | None = None
if kind == "image" and meta.get("image_path"):
image_url = (
f"/api/rag/images/{meta['document_id']}/{Path(meta['image_path']).name}"
)
out.append(
SearchHit(
chunk_id = hit.chunk_id,

54
studio/backend/tests/test_rag_preview_routes.py
View file

@ -570,57 +570,3 @@ class TestFileRoute:
assert resp.headers.get("x-content-type-options") == "nosniff"
cc = resp.headers.get("cache-control", "")
assert "private" in cc
# ── /images tests ─────────────────────────────────────────────────────
class TestImageRoute:
def test_image_route_wrong_subject_returns_404(self, app, db_env, monkeypatch):
"""Extracted images require the same document authorization as /file."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
images = uploads / "images" / doc_id
images.mkdir(parents = True, exist_ok = True)
image = images / "figure.png"
image.write_bytes(b"\x89PNG\r\n\x1a\n")
stored = uploads / "report.pdf"
stored.write_bytes(b"%PDF-1.4")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner = "alice")
_insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf")
client = _make_client(app, "mallory")
try:
resp = client.get(f"/api/rag/images/{doc_id}/figure.png")
finally:
_clear_overrides(app)
assert resp.status_code == 404
def test_image_route_authorized_subject_gets_image(self, app, db_env, monkeypatch):
"""Authorized subject can still fetch an extracted image."""
doc_id, kb_id = _uid(), _uid()
uploads = db_env / "rag" / "uploads"
images = uploads / "images" / doc_id
images.mkdir(parents = True, exist_ok = True)
image = images / "figure.png"
image.write_bytes(b"\x89PNG\r\n\x1a\n")
stored = uploads / "report.pdf"
stored.write_bytes(b"%PDF-1.4")
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(db_env))
with studio_db.get_connection() as conn:
_insert_kb(conn, kb_id, owner = "alice")
_insert_doc(conn, doc_id, kb_id, str(stored), "report.pdf")
client = _make_client(app, "alice")
try:
resp = client.get(f"/api/rag/images/{doc_id}/figure.png")
finally:
_clear_overrides(app)
assert resp.status_code == 200
assert resp.content.startswith(b"\x89PNG")

26
studio/backend/utils/rag/config.py
View file

@ -31,26 +31,12 @@ RAG_EMBEDDING_MODEL: str = (
or "BAAI/bge-small-en-v1.5"
)
# Default embedder per mode.
#
# Text mode is the default: PDF figures are captioned at ingest (chat VLM or
# helper gemma-3n fallback) and spliced into the page markdown before chunking,
# so a single 384-d text embedder handles retrieval. Multimodal adds image-vector
# rows on top via Qwen3-VL-Embedding-2B (2 B, 2048-d, no CLIP text cap — full
# 512-token chunks embed losslessly).
#
# Alternative multimodal embedders kept for manual override:
# - "BAAI/BGE-VL-large" — smaller (~400 M / 768-d) but CLIP-family with a
# 77-token text cap; routed via `_BGEVLAdapter` in core/rag/embeddings.py.
RAG_EMBEDDER_MATRIX: dict[str, str] = {
"text": "BAAI/bge-small-en-v1.5",
"multimodal": "Qwen/Qwen3-VL-Embedding-2B",
}
def resolve_embedder(mode: str) -> str:
"""Embedder for the given mode; unknown modes fall back to RAG_EMBEDDING_MODEL."""
return RAG_EMBEDDER_MATRIX.get(mode, RAG_EMBEDDING_MODEL)
# A single text embedder handles retrieval. PDF figures are captioned at ingest
# (chat VLM or helper gemma-3n fallback) and spliced into the page markdown
# before chunking, so the 384-d text embedder covers figure content too.
def resolve_embedder() -> str:
"""The configured RAG embedder."""
return RAG_EMBEDDING_MODEL
RAG_CHUNK_SIZE: int = _env_int("UNSLOTH_RAG_CHUNK_SIZE", 512)

5
studio/frontend/src/__tests__/knowledge-bases-tab.test.tsx
View file

@ -48,11 +48,6 @@ vi.mock("@/features/rag/components/thread-index-list", () => ({
React.createElement("div", { "data-testid": "thread-index-list" }),
}));
vi.mock("@/features/rag/components/rag-defaults-section", () => ({
RagDefaultsSection: () =>
React.createElement("div", { "data-testid": "rag-defaults-section" }),
}));
beforeEach(() => {
mockUsePreviewStore.__state.target = null;
mockUsePreviewStore.__state.status = "idle";

1
studio/frontend/src/__tests__/preview-a11y.test.tsx
View file

@ -47,7 +47,6 @@ function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
targetPage: 1,
snippet: "safe extracted text",
kind: "text",
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,

1
studio/frontend/src/__tests__/preview-panel.test.tsx
View file

@ -163,7 +163,6 @@ function makeTarget(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
targetPage: null,
snippet: null,
kind: null,
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,

1
studio/frontend/src/__tests__/preview-pdf-smoke.test.tsx
View file

@ -86,7 +86,6 @@ function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
targetPage: 1,
snippet: "target phrase appears here",
kind: "text",
imageUrl: null,
sourcePageIndex: 0,
pageCharStart: 0,
pageCharEnd: 13,

1
studio/frontend/src/__tests__/preview-store.test.ts
View file

@ -100,7 +100,6 @@ function makeTarget(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
targetPage: null,
snippet: null,
kind: null,
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,

1
studio/frontend/src/__tests__/preview-target-locator.test.tsx
View file

@ -23,7 +23,6 @@ function target(overrides: Partial<PreviewTarget> = {}): PreviewTarget {
targetPage: 2,
snippet: "alpha\nhighlighted line\nomega",
kind: "text",
imageUrl: null,
sourcePageIndex: null,
pageCharStart: null,
pageCharEnd: null,

1
studio/frontend/src/__tests__/rag-api.test.ts
View file

@ -66,7 +66,6 @@ function target(): PreviewTarget {
targetPage: 1,
snippet: "excerpt",
kind: "text",
imageUrl: null,
sourcePageIndex: 0,
pageCharStart: 0,
pageCharEnd: 7,

67
studio/frontend/src/components/assistant-ui/tool-ui-search-knowledge-base.tsx
View file

@ -3,14 +3,13 @@
"use client";
import { authFetch } from "@/features/auth";
import { usePreviewStore } from "@/features/rag/stores/preview-store";
import { cn } from "@/lib/utils";
import {
type ToolCallMessagePartComponent,
useAuiState,
} from "@assistant-ui/react";
import { FileTextIcon, ImageIcon, LoaderIcon } from "lucide-react";
import { FileTextIcon, LoaderIcon } from "lucide-react";
import { memo, useCallback, useEffect, useState } from "react";
import {
ToolFallbackContent,
@ -31,7 +30,6 @@ export interface ParsedChunk {
lineStart?: string;
lineEnd?: string;
kind?: string;
imageUrl?: string;
text: string;
/** Durable `rag_documents.id` from tool XML `document_id=`. Absent on
* legacy tool output. */
@ -79,7 +77,6 @@ export function parseChunks(raw: string): ParsedChunk[] {
lineStart: attrs.line_start,
lineEnd: attrs.line_end,
kind: attrs.kind,
imageUrl: attrs.image_url,
text,
// Durable backend ids (legacy XML omits both → preview gated off).
...(attrs.document_id ? { documentId: attrs.document_id } : {}),
@ -93,59 +90,6 @@ export function parseChunks(raw: string): ParsedChunk[] {
return out;
}
/** Fetch a backend image via bearer-authed `authFetch`, expose it as a
* blob URL for `<img src>`. Revokes the object URL on unmount. */
function useAuthedImageUrl(path: string | undefined): string | undefined {
const [url, setUrl] = useState<string | undefined>(undefined);
useEffect(() => {
if (!path) {
setUrl(undefined);
return;
}
let cancelled = false;
let objectUrl: string | undefined;
authFetch(path)
.then((response) => {
if (!response.ok) {
throw new Error(`image fetch ${response.status}`);
}
return response.blob();
})
.then((blob) => {
if (cancelled) return;
objectUrl = URL.createObjectURL(blob);
setUrl(objectUrl);
})
.catch(() => {
if (!cancelled) setUrl(undefined);
});
return () => {
cancelled = true;
if (objectUrl) URL.revokeObjectURL(objectUrl);
};
}, [path]);
return url;
}
function ChunkImage({ url, alt }: { url: string; alt: string }) {
const blobUrl = useAuthedImageUrl(url);
if (!blobUrl) {
return (
<div className="mb-2 flex h-32 items-center justify-center rounded-md bg-muted/60 text-[10px] text-muted-foreground">
<ImageIcon className="mr-1.5 size-3" />
Loading image…
</div>
);
}
return (
<img
src={blobUrl}
alt={alt}
className="mb-2 max-h-64 w-full rounded-md object-contain"
/>
);
}
function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
const openPreview = usePreviewStore((s) => s.open);
const meta: string[] = [];
@ -174,11 +118,7 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
<span className="rounded bg-foreground/10 px-1.5 py-0.5 font-mono text-[10px] font-semibold">
[{chunk.id}]
</span>
{chunk.kind === "image" ? (
<ImageIcon className="size-3 shrink-0 text-muted-foreground" />
) : (
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
)}
<FileTextIcon className="size-3 shrink-0 text-muted-foreground" />
<span className="truncate font-medium" title={chunk.source}>
{chunk.source}
</span>
@ -213,9 +153,6 @@ function ChunkCard({ chunk }: { chunk: ParsedChunk }) {
</span>
) : null}
</div>
{chunk.kind === "image" && chunk.imageUrl ? (
<ChunkImage url={chunk.imageUrl} alt={chunk.source} />
) : null}
{chunk.text ? (
<pre className="max-h-48 overflow-auto whitespace-pre-wrap break-words text-[11px] leading-relaxed text-foreground/80">
{chunk.text}

63
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -45,7 +45,6 @@ import {
TooltipContent,
TooltipTrigger,
} from "@/components/ui/tooltip";
import { type KBMode } from "@/features/rag/api/rag-api";
import { DocumentRow } from "@/features/rag/components/document-row";
import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog";
import { PreviewPanel } from "@/features/rag/components/preview-panel";
@ -559,9 +558,6 @@ export function ChatSettingsPanel({
}
}, [ragSource.kind, activeThreadId, loadThreadSettings]);
const effectiveThreadMode: KBMode =
threadSettings?.mode ?? ragDefaults?.mode ?? "text";
const aui = useAui();
// Brand-new chat has no backend thread yet — initialize the local
// assistant-ui thread to mint a remoteId so per-thread RAG settings
@ -584,31 +580,6 @@ export function ChatSettingsPanel({
}
};
const applyThreadSettingChange = (patch: {
mode?: KBMode;
}) => {
void (async () => {
const threadId = await ensureThreadId();
if (!threadId) return;
if (threadDocs.length === 0) {
void updateThreadSettings(threadId, patch);
return;
}
const ok = window.confirm(
`Re-index ${threadDocs.length} document${threadDocs.length === 1 ? "" : "s"} ` +
`with the new settings? Existing chunks will be deleted and rebuilt.`,
);
if (ok) {
void reingestThread(threadId, {
...patch,
caption_images: ragCaptionImages,
});
} else {
// User declined: refresh so the select snaps back.
void loadThreadSettings(threadId);
}
})();
};
const [kbCreateOpen, setKbCreateOpen] = useState(false);
const ragEnabled = ragSource.kind !== "off";
const activeKbId = ragSource.kind === "kb" ? ragSource.kbId : null;
@ -1414,7 +1385,6 @@ export function ChatSettingsPanel({
) : null}
{knowledgeBases.map((kb) => {
const isActive = kb.id === activeKbId;
const isMultimodal = kb.mode === "multimodal";
return (
<DropdownMenuItem
key={kb.id}
@ -1428,14 +1398,6 @@ export function ChatSettingsPanel({
>
<span className="flex min-w-0 items-center gap-1.5 truncate">
<span className="truncate">{kb.name}</span>
{isMultimodal ? (
<span
className="rounded-sm bg-violet-500/15 px-1 text-[10px] font-medium text-violet-700 dark:text-violet-300"
title="Multimodal mode — text + image embeddings"
>
🖼️ MM
</span>
) : null}
</span>
<button
type="button"
@ -1530,31 +1492,6 @@ export function ChatSettingsPanel({
/>
{ragSource.kind === "thread" ? (
<>
<div className="flex flex-col gap-1">
<label className="text-[11px] font-medium text-muted-foreground">
Mode
</label>
<Select
value={effectiveThreadMode}
onValueChange={(v) => {
const next = v as KBMode;
if (next === effectiveThreadMode) return;
applyThreadSettingChange({ mode: next });
}}
>
<SelectTrigger className="h-8 text-xs">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="text">Text only</SelectItem>
<SelectItem value="multimodal">Multimodal</SelectItem>
</SelectContent>
</Select>
</div>
<p className="text-[11px] text-muted-foreground">
Changing the mode will re-index this thread's existing
documents.
</p>
<div className="flex flex-col gap-1.5">
<label className="text-[12px] font-medium text-muted-foreground">
Documents in this thread

10
studio/frontend/src/features/rag/api/rag-api.ts
View file

@ -6,14 +6,12 @@ import { apiUrl } from "@/lib/api-base";
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
import { EventSourcePolyfill } from "event-source-polyfill";
export type KBMode = "text" | "multimodal";
export interface KnowledgeBase {
id: string;
name: string;
description: string | null;
embedding_model: string;
mode: KBMode;
created_at: number;
}
@ -47,7 +45,6 @@ export interface SearchHit {
page_number: number | null;
filename: string | null;
kind?: "text" | "image" | "caption";
image_url?: string | null;
source_page_index?: number | null;
page_char_start?: number | null;
page_char_end?: number | null;
@ -92,7 +89,6 @@ export interface PreviewTarget {
targetPage: number | null;
snippet: string | null;
kind: PreviewChunkKind | null;
imageUrl: string | null;
sourcePageIndex: number | null;
pageCharStart: number | null;
pageCharEnd: number | null;
@ -181,7 +177,6 @@ export interface CreateKnowledgeBaseRequest {
name: string;
description?: string;
embedding_model?: string;
mode?: KBMode;
}
export async function createKnowledgeBase(
@ -288,7 +283,6 @@ export interface ReingestResponse {
}
export interface ReingestKBOptions {
mode?: KBMode;
embedding_model?: string;
caption_images?: boolean;
}
@ -309,12 +303,10 @@ export async function reingestKnowledgeBase(
}
export interface ThreadRagSettings {
mode: KBMode;
embedding_model: string | null;
}
export interface UpdateThreadRagSettingsRequest {
mode?: KBMode;
embedding_model?: string | null;
// Only consulted by reingest (not persisted as a thread setting).
caption_images?: boolean;
@ -360,7 +352,6 @@ export async function reingestThreadDocuments(
}
export interface RagDefaults {
mode: KBMode;
embedding_model: string | null;
}
@ -370,7 +361,6 @@ export async function getRagDefaults(): Promise<RagDefaults> {
}
export interface UpdateRagDefaultsRequest {
mode?: KBMode;
embedding_model?: string | null;
}

44
studio/frontend/src/features/rag/components/kb-create-dialog.tsx
View file

@ -12,15 +12,8 @@ import {
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { useEffect, useState } from "react";
import type { KBMode, KnowledgeBase } from "../api/rag-api";
import type { KnowledgeBase } from "../api/rag-api";
import { useKnowledgeBases } from "../hooks/use-knowledge-bases";
import { useRagStore } from "../stores/rag-store";
@ -37,13 +30,11 @@ export function KBCreateDialog({
const defaults = useRagStore((s) => s.defaults);
const loadDefaults = useRagStore((s) => s.loadDefaults);
const initialMode: KBMode = defaults?.mode ?? "text";
const initialEmbedder = defaults?.embedding_model ?? "";
const [name, setName] = useState("");
const [description, setDescription] = useState("");
const [embeddingModel, setEmbeddingModel] = useState(initialEmbedder);
const [mode, setMode] = useState<KBMode>(initialMode);
const [submitting, setSubmitting] = useState(false);
const [error, setError] = useState<string | null>(null);
@ -55,7 +46,6 @@ export function KBCreateDialog({
useEffect(() => {
if (open && defaults) {
setMode(defaults.mode);
setEmbeddingModel(defaults.embedding_model ?? "");
}
// Only on open-flip, not on every defaults change.
@ -66,15 +56,11 @@ export function KBCreateDialog({
setName("");
setDescription("");
setEmbeddingModel(defaults?.embedding_model ?? "");
setMode(defaults?.mode ?? "text");
setError(null);
setSubmitting(false);
};
const placeholderEmbedder =
mode === "multimodal"
? "Defaults to BAAI/BGE-VL-base"
: "Defaults to BAAI/bge-small-en-v1.5";
const placeholderEmbedder = "Defaults to BAAI/bge-small-en-v1.5";
const handleSubmit = async (e: React.FormEvent) => {
e.preventDefault();
@ -86,7 +72,6 @@ export function KBCreateDialog({
name: name.trim(),
description: description.trim() || undefined,
embedding_model: embeddingModel.trim() || undefined,
mode,
});
onCreated?.(kb);
reset();
@ -135,31 +120,6 @@ export function KBCreateDialog({
placeholder="What's in this KB?"
/>
</div>
<div className="flex flex-col gap-2">
<Label htmlFor="kb-mode">Mode</Label>
<Select
value={mode}
onValueChange={(v) => setMode(v as KBMode)}
>
<SelectTrigger id="kb-mode">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="text">
Text only — embed text chunks
</SelectItem>
<SelectItem value="multimodal">
Multimodal — also embed images alongside text
</SelectItem>
</SelectContent>
</Select>
<p className="text-[11px] text-muted-foreground">
Multimodal mode extracts figures from your documents and
embeds them in a shared text + image vector space (BGE-VL),
so retrieval can match visual content. Larger embedder
(~1.5&nbsp;GB VRAM).
</p>
</div>
<div className="flex flex-col gap-2">
<Label htmlFor="kb-model">Embedding model (optional)</Label>
<Input

1
studio/frontend/src/features/rag/components/kb-detail-panel.tsx
View file

@ -133,7 +133,6 @@ export function KBDetailPanel({
{/* Own full-width row so the embedder id fits on one line instead of
wrapping next to the action buttons. */}
<p className="truncate text-xs text-muted-foreground">
{kb.mode === "multimodal" ? "🖼️ Multimodal · " : ""}
Embedder: <code>{kb.embedding_model}</code>
</p>
</div>

8
studio/frontend/src/features/rag/components/kb-list.tsx
View file

@ -50,14 +50,6 @@ export function KBList({
<div className="flex min-w-0 flex-col">
<span className="flex min-w-0 items-center gap-1.5 truncate text-sm font-medium">
<span className="truncate">{kb.name}</span>
{kb.mode === "multimodal" ? (
<span
className="shrink-0 rounded-sm bg-violet-500/15 px-1 text-[10px] font-medium text-violet-700 dark:text-violet-300"
title="Multimodal — text + image embeddings"
>
🖼️ MM
</span>
) : null}
</span>
{kb.description ? (
<span className="truncate text-xs text-muted-foreground">

36
studio/frontend/src/features/rag/components/kb-reconfigure-dialog.tsx
View file

@ -12,15 +12,8 @@ import {
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { useEffect, useState } from "react";
import type { KBMode, KnowledgeBase } from "../api/rag-api";
import type { KnowledgeBase } from "../api/rag-api";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { useRagStore } from "../stores/rag-store";
@ -36,7 +29,6 @@ export function KBReconfigureDialog({
documentCount: number;
}) {
const reingestKB = useRagStore((s) => s.reingestKB);
const [mode, setMode] = useState<KBMode>(kb.mode);
const [embeddingModel, setEmbeddingModel] = useState("");
const [submitting, setSubmitting] = useState(false);
const [error, setError] = useState<string | null>(null);
@ -44,17 +36,15 @@ export function KBReconfigureDialog({
// Re-sync when the dialog opens against a different KB.
useEffect(() => {
if (open) {
setMode(kb.mode);
setEmbeddingModel("");
setError(null);
setSubmitting(false);
}
}, [open, kb.id, kb.mode]);
}, [open, kb.id]);
const placeholderEmbedder = `Current: ${kb.embedding_model}`;
const changedSettings =
mode !== kb.mode || embeddingModel.trim() !== "";
const changedSettings = embeddingModel.trim() !== "";
const handleSubmit = async (e: React.FormEvent) => {
e.preventDefault();
@ -73,7 +63,6 @@ export function KBReconfigureDialog({
setError(null);
try {
await reingestKB(kb.id, {
mode,
embedding_model: embeddingModel.trim() || undefined,
caption_images: useChatRuntimeStore.getState().ragCaptionImages,
});
@ -91,29 +80,12 @@ export function KBReconfigureDialog({
<DialogHeader>
<DialogTitle>Reconfigure “{kb.name}”</DialogTitle>
<DialogDescription>
Change the mode or embedder for this KB.
Change the embedder for this KB.
All {documentCount} document{documentCount === 1 ? "" : "s"}{" "}
will be re-ingested from the originals on disk.
</DialogDescription>
</DialogHeader>
<div className="flex flex-col gap-4 py-4">
<div className="flex flex-col gap-2">
<Label htmlFor="reconf-mode">Mode</Label>
<Select
value={mode}
onValueChange={(v) => setMode(v as KBMode)}
>
<SelectTrigger id="reconf-mode">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="text">Text only</SelectItem>
<SelectItem value="multimodal">
Multimodal — text + images
</SelectItem>
</SelectContent>
</Select>
</div>
<div className="flex flex-col gap-2">
<Label htmlFor="reconf-model">Embedding model (optional)</Label>
<Input

78
studio/frontend/src/features/rag/components/rag-defaults-section.tsx
View file

@ -1,78 +0,0 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Label } from "@/components/ui/label";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import { useEffect, useState } from "react";
import type { KBMode } from "../api/rag-api";
import { useRagStore } from "../stores/rag-store";
/** Defaults pre-fill the KB create dialog. */
export function RagDefaultsSection() {
const defaults = useRagStore((s) => s.defaults);
const loadDefaults = useRagStore((s) => s.loadDefaults);
const updateDefaults = useRagStore((s) => s.updateDefaults);
const [mode, setMode] = useState<KBMode>("text");
const [error, setError] = useState<string | null>(null);
useEffect(() => {
void loadDefaults();
}, [loadDefaults]);
useEffect(() => {
if (defaults) {
setMode(defaults.mode);
}
}, [defaults]);
const persist = (patch: {
mode?: KBMode;
embedding_model?: string | null;
}) => {
setError(null);
void updateDefaults(patch).catch((err) => {
setError(err instanceof Error ? err.message : String(err));
});
};
return (
<div className="flex flex-col gap-3">
<div>
<h3 className="text-sm font-medium">Defaults for new knowledge bases</h3>
<p className="text-xs text-muted-foreground">
Pre-fills the KB create dialog. Existing KBs keep their own
settings — use the Reconfigure button to change those.
</p>
</div>
<div className="grid grid-cols-1 gap-3 md:grid-cols-2">
<div className="flex flex-col gap-1.5">
<Label htmlFor="defaults-mode">Mode</Label>
<Select
value={mode}
onValueChange={(v) => {
const next = v as KBMode;
setMode(next);
persist({ mode: next });
}}
>
<SelectTrigger id="defaults-mode">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="text">Text only</SelectItem>
<SelectItem value="multimodal">Multimodal</SelectItem>
</SelectContent>
</Select>
</div>
</div>
{error ? <div className="text-xs text-destructive">{error}</div> : null}
</div>
);
}

3
studio/frontend/src/features/settings/tabs/knowledge-bases-tab.tsx
View file

@ -8,7 +8,6 @@ import { KBCreateDialog } from "@/features/rag/components/kb-create-dialog";
import { KBDetailPanel } from "@/features/rag/components/kb-detail-panel";
import { KBList, type KBPanel } from "@/features/rag/components/kb-list";
import { PreviewPanel } from "@/features/rag/components/preview-panel";
import { RagDefaultsSection } from "@/features/rag/components/rag-defaults-section";
import { ThreadIndexList } from "@/features/rag/components/thread-index-list";
import { useResizablePanelWidth } from "@/features/rag/hooks/use-resizable-width";
import { usePreviewStore } from "@/features/rag/stores/preview-store";
@ -193,8 +192,6 @@ export function KnowledgeBasesTab() {
<Separator />
<ThreadIndexList />
<Separator />
<RagDefaultsSection />
</div>
);
}

57
tests/python/test_rag_multimodal.py
View file

@ -52,57 +52,8 @@ def test_html_parser_returns_images_when_requested(tmp_path):
assert img.nearest_caption == "A tiny figure"
def test_rag_embedder_matrix_is_keyed_by_mode():
from utils.rag.config import RAG_EMBEDDER_MATRIX, resolve_embedder
def test_rag_resolve_embedder_returns_default():
from utils.rag.config import RAG_EMBEDDING_MODEL, resolve_embedder
assert "text" in RAG_EMBEDDER_MATRIX
assert "multimodal" in RAG_EMBEDDER_MATRIX
# Unknown modes fall back to the default, not KeyError.
fallback = resolve_embedder("unknown-mode")
assert isinstance(fallback, str) and fallback
def test_image_path_url_construction():
"""Sanity-check the URL shape served back to the frontend.
The image URL is built relative to /api/rag/images/<doc>/<filename>
purely from the stored image_path (filename only — directory
structure is fixed). Verify the rule.
"""
from pathlib import Path as P
image_path = "/var/data/rag/images/doc-123/img-0042.png"
document_id = "doc-123"
expected = f"/api/rag/images/{document_id}/{P(image_path).name}"
assert expected == "/api/rag/images/doc-123/img-0042.png"
@pytest.mark.server
def test_multimodal_encode_image_returns_vector(tmp_path, monkeypatch):
pytest.importorskip("sentence_transformers")
pytest.importorskip("PIL")
monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base")
# Reset the embedder singleton so the env var applies.
from core.rag import embeddings as embeddings_module
embeddings_module._model = None
embeddings_module._model_name = None
from io import BytesIO
from PIL import Image
img = Image.new("RGB", (32, 32), (200, 100, 50))
buf = BytesIO()
img.save(buf, format = "PNG")
image_bytes = buf.getvalue()
vectors = embeddings_module.encode_images([image_bytes])
assert len(vectors) == 1
dim = vectors[0].shape[0]
assert dim > 0
# Text shares the same dim — the point of a multimodal embedder.
text_vec = embeddings_module.encode(["a red square"])[0]
assert text_vec.shape[0] == dim
assert resolve_embedder() == RAG_EMBEDDING_MODEL
assert isinstance(resolve_embedder(), str) and resolve_embedder()

168
tests/python/test_rag_multimodal_integration.py
View file

@ -1,168 +0,0 @@
"""End-to-end multimodal RAG integration test.
Marked `server` so default pytest runs skip it — downloads BGE-VL-base
(~600 MB) on first run and exercises the real embedding stack. Run
explicitly with:
~/.unsloth/studio/unsloth_studio/bin/python -m pytest \
tests/python/test_rag_multimodal_integration.py -v -m server
Exercises the ingestion subprocess worker in-process (with a regular
queue rather than mp.Queue) so we cover the parse → chunk → load
embedder → encode_images → emit chunks_batch path without spinning up
a child process. The parent-side chunk insertion is covered separately
by test_rag_multimodal.py.
"""
import os
import queue as queue_module
import sys
from io import BytesIO
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
@pytest.mark.server
def test_multimodal_subprocess_emits_image_and_caption_chunks(
tmp_path,
monkeypatch,
):
pymupdf = pytest.importorskip("pymupdf")
pytest.importorskip("pymupdf4llm")
pytest.importorskip("sentence_transformers")
pytest.importorskip("PIL")
pytest.importorskip("torch")
# tmp studio root isolates the subprocess's image writes.
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base")
monkeypatch.setenv("UNSLOTH_RAG_CHUNK_SIZE", "200")
monkeypatch.setenv("UNSLOTH_RAG_CHUNK_OVERLAP", "20")
# Reset module caches so the new env vars apply.
import importlib
import utils.rag.config as rag_config
importlib.reload(rag_config)
from core.rag import embeddings as embeddings_module
embeddings_module._model = None
embeddings_module._model_name = None
# Small PDF: text + one embedded image.
from PIL import Image
img = Image.new("RGB", (96, 64), (200, 100, 50))
img_buf = BytesIO()
img.save(img_buf, format = "PNG")
img_bytes = img_buf.getvalue()
doc = pymupdf.open()
page = doc.new_page(width = 612, height = 792)
page.insert_text(
(72, 100),
"Architecture overview\n\nThe following diagram shows our system.",
fontsize = 11,
)
image_rect = pymupdf.Rect(72, 200, 168, 264)
page.insert_image(image_rect, stream = img_bytes)
page.insert_text(
(72, 290),
"Figure 1: the architecture diagram described above.",
fontsize = 11,
)
pdf_path = tmp_path / "sample.pdf"
doc.save(str(pdf_path))
doc.close()
# Drive the worker in-process via a regular queue.
from core.rag.ingestion import _subprocess_worker
out_queue: "queue_module.Queue[dict]" = queue_module.Queue()
_subprocess_worker(
stored_path = str(pdf_path),
model_name = "BAAI/BGE-VL-base",
chunk_size = 200,
overlap = 20,
batch_size = 4,
out_queue = out_queue,
mode = "multimodal",
document_id = "test-doc-1",
)
# Drain all events (in-process queue, stable order).
events: list[dict] = []
while not out_queue.empty():
events.append(out_queue.get_nowait())
# Expect >=1 chunks_batch and exactly one terminal complete/error.
assert any(e["type"] == "chunks_batch" for e in events)
terminals = [e for e in events if e["type"] in ("complete", "error")]
assert len(terminals) == 1, terminals
assert terminals[0]["type"] == "complete"
# Collect chunks across batches.
all_chunks: list[dict] = []
for e in events:
if e["type"] == "chunks_batch":
all_chunks.extend(e["chunks"])
kinds = [c.get("kind") for c in all_chunks]
assert "text" in kinds, "expected at least one text chunk"
assert "image" in kinds, "expected at least one image chunk"
# Paragraph right after the image triggers caption pairing.
assert "caption" in kinds, "expected at least one caption chunk"
# Image chunks carry a path on disk under the tmp studio root.
image_chunks = [c for c in all_chunks if c.get("kind") == "image"]
for chunk in image_chunks:
assert chunk.get("image_path"), chunk
path_on_disk = Path(chunk["image_path"])
assert path_on_disk.is_file()
assert str(path_on_disk).startswith(str(tmp_path))
# Paired image + caption chunks share a pair_group.
pair_groups: dict[str, list[str]] = {}
for chunk in all_chunks:
group = chunk.get("pair_group")
if group:
pair_groups.setdefault(group, []).append(chunk.get("kind", ""))
paired = [
kinds
for kinds in pair_groups.values()
if "image" in kinds and "caption" in kinds
]
assert paired, f"expected an image/caption pair, got groups={pair_groups}"
@pytest.mark.server
def test_text_and_image_vectors_share_dimension(monkeypatch):
"""BGE-VL is a shared-space embedder — sanity-check before relying on it."""
pytest.importorskip("sentence_transformers")
pytest.importorskip("PIL")
pytest.importorskip("torch")
monkeypatch.setenv("UNSLOTH_RAG_EMBEDDING_MODEL", "BAAI/BGE-VL-base")
from core.rag import embeddings as embeddings_module
embeddings_module._model = None
embeddings_module._model_name = None
from PIL import Image
img = Image.new("RGB", (32, 32), (50, 150, 200))
buf = BytesIO()
img.save(buf, format = "PNG")
image_vectors = embeddings_module.encode_images([buf.getvalue()])
text_vectors = embeddings_module.encode(["a blue square"])
assert (
image_vectors[0].shape == text_vectors[0].shape
), f"text dim {text_vectors[0].shape} != image dim {image_vectors[0].shape}"

17
tests/python/test_rag_reingest.py
View file

@ -9,8 +9,6 @@ import importlib.util
import sys
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
@ -42,18 +40,7 @@ def test_reingest_request_accepts_all_optional_fields():
ReingestKBRequest = _rag_route().ReingestKBRequest
empty = ReingestKBRequest()
assert empty.mode is None
assert empty.embedding_model is None
partial = ReingestKBRequest(mode = "multimodal")
assert partial.mode == "multimodal"
assert partial.embedding_model is None
def test_reingest_request_rejects_unknown_mode():
from pydantic import ValidationError
ReingestKBRequest = _rag_route().ReingestKBRequest
with pytest.raises(ValidationError):
ReingestKBRequest(mode = "augmented")
partial = ReingestKBRequest(embedding_model = "BAAI/bge-small-en-v1.5")
assert partial.embedding_model == "BAAI/bge-small-en-v1.5"

19
tests/python/test_rag_tool_handler.py
View file

@ -146,25 +146,6 @@ def test_format_hits_offsets_ids_by_start_id():
assert '<chunk id="1"' not in result
def test_format_hits_emits_image_url_for_image_kind():
from core.rag.tool import _format_hits_for_llm
hits = [
{
"filename": "paper.pdf",
"text": "Figure 1 shows a bar chart of X over time.",
"kind": "image",
"image_path": "/abs/path/images/doc-123/img-0007.png",
"document_id": "doc-123",
"page_number": 5,
}
]
result = _format_hits_for_llm(hits)
assert 'kind="image"' in result
assert 'image_url="/api/rag/images/doc-123/img-0007.png"' in result
assert "Figure 1 shows a bar chart" in result
def test_format_hits_escapes_xml_in_source():
from core.rag.tool import _format_hits_for_llm