0d7810c78a
`format_for_agent` named the database whenever one was named, so a search over a single named database carried a line with nothing to distinguish. It now takes `include_collection` from the caller, which decides from the search selection rather than from the hits: a search that could have drawn on two collections names them even when everything came back from one. `Collection:` at the model boundary, database in configuration and administration. `source` on results, documents, citations and analysis dictionaries is unchanged.
182 lines
6.9 KiB
Python
182 lines
6.9 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],
|
|
) -> 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)
|
|
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)
|
|
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
|