diff --git a/CHANGELOG.md b/CHANGELOG.md index cba408f5..46a76ca6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index a90ddcbc..06e6f2f6 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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") diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 897745cc..cc26226d 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/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 + } diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index 247e4223..cee84196 100644 --- a/tests/store/test_document_items.py +++ b/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: diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 173e1c7a..6c7ab2ea 100644 --- a/tests/test_picture_in_context.py +++ b/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."""