haiku.rag/haiku_rag_slim/haiku/rag/tools/search.py

import base64
from collections.abc import Callable
from io import BytesIO

from PIL import Image
from pydantic_ai import FunctionToolset, RunContext, ToolFailed
from pydantic_ai.messages import BinaryContent, ToolReturn

from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import RAGDeps

RETRIEVED_IMAGE_TAG = "[haiku.rag/retrieved-image]"
"""Tag every label we attach to a retrieved picture ends with.

Identifies our own pictures on the wire without inferring ownership from position,
which is wrong as soon as two tools' results arrive in one request. Deliberately not
a phrase: a user writing "retrieved from the knowledge base for my report" above
their own picture had it removed, along with their text.
"""


def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
    """Wrap picture bytes for the wire, or return nothing if they will not decode.

    The model adapter renders one vision placeholder per ``BinaryContent``, so
    emitting one for an image the server cannot decode leaves the processor with an
    off-by-one count.
    """
    try:
        with Image.open(BytesIO(data)) as image:
            image.verify()
    except Exception:
        return None
    return BinaryContent(data=data, media_type="image/png", identifier=self_ref)


def build_image_content_from_results(
    results: list[SearchResult],
) -> list[str | BinaryContent]:
    """Decode and validate picture bytes attached to search results, labelled.

    Dedup keyed on ``(document_id, self_ref)`` so the same picture in
    different chunks is sent once. Pictures that fail
    ``PIL.Image.verify()`` are skipped — the model adapter renders one
    vision placeholder per ``BinaryContent``, so emitting one for an
    image the server can't decode leaves the processor with an
    off-by-one count.

    Every picture is preceded by a line naming the result it belongs to.
    ``ToolReturn.content`` reaches the model as a user-role message, so
    retrieved pictures are otherwise indistinguishable from ones the user
    attached, and models narrate them as part of the question: unlabelled,
    gemma4-26b answered about a figure from an unrelated document, and with a
    single note ahead of the batch it still called them "images in the prompt".
    The label also names the chunk to cite for a figure, which
    ``BinaryContent.identifier`` cannot do — it does not survive serialization
    to the vision API.
    """
    collected: list[tuple[str | None, str, BinaryContent]] = []
    seen: set[tuple[str | None, str]] = set()
    for result in results:
        if not result.image_data:
            continue
        for self_ref, b64 in result.image_data.items():
            key = (result.document_id, self_ref)
            if key in seen:
                continue
            picture = decode_picture(base64.b64decode(b64), self_ref)
            if picture is None:
                continue
            collected.append((result.chunk_id, self_ref, picture))
            seen.add(key)

    content: list[str | BinaryContent] = []
    total = len(collected)
    for position, (chunk_id, self_ref, picture) in enumerate(collected, 1):
        content.append(
            f"Page image {position} of {total}, retrieved from the knowledge base "
            f"for search result [{chunk_id}] ({self_ref}). "
            f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
        )
        content.append(picture)
    return content


def create_search_toolset(
    config: AppConfig,
    expand_context: bool = True,
    base_filter: str | None = None,
    tool_name: str = "search",
    on_results: Callable[[list[SearchResult]], None] | None = None,
    max_searches: int | None = None,
) -> FunctionToolset[RAGDeps]:
    """Create a toolset with search capabilities.

    Args:
        config: Application configuration.
        expand_context: Whether to expand search results with surrounding context.
            Defaults to True.
        base_filter: Optional base SQL WHERE clause applied to all searches.
            Combined with any filter passed to the search tool.
        tool_name: Name for the search tool. Defaults to "search".
        on_results: Optional callback invoked with search results after each search.
            Useful for accumulating results externally (e.g., for citation resolution).
        max_searches: Maximum number of searches allowed. When exceeded, the
            tool fails with a message directing the agent to answer with
            existing results.

    Returns:
        FunctionToolset with a search tool.
    """
    # Per-run search counter keyed by run_id. Safe for concurrent runs
    # and reuse across sequential agent.run() calls.
    search_counts: dict[str, int] = {}

    async def search(
        ctx: RunContext[RAGDeps],
        query: str,
        limit: int | None = None,
    ) -> str | ToolReturn:
        """Search the knowledge base for relevant documents.

        Args:
            query: The search query (what to search for).
            limit: Number of results to return (default: from config).

        Returns:
            Formatted search results with content and metadata. When a
            picture-labeled chunk is in the result set, returns a
            ``pydantic_ai.messages.ToolReturn`` whose ``content`` carries the
            corresponding ``BinaryContent`` parts so a vision-capable model
            sees the figures alongside the text.
        """
        rid = ctx.run_id or ""
        search_counts[rid] = search_counts.get(rid, 0) + 1
        if max_searches is not None and search_counts[rid] > max_searches:
            raise ToolFailed(
                "Search limit reached. "
                "Answer the question using the results you already have."
            )

        client = ctx.deps.client

        effective_filter = base_filter
        effective_limit = limit or config.search.limit
        results = await client.search(
            query, limit=effective_limit, filter=effective_filter
        )

        if expand_context:
            results = await client.expand_context(results)

        results_list = list(results)

        if on_results:
            on_results(results_list)

        if not results_list:
            return "No results found."

        total = len(results_list)
        formatted = [
            r.format_for_agent(rank=i + 1, total=total)
            for i, r in enumerate(results_list)
        ]
        text = "\n\n".join(formatted)

        if not config.qa.model.vision:
            return text

        image_content = build_image_content_from_results(results_list)
        if image_content:
            return ToolReturn(return_value=text, content=image_content)
        return text

    toolset: FunctionToolset[RAGDeps] = FunctionToolset()
    toolset.add_function(search, name=tool_name, retries=3)
    return toolset