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.
This commit is contained in:
Yiorgis Gozadinos 2026-07-24 09:39:07 +03:00
parent 29ccb9a035
commit 543aba7547
No known key found for this signature in database
9 changed files with 232 additions and 2 deletions

View file

@ -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.

View file

@ -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:

View file

@ -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.

View file

@ -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``.

View file

@ -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):

View file

@ -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

View file

@ -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",

View file

@ -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

View file

@ -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