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

Name the search result each page image belongs to

Browse source 
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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-08 17:51:06 +03:00
parent d9bd3a701f
commit 46f7ab8d97
No known key found for this signature in database
6 changed files with 109 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -7,6 +7,7 @@
- `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`.
- `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`.
- `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized.
- Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures.
- Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered.
- Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images.
- `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice.

4
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -27,7 +27,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation, resolve_citations
from haiku.rag.tools.search import build_binary_parts_from_results
from haiku.rag.tools.search import build_image_content_from_results
CITATION_GRACE_REQUESTS = 2
"""Requests calling this capability's tools that its cite tool outlives the rest by.
@ -372,7 +372,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
)
state = cast(Any, self.state)
state.searches[query] = results
if self.vision and (parts := build_binary_parts_from_results(results)):
if self.vision and (parts := build_image_content_from_results(results)):
return ToolReturn(return_value=formatted, content=parts)
return formatted

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

@ -211,7 +211,7 @@ class SearchResult(BaseModel):
parts.append(f"Type: {primary_label}")
# Surface picture captions when present. Order matches the binary
# attachments emitted by build_binary_parts_from_results, so the model
# attachments emitted by build_image_content_from_results, so the model
# can correlate caption ↔ attached image by position (BinaryContent
# identifiers don't survive serialization to the OpenAI vision API).
if self.picture_captions:

46
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -11,10 +11,10 @@ from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import RAGDeps
def build_binary_parts_from_results(
def build_image_content_from_results(
results: list[SearchResult],
) -> list[BinaryContent]:
"""Decode and validate picture bytes attached to search results.
) -> list[str | BinaryContent]:
"""Decode and validate picture bytes attached to search results, labelled.
Dedup keyed on ``(document_id, self_ref)`` so the same picture in
different chunks is sent once. Pictures that fail
@ -22,8 +22,18 @@ def build_binary_parts_from_results(
vision placeholder per ``BinaryContent``, so emitting one for an
image the server can't decode leaves the processor with an
off-by-one count.
Every picture is preceded by a line naming the result it belongs to.
``ToolReturn.content`` reaches the model as a user-role message, so
retrieved pictures are otherwise indistinguishable from ones the user
attached, and models narrate them as part of the question: unlabelled,
gemma4-26b answered about a figure from an unrelated document, and with a
single note ahead of the batch it still called them "images in the prompt".
The label also names the chunk to cite for a figure, which
``BinaryContent.identifier`` cannot do — it does not survive serialization
to the vision API.
"""
parts: list[BinaryContent] = []
collected: list[tuple[str | None, str, bytes]] = []
seen: set[tuple[str | None, str]] = set()
for result in results:
if not result.image_data:
@ -38,15 +48,21 @@ def build_binary_parts_from_results(
img.verify()
except Exception:
continue
parts.append(
BinaryContent(
data=data,
media_type="image/png",
identifier=self_ref,
)
)
collected.append((result.chunk_id, self_ref, data))
seen.add(key)
return parts
content: list[str | BinaryContent] = []
total = len(collected)
for position, (chunk_id, self_ref, data) in enumerate(collected, 1):
content.append(
f"Page image {position} of {total}, retrieved from the knowledge "
f"base for search result [{chunk_id}] ({self_ref}). "
"Not provided by the user."
)
content.append(
BinaryContent(data=data, media_type="image/png", identifier=self_ref)
)
return content
def create_search_toolset(
@ -134,9 +150,9 @@ def create_search_toolset(
if not config.qa.model.vision:
return text
binary_parts = build_binary_parts_from_results(results_list)
if binary_parts:
return ToolReturn(return_value=text, content=binary_parts)
image_content = build_image_content_from_results(results_list)
if image_content:
return ToolReturn(return_value=text, content=image_content)
return text
toolset: FunctionToolset[RAGDeps] = FunctionToolset()

12
tests/test_picture_in_context.py
View file

@ -402,8 +402,9 @@ async def test_search_tool_returns_multimodal_when_picture_present():
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
assert len(result.content) == 1
part = result.content[0]
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"
@ -457,11 +458,12 @@ async def test_search_tool_attaches_same_self_ref_from_different_documents():
assert isinstance(result, ToolReturn)
assert result.content is not None
assert len(result.content) == 2, (
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 result.content} # type: ignore[attr-defined]
payloads = {part.data for part in images}
assert PICTURE_BYTES in payloads
assert other_bytes in payloads
@ -905,7 +907,7 @@ async def test_search_tool_drops_invalid_image_bytes():
assert isinstance(result, ToolReturn)
assert result.content is not None
identifiers = {p.identifier for p in result.content} # type: ignore[attr-defined]
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 "

82
tests/tools/test_search.py
View file

@ -210,30 +210,35 @@ def search_config():
return Config
class TestBuildBinaryPartsFromResults:
"""Picture bytes are attached once per (document, self_ref) pair."""
def _png_b64():
import base64
from io import BytesIO
from PIL import Image as PILImage
buf = BytesIO()
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode()
class TestBuildImageContentFromResults:
"""Picture bytes are attached once per (document, self_ref) pair, and labelled."""
def test_results_without_image_data_contribute_nothing(self):
from haiku.rag.tools.search import build_binary_parts_from_results
from haiku.rag.tools.search import build_image_content_from_results
results = [
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
]
assert build_binary_parts_from_results(results) == []
assert build_image_content_from_results(results) == []
def test_duplicate_document_and_ref_is_attached_once(self):
import base64
from io import BytesIO
from pydantic_ai.messages import BinaryContent
from PIL import Image as PILImage
from haiku.rag.tools.search import build_image_content_from_results
from haiku.rag.tools.search import build_binary_parts_from_results
buf = BytesIO()
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
png = base64.b64encode(buf.getvalue()).decode()
shared = {"#/pictures/0": png}
shared = {"#/pictures/0": _png_b64()}
results = [
SearchResult(
content="a",
@ -251,6 +256,53 @@ class TestBuildBinaryPartsFromResults:
),
]
parts = build_binary_parts_from_results(results)
content = build_image_content_from_results(results)
assert len(parts) == 1
images = [item for item in content if isinstance(item, BinaryContent)]
assert len(images) == 1
def test_each_image_is_labelled_with_the_result_it_belongs_to(self):
"""Label every picture, not just the batch.
``ToolReturn.content`` reaches the model as a user-role message, and one
leading note does not override that: with a single note on the wire,
gemma4-26b still reasoned "the user also provided images in the prompt".
A label adjacent to each picture also names the chunk to cite for it,
which ``BinaryContent.identifier`` cannot do — it does not survive
serialization to the vision API.
"""
from pydantic_ai.messages import BinaryContent
from haiku.rag.tools.search import build_image_content_from_results
results = [
SearchResult(
content="a",
score=0.9,
chunk_id="c1",
document_id="doc-1",
image_data={"#/pictures/0": _png_b64()},
),
SearchResult(
content="b",
score=0.8,
chunk_id="c2",
document_id="doc-2",
image_data={"#/pictures/3": _png_b64()},
),
]
content = build_image_content_from_results(results)
# label, image, label, image — each picture preceded by its own line.
assert [type(item) is str for item in content] == [True, False, True, False]
assert isinstance(content[1], BinaryContent)
assert isinstance(content[3], BinaryContent)
first, second = content[0], content[2]
assert isinstance(first, str) and isinstance(second, str)
assert "c1" in first and "#/pictures/0" in first
assert "c2" in second and "#/pictures/3" in second
assert "1 of 2" in first and "2 of 2" in second
for label in (first, second):
assert "not provided by the user" in label.lower()