Merge pull request #482 from ggozad/fix/attach-picture-from-caption

Attach figure bytes when a result matches its caption

CHANGELOG.md

@@ -5,6 +5,10 @@
 
 - Custom rerankers override `RerankerBase._rerank` instead of `rerank`; the base `rerank` handles the empty-input short-circuit.
 
+### Fixed
+
+- A search result whose matched refs include a figure's caption now attaches that figure's picture bytes for vision-capable models, resolved through the caption's adjacent picture item.
+
 ## [0.63.1] - 2026-06-29
 
 ### Changed

haiku_rag_slim/haiku/rag/client/search.py

@@ -110,40 +110,54 @@ def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
 async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) -> None:
     """Attach base64 picture bytes to ``SearchResult.image_data`` in-place.
 
+    A result carries a picture when its refs include the picture directly, or
+    when they include the picture's caption — the common case where a prose
+    chunk carrying a figure's caption ranks while the picture is its own chunk.
     Groups results by document_id and batches one picture-bytes lookup per
     document so a result set spanning N documents costs N reads, not one per
-    picture. Only refs starting with ``PICTURE_REF_PREFIX`` are queried.
+    picture.
     """
+    repo = client.document_item_repository
     by_doc: dict[str, list[SearchResult]] = {}
     for r in results:
-        if not r.document_id:
-            continue
-        if not any(ref.startswith(PICTURE_REF_PREFIX) for ref in r.doc_item_refs):
-            continue
-        by_doc.setdefault(r.document_id, []).append(r)
+        if r.document_id and r.doc_item_refs:
+            by_doc.setdefault(r.document_id, []).append(r)
 
     for doc_id, doc_results in by_doc.items():
+        all_refs = {ref for r in doc_results for ref in r.doc_item_refs}
+        caption_to_picture = await repo.get_caption_picture_refs(doc_id, list(all_refs))
+
+        result_pictures: list[tuple[SearchResult, list[str]]] = []
         wanted: list[str] = []
         seen: set[str] = set()
         for r in doc_results:
+            pictures: list[str] = []
             for ref in r.doc_item_refs:
-                if ref.startswith(PICTURE_REF_PREFIX) and ref not in seen:
-                    wanted.append(ref)
-                    seen.add(ref)
+                picture = (
+                    ref
+                    if ref.startswith(PICTURE_REF_PREFIX)
+                    else caption_to_picture.get(ref)
+                )
+                if picture and picture not in pictures:
+                    pictures.append(picture)
+            if pictures:
+                result_pictures.append((r, pictures))
+                for picture in pictures:
+                    if picture not in seen:
+                        wanted.append(picture)
+                        seen.add(picture)
         if not wanted:
             continue
-        bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk(
-            doc_id, wanted
-        )
+        bytes_by_ref = await repo.get_pictures_for_chunk(doc_id, wanted)
         if not bytes_by_ref:
             continue
-        captions_by_ref = await client.document_item_repository.get_text_for_refs(
+        captions_by_ref = await repo.get_text_for_refs(
             doc_id, list(bytes_by_ref.keys())
         )
-        for r in doc_results:
+        for r, pictures in result_pictures:
             attached: dict[str, str] = {}
             captions: dict[str, str] = {}
-            for ref in r.doc_item_refs:
+            for ref in pictures:
                 blob = bytes_by_ref.get(ref)
                 if blob:
                     attached[ref] = base64.b64encode(blob).decode("ascii")

haiku_rag_slim/haiku/rag/store/repositories/document_item.py

@@ -271,3 +271,47 @@ class DocumentItemRepository:
             if text:
                 result[row["self_ref"]] = text
         return result
+
+    async def get_caption_picture_refs(
+        self, document_id: str, refs: list[str]
+    ) -> dict[str, str]:
+        """Map caption refs to the picture item immediately preceding them.
+
+        Docling emits a figure's caption at the position right after its
+        picture, so a caption's picture is the picture item at
+        ``position - 1``. Returns ``{caption_ref: picture_ref}`` for the
+        caption refs among ``refs`` that have a picture predecessor. Non-caption
+        refs, and captions whose predecessor is not a picture (table captions),
+        map to nothing.
+        """
+        if not refs:
+            return {}
+
+        safe_id = escape_sql_string(document_id)
+        refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
+        caption_rows = await (
+            self.store.document_items_table.query()
+            .select(["self_ref", "position"])
+            .where(
+                f"document_id = '{safe_id}' AND label = 'caption' "
+                f"AND self_ref IN ({refs_sql})"
+            )
+            .to_list()
+        )
+        if not caption_rows:
+            return {}
+
+        prev_to_caption = {row["position"] - 1: row["self_ref"] for row in caption_rows}
+        positions_sql = ", ".join(str(p) for p in prev_to_caption)
+        picture_rows = await (
+            self.store.document_items_table.query()
+            .select(["self_ref", "position"])
+            .where(
+                f"document_id = '{safe_id}' AND label = 'picture' "
+                f"AND position IN ({positions_sql})"
+            )
+            .to_list()
+        )
+        return {
+            prev_to_caption[row["position"]]: row["self_ref"] for row in picture_rows
+        }

tests/store/test_document_items.py

@@ -657,6 +657,62 @@ class TestPictureDataStorage:
             }
             assert await repo.get_text_for_refs("doc-1", []) == {}
 
+    async def test_get_caption_picture_refs(self, temp_db_path):
+        """A caption ref resolves to the picture at the immediately preceding
+        position; a table caption (no preceding picture) resolves to nothing."""
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            repo = DocumentItemRepository(rag.store)
+            await repo.create_items(
+                "doc-1",
+                [
+                    DocumentItem(
+                        document_id="doc-1",
+                        position=0,
+                        self_ref="#/pictures/0",
+                        label="picture",
+                        picture_data=b"\x89PNG\r\n\x1a\nfake",
+                    ),
+                    DocumentItem(
+                        document_id="doc-1",
+                        position=1,
+                        self_ref="#/texts/0",
+                        label="caption",
+                        text="Figure 1. A figure caption.",
+                    ),
+                    DocumentItem(
+                        document_id="doc-1",
+                        position=2,
+                        self_ref="#/texts/1",
+                        label="paragraph",
+                        text="Body prose.",
+                    ),
+                    DocumentItem(
+                        document_id="doc-1",
+                        position=3,
+                        self_ref="#/tables/0",
+                        label="table",
+                        text="| a | b |",
+                    ),
+                    DocumentItem(
+                        document_id="doc-1",
+                        position=4,
+                        self_ref="#/texts/2",
+                        label="caption",
+                        text="Table 1. A table caption.",
+                    ),
+                ],
+            )
+
+            # Figure caption resolves to its picture; table caption does not.
+            got = await repo.get_caption_picture_refs(
+                "doc-1", ["#/texts/0", "#/texts/1", "#/texts/2"]
+            )
+            assert got == {"#/texts/0": "#/pictures/0"}
+
+            # A non-caption ref alone yields nothing.
+            assert await repo.get_caption_picture_refs("doc-1", ["#/texts/1"]) == {}
+            assert await repo.get_caption_picture_refs("doc-1", []) == {}
+
     async def test_hot_paths_exclude_picture_data(self, temp_db_path):
         """Light read paths must NOT pull picture_data into memory."""
         async with HaikuRAG(temp_db_path, create=True) as rag:

tests/test_picture_in_context.py

@@ -78,6 +78,48 @@ async def test_populate_image_data_attaches_base64(temp_db_path):
         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."""