fa319596cc
Search results and capsule entries name the collection they came from; the images attached beside them carried only the chunk id and reference. Two collections can return the same picture of the same document, so the two labels were identical and the model could place neither. The decision is the one already made for the text: `covers_multiple` at the generic search tool, and the flag `search_corpus` computed for the capability tools, which it now returns.
186 lines
7.2 KiB
Python
186 lines
7.2 KiB
Python
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],
|
|
include_collection: bool = False,
|
|
) -> list[str | BinaryContent]:
|
|
"""Decode and validate picture bytes attached to search results, labelled.
|
|
|
|
Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in
|
|
different chunks is sent once, and a copy in another collection is its own. 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 | None, str, BinaryContent]] = []
|
|
seen: set[tuple[str | None, 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.source, 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.source, result.chunk_id, self_ref, picture))
|
|
seen.add(key)
|
|
|
|
content: list[str | BinaryContent] = []
|
|
total = len(collected)
|
|
for position, (source, chunk_id, self_ref, picture) in enumerate(collected, 1):
|
|
collection = f"Collection: {source}. " if include_collection and source else ""
|
|
content.append(
|
|
f"Page image {position} of {total}, retrieved from the knowledge base "
|
|
f"for search result [{chunk_id}] ({self_ref}). {collection}"
|
|
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)
|
|
include_collection = client.covers_multiple
|
|
formatted = [
|
|
r.format_for_agent(
|
|
rank=i + 1, total=total, include_collection=include_collection
|
|
)
|
|
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, include_collection=include_collection
|
|
)
|
|
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
|