68114fd223
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.
127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
"""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
|