46f7ab8d97
ToolReturn.content reaches the model as a user-role message and the pictures arrive bare, so nothing connects a figure to the chunk it came from: BinaryContent.identifier does not survive serialization to the vision API, and the captions in the result text correlate only by position. Precede each picture with its position, source chunk id and self_ref. build_binary_parts_from_results becomes build_image_content_from_results and returns the labels interleaved with the pictures, so both attachment sites emit them the same way. This does not stop a model narrating retrieved pictures as user-supplied. Measured on gemma4-26b with a single note ahead of the batch, and again with per-image labels: it quotes the label and still says the user provided them. The message role wins over its text.
1044 lines
37 KiB
Python
1044 lines
37 KiB
Python
"""Picture-bearing search results: image_data attachment, expansion, multimodal ToolReturn."""
|
|
|
|
import base64
|
|
from dataclasses import dataclass
|
|
from io import BytesIO
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from PIL import Image as PILImageModule
|
|
from pydantic_ai import RunContext
|
|
from pydantic_ai.messages import BinaryContent, ToolReturn
|
|
from pydantic_ai.models.test import TestModel
|
|
from pydantic_ai.usage import RunUsage
|
|
|
|
from haiku.rag.capabilities.rag import RAGState, create_capability
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.client.search import _populate_image_data
|
|
from haiku.rag.config import AppConfig, Config
|
|
from haiku.rag.context import expand_with_items
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
from haiku.rag.tools.search import create_search_toolset
|
|
|
|
|
|
def _make_png(color: str = "red", size: tuple[int, int] = (4, 4)) -> bytes:
|
|
buf = BytesIO()
|
|
PILImageModule.new("RGB", size, color).save(buf, "PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
PICTURE_BYTES = _make_png("red")
|
|
PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_populate_image_data_attaches_base64(temp_db_path):
|
|
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="#/texts/0",
|
|
label="paragraph",
|
|
text="Some text",
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=1,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text="",
|
|
picture_data=PICTURE_BYTES,
|
|
),
|
|
],
|
|
)
|
|
|
|
text_only = SearchResult(
|
|
content="Some text",
|
|
score=1.0,
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/texts/0"],
|
|
labels=["paragraph"],
|
|
)
|
|
with_picture = SearchResult(
|
|
content="A figure caption",
|
|
score=0.9,
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/texts/0", "#/pictures/0"],
|
|
labels=["paragraph", "picture"],
|
|
)
|
|
|
|
await _populate_image_data(rag, [text_only, with_picture])
|
|
|
|
# Text-only result is unchanged
|
|
assert text_only.image_data is None
|
|
# Picture-bearing result has the bytes attached
|
|
assert with_picture.image_data == {"#/pictures/0": PICTURE_B64}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_populate_image_data_attaches_picture_via_caption(temp_db_path):
|
|
"""A result whose matched refs include a figure's caption (but not the
|
|
picture itself) gets the picture bytes attached, resolved through the
|
|
caption's adjacent picture. This is the common case: the prose chunk
|
|
carrying a figure's caption ranks, while the picture is its own chunk."""
|
|
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="Figure 1. The layout.",
|
|
picture_data=PICTURE_BYTES,
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=1,
|
|
self_ref="#/texts/0",
|
|
label="caption",
|
|
text="Figure 1. The layout.",
|
|
),
|
|
],
|
|
)
|
|
|
|
via_caption = SearchResult(
|
|
content="Figure 1. The layout.",
|
|
score=1.0,
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/texts/0"],
|
|
labels=["caption"],
|
|
)
|
|
|
|
await _populate_image_data(rag, [via_caption])
|
|
|
|
assert via_caption.image_data == {"#/pictures/0": PICTURE_B64}
|
|
assert via_caption.picture_captions == {"#/pictures/0": "Figure 1. The layout."}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_search_include_images_false_skips_lookup(temp_db_path):
|
|
"""include_images=False must short-circuit the picture-bytes lookup."""
|
|
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",
|
|
picture_data=PICTURE_BYTES,
|
|
),
|
|
],
|
|
)
|
|
# Spy that we never reach the picture-bytes accessor
|
|
rag.document_item_repository.get_pictures_for_chunk = AsyncMock( # type: ignore[method-assign]
|
|
wraps=rag.document_item_repository.get_pictures_for_chunk
|
|
)
|
|
|
|
from haiku.rag.client.search import search
|
|
|
|
# Stub the chunk-search results so we don't depend on embeddings/FTS
|
|
async def fake_chunk_search(*args, **kwargs):
|
|
return []
|
|
|
|
rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign]
|
|
|
|
await search(rag, "anything", include_images=False)
|
|
rag.document_item_repository.get_pictures_for_chunk.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_path):
|
|
"""A picture item with empty text must keep its self_ref through expansion."""
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
# Build an items table with a section header + a paragraph match + an
|
|
# adjacent picture row that has no text. The expansion must keep
|
|
# picture self_refs even when item.text is empty so picture bytes
|
|
# are still attached downstream.
|
|
await rag.document_item_repository.create_items(
|
|
"doc-1",
|
|
[
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=0,
|
|
self_ref="#/texts/0",
|
|
label="section_header",
|
|
text="Methods",
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=1,
|
|
self_ref="#/texts/1",
|
|
label="paragraph",
|
|
text="The figure below shows the architecture.",
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=2,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text="",
|
|
picture_data=PICTURE_BYTES,
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=3,
|
|
self_ref="#/texts/2",
|
|
label="paragraph",
|
|
text="More commentary follows.",
|
|
),
|
|
],
|
|
)
|
|
|
|
# Match on the paragraph that mentions the figure.
|
|
seed = SearchResult(
|
|
content="The figure below shows the architecture.",
|
|
score=1.0,
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/texts/1"],
|
|
labels=["paragraph"],
|
|
)
|
|
expanded = await expand_with_items(
|
|
rag.document_item_repository,
|
|
"doc-1",
|
|
[seed],
|
|
max_chars=10_000,
|
|
)
|
|
assert len(expanded) == 1
|
|
out = expanded[0]
|
|
assert "#/pictures/0" in out.doc_item_refs, (
|
|
"Picture self_ref should survive expansion even with empty text"
|
|
)
|
|
assert "picture" in out.labels
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
@pytest.mark.asyncio
|
|
async def test_rechunk_preserves_picture_data(temp_db_path):
|
|
"""``rebuild --rechunk`` keeps ``picture_data`` for every picture row."""
|
|
from haiku.rag.client import RebuildMode
|
|
from haiku.rag.client.documents import _store_document_with_chunks
|
|
from haiku.rag.store.models.document import Document
|
|
from tests.store.test_document_items import _docling_doc_with_picture
|
|
|
|
docling_doc = _docling_doc_with_picture()
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
document = Document(content="x", uri="test://doc")
|
|
document.set_docling(docling_doc)
|
|
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
|
assert created.id is not None
|
|
before = await rag.document_item_repository.get_all_picture_data(created.id)
|
|
assert before.get("#/pictures/0") is not None
|
|
|
|
async for _ in rag.rebuild_database(mode=RebuildMode.RECHUNK):
|
|
pass
|
|
|
|
after = await rag.document_item_repository.get_all_picture_data(created.id)
|
|
assert after.get("#/pictures/0") == before.get("#/pictures/0")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_embed_only_preserves_picture_vectors(temp_db_path, monkeypatch):
|
|
"""``rebuild --embed-only`` must re-embed picture chunks through the image
|
|
path. With a multimodal embedder, picture vectors must survive the rebuild
|
|
instead of being overwritten by a text embedding of the caption."""
|
|
from haiku.rag.client import RebuildMode
|
|
from haiku.rag.client.documents import _store_document_with_chunks
|
|
from haiku.rag.config import EmbeddingModelConfig, EmbeddingsConfig
|
|
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
|
|
|
|
TEXT_VEC = [0.1, 0.1, 0.1, 0.1]
|
|
IMAGE_VEC = [0.9, 0.9, 0.9, 0.9]
|
|
|
|
class StubMultimodalEmbedder(EmbedderWrapper):
|
|
supports_images = True
|
|
|
|
def __init__(self):
|
|
super().__init__(embedder=None, vector_dim=4)
|
|
|
|
async def embed_documents(self, texts):
|
|
return [list(TEXT_VEC) for _ in texts]
|
|
|
|
async def embed_image(self, image):
|
|
return list(IMAGE_VEC)
|
|
|
|
monkeypatch.setattr(
|
|
"haiku.rag.store.engine.get_embedder",
|
|
lambda *a, **kw: StubMultimodalEmbedder(),
|
|
)
|
|
|
|
config = AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
|
|
)
|
|
)
|
|
# The fixture picture is 8x8; disable the size filter so it still chunks.
|
|
config.processing.min_picture_size = 0
|
|
|
|
docling_doc = _docling_doc_with_picture()
|
|
|
|
async def _picture_chunk_row(rag):
|
|
rows = await rag.chunk_repository.store.chunks_table.query().to_list()
|
|
picture_rows = [r for r in rows if "#/pictures/0" in (r.get("metadata") or "")]
|
|
assert len(picture_rows) == 1
|
|
return picture_rows[0]
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
|
chunks = await rag.chunk(docling_doc)
|
|
embedded = await embed_chunks(chunks, rag.embedder, rag._config)
|
|
document = Document(content="x", uri="test://doc")
|
|
document.set_docling(docling_doc)
|
|
await _store_document_with_chunks(rag, document, embedded, docling_doc)
|
|
|
|
before = await _picture_chunk_row(rag)
|
|
assert list(before["vector"]) == pytest.approx(IMAGE_VEC)
|
|
|
|
async for _ in rag.rebuild_database(mode=RebuildMode.EMBED_ONLY):
|
|
pass
|
|
|
|
after = await _picture_chunk_row(rag)
|
|
assert list(after["vector"]) == pytest.approx(IMAGE_VEC)
|
|
assert after["id"] == before["id"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expand_context_does_not_attach_expansion_added_pictures(temp_db_path):
|
|
"""expand_context preserves picture bytes from the pre-expansion result and
|
|
does NOT re-fetch bytes for picture self_refs swept in by section
|
|
expansion. The expansion-added picture ref still rides along in
|
|
doc_item_refs for cross-referencing, but image_data stays empty so the
|
|
multimodal payload is bounded by what search originally returned."""
|
|
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="#/texts/0",
|
|
label="section_header",
|
|
text="Methods",
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=1,
|
|
self_ref="#/texts/1",
|
|
label="paragraph",
|
|
text="The figure below shows the architecture.",
|
|
),
|
|
DocumentItem(
|
|
document_id="doc-1",
|
|
position=2,
|
|
self_ref="#/pictures/0",
|
|
label="picture",
|
|
text="",
|
|
picture_data=PICTURE_BYTES,
|
|
),
|
|
],
|
|
)
|
|
seed = SearchResult(
|
|
content="The figure below shows the architecture.",
|
|
score=1.0,
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/texts/1"],
|
|
labels=["paragraph"],
|
|
image_data=None,
|
|
)
|
|
expanded = await rag.expand_context([seed])
|
|
assert len(expanded) == 1
|
|
out = expanded[0]
|
|
assert "#/pictures/0" in out.doc_item_refs
|
|
assert out.image_data is None
|
|
|
|
|
|
@dataclass
|
|
class _Deps:
|
|
client: object
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_tool_returns_multimodal_when_picture_present():
|
|
"""The agent-facing search tool must wrap text + BinaryContent in ToolReturn
|
|
whenever a result carries picture image_data AND the QA model is vision-capable."""
|
|
|
|
picture_result = SearchResult(
|
|
content="A diagram of the layout",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": PICTURE_B64},
|
|
)
|
|
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[picture_result])
|
|
fake_client.expand_context = AsyncMock(return_value=[picture_result])
|
|
|
|
config = AppConfig()
|
|
config.qa.model.vision = True
|
|
toolset = create_search_toolset(config, expand_context=False)
|
|
func = toolset.tools["search"].function
|
|
|
|
ctx = RunContext(
|
|
deps=_Deps(client=fake_client), # type: ignore[arg-type]
|
|
model=TestModel(),
|
|
usage=RunUsage(),
|
|
run_id="run-1",
|
|
)
|
|
result = await func(ctx, "anything")
|
|
|
|
assert isinstance(result, ToolReturn)
|
|
assert isinstance(result.return_value, str)
|
|
assert "Type: picture" in result.return_value or "rank 1" in result.return_value
|
|
assert result.content is not None
|
|
images = [c for c in result.content if isinstance(c, BinaryContent)]
|
|
assert len(images) == 1
|
|
part = images[0]
|
|
assert isinstance(part, BinaryContent)
|
|
assert part.media_type == "image/png"
|
|
assert part.identifier == "#/pictures/0"
|
|
assert part.data == PICTURE_BYTES
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_tool_attaches_same_self_ref_from_different_documents():
|
|
"""Two different documents both have ``#/pictures/0`` — the dedup must
|
|
key on ``(document_id, self_ref)`` so each document's figure reaches
|
|
the model. Keying on ``self_ref`` alone silently drops the second
|
|
document's bytes, leaving the model with text only for that result."""
|
|
other_bytes = _make_png("blue")
|
|
other_b64 = base64.b64encode(other_bytes).decode("ascii")
|
|
|
|
doc_a = SearchResult(
|
|
content="Figure from doc A",
|
|
score=1.0,
|
|
chunk_id="chunk-a",
|
|
document_id="doc-A",
|
|
doc_item_refs=["#/pictures/0"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": PICTURE_B64},
|
|
)
|
|
doc_b = SearchResult(
|
|
content="Figure from doc B (same self_ref, different bytes)",
|
|
score=0.9,
|
|
chunk_id="chunk-b",
|
|
document_id="doc-B",
|
|
doc_item_refs=["#/pictures/0"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": other_b64},
|
|
)
|
|
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[doc_a, doc_b])
|
|
fake_client.expand_context = AsyncMock(return_value=[doc_a, doc_b])
|
|
|
|
config = AppConfig()
|
|
config.qa.model.vision = True
|
|
toolset = create_search_toolset(config, expand_context=False)
|
|
func = toolset.tools["search"].function
|
|
|
|
ctx = RunContext(
|
|
deps=_Deps(client=fake_client), # type: ignore[arg-type]
|
|
model=TestModel(),
|
|
usage=RunUsage(),
|
|
run_id="run-1",
|
|
)
|
|
result = await func(ctx, "anything")
|
|
|
|
assert isinstance(result, ToolReturn)
|
|
assert result.content is not None
|
|
images = [c for c in result.content if isinstance(c, BinaryContent)]
|
|
assert len(images) == 2, (
|
|
"Both documents' figures must reach the model — dedup keyed on "
|
|
"self_ref alone would drop doc-B's bytes."
|
|
)
|
|
payloads = {part.data for part in images}
|
|
assert PICTURE_BYTES in payloads
|
|
assert other_bytes in payloads
|
|
|
|
|
|
# Synthetic picture chunks at ingest
|
|
|
|
|
|
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 == []
|
|
|
|
|
|
def _doc_with_picture_images(*images):
|
|
"""DoclingDocument with one paragraph and one PictureItem per PIL image."""
|
|
from docling_core.types.doc.document import DoclingDocument, ImageRef
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
|
|
doc = DoclingDocument(name="pics")
|
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
|
|
for img in images:
|
|
doc.add_picture(image=ImageRef.from_pil(img, dpi=72))
|
|
return doc
|
|
|
|
|
|
def test_build_picture_chunks_dedupes_identical_bytes():
|
|
"""Identical picture bytes within a document produce one chunk — the
|
|
first occurrence. A watermark repeated on every page embeds once."""
|
|
from PIL import Image as PILImageModule
|
|
|
|
from haiku.rag.client.processing import build_picture_chunks
|
|
|
|
red = PILImageModule.new("RGB", (100, 100), "red")
|
|
blue = PILImageModule.new("RGB", (100, 100), "blue")
|
|
doc = _doc_with_picture_images(red, red, blue, red)
|
|
|
|
chunks = build_picture_chunks(doc, document_id="doc-1")
|
|
|
|
refs = [c.metadata["doc_item_refs"][0] for c in chunks]
|
|
assert refs == ["#/pictures/0", "#/pictures/2"]
|
|
|
|
|
|
def test_build_picture_chunks_skips_small_pictures():
|
|
"""Pictures whose smaller side is under min_picture_size are not chunked."""
|
|
from PIL import Image as PILImageModule
|
|
|
|
from haiku.rag.client.processing import build_picture_chunks
|
|
|
|
icon = PILImageModule.new("RGB", (16, 16), "red")
|
|
figure = PILImageModule.new("RGB", (100, 100), "blue")
|
|
banner = PILImageModule.new("RGB", (200, 16), "green")
|
|
doc = _doc_with_picture_images(icon, figure, banner)
|
|
|
|
chunks = build_picture_chunks(doc, document_id="doc-1", min_picture_size=64)
|
|
|
|
assert [c.metadata["doc_item_refs"][0] for c in chunks] == ["#/pictures/1"]
|
|
|
|
|
|
def test_build_picture_chunks_measures_snapshot_bytes():
|
|
"""Rebuild path: picture.image is None, so size comes from a PIL header
|
|
read of the snapshot bytes — existing DBs shed small pictures on rebuild."""
|
|
import io
|
|
|
|
from PIL import Image as PILImageModule
|
|
|
|
from haiku.rag.client.processing import build_picture_chunks
|
|
|
|
icon_png = io.BytesIO()
|
|
PILImageModule.new("RGB", (16, 16), "red").save(icon_png, format="PNG")
|
|
doc = _doc_with_picture_images(PILImageModule.new("RGB", (16, 16), "red"))
|
|
for picture in doc.pictures:
|
|
picture.image = None
|
|
|
|
chunks = build_picture_chunks(
|
|
doc,
|
|
document_id="doc-1",
|
|
existing_picture_data={"#/pictures/0": icon_png.getvalue()},
|
|
min_picture_size=64,
|
|
)
|
|
assert chunks == []
|
|
|
|
|
|
def test_build_picture_chunks_keeps_unmeasurable_bytes():
|
|
"""Bytes PIL can't parse are kept — the filter only drops what it can
|
|
measure."""
|
|
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"not-an-image"},
|
|
min_picture_size=64,
|
|
)
|
|
assert len(chunks) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chunk_filters_small_pictures_by_config(monkeypatch):
|
|
"""``chunk()`` applies ``processing.min_picture_size`` — with the default
|
|
config, icon-sized pictures don't become picture chunks."""
|
|
from haiku.rag.client.processing import chunk
|
|
from haiku.rag.embeddings import EmbedderWrapper
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
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)
|
|
|
|
class StubChunker:
|
|
async def chunk(self, document):
|
|
return [Chunk(content="text", metadata={"doc_item_refs": ["#/texts/0"]})]
|
|
|
|
monkeypatch.setattr(
|
|
"haiku.rag.chunkers.get_chunker", lambda *a, **kw: StubChunker()
|
|
)
|
|
|
|
doc = _docling_doc_with_picture() # 8x8 picture, below the 64px default
|
|
chunks = await chunk(AppConfig(), doc, embedder=StubMultimodalEmbedder())
|
|
|
|
assert [c.content for c in chunks] == ["text"]
|
|
assert not any("picture" in (c.metadata or {}).get("labels", []) for c in 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.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.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", (64, 64), "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, embedder=StubMultimodalEmbedder())
|
|
|
|
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():
|
|
"""embed_chunks routes text chunks through embed_documents (batched) and
|
|
picture chunks through embed_image (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(self, image):
|
|
image_calls.append(image)
|
|
return [0.9, 0.8, 0.7, 0.6]
|
|
|
|
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()], StubEmbedder()
|
|
)
|
|
|
|
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():
|
|
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]
|
|
|
|
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], TextOnlyEmbedder())
|
|
|
|
|
|
@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(self, image):
|
|
return [0.9] * 4
|
|
|
|
monkeypatch.setattr(
|
|
"haiku.rag.store.engine.get_embedder",
|
|
lambda *a, **kw: StubMultimodalEmbedder(),
|
|
)
|
|
|
|
docling_doc = _docling_doc_with_picture()
|
|
|
|
from haiku.rag.config import EmbeddingModelConfig, EmbeddingsConfig
|
|
|
|
config = AppConfig(
|
|
embeddings=EmbeddingsConfig(
|
|
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
|
|
)
|
|
)
|
|
# The fixture picture is 8x8; disable the size filter so it still chunks.
|
|
config.processing.min_picture_size = 0
|
|
|
|
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
|
chunks = await rag.chunk(docling_doc)
|
|
embedded = await embed_chunks(chunks, rag.embedder, 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_skips_binary_content_when_qa_model_is_text_only():
|
|
"""The agent search tool must NOT attach picture bytes when the QA model
|
|
is text-only (``qa.model.vision = False``, the default). Sending image
|
|
parts to a text-only model would cause it to hallucinate confidently —
|
|
Ollama silently accepts the bytes and the model guesses."""
|
|
|
|
picture_result = SearchResult(
|
|
content="A diagram of the layout",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": PICTURE_B64},
|
|
)
|
|
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[picture_result])
|
|
fake_client.expand_context = AsyncMock(return_value=[picture_result])
|
|
|
|
config = AppConfig()
|
|
# vision defaults to False; assert anyway so the test reads explicitly.
|
|
assert config.qa.model.vision is False
|
|
toolset = create_search_toolset(config, expand_context=False)
|
|
func = toolset.tools["search"].function
|
|
|
|
ctx = RunContext(
|
|
deps=_Deps(client=fake_client), # type: ignore[arg-type]
|
|
model=TestModel(),
|
|
usage=RunUsage(),
|
|
run_id="run-1",
|
|
)
|
|
result = await func(ctx, "anything")
|
|
|
|
assert isinstance(result, str)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_tool_drops_invalid_image_bytes():
|
|
"""A picture whose bytes cannot be decoded by PIL must not produce a
|
|
BinaryContent part. Otherwise the model adapter emits a vision
|
|
placeholder for an image the server can't decode, leaving Qwen3-VL's
|
|
processor with an off-by-one count and a 400 from
|
|
``Qwen3VLProcessor``."""
|
|
from io import BytesIO
|
|
|
|
from PIL import Image as PILImageModule
|
|
|
|
buf = BytesIO()
|
|
PILImageModule.new("RGB", (4, 4), "red").save(buf, "PNG")
|
|
valid_png = buf.getvalue()
|
|
valid_b64 = base64.b64encode(valid_png).decode("ascii")
|
|
|
|
# Truthy bytes (passes ``if blob`` guards) but not a decodable PNG.
|
|
invalid_bytes = b"\x89PNG\r\n\x1a\ngarbage"
|
|
invalid_b64 = base64.b64encode(invalid_bytes).decode("ascii")
|
|
|
|
picture_result = SearchResult(
|
|
content="Two figures",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0", "#/pictures/1"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": valid_b64, "#/pictures/1": invalid_b64},
|
|
)
|
|
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[picture_result])
|
|
fake_client.expand_context = AsyncMock(return_value=[picture_result])
|
|
|
|
config = AppConfig()
|
|
config.qa.model.vision = True
|
|
toolset = create_search_toolset(config, expand_context=False)
|
|
func = toolset.tools["search"].function
|
|
|
|
ctx = RunContext(
|
|
deps=_Deps(client=fake_client), # type: ignore[arg-type]
|
|
model=TestModel(),
|
|
usage=RunUsage(),
|
|
run_id="run-1",
|
|
)
|
|
result = await func(ctx, "anything")
|
|
|
|
assert isinstance(result, ToolReturn)
|
|
assert result.content is not None
|
|
identifiers = {p.identifier for p in result.content if isinstance(p, BinaryContent)}
|
|
assert identifiers == {"#/pictures/0"}, (
|
|
"Only the decodable PNG should reach the model — the corrupt "
|
|
"ref must be dropped so we don't emit a placeholder for an "
|
|
"image the server can't decode."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_tool_drops_all_invalid_returns_plain_text():
|
|
"""If every picture in the result set fails decode, fall back to a
|
|
plain string return — there's nothing to attach, so wrapping in
|
|
``ToolReturn`` with an empty ``content`` list would surface an empty
|
|
user message downstream."""
|
|
bad_b64 = base64.b64encode(b"\x89PNGnope").decode("ascii")
|
|
picture_result = SearchResult(
|
|
content="One broken figure",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": bad_b64},
|
|
)
|
|
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[picture_result])
|
|
fake_client.expand_context = AsyncMock(return_value=[picture_result])
|
|
|
|
config = AppConfig()
|
|
config.qa.model.vision = True
|
|
toolset = create_search_toolset(config, expand_context=False)
|
|
func = toolset.tools["search"].function
|
|
|
|
ctx = RunContext(
|
|
deps=_Deps(client=fake_client), # type: ignore[arg-type]
|
|
model=TestModel(),
|
|
usage=RunUsage(),
|
|
run_id="run-1",
|
|
)
|
|
result = await func(ctx, "anything")
|
|
|
|
assert isinstance(result, str)
|
|
|
|
|
|
@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
|
|
ToolReturn wrapper) so non-vision flows are unaffected."""
|
|
text_result = SearchResult(
|
|
content="Some text",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/texts/0"],
|
|
labels=["paragraph"],
|
|
)
|
|
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[text_result])
|
|
fake_client.expand_context = AsyncMock(return_value=[text_result])
|
|
|
|
toolset = create_search_toolset(Config, expand_context=False)
|
|
func = toolset.tools["search"].function
|
|
|
|
ctx = RunContext(
|
|
deps=_Deps(client=fake_client), # type: ignore[arg-type]
|
|
model=TestModel(),
|
|
usage=RunUsage(),
|
|
run_id="run-1",
|
|
)
|
|
result = await func(ctx, "anything")
|
|
|
|
assert isinstance(result, str)
|
|
assert "rank 1" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rag_capability_attaches_images_for_vision_model(temp_db_path):
|
|
picture_result = SearchResult(
|
|
content="A figure",
|
|
score=1.0,
|
|
chunk_id="chunk-1",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0"],
|
|
labels=["picture"],
|
|
image_data={"#/pictures/0": PICTURE_B64},
|
|
)
|
|
fake_client = AsyncMock()
|
|
fake_client.search = AsyncMock(return_value=[picture_result])
|
|
fake_client.expand_context = AsyncMock(return_value=[picture_result])
|
|
config = AppConfig()
|
|
config.qa.model.vision = True
|
|
capability = create_capability(
|
|
db_path=temp_db_path,
|
|
config=config,
|
|
defer_loading=False,
|
|
)
|
|
capability.state = RAGState()
|
|
capability.rag = fake_client
|
|
|
|
result = await capability._search("anything", None)
|
|
|
|
assert isinstance(result, ToolReturn)
|
|
assert result.content is not None
|
|
assert any(isinstance(part, BinaryContent) for part in result.content)
|
|
|
|
|
|
def test_build_picture_chunks_records_provenance_pages():
|
|
"""A picture with provenance contributes its page numbers to the chunk."""
|
|
from docling_core.types.doc.base import BoundingBox
|
|
from docling_core.types.doc.document import ProvenanceItem
|
|
|
|
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()
|
|
picture = doc.pictures[0]
|
|
picture.prov = [
|
|
ProvenanceItem(
|
|
page_no=3,
|
|
bbox=BoundingBox(l=0, t=10, r=10, b=0),
|
|
charspan=(0, 0),
|
|
),
|
|
# A repeat of the same page must not be counted twice.
|
|
ProvenanceItem(
|
|
page_no=3,
|
|
bbox=BoundingBox(l=0, t=20, r=10, b=10),
|
|
charspan=(0, 0),
|
|
),
|
|
]
|
|
|
|
chunks = build_picture_chunks(doc, document_id="doc-1")
|
|
|
|
assert chunks[0].metadata["page_numbers"] == [3]
|