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.
447 lines
16 KiB
Python
447 lines
16 KiB
Python
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from pydantic_ai import ToolFailed
|
|
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.tools.search import create_search_toolset
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vcr_cassette_dir():
|
|
return str(Path(__file__).parent.parent / "cassettes" / "test_search_tools")
|
|
|
|
|
|
def make_ctx(client, run_id="test-run"):
|
|
"""Create a lightweight RunContext-like object for direct tool function calls."""
|
|
return SimpleNamespace(deps=SimpleNamespace(client=client), run_id=run_id)
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
class TestSearchToolset:
|
|
"""Tests for create_search_toolset."""
|
|
|
|
def test_create_search_toolset_returns_function_toolset(self, search_config):
|
|
"""create_search_toolset returns a FunctionToolset."""
|
|
from pydantic_ai import FunctionToolset
|
|
|
|
toolset = create_search_toolset(search_config)
|
|
assert isinstance(toolset, FunctionToolset)
|
|
|
|
def test_search_toolset_has_search_tool(self, search_config):
|
|
"""The toolset includes a 'search' tool."""
|
|
toolset = create_search_toolset(search_config)
|
|
|
|
# toolset.tools is a dict with tool names as keys
|
|
assert "search" in toolset.tools
|
|
|
|
|
|
class TestNamingTheCollection:
|
|
"""The generic tool has no source selector, so what the client covers is
|
|
what the search spans."""
|
|
|
|
@staticmethod
|
|
def _client(covers_multiple: bool, source: str | None, *, pictures: bool = False):
|
|
from unittest.mock import AsyncMock
|
|
|
|
results = [
|
|
SearchResult(
|
|
content="body",
|
|
score=0.9,
|
|
source=source,
|
|
chunk_id="c1",
|
|
document_id="d1",
|
|
document_title="Report",
|
|
image_data={"#/pictures/0": _png_b64()} if pictures else None,
|
|
)
|
|
]
|
|
return SimpleNamespace(
|
|
covers_multiple=covers_multiple,
|
|
search=AsyncMock(return_value=results),
|
|
expand_context=AsyncMock(return_value=results),
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_client_covering_a_set_names_each_result(self, search_config):
|
|
toolset = create_search_toolset(search_config)
|
|
client = self._client(covers_multiple=True, source="alpha")
|
|
|
|
text = await toolset.tools["search"].function(make_ctx(client), "cats")
|
|
|
|
assert "Collection: alpha" in text
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_named_collection_is_not_named(self, search_config):
|
|
toolset = create_search_toolset(search_config)
|
|
client = self._client(covers_multiple=False, source="alpha")
|
|
|
|
text = await toolset.tools["search"].function(make_ctx(client), "cats")
|
|
|
|
assert "Collection" not in text
|
|
|
|
@staticmethod
|
|
def _seeing_config(search_config):
|
|
config = search_config.model_copy(deep=True)
|
|
config.qa.model.vision = True
|
|
return config
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_client_covering_a_set_names_each_image(self, search_config):
|
|
"""Images travel beside the results and are labelled the same way."""
|
|
from pydantic_ai.messages import ToolReturn
|
|
|
|
toolset = create_search_toolset(self._seeing_config(search_config))
|
|
client = self._client(covers_multiple=True, source="alpha", pictures=True)
|
|
|
|
returned = await toolset.tools["search"].function(make_ctx(client), "cats")
|
|
|
|
assert isinstance(returned, ToolReturn)
|
|
labels = [item for item in returned.content if isinstance(item, str)]
|
|
assert "Collection: alpha." in labels[0]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_named_collection_is_not_named_on_an_image(self, search_config):
|
|
from pydantic_ai.messages import ToolReturn
|
|
|
|
toolset = create_search_toolset(self._seeing_config(search_config))
|
|
client = self._client(covers_multiple=False, source="alpha", pictures=True)
|
|
|
|
returned = await toolset.tools["search"].function(make_ctx(client), "cats")
|
|
|
|
assert isinstance(returned, ToolReturn)
|
|
assert not [
|
|
item
|
|
for item in returned.content
|
|
if isinstance(item, str) and "Collection" in item
|
|
]
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
class TestSearchToolExecution:
|
|
"""Tests for search tool execution."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_returns_formatted_results(self, search_client, search_config):
|
|
"""Search tool returns formatted results."""
|
|
toolset = create_search_toolset(search_config)
|
|
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
result = await search_tool.function(ctx, "Python")
|
|
|
|
assert "Python" in result or "programming" in result
|
|
assert "No results found" not in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_with_no_results(self, temp_db_path, search_config):
|
|
"""Search tool returns appropriate message when no results."""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
# Use empty database
|
|
async with HaikuRAG(temp_db_path, create=True) as empty_client:
|
|
toolset = create_search_toolset(search_config)
|
|
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(empty_client)
|
|
result = await search_tool.function(ctx, "anything")
|
|
|
|
assert result == "No results found."
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_with_base_filter(self, search_client, search_config):
|
|
"""Search toolset respects base_filter parameter."""
|
|
accumulated: list[SearchResult] = []
|
|
toolset = create_search_toolset(
|
|
search_config,
|
|
base_filter="title LIKE '%Python%'",
|
|
on_results=accumulated.extend,
|
|
)
|
|
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
await search_tool.function(ctx, "programming")
|
|
|
|
assert len(accumulated) > 0
|
|
for r in accumulated:
|
|
assert "JavaScript" not in (r.document_title or "")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_on_results_callback(self, search_client, search_config):
|
|
"""on_results callback receives search results."""
|
|
accumulated: list[SearchResult] = []
|
|
toolset = create_search_toolset(search_config, on_results=accumulated.extend)
|
|
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
await search_tool.function(ctx, "Python")
|
|
|
|
assert len(accumulated) > 0
|
|
assert any("Python" in r.content for r in accumulated)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_on_results_accumulates_across_calls(
|
|
self, search_client, search_config
|
|
):
|
|
"""Multiple searches accumulate results via on_results callback."""
|
|
accumulated: list[SearchResult] = []
|
|
toolset = create_search_toolset(search_config, on_results=accumulated.extend)
|
|
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
await search_tool.function(ctx, "Python")
|
|
first_count = len(accumulated)
|
|
|
|
await search_tool.function(ctx, "JavaScript")
|
|
assert len(accumulated) > first_count
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_without_on_results(self, search_client, search_config):
|
|
"""Search works without on_results callback."""
|
|
toolset = create_search_toolset(search_config)
|
|
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
result = await search_tool.function(ctx, "Python")
|
|
|
|
assert "Python" in result or "programming" in result
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
class TestSearchMaxSearches:
|
|
"""Tests for max_searches cap on search toolset."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_searches_within_limit_return_results(
|
|
self, search_client, search_config
|
|
):
|
|
"""Searches within max_searches return normal results."""
|
|
toolset = create_search_toolset(search_config, max_searches=2)
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
|
|
assert await search_tool.function(ctx, "Python")
|
|
assert await search_tool.function(ctx, "JavaScript")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_searches_beyond_limit_fail_the_tool(
|
|
self, search_client, search_config
|
|
):
|
|
"""Searches beyond max_searches fail with the limit message."""
|
|
toolset = create_search_toolset(search_config, max_searches=1)
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
|
|
assert await search_tool.function(ctx, "Python")
|
|
|
|
with pytest.raises(ToolFailed, match="Search limit reached"):
|
|
await search_tool.function(ctx, "JavaScript")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_counter_resets_across_runs(self, search_client, search_config):
|
|
"""Search counter resets when run_id changes (new agent run)."""
|
|
toolset = create_search_toolset(search_config, max_searches=1)
|
|
search_tool = toolset.tools["search"]
|
|
|
|
ctx_run1 = make_ctx(search_client, run_id="run-1")
|
|
assert await search_tool.function(ctx_run1, "Python")
|
|
|
|
with pytest.raises(ToolFailed, match="Search limit reached"):
|
|
await search_tool.function(ctx_run1, "JavaScript")
|
|
|
|
ctx_run2 = make_ctx(search_client, run_id="run-2")
|
|
assert await search_tool.function(ctx_run2, "Python")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_limit_by_default(self, search_client, search_config):
|
|
"""Without max_searches, searches are unlimited."""
|
|
toolset = create_search_toolset(search_config)
|
|
search_tool = toolset.tools["search"]
|
|
ctx = make_ctx(search_client)
|
|
|
|
for _ in range(5):
|
|
assert await search_tool.function(ctx, "Python")
|
|
|
|
|
|
@pytest.fixture
|
|
async def search_client(temp_db_path):
|
|
"""Create a HaikuRAG client with test data for search tests."""
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
await rag.create_document(
|
|
"Python is a programming language. It is widely used for web development.",
|
|
uri="test://python",
|
|
title="Python Guide",
|
|
)
|
|
await rag.create_document(
|
|
"JavaScript runs in the browser. It powers interactive web pages.",
|
|
uri="test://javascript",
|
|
title="JavaScript Guide",
|
|
)
|
|
yield rag
|
|
|
|
|
|
@pytest.fixture
|
|
def search_config():
|
|
"""Default AppConfig for search tests."""
|
|
from haiku.rag.config import get_config
|
|
|
|
return get_config()
|
|
|
|
|
|
def _png_b64():
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
from PIL import Image as PILImage
|
|
|
|
buf = BytesIO()
|
|
PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode()
|
|
|
|
|
|
class TestBuildImageContentFromResults:
|
|
"""Picture bytes are attached once per (document, self_ref) pair, and labelled."""
|
|
|
|
def test_results_without_image_data_contribute_nothing(self):
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
results = [
|
|
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
|
|
]
|
|
|
|
assert build_image_content_from_results(results) == []
|
|
|
|
def test_duplicate_document_and_ref_is_attached_once(self):
|
|
from pydantic_ai.messages import BinaryContent
|
|
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
shared = {"#/pictures/0": _png_b64()}
|
|
results = [
|
|
SearchResult(
|
|
content="a",
|
|
score=0.9,
|
|
chunk_id="c1",
|
|
document_id="doc-1",
|
|
image_data=shared,
|
|
),
|
|
SearchResult(
|
|
content="b",
|
|
score=0.8,
|
|
chunk_id="c2",
|
|
document_id="doc-1",
|
|
image_data=shared,
|
|
),
|
|
]
|
|
|
|
content = build_image_content_from_results(results)
|
|
|
|
images = [item for item in content if isinstance(item, BinaryContent)]
|
|
assert len(images) == 1
|
|
|
|
@staticmethod
|
|
def _one_picture_in_two_collections():
|
|
"""A document copied into another collection keeps its ids and refs."""
|
|
shared = {"#/pictures/0": _png_b64()}
|
|
return [
|
|
SearchResult(
|
|
content="a",
|
|
score=0.9,
|
|
chunk_id="c1",
|
|
document_id="doc-1",
|
|
source="papers",
|
|
image_data=shared,
|
|
),
|
|
SearchResult(
|
|
content="a",
|
|
score=0.8,
|
|
chunk_id="c1",
|
|
document_id="doc-1",
|
|
source="wiki",
|
|
image_data=shared,
|
|
),
|
|
]
|
|
|
|
def test_the_same_picture_in_two_collections_is_attached_from_each(self):
|
|
from pydantic_ai.messages import BinaryContent
|
|
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
content = build_image_content_from_results(
|
|
self._one_picture_in_two_collections()
|
|
)
|
|
|
|
images = [item for item in content if isinstance(item, BinaryContent)]
|
|
assert len(images) == 2
|
|
|
|
def test_each_image_is_labelled_with_the_collection_it_came_from(self):
|
|
"""Nothing else tells the two apart: same chunk id, same document, same
|
|
reference."""
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
content = build_image_content_from_results(
|
|
self._one_picture_in_two_collections(), include_collection=True
|
|
)
|
|
|
|
labels = [item for item in content if isinstance(item, str)]
|
|
assert "Collection: papers." in labels[0]
|
|
assert "Collection: wiki." in labels[1]
|
|
|
|
def test_an_unasked_for_collection_is_not_named_on_an_image(self):
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
content = build_image_content_from_results(
|
|
self._one_picture_in_two_collections()
|
|
)
|
|
|
|
assert not [
|
|
item for item in content if isinstance(item, str) and "Collection" in item
|
|
]
|
|
|
|
def test_each_image_is_labelled_with_the_result_it_belongs_to(self):
|
|
"""Label every picture, not just the batch.
|
|
|
|
``ToolReturn.content`` reaches the model as a user-role message, and one
|
|
leading note does not override that: with a single note on the wire,
|
|
gemma4-26b still reasoned "the user also provided images in the prompt".
|
|
A label adjacent to each picture also names the chunk to cite for it,
|
|
which ``BinaryContent.identifier`` cannot do — it does not survive
|
|
serialization to the vision API.
|
|
"""
|
|
from pydantic_ai.messages import BinaryContent
|
|
|
|
from haiku.rag.tools.search import build_image_content_from_results
|
|
|
|
results = [
|
|
SearchResult(
|
|
content="a",
|
|
score=0.9,
|
|
chunk_id="c1",
|
|
document_id="doc-1",
|
|
image_data={"#/pictures/0": _png_b64()},
|
|
),
|
|
SearchResult(
|
|
content="b",
|
|
score=0.8,
|
|
chunk_id="c2",
|
|
document_id="doc-2",
|
|
image_data={"#/pictures/3": _png_b64()},
|
|
),
|
|
]
|
|
|
|
content = build_image_content_from_results(results)
|
|
|
|
# label, image, label, image — each picture preceded by its own line.
|
|
assert [type(item) is str for item in content] == [True, False, True, False]
|
|
assert isinstance(content[1], BinaryContent)
|
|
assert isinstance(content[3], BinaryContent)
|
|
|
|
first, second = content[0], content[2]
|
|
assert isinstance(first, str) and isinstance(second, str)
|
|
assert "c1" in first and "#/pictures/0" in first
|
|
assert "c2" in second and "#/pictures/3" in second
|
|
assert "1 of 2" in first and "2 of 2" in second
|
|
for label in (first, second):
|
|
assert "not provided by the user" in label.lower()
|