460215158d
`_attach_picture_data` fetched picture bytes once per document, over the `limit * 10` candidates reranking asks for, so it was the per-document fetch with the most candidates behind it. It now issues one query however many documents the candidates span: one for ten documents, as for one. Removes `get_text_for_refs`, whose only caller now gets the text back with the bytes from `get_pictures_grouped`. `test_client_search_include_images_false_skips_lookup` returned no search results, so asserting the picture accessor went uncalled held whatever the code did. It now returns a picture-carrying result, making "did not fetch" the assertion rather than "had nothing to fetch".
1055 lines
37 KiB
Python
1055 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.store.models.chunk import Chunk, SearchResult
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
from haiku.rag.tools.search import create_search_toolset
|
|
from tests.test_context import _fetch_and_expand
|
|
|
|
|
|
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_grouped = AsyncMock( # type: ignore[method-assign]
|
|
wraps=rag.document_item_repository.get_pictures_grouped
|
|
)
|
|
|
|
from haiku.rag.client.search import search
|
|
|
|
# A real picture-carrying result, so not fetching is the assertion
|
|
# rather than there being nothing to fetch.
|
|
async def fake_chunk_search(*args, **kwargs):
|
|
return [
|
|
(
|
|
Chunk(
|
|
id="chunk-1",
|
|
document_id="doc-1",
|
|
content="body",
|
|
metadata={"doc_item_refs": ["#/pictures/0"]},
|
|
),
|
|
0.9,
|
|
)
|
|
]
|
|
|
|
rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign]
|
|
|
|
results = await search(rag, "anything", include_images=False)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].image_data is None
|
|
rag.document_item_repository.get_pictures_grouped.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 _fetch_and_expand(
|
|
rag.document_item_repository, "doc-1", [seed], 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]
|