danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

emit synthetic picture chunks at ingest under multimodal embedders.

Browse source 
processing.chunk() merges text chunks with one synthetic Chunk per PictureItem-with-bytes,
sorted by iterate_items() position so chunk.order is structural.
embed_chunks dispatches on a Chunk._picture_data PrivateAttr
(text through embed_documents, picture through embed_image_query)
This commit is contained in:
Yiorgis Gozadinos 2026-05-03 17:03:56 +03:00
parent a36f920f5d
commit c9227c649f
No known key found for this signature in database
9 changed files with 449 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -3,6 +3,7 @@
### Added
- **Synthetic picture chunks at ingest under multimodal embedders.** `build_picture_chunks` (in `client/processing.py`) walks a `DoclingDocument`'s `pictures` and emits one synthetic `Chunk` per `PictureItem` with available bytes. Bytes ride on a `Chunk._picture_data` PrivateAttr (not serialized) so `embed_chunks` can route them through `embed_images` while text chunks keep going through `embed_documents`. Wired into the three ingest paths (`create_document`, `_create_document_from_file`, `_create_or_update_document_from_url`, `_update_document_with_chunks`, and `_rebuild_rechunk`) — guarded by `embedder.supports_images` so text-only configurations are unaffected. Snapshot/merge with `existing_picture_data` keeps picture chunks alive across rebuild round-trips. Picture chunks land in the same `chunks` table with the same vector dim as text chunks, so cross-modal search reuses the existing hybrid+RRF pipeline.
- **Multimodal embedder support (`provider="mlx"` and `provider="vllm"`).** `EmbedderWrapper` gains `supports_images: bool`, `embed_image_query`, and `embed_images`. Two pluggable paths:
- `provider="mlx"` — Apple Silicon, in-process via the new `[mlx]` optional extra (env-marker-guarded so `uv sync --all-extras` works on Linux/Windows/Intel-Mac without resolver errors). Loads any HF repo that ships an MLX `load_model.py` (default tested model: `jinaai/jina-embeddings-v4-mlx-8bit`, 2048-dim).
- `provider="vllm"` — cross-platform, talks HTTP to a vLLM server's OpenAI-compatible `/v1/embeddings` endpoint with vLLM's `messages` superset (text or `image_url` content parts, base64 data URIs). Works with `Qwen/Qwen3-VL-Embedding-8B` and `jinaai/jina-embeddings-v4`. No Python ML deps added — uses `httpx`.

15
haiku_rag_slim/haiku/rag/client/__init__.py
View file

@ -150,10 +150,21 @@ class HaikuRAG:
return await convert(self._config, source, format=format)
async def chunk(self, docling_document: "DoclingDocument") -> list[Chunk]:
async def chunk(
self,
docling_document: "DoclingDocument",
*,
existing_picture_data: dict[str, bytes] | None = None,
document_id: str | None = None,
) -> list[Chunk]:
from haiku.rag.client.processing import chunk
return await chunk(self._config, docling_document)
return await chunk(
self._config,
docling_document,
existing_picture_data=existing_picture_data,
document_id=document_id,
)
# =========================================================================
# Title Generation

22
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -79,6 +79,15 @@ async def _update_document_with_chunks(
"""
assert document.id is not None, "Document ID is required for update"
# Snapshot existing picture bytes before deleting items so the post-delete
# extract_items can merge them back. Skip under pictures="none" so updates
# reclaim storage.
existing_picture_data: dict[str, bytes] | None = None
if docling_document is not None and client._config.processing.pictures != "none":
existing_picture_data = (
await client.document_item_repository.get_all_picture_data(document.id)
)
chunks = await ensure_chunks_embedded(client._config, chunks)
versions = await client.store.current_table_versions()
@ -96,20 +105,7 @@ async def _update_document_with_chunks(
await client.chunk_repository.create(chunks)
# Replace document items when a new DoclingDocument is provided.
# Snapshot existing picture bytes first so they survive the
# delete-and-re-extract cycle when the live docling has already had
# its picture URIs stripped (rebuild / round-trip scenarios). Under
# `pictures="none"` we skip the snapshot so updates reclaim storage.
if docling_document is not None:
keep_picture_data = client._config.processing.pictures != "none"
existing_picture_data = (
await client.document_item_repository.get_all_picture_data(
updated_doc.id
)
if keep_picture_data
else None
)
await client.document_item_repository.delete_by_document_id(updated_doc.id)
items = extract_items(
updated_doc.id,

110
haiku_rag_slim/haiku/rag/client/processing.py
View file

@ -91,16 +91,118 @@ async def convert(
return await converter.convert_text(source, format=format)
async def chunk(config: AppConfig, docling_document: "DoclingDocument") -> list[Chunk]:
async def chunk(
config: AppConfig,
docling_document: "DoclingDocument",
*,
existing_picture_data: dict[str, bytes] | None = None,
document_id: str | None = None,
) -> list[Chunk]:
"""Chunk a DoclingDocument into Chunks.
Returns chunks without embeddings or document_id. Each chunk's `order`
field is set to its position in the list.
When the configured embedder supports images, also emit one synthetic
Chunk per ``PictureItem`` with available bytes (see ``build_picture_chunks``)
and merge them with text chunks in structural (``iterate_items()``) order.
``chunk.order`` is the index in the merged list.
``existing_picture_data`` (snapshot keyed by ``self_ref``) supplies bytes
for pictures whose ``image.uri`` has been stripped — used by the rebuild
path where the docling is loaded from the stored blob.
"""
from haiku.rag.chunkers import get_chunker
from haiku.rag.embeddings import get_embedder
chunker = get_chunker(config)
return await chunker.chunk(docling_document)
text_chunks = await chunker.chunk(docling_document)
if not get_embedder(config).supports_images:
for i, c in enumerate(text_chunks):
c.order = i
return text_chunks
picture_chunks = build_picture_chunks(
docling_document,
document_id=document_id,
existing_picture_data=existing_picture_data,
)
if not picture_chunks:
for i, c in enumerate(text_chunks):
c.order = i
return text_chunks
positions = {
item.self_ref: pos
for pos, (item, _level) in enumerate(docling_document.iterate_items())
}
def first_pos(c: Chunk) -> int:
refs = (c.metadata or {}).get("doc_item_refs") or []
return positions.get(refs[0], len(positions)) if refs else len(positions)
merged = sorted(text_chunks + picture_chunks, key=first_pos)
for i, c in enumerate(merged):
c.order = i
return merged
def build_picture_chunks(
docling_document: "DoclingDocument",
*,
document_id: str | None = None,
existing_picture_data: dict[str, bytes] | None = None,
) -> list[Chunk]:
"""Emit one synthetic ``Chunk`` per ``PictureItem`` with available bytes.
Bytes come from ``picture.image.uri`` (live data URI on a freshly-converted
docling) or from ``existing_picture_data`` keyed by ``self_ref`` (snapshot
taken before a delete-and-re-extract cycle, when the live docling has had
its picture URIs stripped). Pictures with no available bytes are skipped.
The bytes ride on ``Chunk._picture_data`` (a PrivateAttr — not serialized)
so ``embed_chunks`` can route them through ``embed_image_query``. The
``order`` field is left at its default (0); the caller (``chunk()``)
reassigns it after merging with text chunks in structural order.
"""
from haiku.rag.store.models.document_item import (
_decode_picture_bytes,
extract_item_text,
)
existing = existing_picture_data or {}
chunks: list[Chunk] = []
for picture in docling_document.pictures:
picture_data = _decode_picture_bytes(picture)
if picture_data is None:
picture_data = existing.get(picture.self_ref)
if picture_data is None:
continue
text = extract_item_text(picture, docling_document) or ""
page_numbers: list[int] = []
if prov := getattr(picture, "prov", None):
for p in prov:
page_no = getattr(p, "page_no", None)
if page_no is not None and page_no not in page_numbers:
page_numbers.append(page_no)
metadata = {
"doc_item_refs": [picture.self_ref],
"labels": ["picture"],
"page_numbers": sorted(page_numbers),
"headings": None,
}
chunk = Chunk(
document_id=document_id,
content=text,
metadata=metadata,
)
chunk._picture_data = picture_data
chunks.append(chunk)
return chunks
async def ensure_chunks_embedded(config: AppConfig, chunks: list[Chunk]) -> list[Chunk]:

17
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -224,11 +224,12 @@ async def _rebuild_rechunk(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Re-chunk and re-embed each document from its stored docling blob."""
from haiku.rag.embeddings import embed_chunks
from haiku.rag.embeddings import embed_chunks, get_embedder
pending_chunks: list[Chunk] = []
pending_docs: list[Document] = []
pending_doc_ids: list[str] = []
embedder = get_embedder(client._config)
for doc in documents:
assert doc.id is not None
@ -240,8 +241,18 @@ async def _rebuild_rechunk(
"requires it. Run a full rebuild (without --rechunk) instead."
)
# Chunk and embed
chunks = await client.chunk(docling_document)
# Stored blob has stripped picture URIs; pass the snapshot so
# build_picture_chunks (inside chunk()) can recover the bytes.
existing_picture_data = (
await client.document_item_repository.get_all_picture_data(doc.id)
if embedder.supports_images
else None
)
chunks = await client.chunk(
docling_document,
existing_picture_data=existing_picture_data,
document_id=doc.id,
)
embedded_chunks = await embed_chunks(chunks, client._config)
# Prepare chunks with document_id and order

61
haiku_rag_slim/haiku/rag/embeddings/__init__.py
View file

@ -90,20 +90,12 @@ EMBEDDING_BATCH_SIZE = 512
async def embed_chunks(
chunks: list["Chunk"], config: AppConfig = Config
) -> list["Chunk"]:
"""Generate embeddings for chunks.
"""Generate embeddings for chunks, dispatching text vs picture variants.
Contextualizes chunks (prepends headings) before embedding for better
semantic search. Returns new Chunk objects with embeddings set.
Embeddings are generated in batches to avoid request size limits
and timeouts with large document sets.
Args:
chunks: List of chunks to embed.
config: Configuration for embedder selection.
Returns:
New list of Chunk objects with embedding field populated.
Text chunks are contextualized (headings prepended) and routed through
``embed_documents``. Picture chunks (those carrying ``_picture_data``)
are routed through ``embed_images`` and require a multimodal embedder.
Vectors land in the original chunk order.
"""
if not chunks:
return []
@ -111,15 +103,36 @@ async def embed_chunks(
from haiku.rag.store.models.chunk import Chunk
embedder = get_embedder(config)
texts = contextualize(chunks)
# Batch embedding calls to avoid request size limits
all_embeddings: list[list[float]] = []
for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
batch = texts[i : i + EMBEDDING_BATCH_SIZE]
batch_embeddings = await embedder.embed_documents(batch)
all_embeddings.extend(batch_embeddings)
text_chunks: list[Chunk] = []
picture_chunks: list[Chunk] = []
for chunk in chunks:
if chunk._picture_data is not None:
picture_chunks.append(chunk)
else:
text_chunks.append(chunk)
text_embeddings: list[list[float]] = []
if text_chunks:
texts = contextualize(text_chunks)
for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
batch = texts[i : i + EMBEDDING_BATCH_SIZE]
text_embeddings.extend(await embedder.embed_documents(batch))
picture_embeddings: list[list[float]] = []
if picture_chunks:
if not embedder.supports_images:
raise ValueError(
"Picture chunks require a multimodal embedder. Configure "
"provider='mlx' or provider='vllm', or omit picture chunks."
)
for chunk in picture_chunks:
picture_embeddings.append(
await embedder.embed_image_query(chunk._picture_data)
)
text_iter = iter(text_embeddings)
picture_iter = iter(picture_embeddings)
return [
Chunk(
id=chunk.id,
@ -130,9 +143,13 @@ async def embed_chunks(
document_uri=chunk.document_uri,
document_title=chunk.document_title,
document_meta=chunk.document_meta,
embedding=embedding,
embedding=(
next(picture_iter)
if chunk._picture_data is not None
else next(text_iter)
),
)
for chunk, embedding in zip(chunks, all_embeddings)
for chunk in chunks
]

7
haiku_rag_slim/haiku/rag/store/models/chunk.py
View file

@ -1,6 +1,6 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel
from pydantic import BaseModel, PrivateAttr
if TYPE_CHECKING:
from docling_core.types.doc.document import DocItem, DoclingDocument
@ -103,6 +103,11 @@ class Chunk(BaseModel):
document_meta: dict = {}
embedding: list[float] | None = None
# Transient: picture bytes for synthetic picture chunks. Set by
# build_picture_chunks; consumed by embed_chunks to route through
# embed_images. Excluded from serialization (PrivateAttr).
_picture_data: bytes | None = PrivateAttr(default=None)
def get_chunk_metadata(self) -> ChunkMetadata:
"""Parse metadata dict into structured ChunkMetadata."""
return ChunkMetadata.model_validate(self.metadata)

23
tests/cassettes/test_rebuild/test_rebuild_rechunk.yaml
View file

File diff suppressed because one or more lines are too long

246
tests/test_picture_in_context.py
View file

@ -322,6 +322,252 @@ async def test_search_tool_returns_multimodal_when_picture_present():
assert part.data == PICTURE_BYTES
# B2: synthetic picture chunks at ingestion
def test_build_picture_chunks_uses_live_uri():
from haiku.rag.client.processing import build_picture_chunks
from tests.store.test_document_items import _docling_doc_with_picture
doc = _docling_doc_with_picture()
chunks = build_picture_chunks(doc, document_id="doc-1")
assert len(chunks) == 1
chunk = chunks[0]
assert chunk.metadata["doc_item_refs"] == ["#/pictures/0"]
assert chunk.metadata["labels"] == ["picture"]
assert chunk._picture_data is not None
assert chunk._picture_data.startswith(b"\x89PNG")
assert chunk.document_id == "doc-1"
def test_build_picture_chunks_falls_back_to_existing_picture_data():
"""When the live docling has its picture URIs stripped, the snapshot
fills the gap so rebuild round-trips don't lose picture chunks."""
from haiku.rag.client.processing import build_picture_chunks
from tests.store.test_document_items import _docling_doc_with_picture
doc = _docling_doc_with_picture()
for picture in doc.pictures:
picture.image = None
chunks = build_picture_chunks(
doc,
document_id="doc-1",
existing_picture_data={"#/pictures/0": b"snapshot-bytes"},
)
assert len(chunks) == 1
assert chunks[0]._picture_data == b"snapshot-bytes"
def test_build_picture_chunks_skips_pictures_without_bytes():
from haiku.rag.client.processing import build_picture_chunks
from tests.store.test_document_items import _docling_doc_with_picture
doc = _docling_doc_with_picture()
for picture in doc.pictures:
picture.image = None
chunks = build_picture_chunks(doc, document_id="doc-1")
assert chunks == []
@pytest.mark.asyncio
async def test_chunk_interleaves_picture_in_structural_order(monkeypatch):
"""``chunk()`` merges text and picture chunks by their first
``doc_item_ref``'s position in ``iterate_items()``, so picture chunks
sit where they appear in the document, not appended at the end.
"""
from haiku.rag.client.processing import chunk
from haiku.rag.config import AppConfig
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.store.models.chunk import Chunk
class StubMultimodalEmbedder(EmbedderWrapper):
supports_images = True
def __init__(self):
super().__init__(embedder=None, vector_dim=4)
class StubChunker:
async def chunk(self, document):
# Two text chunks straddling the picture's structural position.
# iterate_items order on the fixture below: texts/0, texts/1,
# pictures/0, texts/2 — positions 0,1,2,3.
return [
Chunk(
content="before",
metadata={"doc_item_refs": ["#/texts/0", "#/texts/1"]},
),
Chunk(
content="after",
metadata={"doc_item_refs": ["#/texts/2"]},
),
]
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder",
lambda *a, **kw: StubMultimodalEmbedder(),
)
monkeypatch.setattr(
"haiku.rag.chunkers.get_chunker", lambda *a, **kw: StubChunker()
)
from docling_core.types.doc.document import DoclingDocument, ImageRef
from docling_core.types.doc.labels import DocItemLabel
from PIL import Image as PILImageModule
img = PILImageModule.new("RGB", (8, 8), "blue")
doc = DoclingDocument(name="ordered")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="A")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="B")
doc.add_picture(image=ImageRef.from_pil(img, dpi=72))
doc.add_text(label=DocItemLabel.PARAGRAPH, text="C")
chunks = await chunk(AppConfig(), doc)
contents = [c.content for c in chunks]
assert contents == ["before", "", "after"], (
f"expected [before, picture, after], got {contents}"
)
assert chunks[1].metadata["labels"] == ["picture"]
assert chunks[1].metadata["doc_item_refs"] == ["#/pictures/0"]
assert chunks[1]._picture_data is not None
# chunk.order matches list index after the merge sort.
for i, c in enumerate(chunks):
assert c.order == i
@pytest.mark.asyncio
async def test_embed_chunks_dispatches_text_vs_picture(monkeypatch):
"""embed_chunks routes text chunks through embed_documents (batched) and
picture chunks through embed_image_query (one at a time), reassembling
in original order."""
from haiku.rag.embeddings import EmbedderWrapper, embed_chunks
from haiku.rag.store.models.chunk import Chunk
text_calls: list[list[str]] = []
image_calls: list[bytes] = []
class StubEmbedder(EmbedderWrapper):
supports_images = True
def __init__(self):
super().__init__(embedder=None, vector_dim=4)
async def embed_documents(self, texts):
text_calls.append(list(texts))
return [[0.1, 0.2, 0.3, 0.4] for _ in texts]
async def embed_image_query(self, image):
image_calls.append(image)
return [0.9, 0.8, 0.7, 0.6]
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder", lambda *a, **kw: StubEmbedder()
)
text_chunk = Chunk(content="hello", order=0)
pic_chunk = Chunk(
content="figure 1",
metadata={"labels": ["picture"], "doc_item_refs": ["#/pictures/0"]},
order=1,
)
pic_chunk._picture_data = b"PNGBYTES"
embedded = await embed_chunks([text_chunk, pic_chunk, text_chunk.model_copy()])
assert len(embedded) == 3
assert embedded[0].embedding == [0.1, 0.2, 0.3, 0.4]
assert embedded[1].embedding == [0.9, 0.8, 0.7, 0.6]
assert embedded[2].embedding == [0.1, 0.2, 0.3, 0.4]
assert text_calls == [["hello", "hello"]]
assert image_calls == [b"PNGBYTES"]
@pytest.mark.asyncio
async def test_embed_chunks_raises_on_picture_chunks_with_text_only_embedder(
monkeypatch,
):
from haiku.rag.embeddings import EmbedderWrapper, embed_chunks
from haiku.rag.store.models.chunk import Chunk
class TextOnlyEmbedder(EmbedderWrapper):
def __init__(self):
super().__init__(embedder=None, vector_dim=4)
async def embed_documents(self, texts):
return [[0.0] * 4 for _ in texts]
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder", lambda *a, **kw: TextOnlyEmbedder()
)
pic_chunk = Chunk(content="x", metadata={"labels": ["picture"]}, order=0)
pic_chunk._picture_data = b"PNG"
with pytest.raises(ValueError, match="multimodal embedder"):
await embed_chunks([pic_chunk])
@pytest.mark.asyncio
async def test_ingest_emits_picture_chunks_with_multimodal_embedder(
temp_db_path, monkeypatch
):
"""End-to-end: ingest a docling doc with one picture under a stub
multimodal embedder; chunks_table contains a picture-labelled chunk
pointing at the picture's self_ref."""
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.embeddings import EmbedderWrapper, embed_chunks
from haiku.rag.store.models.document import Document
from tests.store.test_document_items import _docling_doc_with_picture
class StubMultimodalEmbedder(EmbedderWrapper):
supports_images = True
def __init__(self):
super().__init__(embedder=None, vector_dim=4)
async def embed_documents(self, texts):
return [[0.1] * 4 for _ in texts]
async def embed_image_query(self, image):
return [0.9] * 4
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder",
lambda *a, **kw: StubMultimodalEmbedder(),
)
docling_doc = _docling_doc_with_picture()
from haiku.rag.config import AppConfig, EmbeddingModelConfig, EmbeddingsConfig
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
)
)
config.processing.pictures = "image"
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
chunks = await rag.chunk(docling_doc)
embedded = await embed_chunks(chunks, rag._config)
document = Document(content="x", uri="test://doc")
document.set_docling(docling_doc)
await _store_document_with_chunks(rag, document, embedded, docling_doc)
all_db_chunks = await rag.chunk_repository.store.chunks_table.query().to_list()
picture_db_chunks = [
c for c in all_db_chunks if "picture" in (c.get("metadata") or "")
]
assert len(picture_db_chunks) >= 1
assert any(
"#/pictures/0" in (c.get("metadata") or "") for c in picture_db_chunks
)
@pytest.mark.asyncio
async def test_search_tool_returns_plain_string_when_no_pictures():
"""When no result carries image_data the tool returns a plain str (no