diff --git a/packages/shared-python/shared/tests/test_asset_inline.py b/packages/shared-python/shared/tests/test_asset_inline.py index 888fe2e0d..7da82b1d7 100644 --- a/packages/shared-python/shared/tests/test_asset_inline.py +++ b/packages/shared-python/shared/tests/test_asset_inline.py @@ -2,6 +2,11 @@ from __future__ import annotations +import base64 +import tempfile +from pathlib import Path +from unittest.mock import Mock + import pytest from shared.services.retrieval.hydration.asset_inline import ( @@ -10,6 +15,7 @@ from shared.services.retrieval.hydration.result_assembly import ( assemble_retrieval_results, ) +from shared.services.storage.result_storage import JobResultStorage from shared.services.retrieval.scoring.hierarchy import ProviderToolSpace from shared.services.retrieval.scoring.knowhere_provider import ( KnowhereProvider, @@ -167,6 +173,119 @@ async def hydrate_connected_rows(**_kwargs: object) -> list[dict[str, object]]: assert display_marker not in assembled[0]["content"] +@pytest.mark.asyncio +async def test_assemble_skips_connected_image_without_placeholder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + storage_factory = Mock() + monkeypatch.setattr( + "shared.services.retrieval.hydration.evidence_compose.get_result_storage", + storage_factory, + ) + rows = [ + { + "chunk_id": "text-1", + "chunk_type": "text", + "content": "The process is illustrated below.", + "chunk_metadata": { + "connect_to": [ + { + "target": "image-1", + "relation": "embeds", + "ref": "[images/flow.png]", + } + ] + }, + }, + { + "chunk_id": "image-1", + "chunk_type": "image", + "content": "Flowchart of the ingestion pipeline.", + "file_path": "images/flow.png", + "job_id": "job-synth", + }, + ] + assembled = await assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + ) + + assert [row["chunk_id"] for row in assembled] == ["text-1"] + assert assembled[0]["content"] == rows[0]["content"] + assert assembled[0]["composed"] == [ + {"type": "text", "text": "The process is illustrated below."} + ] + storage_factory.assert_not_called() + + +@pytest.mark.asyncio +async def test_assemble_inserts_image_at_placeholder( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_bytes = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGD4" + "DwABBAEAX+XDSwAAAABJRU5ErkJggg==" + ) + image_path = tmp_path / "a.png" + image_path.write_bytes(image_bytes) + storage = Mock(spec=JobResultStorage) + storage.normalize_artifact_ref.return_value = "images/a.png" + storage.download_raw_to_temp.return_value = str(image_path) + monkeypatch.setattr( + "shared.services.retrieval.hydration.evidence_compose.get_result_storage", + lambda: storage, + ) + rows = [ + { + "chunk_id": "text-1", + "chunk_type": "text", + "content": "see [images/a.png] end", + "chunk_metadata": { + "connect_to": [ + { + "target": "img-1", + "relation": "embeds", + "ref": "[images/a.png]", + } + ] + }, + }, + { + "chunk_id": "img-1", + "chunk_type": "image", + "content": "chart summary", + "file_path": "images/a.png", + "job_id": "job-synth", + }, + ] + assembled = await assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + ) + + assert [row["chunk_id"] for row in assembled] == ["text-1"] + assert assembled[0]["content"] == "see [images/a.png] end" + assert assembled[0]["composed"] == [ + {"type": "text", "text": "see \n"}, + { + "type": "image", + "media_type": "image/png", + "data": base64.b64encode(image_bytes).decode("ascii"), + }, + {"type": "text", "text": "\n"}, + {"type": "text", "text": " end"}, + ] + storage.download_raw_to_temp.assert_called_once_with( + job_id="job-synth", + relative_path="images/a.png", + suffix=".png", + temp_dir=tempfile.gettempdir(), + ) + + def test_node_unit_span_inlines_section_assets() -> None: provider = KnowhereProvider( doc_id="doc-1",