Rename get_captions_for_chunk() to get_text_for_refs()

This commit is contained in:
Yiorgis Gozadinos 2026-05-18 15:32:42 +03:00
parent bdc77bd467
commit 3b6bef3bfb
No known key found for this signature in database
3 changed files with 14 additions and 12 deletions

View file

@ -139,7 +139,7 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult])
) )
if not bytes_by_ref: if not bytes_by_ref:
continue continue
captions_by_ref = await client.document_item_repository.get_captions_for_chunk( captions_by_ref = await client.document_item_repository.get_text_for_refs(
doc_id, list(bytes_by_ref.keys()) doc_id, list(bytes_by_ref.keys())
) )
for r in doc_results: for r in doc_results:

View file

@ -207,16 +207,17 @@ class DocumentItemRepository:
result[row["self_ref"]] = data result[row["self_ref"]] = data
return result return result
async def get_captions_for_chunk( async def get_text_for_refs(
self, document_id: str, refs: list[str] self, document_id: str, refs: list[str]
) -> dict[str, str]: ) -> dict[str, str]:
"""Fetch caption text for multiple self_refs within a single document. """Fetch the ``text`` field for multiple self_refs within a single document.
Returns ``{self_ref: text}`` for refs that have non-empty text. Used Returns ``{self_ref: text}`` for refs whose text is non-empty. Used
alongside ``get_pictures_for_chunk`` to label figures in agent-facing alongside ``get_pictures_for_chunk`` to label figures in agent-facing
search results the OpenAI vision message format has no identifier search results: picture items carry their VLM-generated caption in
field for binary parts, so the caption is the only signal a model can the ``text`` field, and the OpenAI vision message format has no
use to correlate a description with the picture it sees. identifier on binary parts, so the caption text is the only signal a
model can use to correlate a description with the picture it sees.
""" """
if not refs: if not refs:
return {} return {}

View file

@ -502,13 +502,14 @@ class TestPictureDataStorage:
# Empty refs returns empty dict # Empty refs returns empty dict
assert await repo.get_pictures_for_chunk("doc-1", []) == {} assert await repo.get_pictures_for_chunk("doc-1", []) == {}
async def test_get_captions_for_chunk(self, temp_db_path): async def test_get_text_for_refs(self, temp_db_path):
"""Captions are returned for refs whose text is non-empty. """Text is returned for any ref with non-empty ``text``, regardless of label.
In practice pictures carry their caption in the ``text`` field In practice pictures carry their caption in the ``text`` field
(populated by the VLM picture-description pass during ingest); this (populated by the VLM picture-description pass during ingest); this
method surfaces that text alongside the picture bytes so the model can method surfaces that text alongside the picture bytes so the model can
correlate a description with the binary it sees. correlate a description with the binary it sees. The same method also
returns text for non-picture refs callers filter by label.
""" """
async with HaikuRAG(temp_db_path, create=True) as rag: async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store) repo = DocumentItemRepository(rag.store)
@ -541,7 +542,7 @@ class TestPictureDataStorage:
], ],
) )
captions = await repo.get_captions_for_chunk( captions = await repo.get_text_for_refs(
"doc-1", "doc-1",
["#/pictures/0", "#/pictures/1", "#/texts/0", "#/pictures/999"], ["#/pictures/0", "#/pictures/1", "#/texts/0", "#/pictures/999"],
) )
@ -549,7 +550,7 @@ class TestPictureDataStorage:
"#/pictures/0": "Figure 1. CCS generation over time.", "#/pictures/0": "Figure 1. CCS generation over time.",
"#/texts/0": "Inline prose.", "#/texts/0": "Inline prose.",
} }
assert await repo.get_captions_for_chunk("doc-1", []) == {} assert await repo.get_text_for_refs("doc-1", []) == {}
async def test_hot_paths_exclude_picture_data(self, temp_db_path): async def test_hot_paths_exclude_picture_data(self, temp_db_path):
"""Light read paths must NOT pull picture_data into memory.""" """Light read paths must NOT pull picture_data into memory."""