attach picture bytes through skill search and document the schema

haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2

@@ -11,7 +11,7 @@ description: {{ description }}
 {% if "search" in tool_names %}
 
 ### search
-Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content. Use for answering questions, finding passages, exploring topics.
+Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content. Use for answering questions, finding passages, exploring topics. Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots.
 {% endif %}
 {% if "list_documents" in tool_names %}
 

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -1,8 +1,10 @@
+import base64
 from pathlib import Path
 from typing import Any
 
 from pydantic import BaseModel
 from pydantic_ai import RunContext
+from pydantic_ai.messages import BinaryContent, ToolReturn
 
 from haiku.rag.agents.research.models import Citation
 from haiku.rag.client import HaikuRAG
@@ -167,10 +169,14 @@ def create_skill_tools(
 
         async def search(
             ctx: RunContext[RAGRunDeps], query: str, limit: int | None = None
-        ) -> str:
+        ) -> str | ToolReturn:
             """Search the knowledge base using hybrid search (vector + full-text).
 
-            Returns ranked results with content and metadata.
+            Returns ranked results with content and metadata. When picture
+            content is in the result set and the configured QA model is
+            vision-capable (``qa.model.vision = true``), picture bytes are
+            attached as ``BinaryContent`` parts so the model sees figures
+            alongside text.
 
             Args:
                 query: The search query.
@@ -192,6 +198,29 @@ def create_skill_tools(
             )
             if state:
                 state.searches[query] = results
+
+            if not config.qa.model.vision:
+                return formatted
+
+            binary_parts: list[BinaryContent] = []
+            seen: set[str] = set()
+            for result in results:
+                if not result.image_data:
+                    continue
+                for self_ref, b64 in result.image_data.items():
+                    if self_ref in seen:
+                        continue
+                    binary_parts.append(
+                        BinaryContent(
+                            data=base64.b64decode(b64),
+                            media_type="image/png",
+                            identifier=self_ref,
+                        )
+                    )
+                    seen.add(self_ref)
+
+            if binary_parts:
+                return ToolReturn(return_value=formatted, content=binary_parts)
             return formatted
 
         tools["search"] = search

haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md

@@ -26,7 +26,7 @@ Available modules: `json`, `re`, `math`, `pathlib`
 Not supported: class definitions, generators/yield, match statements, decorators, `with` statements
 
 ### search
-Search the knowledge base directly (outside code execution). Use for initial exploration before writing code.
+Search the knowledge base directly (outside code execution). Use for initial exploration before writing code. Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots.
 
 ### list_documents
 List available documents. Use to discover what's in the knowledge base.

haiku_rag_slim/haiku/rag/skills/rag/SKILL.md

@@ -16,9 +16,11 @@ Search the knowledge base using hybrid search (vector + full-text). Returns rank
 Each result includes:
 - `chunk_id` in brackets and rank position (rank 1 = most relevant)
 - Source: document title and section hierarchy
-- Type: content type (paragraph, table, code, list_item)
+- Type: content type (paragraph, table, code, list_item, picture)
 - Content: the actual text
 
+When a result's Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text. Use the image directly to answer questions about figures, diagrams, charts, screenshots.
+
 ### list_documents
 List available documents in the knowledge base. Use when the user wants to browse what's available.
 

tests/test_skill_tools.py

@@ -0,0 +1,187 @@
+"""Tests for skill tool closures from ``haiku.rag.skills._tools.create_skill_tools``.
+
+These cover the vision toggle on the skill ``search`` tool: when the configured
+QA model is vision-capable, picture bytes from search results must reach the
+sub-agent as ``BinaryContent`` parts (so a vision model can read figures).
+When the QA model is not vision-capable, the same search must return plain
+text only.
+"""
+
+import base64
+from pathlib import Path
+from unittest.mock import AsyncMock
+
+import pytest
+from pydantic_ai import RunContext
+from pydantic_ai.messages import BinaryContent, ToolReturn
+from pydantic_ai.models.test import TestModel
+from pydantic_ai.usage import RunUsage
+
+from haiku.rag.config import AppConfig
+from haiku.rag.skills._deps import RAGRunDeps
+from haiku.rag.skills._tools import create_skill_tools
+from haiku.rag.skills.rag import RAGState
+from haiku.rag.store.models.chunk import SearchResult
+
+PICTURE_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes"
+PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii")
+
+
+def _picture_result() -> SearchResult:
+    return SearchResult(
+        content="A diagram of the layout",
+        score=1.0,
+        chunk_id="chunk-1",
+        document_id="doc-1",
+        doc_item_refs=["#/pictures/0"],
+        labels=["picture"],
+        image_data={"#/pictures/0": PICTURE_B64},
+    )
+
+
+def _text_result() -> SearchResult:
+    return SearchResult(
+        content="Some surrounding paragraph text",
+        score=0.9,
+        chunk_id="chunk-2",
+        document_id="doc-1",
+        doc_item_refs=["#/texts/3"],
+        labels=["paragraph"],
+        image_data=None,
+    )
+
+
+def _make_ctx(rag, state: RAGState) -> RunContext[RAGRunDeps]:
+    deps = RAGRunDeps(state=state, rag=rag)
+    return RunContext(
+        deps=deps,
+        model=TestModel(),
+        usage=RunUsage(),
+        run_id="run-1",
+    )
+
+
+def _build_search_tool(config: AppConfig):
+    tools = create_skill_tools(
+        db_path=Path("/tmp/unused.lancedb"),
+        config=config,
+        state_type=RAGState,
+        tool_names=["search"],
+    )
+    return tools["search"]
+
+
+def _fake_rag(results: list[SearchResult]) -> AsyncMock:
+    rag = AsyncMock()
+    rag.search = AsyncMock(return_value=results)
+    rag.expand_context = AsyncMock(return_value=results)
+    return rag
+
+
+@pytest.mark.asyncio
+async def test_skill_search_attaches_binary_content_when_vision_capable():
+    """vision=True + picture in results → ToolReturn carries BinaryContent."""
+    config = AppConfig()
+    config.qa.model.vision = True
+
+    search = _build_search_tool(config)
+    rag = _fake_rag([_picture_result()])
+    ctx = _make_ctx(rag, RAGState())
+
+    result = await search(ctx, "diagram")
+
+    assert isinstance(result, ToolReturn)
+    assert isinstance(result.return_value, str)
+    assert "rank 1" in result.return_value
+    assert result.content is not None
+    assert len(result.content) == 1
+    part = result.content[0]
+    assert isinstance(part, BinaryContent)
+    assert part.data == PICTURE_BYTES
+    assert part.media_type == "image/png"
+    assert part.identifier == "#/pictures/0"
+
+
+@pytest.mark.asyncio
+async def test_skill_search_returns_plain_string_when_not_vision_capable():
+    """vision=False (the default) + picture in results → plain text only.
+    The picture bytes must not reach a text-only model — providers behave
+    inconsistently with image content (Ollama silently accepts and the
+    model hallucinates)."""
+    config = AppConfig()
+    assert config.qa.model.vision is False
+
+    search = _build_search_tool(config)
+    rag = _fake_rag([_picture_result()])
+    ctx = _make_ctx(rag, RAGState())
+
+    result = await search(ctx, "diagram")
+
+    assert isinstance(result, str)
+    assert "rank 1" in result
+
+
+@pytest.mark.asyncio
+async def test_skill_search_returns_plain_string_when_no_pictures():
+    """vision=True + no pictures in results → no ToolReturn wrapper, just
+    text. The wrapper is only needed when there's actually image content
+    to carry."""
+    config = AppConfig()
+    config.qa.model.vision = True
+
+    search = _build_search_tool(config)
+    rag = _fake_rag([_text_result()])
+    ctx = _make_ctx(rag, RAGState())
+
+    result = await search(ctx, "paragraph")
+
+    assert isinstance(result, str)
+    assert "rank 1" in result
+
+
+@pytest.mark.asyncio
+async def test_skill_search_records_results_into_state():
+    """Whether or not the QA model is vision-capable, the SearchResult
+    list must land in state.searches[query] so cite/visualize_chunk can
+    look chunks up later."""
+    config = AppConfig()
+
+    search = _build_search_tool(config)
+    rag = _fake_rag([_picture_result(), _text_result()])
+    state = RAGState()
+    ctx = _make_ctx(rag, state)
+
+    await search(ctx, "anything")
+
+    assert "anything" in state.searches
+    assert len(state.searches["anything"]) == 2
+
+
+@pytest.mark.asyncio
+async def test_skill_search_dedups_picture_bytes_by_self_ref():
+    """When two search results reference the same picture self_ref (e.g. a
+    text chunk and a synthetic picture chunk), the BinaryContent list
+    must include that picture exactly once. Otherwise the model receives
+    duplicate image content and pays double the image-token cost."""
+    config = AppConfig()
+    config.qa.model.vision = True
+
+    other = SearchResult(
+        content="Surrounding text mentioning the figure",
+        score=0.8,
+        chunk_id="chunk-3",
+        document_id="doc-1",
+        doc_item_refs=["#/texts/2", "#/pictures/0"],
+        labels=["paragraph", "picture"],
+        image_data={"#/pictures/0": PICTURE_B64},
+    )
+    search = _build_search_tool(config)
+    rag = _fake_rag([_picture_result(), other])
+    ctx = _make_ctx(rag, RAGState())
+
+    result = await search(ctx, "figure")
+
+    assert isinstance(result, ToolReturn)
+    assert result.content is not None
+    assert len(result.content) == 1
+    assert result.content[0].identifier == "#/pictures/0"  # type: ignore[attr-defined]