From 46b81f3fa1be1d9b7b2fab10fe7db6952bb4342f Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 5 May 2026 14:20:37 +0300
Subject: [PATCH] cover picture-description provider and rebuild edge cases

---
 tests/test_picture_description.py | 143 +++++++++++++++++++++++++++++
 tests/test_rebuild.py             | 146 ++++++++++++++++++++++++++++++
 2 files changed, 289 insertions(+)
 create mode 100644 tests/test_picture_description.py

diff --git a/tests/test_picture_description.py b/tests/test_picture_description.py
new file mode 100644
index 00000000..5fed0fc8
--- /dev/null
+++ b/tests/test_picture_description.py
@@ -0,0 +1,143 @@
+"""Tests for the direct VLM client used by ``rebuild --descriptions``."""
+
+import logging
+from typing import Any
+
+import pytest
+
+from haiku.rag.config import AppConfig
+from haiku.rag.providers.picture_description import describe_pictures
+
+
+class _StubAgent:
+    """Minimal stand-in for ``pydantic_ai.Agent`` that returns a queue of
+    pre-baked responses (or raises) for each ``run`` call."""
+
+    def __init__(self, outputs: list[Any]):
+        self._outputs = list(outputs)
+        self.calls: list[Any] = []
+
+    async def run(self, prompt: list) -> Any:
+        self.calls.append(prompt)
+        out = self._outputs.pop(0)
+        if isinstance(out, BaseException):
+            raise out
+
+        class _Result:
+            def __init__(self, output: str):
+                self.output = output
+
+        return _Result(out)
+
+
+def _patch_agent(monkeypatch, outputs: list[Any]) -> _StubAgent:
+    """Replace pydantic_ai.Agent in our module with a constructor that
+    returns a single shared StubAgent."""
+    stub = _StubAgent(outputs)
+    monkeypatch.setattr(
+        "haiku.rag.providers.picture_description.Agent",
+        lambda **kwargs: stub,
+    )
+    # Skip real model construction — we don't use the returned model anyway.
+    monkeypatch.setattr(
+        "haiku.rag.providers.picture_description.get_model",
+        lambda model_config, app_config: object(),
+    )
+    return stub
+
+
+@pytest.mark.asyncio
+async def test_describe_pictures_returns_text_per_self_ref(monkeypatch):
+    """Happy path: each picture gets one VLM call and the response text
+    lands in the result map keyed by self_ref."""
+    stub = _patch_agent(
+        monkeypatch,
+        outputs=["A red square.", "A blue triangle."],
+    )
+
+    config = AppConfig()
+    out = await describe_pictures(
+        {"#/pictures/0": b"red-bytes", "#/pictures/1": b"blue-bytes"},
+        config=config,
+    )
+
+    assert out == {
+        "#/pictures/0": "A red square.",
+        "#/pictures/1": "A blue triangle.",
+    }
+    assert len(stub.calls) == 2
+
+
+@pytest.mark.asyncio
+async def test_describe_pictures_drops_empty_output(monkeypatch):
+    """Pictures whose VLM response is empty/whitespace are dropped from
+    the result map. Caller can decide whether the partial result is
+    acceptable."""
+    _patch_agent(monkeypatch, outputs=["A real description.", "   ", ""])
+
+    out = await describe_pictures(
+        {
+            "#/pictures/0": b"a",
+            "#/pictures/1": b"b",
+            "#/pictures/2": b"c",
+        },
+        config=AppConfig(),
+    )
+
+    assert out == {"#/pictures/0": "A real description."}
+
+
+@pytest.mark.asyncio
+async def test_describe_pictures_swallows_exceptions(monkeypatch, caplog):
+    """A failing VLM call is logged as a warning and the picture is
+    skipped — the rest of the batch still gets described."""
+    _patch_agent(
+        monkeypatch,
+        outputs=[
+            RuntimeError("boom"),
+            "After the failure.",
+        ],
+    )
+
+    with caplog.at_level(
+        logging.WARNING, logger="haiku.rag.providers.picture_description"
+    ):
+        out = await describe_pictures(
+            {"#/pictures/0": b"a", "#/pictures/1": b"b"},
+            config=AppConfig(),
+        )
+
+    assert out == {"#/pictures/1": "After the failure."}
+    # caplog may not catch records due to project-wide propagate=False on the
+    # haiku.rag logger; fall back to checking the result reflects the skip.
+
+
+@pytest.mark.asyncio
+async def test_describe_pictures_empty_input(monkeypatch):
+    """No pictures means no VLM calls and an empty result."""
+    stub = _patch_agent(monkeypatch, outputs=[])
+
+    out = await describe_pictures({}, config=AppConfig())
+
+    assert out == {}
+    assert stub.calls == []
+
+
+@pytest.mark.asyncio
+async def test_describe_pictures_passes_binary_content(monkeypatch):
+    """The VLM call receives the picture bytes as a BinaryContent part with
+    media_type=image/png so model providers route the request correctly."""
+    from pydantic_ai.messages import BinaryContent
+
+    stub = _patch_agent(monkeypatch, outputs=["ok"])
+
+    await describe_pictures(
+        {"#/pictures/0": b"\x89PNG\r\n\x1a\nfake"}, config=AppConfig()
+    )
+
+    assert len(stub.calls) == 1
+    parts = stub.calls[0]
+    assert isinstance(parts, list) and len(parts) == 1
+    assert isinstance(parts[0], BinaryContent)
+    assert parts[0].data == b"\x89PNG\r\n\x1a\nfake"
+    assert parts[0].media_type == "image/png"
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
index cec7d57f..51c0b90a 100644
--- a/tests/test_rebuild.py
+++ b/tests/test_rebuild.py
@@ -557,3 +557,149 @@ async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkey
             getattr(getattr(meta, "description", None), "text", None) if meta else None
         )
         assert text == "Pre-existing description."
+
+
+@pytest.mark.asyncio
+async def test_patch_picture_descriptions_returns_zero_for_doc_without_pictures(
+    temp_db_path,
+):
+    """A document with no pictures returns 0 without ever calling the VLM."""
+    from haiku.rag.client.rebuild import _patch_picture_descriptions
+    from haiku.rag.config import AppConfig
+
+    config = AppConfig()
+    config.processing.conversion_options.picture_description.enabled = True
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
+        doc = await rag.create_document(content="Just text, no pictures.")
+        assert doc.id is not None
+        n = await _patch_picture_descriptions(rag, doc)
+        assert n == 0
+
+
+@pytest.mark.asyncio
+async def test_patch_picture_descriptions_warns_on_missing_bytes(temp_db_path, caplog):
+    """When the docling blob has pictures but document_items.picture_data is
+    empty (e.g. legacy DB ingested before A2b), the helper logs a warning
+    and returns 0 instead of trying to drive the VLM with no input."""
+    import logging
+
+    from haiku.rag.client.documents import _store_document_with_chunks
+    from haiku.rag.client.rebuild import _patch_picture_descriptions
+    from haiku.rag.config import AppConfig
+    from haiku.rag.store.models.document import Document
+    from tests.store.test_document_items import _docling_doc_with_picture
+
+    docling_doc = _docling_doc_with_picture()
+    config = AppConfig()
+    config.processing.conversion_options.picture_description.enabled = True
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
+        document = Document(content="x", uri="test://doc")
+        document.set_docling(docling_doc)
+        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        assert created.id is not None
+
+        # Wipe the stored picture bytes to simulate a doc that knows about
+        # pictures but doesn't have them on disk.
+        await rag.store.document_items_table.update(
+            {"picture_data": None},
+            where=f"document_id = '{created.id}' AND label = 'picture'",
+        )
+
+        # Capture warnings directly off the rebuild module logger — the
+        # haiku.rag parent logger is configured non-propagating elsewhere
+        # in the suite so caplog can miss records.
+        from haiku.rag.client import rebuild as rebuild_module
+
+        records: list[logging.LogRecord] = []
+
+        class _ListHandler(logging.Handler):
+            def emit(self, record: logging.LogRecord) -> None:
+                records.append(record)
+
+        handler = _ListHandler(level=logging.WARNING)
+        rebuild_module.logger.addHandler(handler)
+        try:
+            n = await _patch_picture_descriptions(rag, created)
+        finally:
+            rebuild_module.logger.removeHandler(handler)
+
+        assert n == 0
+        assert any("no stored picture bytes" in r.getMessage() for r in records)
+
+
+@pytest.mark.asyncio
+async def test_patch_picture_descriptions_skips_when_all_already_described(
+    temp_db_path, monkeypatch
+):
+    """If every picture already has meta.description.text, the helper does
+    not call the VLM and returns 0."""
+    from docling_core.types.doc.document import DescriptionMetaField, PictureMeta
+
+    from haiku.rag.client.documents import _store_document_with_chunks
+    from haiku.rag.client.rebuild import _patch_picture_descriptions
+    from haiku.rag.config import AppConfig
+    from haiku.rag.store.models.document import Document
+    from tests.store.test_document_items import _docling_doc_with_picture
+
+    docling_doc = _docling_doc_with_picture()
+    docling_doc.pictures[0].meta = PictureMeta(
+        description=DescriptionMetaField(text="Pre-described.")
+    )
+
+    config = AppConfig()
+    config.processing.conversion_options.picture_description.enabled = True
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
+        document = Document(content="x", uri="test://doc")
+        document.set_docling(docling_doc)
+        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        assert created.id is not None
+
+        called = False
+
+        async def fake_describe(*args, **kwargs):
+            nonlocal called
+            called = True
+            return {}
+
+        monkeypatch.setattr(
+            "haiku.rag.providers.picture_description.describe_pictures",
+            fake_describe,
+        )
+
+        n = await _patch_picture_descriptions(rag, created)
+        assert n == 0
+        assert called is False
+
+
+@pytest.mark.asyncio
+async def test_rebuild_descriptions_raises_when_blob_is_missing(
+    temp_db_path, monkeypatch
+):
+    """Documents without a stored docling blob can't be re-described —
+    surface a clear error pointing the user at full rebuild instead."""
+    from haiku.rag.client.documents import _store_document_with_chunks
+    from haiku.rag.config import AppConfig
+    from haiku.rag.store.models.document import Document
+    from tests.store.test_document_items import _docling_doc_with_picture
+
+    docling_doc = _docling_doc_with_picture()
+    config = AppConfig()
+    config.processing.conversion_options.picture_description.enabled = True
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
+        document = Document(content="x", uri="test://doc")
+        document.set_docling(docling_doc)
+        created = await _store_document_with_chunks(rag, document, [], docling_doc)
+        assert created.id is not None
+
+        # Force the stored doc to come back without a docling blob.
+        await rag.store.documents_table.update(
+            {"docling_document": None}, where=f"id = '{created.id}'"
+        )
+
+        with pytest.raises(ValueError, match="rebuild --descriptions requires"):
+            async for _ in rag.rebuild_database(mode=RebuildMode.DESCRIPTIONS):
+                pass