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

Studio: multimodal RAG mode (Phase 3B-multimodal)

Browse source 
When a KB has mode = 'multimodal', ingestion extracts images alongside
text and embeds both into a shared 512-d vector space via BGE-VL-base.
Image hits become first-class search results — useful for slides,
reports, and diagrams where text-only retrieval loses ~30-50% of the
content.

Backend
- embeddings.py: new encode_images(image_bytes_list) — opens bytes via
  PIL and routes to the SentenceTransformer (BGE-VL accepts PIL images
  in the same encode call as text).
- ingestion.py: _subprocess_worker gains document_id arg and a new
  _stream_image_chunks() helper. For multimodal KBs the standard text
  chunking runs first, then images are saved to
  rag_uploads_root() / 'images' / <document_id> / img-NNNN.<ext> and
  embedded; for each image with an adjacent caption, both an
  'image'-kind chunk (vector = encoded image) and a 'caption'-kind
  chunk (vector = encoded caption text) are streamed back with a
  shared pair_group field.
- ingestion.py parent: _insert_chunks_and_collect_for_bm25 now reads
  kind / image_path / pair_group from the subprocess message,
  populates the new rag_chunks columns, and runs a second pass that
  sets linked_chunk_id for each image ↔ caption pair. BM25 indexes
  text + caption chunks only — image chunks have no tokenisable body.
- retrieval.py: Hit gains a `kind` field plumbed through bm25, dense,
  RRF, and rerank paths.
- reranker.py: image-kind hits skip CrossEncoder rerank (text-only
  model) but are appended back in their original relative position
  rather than dropped.
- routes/rag.py: new GET /api/rag/images/{document_id}/{filename}
  static-file route with realpath containment check. SearchHit gains
  `kind` and `image_url` fields so the chat UI can render image
  thumbnails alongside text hits. KB-doc upload threads kind/mode
  through to ingestion.

Frontend
- rag-api.ts: SearchHit gains optional `kind` and `image_url`.
- kb-create-dialog.tsx: new Mode select (Text / Multimodal) alongside
  the existing Chunking strategy select. The forbidden
  (multimodal + late) combo is enforced in the UI — each side
  disables the conflicting option on the other side with a tooltip
  explaining why. Embedding-model placeholder cycles through the
  three valid defaults (bge-small / nomic / BGE-VL).
- kb-list.tsx + chat-settings-sheet.tsx: 🖼️ MM badge alongside the
  ⚡ Late one so multimodal KBs are obvious at a glance.

Tests
- test_rag_multimodal.py: parser returns images when want_images=True
  and skips them when False; _validate_mode_combo rejects the
  forbidden (multimodal, late) pair with 400; RAG_EMBEDDER_MATRIX
  contains the three valid combos and excludes the forbidden one;
  image URL construction shape is verified. A server-marked test
  loads BGE-VL-base end-to-end and confirms image + text vectors
  share the same dimension.

Phase 3 of the plan is now feature-complete on the backend; the
remaining items (re-ingest UX for changing strategy on existing KBs)
are tracked under "Backfill UX" and can land separately.
This commit is contained in:
Roland Tannous 2026-05-24 12:33:08 +04:00
commit 68114fd223
10 changed files with 526 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -81,6 +81,38 @@ def encode(
)
def encode_images(
image_bytes_list: list[bytes],
*,
model_name: str | None = None,
batch_size: int | None = None,
normalize: bool = True,
):
"""Embed raw image bytes via a multimodal SentenceTransformer.
Works with CLIP-family models (BGE-VL, openai/clip-*) whose
`encode` accepts PIL.Image objects in the same call as text. The
returned vectors live in the same 512-d (or model-specific) space
as text vectors from this model, so a single Qdrant collection
holds both kinds.
"""
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 ``len(tokenize(text))`` callable using the embedder's tokenizer.

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

@ -47,6 +47,18 @@ _QUEUE_TIMEOUT_SECONDS = 300
# Subprocess worker
# ------------------------------------------------------------------
_MIME_TO_EXT = {
"image/png": ".png",
"image/jpeg": ".jpg",
"image/jpg": ".jpg",
"image/gif": ".gif",
"image/webp": ".webp",
"image/bmp": ".bmp",
"image/tiff": ".tiff",
"image/svg+xml": ".svg",
}
def _subprocess_worker(
stored_path: str,
model_name: str,
@ -56,20 +68,17 @@ def _subprocess_worker(
out_queue: Any,
chunking_strategy: str = "standard",
mode: str = "text",
document_id: str = "",
) -> None:
try:
from core.rag.chunking import chunk_pages, chunk_pages_with_spans
from core.rag.chunking import chunk_pages
from core.rag.parsers import parse
out_queue.put({"type": "progress", "stage": "parse", "progress": 0.05})
# want_images is True only for multimodal KBs. The image side of
# the pipeline lands in Phase 3B-multimodal; for now the parser
# collects the bytes anyway in case we want them later, but only
# the text pages are consumed.
parsed = parse(Path(stored_path), want_images = (mode == "multimodal"))
pages = parsed.pages
if not pages:
out_queue.put({"type": "error", "error": "no extractable text in document"})
if not pages and not parsed.images:
out_queue.put({"type": "error", "error": "no extractable content in document"})
return
out_queue.put({"type": "progress", "stage": "load_model", "progress": 0.1})
@ -94,17 +103,30 @@ def _subprocess_worker(
late_chunk_encode = late_chunk_encode,
out_queue = out_queue,
)
else:
_run_standard_chunking(
pages = pages,
chunk_size = chunk_size,
overlap = overlap,
counter = counter,
batch_size = batch_size,
model = model,
chunk_pages = chunk_pages,
return
# Standard chunking path (text + optional images for multimodal mode).
text_count = _run_standard_chunking(
pages = pages,
chunk_size = chunk_size,
overlap = overlap,
counter = counter,
batch_size = batch_size,
model = model,
chunk_pages = chunk_pages,
out_queue = out_queue,
send_complete = False,
)
image_count = 0
if mode == "multimodal" and parsed.images and document_id:
image_count = _stream_image_chunks(
images = parsed.images,
document_id = document_id,
model_name = model_name,
out_queue = out_queue,
first_index = text_count,
)
out_queue.put({"type": "complete", "num_chunks": text_count + image_count})
except Exception as exc: # noqa: BLE001
logger.exception("ingestion subprocess failed")
out_queue.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"})
@ -120,7 +142,14 @@ def _run_standard_chunking(
model,
chunk_pages,
out_queue,
) -> None:
send_complete: bool = True,
) -> int:
"""Stream text chunks back to the parent. Returns the number streamed.
When called as part of the multimodal pipeline, `send_complete` is
False because image chunks still need to be streamed before the
document is marked complete.
"""
out_queue.put({"type": "progress", "stage": "chunk", "progress": 0.2})
chunks = chunk_pages(
pages,
@ -129,8 +158,9 @@ def _run_standard_chunking(
token_counter = counter,
)
if not chunks:
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
return
if send_complete:
out_queue.put({"type": "error", "error": "chunker produced no chunks"})
return 0
total = len(chunks)
for i in range(0, total, batch_size):
@ -151,6 +181,7 @@ def _run_standard_chunking(
"text": c.text,
"token_count": c.token_count,
"page_number": c.page_number,
"kind": "text",
}
for c in batch
],
@ -160,7 +191,121 @@ def _run_standard_chunking(
progress = 0.3 + 0.65 * min(1.0, (i + len(batch)) / total)
out_queue.put({"type": "progress", "stage": "embed", "progress": progress})
out_queue.put({"type": "complete", "num_chunks": total})
if send_complete:
out_queue.put({"type": "complete", "num_chunks": total})
return total
def _stream_image_chunks(
*,
images,
document_id: str,
model_name: str,
out_queue,
first_index: int,
) -> int:
"""Save extracted images to disk and stream image+caption chunks.
Each image becomes an `image`-kind chunk; if the parser found an
adjacent caption, a paired `caption`-kind chunk is also emitted.
Pairs share a ``pair_group`` so the parent can link them via
rag_chunks.linked_chunk_id.
Images are written to ``rag_uploads_root() / 'images' /
<document_id> / img-NNN.<ext>`` so the parent can serve them via
the static-image route without holding bytes in memory.
"""
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)
# Persist bytes, build parallel lists for encoding.
paths: list[str] = []
bytes_for_encoding: list[bytes] = []
captions: list[str] = []
pages: list[int | None] = []
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 %s; skipping", path)
continue
paths.append(str(path))
bytes_for_encoding.append(img.image_bytes)
captions.append(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)
# Embed only the non-empty captions; track which images they map to.
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())
# Emit the caption chunk right after its image so the parent
# sees them adjacent (simplifies pair linking).
if next_cap is not None and next_cap[0] == idx:
_cap_index, cap_vec = next_cap
out_chunks.append(
{
"text": caption,
"token_count": max(1, len(caption.split())),
"page_number": page,
"kind": "caption",
"image_path": None,
"pair_group": group_id,
}
)
out_vectors.append(cap_vec)
next_cap = next(cap_iter, None)
out_queue.put(
{
"type": "chunks_batch",
"first_index": first_index,
"chunks": out_chunks,
"vectors": out_vectors,
}
)
out_queue.put({"type": "progress", "stage": "extract_images", "progress": 0.95})
return len(out_chunks)
def _run_late_chunking(
@ -303,13 +448,26 @@ def _insert_chunks_and_collect_for_bm25(
chunks_meta: list[dict],
vectors: list[list[float]],
) -> list[dict]:
"""Insert chunks into sqlite + Qdrant; return [{id, text}] for BM25."""
"""Insert chunks into sqlite + Qdrant; return [{id, text}] for BM25.
Image-kind chunks ship a stable image_path and skip BM25 (no text
body to tokenise). Paired image/caption chunks share a pair_group
field — the second pass links them via rag_chunks.linked_chunk_id
so retrieval can dereference an image hit to its caption.
"""
rows: list[tuple] = []
points: list[dict] = []
bm25_rows: list[dict] = []
pair_groups: dict[str, list[str]] = {}
for offset, (meta, vec) in enumerate(zip(chunks_meta, vectors)):
chunk_index = first_index + offset
chunk_id = str(uuid4())
kind = meta.get("kind", "text")
image_path = meta.get("image_path")
pair_group = meta.get("pair_group")
if pair_group:
pair_groups.setdefault(pair_group, []).append(chunk_id)
rows.append(
(
chunk_id,
@ -318,6 +476,8 @@ def _insert_chunks_and_collect_for_bm25(
meta["text"],
meta["token_count"],
meta["page_number"],
kind,
image_path,
)
)
points.append(
@ -329,19 +489,41 @@ def _insert_chunks_and_collect_for_bm25(
"chunk_index": chunk_index,
"text": meta["text"],
"page_number": meta["page_number"],
"kind": kind,
"image_path": image_path,
},
}
)
bm25_rows.append({"id": chunk_id, "text": meta["text"]})
# BM25 indexes text + caption chunks. Image chunks have no
# tokenisable body — their caption (if any) is in a separate
# caption-kind chunk that BM25 will index.
if kind in ("text", "caption") and meta["text"]:
bm25_rows.append({"id": chunk_id, "text": meta["text"]})
with get_connection() as conn:
conn.executemany(
"""
INSERT INTO rag_chunks
(id, document_id, chunk_index, text, token_count, page_number)
VALUES (?, ?, ?, ?, ?, ?)
(id, document_id, chunk_index, text, token_count, page_number,
kind, image_path)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
rows,
)
# Link image ↔ caption pairs. We only set linked_chunk_id when
# a pair_group has exactly two members; lone images stay
# unlinked (no caption was paired).
for ids in pair_groups.values():
if len(ids) != 2:
continue
id_a, id_b = ids
conn.execute(
"UPDATE rag_chunks SET linked_chunk_id = ? WHERE id = ?",
(id_b, id_a),
)
conn.execute(
"UPDATE rag_chunks SET linked_chunk_id = ? WHERE id = ?",
(id_a, id_b),
)
conn.commit()
vector_store.upsert_chunks(scope, points)
return bm25_rows
@ -536,6 +718,7 @@ def enqueue_ingestion(
out_queue,
chunking_strategy,
mode,
document_id,
),
daemon = True,
)

50
studio/backend/core/rag/reranker.py
View file

@ -86,28 +86,38 @@ def rerank(
"""
if not pairs:
return []
# CrossEncoder is text-only; image-kind hits get appended at the end
# in their original relative order so they're never dropped, just
# never reranked. Caption-kind hits are eligible (they carry text).
text_pairs = [(h, t) for h, t in pairs if h.kind != "image"]
image_hits = [h for h, _t in pairs if h.kind == "image"]
model = get_reranker(model_name)
inputs = [(query, text) for _, text in pairs]
scores = model.predict(
inputs,
batch_size = RAG_RERANK_BATCH_SIZE,
show_progress_bar = False,
)
ranked = sorted(
zip(pairs, scores),
key = lambda item: float(item[1]),
reverse = True,
)
out: list[Hit] = []
for (hit, _text), score in ranked:
out.append(
Hit(
chunk_id = hit.chunk_id,
score = float(score),
document_id = hit.document_id,
chunk_index = hit.chunk_index,
)
if text_pairs:
inputs = [(query, text) for _, text in text_pairs]
scores = model.predict(
inputs,
batch_size = RAG_RERANK_BATCH_SIZE,
show_progress_bar = False,
)
ranked = sorted(
zip(text_pairs, scores),
key = lambda item: float(item[1]),
reverse = True,
)
reranked_text = [
Hit(
chunk_id = h.chunk_id,
score = float(s),
document_id = h.document_id,
chunk_index = h.chunk_index,
kind = h.kind,
)
for (h, _t), s in ranked
]
else:
reranked_text = []
out: list[Hit] = reranked_text + image_hits
if top_k is not None:
out = out[:top_k]
return out

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

@ -29,6 +29,7 @@ class Hit:
score: float
document_id: str | None = None
chunk_index: int | None = None
kind: str = "text"
def retrieve_bm25(scope: str, query: str, k: int | None = None) -> list[Hit]:
@ -60,6 +61,7 @@ def retrieve_dense(
score = r["score"],
document_id = payload.get("document_id"),
chunk_index = payload.get("chunk_index"),
kind = payload.get("kind", "text"),
)
)
return out
@ -85,6 +87,7 @@ def _rrf_fuse(
score = score,
document_id = seen[cid].document_id,
chunk_index = seen[cid].chunk_index,
kind = seen[cid].kind,
)
for cid, score in ordered
]

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

@ -29,7 +29,7 @@ from typing import Any, Literal, Optional
from uuid import uuid4
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, UploadFile
from fastapi.responses import StreamingResponse
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject, get_current_subject_sse
@ -141,6 +141,8 @@ class SearchHit(BaseModel):
score: float
page_number: int | None = None
filename: str | None = None
kind: str = "text"
image_url: str | None = None
class SearchResponse(BaseModel):
@ -513,6 +515,33 @@ 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 image extracted during multimodal ingestion.
Files live under ``rag_uploads_root() / 'images' / <document_id>``.
We realpath-check the resolved file against that root to refuse
path-traversal attempts (``..`` segments, symlinks pointing
elsewhere). Filenames are constrained to a single path component.
"""
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,
@ -706,7 +735,8 @@ def search(
rows = conn.execute(
f"""
SELECT c.id AS chunk_id, c.document_id, c.chunk_index, c.text,
c.page_number, d.filename
c.page_number, c.kind, c.image_path, c.linked_chunk_id,
d.filename
FROM rag_chunks c
JOIN rag_documents d ON d.id = c.document_id
WHERE c.id IN ({placeholders})
@ -736,15 +766,24 @@ def search(
meta = chunk_lookup.get(hit.chunk_id)
if not meta:
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']}/"
f"{Path(meta['image_path']).name}"
)
out.append(
SearchHit(
chunk_id = hit.chunk_id,
document_id = meta["document_id"],
chunk_index = meta["chunk_index"],
text = meta["text"],
text = meta["text"] or "",
score = hit.score,
page_number = meta.get("page_number"),
filename = meta.get("filename"),
kind = kind,
image_url = image_url,
)
)
return SearchResponse(hits = out)

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

@ -1233,6 +1233,7 @@ export function ChatSettingsPanel({
{knowledgeBases.map((kb) => {
const isActive = kb.id === activeKbId;
const isLate = kb.chunking_strategy === "late";
const isMultimodal = kb.mode === "multimodal";
return (
<DropdownMenuItem
key={kb.id}
@ -1254,6 +1255,14 @@ export function ChatSettingsPanel({
⚡ Late
</span>
) : null}
{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"

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

@ -45,6 +45,11 @@ export interface SearchHit {
score: number;
page_number: number | null;
filename: string | null;
// Phase 3B-multimodal: "image" hits have an image_url (served by
// GET /api/rag/images/{document_id}/{filename}); "caption" hits
// carry the paired image's caption text and are LLM-friendly.
kind?: "text" | "image" | "caption";
image_url?: string | null;
}
export interface SearchRequest {

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

@ -20,7 +20,11 @@ import {
SelectValue,
} from "@/components/ui/select";
import { useState } from "react";
import type { ChunkingStrategy, KnowledgeBase } from "../api/rag-api";
import type {
ChunkingStrategy,
KBMode,
KnowledgeBase,
} from "../api/rag-api";
import { useKnowledgeBases } from "../hooks/use-knowledge-bases";
export function KBCreateDialog({
@ -38,6 +42,7 @@ export function KBCreateDialog({
const [embeddingModel, setEmbeddingModel] = useState("");
const [chunkingStrategy, setChunkingStrategy] =
useState<ChunkingStrategy>("standard");
const [mode, setMode] = useState<KBMode>("text");
const [submitting, setSubmitting] = useState(false);
const [error, setError] = useState<string | null>(null);
@ -46,10 +51,23 @@ export function KBCreateDialog({
setDescription("");
setEmbeddingModel("");
setChunkingStrategy("standard");
setMode("text");
setError(null);
setSubmitting(false);
};
// Forbidden combo: multimodal + late chunking. We disable each side
// when the other side is picking it, with a tooltip explaining why.
const lateDisabled = mode === "multimodal";
const multimodalDisabled = chunkingStrategy === "late";
const placeholderEmbedder =
mode === "multimodal"
? "Defaults to BAAI/BGE-VL-base"
: chunkingStrategy === "late"
? "Defaults to nomic-ai/nomic-embed-text-v1.5"
: "Defaults to BAAI/bge-small-en-v1.5";
const handleSubmit = async (e: React.FormEvent) => {
e.preventDefault();
if (!name.trim() || submitting) return;
@ -61,6 +79,7 @@ export function KBCreateDialog({
description: description.trim() || undefined,
embedding_model: embeddingModel.trim() || undefined,
chunking_strategy: chunkingStrategy,
mode,
});
onCreated?.(kb);
reset();
@ -109,6 +128,40 @@ 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"
disabled={multimodalDisabled}
title={
multimodalDisabled
? "Multimodal cannot be combined with late chunking"
: undefined
}
>
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) and cannot be combined with late
chunking.
</p>
</div>
<div className="flex flex-col gap-2">
<Label htmlFor="kb-strategy">Chunking strategy</Label>
<Select
@ -122,7 +175,15 @@ export function KBCreateDialog({
<SelectItem value="standard">
Standard — heading-aware recursive splitter
</SelectItem>
<SelectItem value="late">
<SelectItem
value="late"
disabled={lateDisabled}
title={
lateDisabled
? "Late chunking cannot be combined with multimodal mode"
: undefined
}
>
Late chunking — single-pass embedder, slower ingest
</SelectItem>
</SelectContent>
@ -141,11 +202,7 @@ export function KBCreateDialog({
id="kb-model"
value={embeddingModel}
onChange={(e) => setEmbeddingModel(e.target.value)}
placeholder={
chunkingStrategy === "late"
? "Defaults to nomic-ai/nomic-embed-text-v1.5"
: "Defaults to BAAI/bge-small-en-v1.5"
}
placeholder={placeholderEmbedder}
/>
</div>
{error ? (

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

@ -68,6 +68,14 @@ export function KBList({
⚡ Late
</span>
) : null}
{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">

127
tests/python/test_rag_multimodal.py Normal file
View file

@ -0,0 +1,127 @@
"""Multimodal RAG tests (Phase 3B-multimodal).
Most of the multimodal pipeline depends on real models (BGE-VL ~1.5 GB
VRAM) and a writable filesystem under rag_uploads_root() — those tests
are gated behind the `server` marker. The pure-python pieces (parser
returns images when asked, route accepts the mode field, constraint
validator rejects illegal combos) run in every test invocation.
"""
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:
sys.path.insert(0, str(STUDIO_BACKEND))
def test_html_parser_returns_images_when_requested(tmp_path):
pytest.importorskip("bs4")
pytest.importorskip("lxml")
pytest.importorskip("markdownify")
from core.rag.parsers import parse
# A tiny 1x1 transparent PNG.
png_bytes = (
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx\x9cc\xfc\xff"
b"\xff?\x00\x05\xfe\x02\xfe\xa3\xb0\xa9\xa8\x00\x00\x00\x00IEND\xaeB`\x82"
)
img_path = tmp_path / "tiny.png"
img_path.write_bytes(png_bytes)
html_path = tmp_path / "sample.html"
html_path.write_text(
f'<html><body><h1>Doc</h1>'
f'<p>Body text.</p>'
f'<img src="tiny.png" alt="A tiny figure">'
f'</body></html>',
encoding = "utf-8",
)
no_images = parse(html_path, want_images = False)
assert no_images.images == []
with_images = parse(html_path, want_images = True)
assert len(with_images.images) == 1
img = with_images.images[0]
assert img.image_bytes == png_bytes
assert img.mime_type == "image/png"
assert img.nearest_caption == "A tiny figure"
def test_multimodal_late_combo_validator():
from fastapi import HTTPException
from routes.rag import _validate_mode_combo
# Allowed combos return None.
assert _validate_mode_combo("text", "standard") is None
assert _validate_mode_combo("text", "late") is None
assert _validate_mode_combo("multimodal", "standard") is None
# Forbidden combo raises 400.
with pytest.raises(HTTPException) as excinfo:
_validate_mode_combo("multimodal", "late")
assert excinfo.value.status_code == 400
def test_rag_embedder_matrix_excludes_multimodal_late():
from utils.rag.config import RAG_EMBEDDER_MATRIX, resolve_embedder
assert ("multimodal", "late") not in RAG_EMBEDDER_MATRIX
assert ("text", "standard") in RAG_EMBEDDER_MATRIX
assert ("text", "late") in RAG_EMBEDDER_MATRIX
assert ("multimodal", "standard") in RAG_EMBEDDER_MATRIX
# Unknown combos fall back to the legacy default rather than KeyError.
fallback = resolve_embedder("multimodal", "late")
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 takes effect.
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 from the same model should also be `dim`-d — shared space is
# the whole point of multimodal embedders.
text_vec = embeddings_module.encode(["a red square"])[0]
assert text_vec.shape[0] == dim