test: lock connected image composition during retrieval assembly

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",