danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #597 from ggozad/feat/search-fanout-units

Browse source 
Price search budget in units and deduplicate fan-out results
This commit is contained in:
Yiorgis Gozadinos 2026-09-03 07:09:13 -05:00 committed by GitHub
commit 4f7f69aaf0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 598 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
CHANGELOG.md
View file

@ -2,6 +2,15 @@
## [Unreleased]
### Changed
- `qa.max_searches` counts search units: searches a model emits in one
response share a unit, up to 3 per unit; sequential searches pay one unit
each.
- Searches in one model response deduplicate their results: evidence a sibling
search already showed collapses to a reference line, and a picture attaches
once per response.
## [0.81.0] - 2026-09-01
### Added

4
docs/configuration/qa.md
View file

@ -30,12 +30,12 @@ qa:
enable_thinking: true
temperature: 0.3 # Default: 0.3
vision: false # Set true for vision-capable models
max_searches: 5 # Maximum search tool calls per question
max_searches: 5 # Maximum search units per question
```
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The capability's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix.
- **max_searches**: Maximum number of search tool calls a capability can make per question (default: 5). Shared by the RAG and analysis capabilities.
- **max_searches**: Maximum number of search units a capability can spend per question (default: 5). Up to three searches emitted in the same model response share one unit, so a model that rephrases its query in one response spends one unit. A search in a later response starts a new unit, as does each further group of three within one response. Shared by the RAG and analysis capabilities. Searches in one response also deduplicate their returns: evidence a sibling search already showed collapses to a reference line, and each picture attaches once per response.
!!! note "Thinking on vLLM"
`enable_thinking` only applies to models with a pydantic-ai reasoning profile (o-series, gpt-5, gpt-oss). For other vLLM-served models such as Qwen3 or the Gemma family, the field is a silent no-op — set the chat template switch via [`extra_body`](providers.md#raw-provider-pass-through) instead.

64
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -14,6 +14,7 @@ from pydantic_ai import (
)
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import (
BinaryContent,
InstructionPart,
ModelMessage,
ModelRequest,
@ -29,6 +30,7 @@ from pydantic_ai.toolsets import AgentToolset
from haiku.rag.capabilities._tools import (
CodeExecutionEntry,
EvidenceKey,
merge_results,
search_corpus,
)
@ -43,7 +45,7 @@ from haiku.rag.store.models.citation import (
ambiguous_citation,
resolve_citations,
)
from haiku.rag.tools.search import build_image_content_from_results
from haiku.rag.tools.search import PictureKey, build_image_content_from_results
CITATION_GRACE_REQUESTS = 2
"""Requests calling this capability's tools that its cite tool outlives the rest by.
@ -61,6 +63,15 @@ Calibration knob. Two unrelated UUID4s reach about 0.5, while dropping or
duplicating a character or a whole group stays above 0.75, so the gap is wide.
"""
FREE_SIBLINGS_PER_ROUND = 3
"""Searches one budget unit covers when emitted in the same model response.
Calibration knob, sized to the measured modal burst. ``qa.max_searches``
counts units, so a model rephrasing its query a few times in one response
spends one unit, while every search of a sequential searcher is a unit of its
own.
"""
def _ambiguous_retry(error: AmbiguousCitationError) -> ModelRetry:
"""The only way out of an id that names a chunk in two databases.
@ -177,6 +188,12 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
search_count: int = field(default=0, repr=False)
search_step: int = field(default=0, repr=False)
"""The run_step whose searches are being priced and deduplicated."""
step_searches: int = field(default=0, repr=False)
step_rejected: bool = field(default=False, repr=False)
step_shown: set[EvidenceKey] = field(default_factory=set, repr=False)
step_pictures: set[PictureKey] = field(default_factory=set, repr=False)
request_count: int = field(default=0, repr=False)
grace_requests_used: int = field(default=0, repr=False)
epoch: int = field(default=0, repr=False)
@ -226,6 +243,11 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
rag_lock=asyncio.Lock(),
resource_lock=asyncio.Lock(),
search_count=0,
search_step=0,
step_searches=0,
step_rejected=False,
step_shown=set(),
step_pictures=set(),
request_count=0,
grace_requests_used=0,
epoch=0,
@ -493,32 +515,54 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
retrieved_now=retrieved,
)
async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
async def _search(
self, query: str, limit: int | None, run_step: int
) -> str | ToolReturn:
assert self.state is not None
self.search_count += 1
if self.search_count > self._max_searches:
if run_step != self.search_step:
self.search_step = run_step
self.step_searches = 0
self.step_rejected = False
self.step_shown = set()
self.step_pictures = set()
self.step_searches += 1
if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0:
self.search_count += 1
if self.step_rejected or self.search_count > self._max_searches:
self.step_rejected = True
raise ToolFailed(
"Search limit reached. Answer the question using "
"the results you already have."
)
async with self.rag_lock:
formatted, results, include_collection = await search_corpus(
formatted, results, rendered, include_collection = await search_corpus(
await self._ensure_rag(),
query,
limit=limit,
document_filter=self.state.document_filter,
sources=self.state.sources,
shown=self.step_shown,
)
parts: list[str | BinaryContent] = []
emitted: set[PictureKey] = set()
if self.vision:
parts, emitted = build_image_content_from_results(
results,
include_collection=include_collection,
exclude=self.step_pictures,
)
# Everything the search produced commits together, after formatting and
# image construction have both succeeded: a search that raises must not
# leave results citable, note evidence the model never received, or
# suppress a later sibling's results.
state = self.state
# A model can search the same query twice with different limits, and the
# narrower return must not drop what the wider one already showed it.
merge_results(state.searches.setdefault(query, []), results)
self._note_evidence()
if self.vision and (
parts := build_image_content_from_results(
results, include_collection=include_collection
)
):
self.step_shown |= rendered
self.step_pictures |= emitted
if parts:
return ToolReturn(return_value=formatted, content=parts)
return formatted

69
haiku_rag_slim/haiku/rag/capabilities/_tools.py
View file

@ -1,9 +1,11 @@
from collections.abc import Iterable
from collections.abc import Set as AbstractSet
from pydantic import BaseModel
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.chunk import SearchResult, qualified_id
from haiku.rag.tools.search import picture_keys
class CodeExecutionEntry(BaseModel):
@ -13,14 +15,48 @@ class CodeExecutionEntry(BaseModel):
success: bool = True
EvidenceKey = tuple[tuple[str | None, str | None], tuple[str, frozenset]]
"""What tells one rendered result from another: qualified id, then signature.
The qualified id comes first because the rendered string alone would conflate
identical renderings of the same chunk id held by two databases.
"""
def evidence_signature(result: SearchResult, include_collection: bool) -> tuple:
"""The rendered evidence a result shows the model, as an equivalence key.
Rank and total are held at neutral values: they vary with a result's
position, and position (like score) must not tell two renderings apart.
"""
return (
result.format_for_agent(rank=0, total=0, include_collection=include_collection),
picture_keys(result),
)
def evidence_key(result: SearchResult, include_collection: bool) -> EvidenceKey:
return (
qualified_id(result.source, result.chunk_id),
evidence_signature(result, include_collection),
)
async def search_corpus(
rag: HaikuRAG,
query: str,
limit: int | None = None,
document_filter: str | None = None,
sources: list[str] | None = None,
) -> tuple[str, list[SearchResult], bool]:
"""Search and context-expand results, and whether they name their collection."""
shown: AbstractSet[EvidenceKey] = frozenset(),
) -> tuple[str, list[SearchResult], set[EvidenceKey], bool]:
"""Search and context-expand results, eliding evidence already shown.
Returns the formatted results, the full result list, the evidence keys the
formatting rendered in full, and whether results name their collection. A
result whose key is in ``shown`` keeps its slot but collapses to one line;
the result list is never filtered.
"""
results = await rag.search(
query, limit=limit, filter=document_filter, sources=sources
)
@ -29,13 +65,25 @@ async def search_corpus(
# two collections names them even when everything came back from one.
selected = rag.source_names if sources is None else sources
include_collection = len(set(selected)) > 1
formatted = "\n\n---\n\n".join(
result.format_for_agent(
rank=index + 1, total=len(results), include_collection=include_collection
)
for index, result in enumerate(results)
)
return formatted or "No results found.", list(results), include_collection
rendered: set[EvidenceKey] = set()
parts: list[str] = []
total = len(results)
for index, result in enumerate(results):
key = evidence_key(result, include_collection)
if key in shown or key in rendered:
parts.append(
f"Also matched, shown above: [{result.chunk_id}] "
f"[rank {index + 1} of {total}]"
)
else:
parts.append(
result.format_for_agent(
rank=index + 1, total=total, include_collection=include_collection
)
)
rendered.add(key)
formatted = "\n\n---\n\n".join(parts)
return formatted or "No results found.", list(results), rendered, include_collection
def merge_results(
@ -56,6 +104,9 @@ def merge_results(
__all__ = [
"CodeExecutionEntry",
"EvidenceKey",
"evidence_key",
"evidence_signature",
"merge_results",
"search_corpus",
]

2
haiku_rag_slim/haiku/rag/capabilities/analysis.py
View file

@ -172,7 +172,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
ctx: RunContext[Any], query: str, limit: int | None = None
) -> str | ToolReturn:
"""Search the knowledge base for evidence to analyze."""
return await self._with_state(self._search(query, limit))
return await self._with_state(self._search(query, limit, ctx.run_step))
async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any:
"""Execute Python against the sandboxed document filesystem."""

2
haiku_rag_slim/haiku/rag/capabilities/rag.py
View file

@ -81,7 +81,7 @@ class RAGCapability(RAGCapabilityBase[RAGState]):
ctx: RunContext[Any], query: str, limit: int | None = None
) -> str | ToolReturn:
"""Search the knowledge base using hybrid vector and full-text search."""
return await self._with_state(self._search(query, limit))
return await self._with_state(self._search(query, limit, ctx.run_step))
async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any:
"""Register exact search-result chunk IDs as citations for the answer."""

34
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -1,5 +1,6 @@
import base64
from collections.abc import Callable
from collections.abc import Set as AbstractSet
from io import BytesIO
from PIL import Image
@ -20,6 +21,22 @@ their own picture had it removed, along with their text.
"""
PictureKey = tuple[str | None, str | None, str]
"""Identity of one attached picture: (source, document_id, self_ref).
``self_ref`` alone collides across documents, and a copy of a document in
another collection carries its own pictures.
"""
def picture_keys(result: SearchResult) -> frozenset[PictureKey]:
"""The identity of every picture this result carries."""
return frozenset(
(result.source, result.document_id, self_ref)
for self_ref in (result.image_data or {})
)
def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
"""Wrap picture bytes for the wire, or return nothing if they will not decode.
@ -38,11 +55,14 @@ def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
def build_image_content_from_results(
results: list[SearchResult],
include_collection: bool = False,
) -> list[str | BinaryContent]:
exclude: AbstractSet[PictureKey] = frozenset(),
) -> tuple[list[str | BinaryContent], set[PictureKey]]:
"""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
Returns the labelled content and the ``PictureKey`` of every picture it
emitted. Dedup keyed on ``PictureKey`` so the same picture in
different chunks is sent once, and a copy in another collection is its
own; ``exclude`` seeds that dedup with pictures already sent. 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
@ -59,7 +79,8 @@ def build_image_content_from_results(
to the vision API.
"""
collected: list[tuple[str | None, str | None, str, BinaryContent]] = []
seen: set[tuple[str | None, str | None, str]] = set()
seen: set[PictureKey] = set(exclude)
emitted: set[PictureKey] = set()
for result in results:
if not result.image_data:
continue
@ -72,6 +93,7 @@ def build_image_content_from_results(
continue
collected.append((result.source, result.chunk_id, self_ref, picture))
seen.add(key)
emitted.add(key)
content: list[str | BinaryContent] = []
total = len(collected)
@ -83,7 +105,7 @@ def build_image_content_from_results(
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
)
content.append(picture)
return content
return content, emitted
def create_search_toolset(
@ -174,7 +196,7 @@ def create_search_toolset(
if not config.qa.model.vision:
return text
image_content = build_image_content_from_results(
image_content, _ = build_image_content_from_results(
results_list, include_collection=include_collection
)
if image_content:

18
tests/capabilities/test_capabilities.py
View file

@ -412,7 +412,7 @@ async def test_a_spent_search_budget_fails_the_tool(temp_db_path):
capability.state = RAGState()
with pytest.raises(ToolFailed, match="Search limit reached"):
await capability._search("anything", None)
await capability._search("anything", None, 1)
def _stub_client(*batches: list[SearchResult]) -> AsyncMock:
@ -452,7 +452,7 @@ async def _labels_of_search(temp_db_path, *sources: str) -> list[str]:
client.source_names = sources
with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)):
returned = await capability._search("cats", None)
returned = await capability._search("cats", None, 1)
assert isinstance(returned, ToolReturn)
assert returned.content is not None
@ -483,7 +483,9 @@ async def test_a_fruitless_search_says_so(temp_db_path):
capability.state = RAGState()
capability.borrowed_rag = _stub_client([])
assert await capability._search("nothing about this", None) == "No results found."
assert (
await capability._search("nothing about this", None, 1) == "No results found."
)
@pytest.mark.asyncio
@ -500,8 +502,8 @@ async def test_a_narrower_repeat_keeps_what_the_wider_search_returned(temp_db_pa
[SearchResult(content="first", score=1.0, chunk_id="chunk-1")],
)
await capability._search("Figure 3-1", 20)
await capability._search("Figure 3-1", None)
await capability._search("Figure 3-1", 20, 1)
await capability._search("Figure 3-1", None, 2)
stored = capability.state.searches["Figure 3-1"]
assert [result.chunk_id for result in stored] == [
@ -525,8 +527,8 @@ async def test_two_databases_holding_one_chunk_id_both_survive(temp_db_path):
],
)
await capability._search("cats", 20)
await capability._search("cats", None)
await capability._search("cats", 20, 1)
await capability._search("cats", None, 2)
stored = capability.state.searches["cats"]
assert [(r.source, r.chunk_id) for r in stored] == [
@ -1105,7 +1107,7 @@ def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord:
return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"])
async def _stub_search(self, query: str, _limit: int | None) -> str:
async def _stub_search(self, query: str, _limit: int | None, _run_step: int) -> str:
"""Record a result the way the real search does, so citing resolves."""
cast(Any, self.state).searches[query] = [
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")

2
tests/capabilities/test_citation_policy.py
View file

@ -29,7 +29,7 @@ class Deps:
state: dict[str, Any] = field(default_factory=dict)
async def stub_search(self, query: str, _limit: int | None) -> str:
async def stub_search(self, query: str, _limit: int | None, _run_step: int) -> str:
cast(Any, self.state).searches[query] = [
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
]

87
tests/capabilities/test_evidence_wire.py
View file

@ -437,7 +437,9 @@ REAL_PNG = base64.b64decode(
)
async def _search_with_a_picture(self, query: str, _limit: int | None) -> str:
async def _search_with_a_picture(
self, query: str, _limit: int | None, _run_step: int
) -> str:
"""Record a result carrying a page image, the way a real search does."""
cast(Any, self.state).searches[query] = [
SearchResult(
@ -515,6 +517,89 @@ async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label(
assert texts_of(wire[-1]) == []
def _burst_result() -> SearchResult:
return SearchResult(
content="evidence",
score=1.0,
chunk_id="chunk-1",
document_id="doc-1",
source="main",
doc_item_refs=["#/pictures/0"],
image_data={"#/pictures/0": base64.b64encode(REAL_PNG).decode()},
)
async def _fanout_question_then_another(temp_db_path, cite: bool) -> list[list[Any]]:
"""Question 1 fans out over one picture chunk; question 2 follows compacted."""
rag = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False, vision=True
)
client = AsyncMock()
client.search.side_effect = [[_burst_result()], [_burst_result()]]
client.expand_context.side_effect = lambda results: results
client.source_names = ["main"]
rag.borrowed_rag = client
citing = (
[[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-3")]]
if cite
else []
)
calls = iter(
[
[
ToolCallPart("rag_search", {"query": "figure"}, "call-1"),
ToolCallPart("rag_search", {"query": "the figure"}, "call-2"),
],
*citing,
[TextPart("first answer")],
[TextPart("second answer")],
]
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return ModelResponse(parts=next(calls))
agent = Agent(
FunctionModel(model),
deps_type=Deps,
capabilities=[rag, create_compaction()],
)
deps = Deps()
with patch.object(
RAGCapability, "get_picture_bytes", AsyncMock(return_value=REAL_PNG)
):
first = await agent.run("what does the figure show?", deps=deps)
await agent.run(
"and what else?", deps=deps, message_history=first.all_messages()
)
return wire
@pytest.mark.asyncio
async def test_a_burst_deduplicated_picture_survives_compaction_when_cited(
temp_db_path,
):
"""Dedup attaches the picture once in its own question; the capsule re-fetches
it for the next. Neither pass may leave the model without it."""
wire = await _fanout_question_then_another(temp_db_path, cite=True)
assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG]
assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG]
@pytest.mark.asyncio
async def test_a_burst_deduplicated_picture_is_dropped_by_compaction_uncited(
temp_db_path,
):
wire = await _fanout_question_then_another(temp_db_path, cite=False)
assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG]
assert images_of(wire[-1]) == []
@pytest.mark.asyncio
async def test_the_capsule_is_built_once_per_request_and_again_for_the_next(
temp_db_path,

329
tests/capabilities/test_search_units.py Normal file
View file

@ -0,0 +1,329 @@
import base64
from dataclasses import dataclass, field
from io import BytesIO
from typing import Any
from unittest.mock import AsyncMock
import pytest
from PIL import Image as PILImage
from pydantic_ai import Agent
from pydantic_ai.messages import (
BinaryContent,
ModelResponse,
TextPart,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.models.function import FunctionModel
from pydantic_ai.run import AgentRunResult
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.capabilities.rag import RAGState
from haiku.rag.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
@dataclass
class Deps:
state: dict[str, Any] = field(default_factory=dict)
def burst_model(bursts: list[list[str]]) -> FunctionModel:
"""Emit one `rag_search` call per query in each burst, then answer."""
responses = 0
def model_function(_messages, _info) -> ModelResponse:
nonlocal responses
responses += 1
if responses <= len(bursts):
return ModelResponse(
parts=[
ToolCallPart("rag_search", {"query": query})
for query in bursts[responses - 1]
]
)
return ModelResponse(parts=[TextPart("done")])
return FunctionModel(model_function)
def burst_agent(
bursts: list[list[str]], db_path, max_searches: int
) -> Agent[Deps, str]:
config = AppConfig()
config.qa.max_searches = max_searches
return Agent(
burst_model(bursts),
deps_type=Deps,
capabilities=[create_rag(db_path=db_path, config=config, defer_loading=False)],
)
def search_returns(result: AgentRunResult[Any]) -> list[ToolReturnPart]:
return [
part
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolReturnPart) and part.tool_name == "rag_search"
]
def outcomes(result: AgentRunResult[Any]) -> list[str]:
return [
"failed" if part.outcome == "failed" else "ok"
for part in search_returns(result)
]
@pytest.mark.asyncio
async def test_a_burst_in_one_response_consumes_one_unit(rag_db):
"""Three searches emitted together cost one unit and run in emission order."""
agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 1)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["ok", "ok", "ok"]
calls = [
part
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolCallPart) and part.tool_name == "rag_search"
]
assert [part.tool_call_id for part in search_returns(result)] == [
part.tool_call_id for part in calls
]
@pytest.mark.asyncio
async def test_sequential_searches_pay_one_unit_each(rag_db):
agent = burst_agent([["ai"], ["machine learning"]], rag_db, 1)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["ok", "failed"]
assert "Search limit reached" in str(search_returns(result)[1].content)
@pytest.mark.asyncio
async def test_max_searches_zero_fails_every_sibling(rag_db):
agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 0)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["failed", "failed", "failed"]
@pytest.mark.asyncio
async def test_a_rejected_round_fails_all_its_siblings(rag_db):
agent = burst_agent([["ai"], ["ml", "deep learning", "supervised"]], rag_db, 1)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["ok", "failed", "failed", "failed"]
@pytest.mark.asyncio
async def test_a_sibling_past_the_allowance_pays_its_own_unit(rag_db):
burst = [["ai", "machine learning", "deep learning", "supervised learning"]]
within = await burst_agent(burst, rag_db, 2).run("question", deps=Deps())
over = await burst_agent(burst, rag_db, 1).run("question", deps=Deps())
assert outcomes(within) == ["ok", "ok", "ok", "ok"]
assert outcomes(over) == ["ok", "ok", "ok", "failed"]
@pytest.mark.asyncio
async def test_unit_tracking_resets_between_runs(rag_db):
"""A second run's opening burst prices like a first run's."""
def model_function(messages, _info) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for part in messages[-1].parts):
return ModelResponse(parts=[TextPart("done")])
return ModelResponse(
parts=[
ToolCallPart("rag_search", {"query": query})
for query in ["ai", "machine learning", "deep learning"]
]
)
config = AppConfig()
config.qa.max_searches = 1
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
capabilities=[create_rag(db_path=rag_db, config=config, defer_loading=False)],
)
deps = Deps()
first = await agent.run("question", deps=deps)
second = await agent.run("another", deps=deps, message_history=first.all_messages())
assert outcomes(first) == ["ok", "ok", "ok"]
assert outcomes(second)[-3:] == ["ok", "ok", "ok"]
def _png() -> str:
buffer = BytesIO()
PILImage.new("RGB", (4, 4), "red").save(buffer, format="PNG")
return base64.b64encode(buffer.getvalue()).decode()
def make_result(**overrides: Any) -> SearchResult:
fields: dict[str, Any] = {
"content": "body",
"score": 0.9,
"source": "main",
"chunk_id": "c1",
"document_id": "d1",
"image_data": {"#/pictures/0": _png()},
}
fields.update(overrides)
return SearchResult(**fields)
def stub_client(
*batches: list[SearchResult], sources: list[str] | None = None
) -> AsyncMock:
client = AsyncMock()
client.search.side_effect = list(batches)
client.expand_context.side_effect = lambda results: results
client.source_names = sources or ["main"]
return client
def dedup_capability(client: AsyncMock, temp_db_path, *, vision: bool = True):
capability = create_rag(db_path=temp_db_path, config=AppConfig(), vision=vision)
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
capability.borrowed_rag = client
return capability
def images_of(returned: Any) -> list[BinaryContent]:
if isinstance(returned, str):
return []
return [item for item in returned.content if isinstance(item, BinaryContent)]
def text_of(returned: Any) -> str:
return returned if isinstance(returned, str) else returned.return_value
@pytest.mark.asyncio
async def test_a_duplicate_sibling_is_elided_and_stays_citable(temp_db_path):
duplicate, novel = make_result(), make_result(chunk_id="c2", content="novel")
client = stub_client([make_result()], [duplicate, novel])
capability = dedup_capability(client, temp_db_path)
first = await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert len(images_of(first)) == 1
assert images_of(second) == []
text = text_of(second)
assert "Also matched, shown above: [c1] [rank 1 of 2]" in text
assert "body" not in text
assert "[rank 2 of 2]" in text and "novel" in text
assert [r.chunk_id for r in capability.state.searches["q rephrased"]] == [
"c1",
"c2",
]
assert await capability._cite(["c1"]) == "Registered 1 citation(s)."
@pytest.mark.asyncio
async def test_a_new_run_step_formats_shown_results_in_full(temp_db_path):
client = stub_client([make_result()], [make_result()])
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q again", None, 2)
assert "body" in text_of(second)
assert len(images_of(second)) == 1
@pytest.mark.asyncio
async def test_same_chunk_id_from_another_collection_is_not_elided(temp_db_path):
client = stub_client(
[make_result(source="alpha")],
[make_result(source="beta")],
sources=["alpha", "beta"],
)
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert "body" in text_of(second)
@pytest.mark.asyncio
async def test_same_anchor_with_new_evidence_formats_in_full(temp_db_path):
shared, extra = _png(), _png()
client = stub_client(
[make_result(content="c1 with c2", image_data={"#/pictures/1": shared})],
[
make_result(
content="c1 with c3",
image_data={"#/pictures/1": shared, "#/pictures/3": extra},
)
],
)
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert "c1 with c3" in text_of(second)
assert len(images_of(second)) == 1
labels = [item for item in second.content if isinstance(item, str)]
assert any("#/pictures/3" in label for label in labels)
@pytest.mark.parametrize(
("overrides", "elided"),
[
({"score": 0.1}, True),
({"content": "different"}, False),
({"document_title": "Other"}, False),
({"headings": ["Heading"]}, False),
({"labels": ["table"]}, False),
({"picture_captions": {"#/pictures/0": "A caption"}}, False),
({"image_data": {"#/pictures/9": _png()}}, False),
],
)
@pytest.mark.asyncio
async def test_equivalence_follows_the_rendered_evidence(
temp_db_path, overrides: dict[str, Any], elided: bool
):
"""Any rendered field or picture identity defeats elision; score alone does not."""
client = stub_client([make_result()], [make_result(**overrides)])
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert ("Also matched, shown above" in text_of(second)) is elided
@pytest.mark.asyncio
async def test_a_failed_sibling_commits_nothing(temp_db_path):
client = stub_client(
[make_result(image_data={"#/pictures/0": "AAA"})],
[make_result()],
)
capability = dedup_capability(client, temp_db_path)
evidence_before = capability.state.evidence.model_dump()
with pytest.raises(Exception):
await capability._search("q", None, 1)
assert capability.state.searches == {}
assert capability.state.evidence.model_dump() == evidence_before
second = await capability._search("q rephrased", None, 1)
assert "body" in text_of(second)
assert len(images_of(second)) == 1

14
tests/multi_db/test_capabilities.py
View file

@ -30,7 +30,7 @@ class TestAskAcrossDatabases:
capability = create_capability(config=config, rag=rag, defer_loading=False)
capability.state = RAGState(sources=["alpha"])
formatted = await capability._search("cats", limit=10)
formatted = await capability._search("cats", 10, 1)
assert isinstance(formatted, str)
assert "alpha" in formatted
@ -47,7 +47,7 @@ class TestAskAcrossDatabases:
capability = create_capability(config=config, rag=rag, defer_loading=False)
capability.state = RAGState()
formatted = await capability._search("cats", limit=10)
formatted = await capability._search("cats", 10, 1)
assert isinstance(formatted, str)
assert "alpha document" in formatted
@ -72,7 +72,7 @@ class TestStandaloneCapabilities:
assert capability.scope.names == ("alpha", "beta")
run = await capability.for_run(make_context(Deps()))
try:
formatted = await run._search("cats", limit=10)
formatted = await run._search("cats", 10, 1)
finally:
await run._close()
@ -135,7 +135,7 @@ class TestAnalyzeAcrossDatabases:
capability = create_analysis(config=config, rag=rag, defer_loading=False)
capability.state = AnalysisState(sources=["alpha"])
formatted = await capability._search("cats", limit=10)
formatted = await capability._search("cats", 10, 1)
sandbox = await capability._ensure_sandbox()
await capability._close()
@ -305,8 +305,10 @@ class TestWhenTheModelIsToldTheCollection:
async with HaikuRAG(config=config) as rag:
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
spanning, _, spans = await search_corpus(rag, "cats")
narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"])
spanning, _, _, spans = await search_corpus(rag, "cats")
narrowed, _, _, narrows = await search_corpus(
rag, "cats", sources=["alpha"]
)
assert "Collection: alpha" in spanning
assert "Collection" not in narrowed

4
tests/multi_db/test_citations.py
View file

@ -342,7 +342,7 @@ class TestCiteFallback:
)
run = await capability.for_run(make_context(deps))
# The search returns the cats chunk, never the aardvark one.
await run._search("cats", limit=10)
await run._search("cats", 10, 1)
await run._cite([aardvark.id])
@ -374,7 +374,7 @@ class TestCiteFallback:
state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
)
run = await capability.for_run(make_context(deps))
await run._search("cats", limit=10)
await run._search("cats", 10, 1)
with pytest.raises(ModelRetry, match="None of the supplied chunk_ids"):
await run._cite([outside.id])

2
tests/test_picture_in_context.py
View file

@ -1022,7 +1022,7 @@ async def test_rag_capability_attaches_images_for_vision_model(temp_db_path):
capability.state = RAGState()
capability.rag = fake_client
result = await capability._search("anything", None)
result = await capability._search("anything", None, 1)
assert isinstance(result, ToolReturn)
assert result.content is not None

12
tests/tools/test_search.py
View file

@ -311,7 +311,7 @@ class TestBuildImageContentFromResults:
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
]
assert build_image_content_from_results(results) == []
assert build_image_content_from_results(results) == ([], set())
def test_duplicate_document_and_ref_is_attached_once(self):
from pydantic_ai.messages import BinaryContent
@ -336,7 +336,7 @@ class TestBuildImageContentFromResults:
),
]
content = build_image_content_from_results(results)
content, _ = build_image_content_from_results(results)
images = [item for item in content if isinstance(item, BinaryContent)]
assert len(images) == 1
@ -369,7 +369,7 @@ class TestBuildImageContentFromResults:
from haiku.rag.tools.search import build_image_content_from_results
content = build_image_content_from_results(
content, _ = build_image_content_from_results(
self._one_picture_in_two_collections()
)
@ -381,7 +381,7 @@ class TestBuildImageContentFromResults:
reference."""
from haiku.rag.tools.search import build_image_content_from_results
content = build_image_content_from_results(
content, _ = build_image_content_from_results(
self._one_picture_in_two_collections(), include_collection=True
)
@ -392,7 +392,7 @@ class TestBuildImageContentFromResults:
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(
content, _ = build_image_content_from_results(
self._one_picture_in_two_collections()
)
@ -431,7 +431,7 @@ class TestBuildImageContentFromResults:
),
]
content = build_image_content_from_results(results)
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]