Studio: multimodal RAG mode (Phase 3B-multimodal)
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:
parent
673b7f86ba
commit
68114fd223
10 changed files with 526 additions and 53 deletions
127
tests/python/test_rag_multimodal.py
Normal file
127
tests/python/test_rag_multimodal.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue