From 543aba7547cd526b94018082eb30aef25211fa05 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 24 Jul 2026 09:39:07 +0300 Subject: [PATCH 1/3] Multimodal reranking: send picture chunks to vllm rerankers as images reranking.multimodal (vllm provider only) attaches picture bytes to synthetic picture chunks before rerank; VLLMReranker sends them as content-parts documents (base64 data URI + description text) in the same /v1/rerank request as plain text documents. --- CHANGELOG.md | 4 + docs/configuration/index.md | 1 + docs/configuration/providers.md | 15 +++ haiku_rag_slim/haiku/rag/client/search.py | 25 +++++ haiku_rag_slim/haiku/rag/config/models.py | 9 ++ .../haiku/rag/reranking/__init__.py | 3 + haiku_rag_slim/haiku/rag/reranking/vllm.py | 22 ++++- tests/test_reranker.py | 92 +++++++++++++++++++ tests/test_search.py | 63 +++++++++++++ 9 files changed, 232 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 146020b8..e3c7c2ad 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/docs/configuration/index.md b/docs/configuration/index.md index a7e37d1f..ba1103ee 100644 --- a/docs/configuration/index.md +++ b/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: diff --git a/docs/configuration/providers.md b/docs/configuration/providers.md index c99a7cbe..39db9e20 100644 --- a/docs/configuration/providers.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index ef814b9d..46752242 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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``. diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index e9d429c6..1bca596c 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/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): diff --git a/haiku_rag_slim/haiku/rag/reranking/__init__.py b/haiku_rag_slim/haiku/rag/reranking/__init__.py index a4b0e535..e02e64db 100644 --- a/haiku_rag_slim/haiku/rag/reranking/__init__.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/reranking/vllm.py b/haiku_rag_slim/haiku/rag/reranking/vllm.py index 67909c8d..4a5e4d22 100644 --- a/haiku_rag_slim/haiku/rag/reranking/vllm.py +++ b/haiku_rag_slim/haiku/rag/reranking/vllm.py @@ -1,9 +1,28 @@ +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 @@ -17,8 +36,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", diff --git a/tests/test_reranker.py b/tests/test_reranker.py index b05e941b..0f8c77bf 100644 --- a/tests/test_reranker.py +++ b/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 diff --git a/tests/test_search.py b/tests/test_search.py index 94c2bd51..b9755add 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -425,6 +425,69 @@ 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"]}, + ) + + async def fake_chunk_search(query, limit, search_type, filter): + return [(text_chunk, 0.9), (picture_chunk, 0.8)] + + 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 = captured["chunks"] + assert reranked_text._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 From f31060e7417597abae3e87081fa9d46b168e1b31 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 24 Jul 2026 12:34:31 +0300 Subject: [PATCH 2/3] Set explicit 120s timeout on the vllm reranker HTTP client --- haiku_rag_slim/haiku/rag/reranking/vllm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/haiku_rag_slim/haiku/rag/reranking/vllm.py b/haiku_rag_slim/haiku/rag/reranking/vllm.py index 4a5e4d22..a2b40087 100644 --- a/haiku_rag_slim/haiku/rag/reranking/vllm.py +++ b/haiku_rag_slim/haiku/rag/reranking/vllm.py @@ -28,7 +28,9 @@ class VLLMReranker(RerankerBase): 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() From 35b8b046ce6b8019acf7223d15b7a32a127bce22 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 24 Jul 2026 13:07:39 +0300 Subject: [PATCH 3/3] Cover the detached-chunk guard in the multimodal rerank search test --- tests/test_search.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_search.py b/tests/test_search.py index b9755add..e1eccf66 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -456,9 +456,13 @@ async def test_search_attaches_picture_bytes_for_multimodal_reranker( 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)] + 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( @@ -480,8 +484,9 @@ async def test_search_attaches_picture_bytes_for_multimodal_reranker( await rag.search("totals", include_images=False) - reranked_text, reranked_picture = captured["chunks"] + 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: