Merge pull request #513 from ggozad/feat/multimodal-reranker

Multimodal reranking: send picture chunks to vllm rerankers as images

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `reranking.multimodal` config flag: picture chunks are sent to a vllm reranker as image documents.
+
 ### Fixed
 
 - `list`, `get`, and `visualize` no longer require an embeddings config matching the database.

docs/configuration/index.md

@@ -99,6 +99,7 @@ reranking:
   model:
     provider: ""  # Empty to disable, or cross-encoder, cohere, zeroentropy, vllm
     name: ""
+  multimodal: false  # vllm only: send picture chunks to the reranker as images
 
 qa:
   model:

docs/configuration/providers.md

@@ -450,6 +450,21 @@ reranking:
 
 **Note:** vLLM reranking uses the `/v1/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded.
 
+#### Multimodal reranking
+
+When serving a vision reranker (for example `nvidia/llama-nemotron-rerank-vl-1b-v2`), set `multimodal: true` to score picture chunks by their image bytes in addition to their description text:
+
+```yaml
+reranking:
+  multimodal: true
+  model:
+    provider: vllm
+    name: nvidia/llama-nemotron-rerank-vl-1b-v2
+    base_url: http://localhost:8001
+```
+
+Picture chunks are sent as image documents (base64 data URIs) alongside plain text documents in the same rerank request. The flag is supported on the vllm provider only, and the served model must accept multimodal inputs.
+
 ### Jina AI
 
 Jina provides high-quality reranking with two deployment options: API mode and local inference.

haiku_rag_slim/haiku/rag/client/search.py

@@ -54,6 +54,8 @@ async def search(
                 query, search_limit, search_type, filter
             )
             chunks = [chunk for chunk, _ in raw_results]
+            if client._config.reranking.multimodal:
+                await _attach_picture_data(client, chunks)
             chunk_results = await reranker.rerank(query, chunks, top_n=limit)
     else:
         embedder = client.embedder
@@ -80,6 +82,29 @@ async def search(
     return results
 
 
+async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None:
+    """Attach picture bytes to synthetic picture chunks in-place, so a
+    multimodal reranker can score the pixels instead of just the chunk's
+    description text. Batches one picture-bytes lookup per document."""
+    by_doc: dict[str, list[tuple[Chunk, str]]] = {}
+    for chunk in chunks:
+        if chunk.document_id is None:
+            continue
+        refs = chunk.get_chunk_metadata().doc_item_refs
+        if len(refs) == 1 and refs[0].startswith(PICTURE_REF_PREFIX):
+            by_doc.setdefault(chunk.document_id, []).append((chunk, refs[0]))
+
+    for doc_id, doc_chunks in by_doc.items():
+        refs = [ref for _, ref in doc_chunks]
+        bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk(
+            doc_id, refs
+        )
+        for chunk, ref in doc_chunks:
+            data = bytes_by_ref.get(ref)
+            if data:
+                chunk._picture_data = data
+
+
 def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
     """Collapse duplicate picture-only chunks to one result per ``self_ref``.
 

haiku_rag_slim/haiku/rag/config/models.py

@@ -74,7 +74,16 @@ class EmbeddingsConfig(BaseModel):
 
 
 class RerankingConfig(BaseModel):
+    """Configuration for reranking search results.
+
+    Attributes:
+        model: Reranker model, or None to disable reranking.
+        multimodal: Whether the reranker scores picture chunks by their image
+            bytes in addition to text. Supported on the vllm provider only.
+    """
+
     model: ModelConfig | None = None
+    multimodal: bool = False
 
 
 class QAConfig(BaseModel):

haiku_rag_slim/haiku/rag/reranking/__init__.py

@@ -9,6 +9,9 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
     if model is None:
         return None
 
+    if config.reranking.multimodal and model.provider != "vllm":
+        raise ValueError("reranking.multimodal is only supported on the vllm provider")
+
     try:
         if model.provider == "cohere":
             from haiku.rag.reranking.cohere import CohereReranker

haiku_rag_slim/haiku/rag/reranking/vllm.py

@@ -1,15 +1,36 @@
+import base64
+
 import httpx
 
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
 
+def _document(chunk: Chunk) -> str | dict:
+    """Rerank document for a chunk: plain text, or content parts carrying the
+    picture bytes as a data URI when the chunk has them (multimodal rerank)."""
+    data = chunk._picture_data
+    if data is None:
+        return chunk.content
+
+    mime = "image/jpeg" if data.startswith(b"\xff\xd8") else "image/png"
+    encoded = base64.b64encode(data).decode("ascii")
+    parts: list[dict] = [
+        {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{encoded}"}}
+    ]
+    if chunk.content:
+        parts.append({"type": "text", "text": chunk.content})
+    return {"content": parts}
+
+
 class VLLMReranker(RerankerBase):
     def __init__(self, model: str, base_url: str):
         self._model = model
         self._base_url = base_url
         # One client reused across rerank calls (connection kept alive).
-        self._client = httpx.AsyncClient()
+        # Multimodal document batches can take far longer than httpx's 5s
+        # default timeout to score.
+        self._client = httpx.AsyncClient(timeout=httpx.Timeout(120.0))
 
     async def aclose(self) -> None:
         await self._client.aclose()
@@ -17,8 +38,7 @@ class VLLMReranker(RerankerBase):
     async def _rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
     ) -> list[tuple[Chunk, float]]:
-        # Prepare documents for reranking
-        documents = [chunk.content for chunk in chunks]
+        documents = [_document(chunk) for chunk in chunks]
 
         response = await self._client.post(
             f"{self._base_url}/v1/rerank",

tests/test_reranker.py

@@ -125,6 +125,32 @@ class TestGetReranker:
         with pytest.raises(ValueError, match="cross-encoder reranker requires name"):
             get_reranker(config)
 
+    def test_multimodal_requires_vllm_provider(self):
+        config = AppConfig(
+            reranking=RerankingConfig(
+                model=ModelConfig(provider="cohere", name="rerank-v3.5"),
+                multimodal=True,
+            )
+        )
+        with pytest.raises(ValueError, match="multimodal"):
+            get_reranker(config)
+
+    def test_multimodal_vllm_provider_builds_reranker(self):
+        pytest.importorskip("haiku.rag.reranking.vllm")
+        from haiku.rag.reranking.vllm import VLLMReranker
+
+        config = AppConfig(
+            reranking=RerankingConfig(
+                model=ModelConfig(
+                    provider="vllm",
+                    name="nvidia/llama-nemotron-rerank-vl-1b-v2",
+                    base_url="http://localhost:8000",
+                ),
+                multimodal=True,
+            )
+        )
+        assert isinstance(get_reranker(config), VLLMReranker)
+
     @pytest.mark.parametrize(
         "provider, model_name, class_module, class_name, extra_model_kwargs, expected_attrs, env_vars",
         [
@@ -287,6 +313,72 @@ async def test_vllm_reranker_reuses_pooled_client(monkeypatch):
     assert stats.closed == 1
 
 
+@pytest.mark.asyncio
+async def test_vllm_reranker_builds_multimodal_documents(monkeypatch):
+    """Chunks carrying picture bytes are sent as content-parts documents
+    (data-URI image plus text when the chunk has content); plain text chunks
+    stay strings in the same request."""
+    import base64
+
+    from haiku.rag.reranking.vllm import VLLMReranker
+
+    captured = {}
+
+    class FakeResponse:
+        def raise_for_status(self):
+            pass
+
+        def json(self):
+            return {
+                "results": [
+                    {"index": 0, "relevance_score": 0.9},
+                    {"index": 1, "relevance_score": 0.8},
+                    {"index": 2, "relevance_score": 0.7},
+                ]
+            }
+
+    class FakeAsyncClient:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        async def post(self, url, json, headers):
+            captured["json"] = json
+            return FakeResponse()
+
+        async def aclose(self):
+            pass
+
+    monkeypatch.setattr("httpx.AsyncClient", FakeAsyncClient)
+
+    png = b"\x89PNG\r\n\x1a\n" + b"png-payload"
+    jpeg = b"\xff\xd8\xff" + b"jpeg-payload"
+
+    text_chunk = Chunk(content="plain text")
+    described = Chunk(content="a described picture")
+    described._picture_data = png
+    undescribed = Chunk(content="")
+    undescribed._picture_data = jpeg
+
+    reranker = VLLMReranker(model="m", base_url="http://localhost:8000")
+    reranked = await reranker.rerank("q", [text_chunk, described, undescribed])
+
+    docs = captured["json"]["documents"]
+    assert docs[0] == "plain text"
+
+    image_part, text_part = docs[1]["content"]
+    prefix = "data:image/png;base64,"
+    assert image_part["type"] == "image_url"
+    assert image_part["image_url"]["url"].startswith(prefix)
+    assert base64.b64decode(image_part["image_url"]["url"].removeprefix(prefix)) == png
+    assert text_part == {"type": "text", "text": "a described picture"}
+
+    (jpeg_part,) = docs[2]["content"]
+    assert jpeg_part["image_url"]["url"].startswith("data:image/jpeg;base64,")
+
+    # Result indices pair back to the right chunks.
+    assert [c for c, _ in reranked] == [text_chunk, described, undescribed]
+
+
 @pytest.mark.asyncio
 async def test_jina_reranker_reuses_pooled_client(monkeypatch):
     """One httpx client is built and reused across rerank calls; aclose

tests/test_search.py

@@ -425,6 +425,74 @@ async def test_reranker_built_once_across_searches(temp_db_path, monkeypatch):
     assert build_count == 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("multimodal", [True, False])
+async def test_search_attaches_picture_bytes_for_multimodal_reranker(
+    temp_db_path, multimodal
+):
+    """With reranking.multimodal on, picture chunks reach the reranker with
+    their picture bytes attached; text chunks and the multimodal-off path are
+    untouched."""
+    from haiku.rag.store.models.chunk import Chunk
+    from haiku.rag.store.models.document_item import DocumentItem
+
+    captured = {}
+
+    class StubReranker:
+        async def rerank(self, query, chunks, top_n):
+            captured["chunks"] = chunks
+            return [(chunk, 1.0) for chunk in chunks][:top_n]
+
+        async def aclose(self):
+            pass
+
+    text_chunk = Chunk(
+        content="prose",
+        document_id="doc-1",
+        metadata={"doc_item_refs": ["#/texts/0"], "labels": ["paragraph"]},
+    )
+    picture_chunk = Chunk(
+        content="a chart of quarterly totals",
+        document_id="doc-1",
+        metadata={"doc_item_refs": ["#/pictures/0"], "labels": ["picture"]},
+    )
+    detached_chunk = Chunk(
+        content="no parent document",
+        metadata={"doc_item_refs": ["#/pictures/1"], "labels": ["picture"]},
+    )
+
+    async def fake_chunk_search(query, limit, search_type, filter):
+        return [(text_chunk, 0.9), (picture_chunk, 0.8), (detached_chunk, 0.7)]
+
+    async with HaikuRAG(temp_db_path, create=True) as rag:
+        await rag.document_item_repository.create_items(
+            "doc-1",
+            [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=0,
+                    self_ref="#/pictures/0",
+                    label="picture",
+                    text="a chart of quarterly totals",
+                    picture_data=b"picture-bytes",
+                ),
+            ],
+        )
+        rag.chunk_repository.search = fake_chunk_search  # type: ignore[method-assign]
+        rag.__dict__["reranker"] = StubReranker()
+        rag._config.reranking.multimodal = multimodal
+
+        await rag.search("totals", include_images=False)
+
+    reranked_text, reranked_picture, reranked_detached = captured["chunks"]
+    assert reranked_text._picture_data is None
+    assert reranked_detached._picture_data is None
+    if multimodal:
+        assert reranked_picture._picture_data == b"picture-bytes"
+    else:
+        assert reranked_picture._picture_data is None
+
+
 @pytest.mark.asyncio
 async def test_search_with_pil_image_works_like_bytes(temp_db_path, monkeypatch):
     from PIL import Image as PILImageModule