diff --git a/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2 b/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2 index fc25efe6..cf83f3fb 100644 --- a/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2 +++ b/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 %} diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 0ce48a4b..d5613e80 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md index 65738896..b23644e1 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md index fafa8c39..701fee96 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md +++ b/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. diff --git a/tests/test_skill_tools.py b/tests/test_skill_tools.py new file mode 100644 index 00000000..49a1ad16 --- /dev/null +++ b/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]