Deduplicate search results within one model response
Sibling searches emitted in one response overlap heavily (40.6% of returned chunk slots on Glimmer ORB fan-out cases). A result whose rendered evidence a sibling already showed keeps its rank slot but collapses to a reference line, and a picture attaches once per response keyed on (source, document_id, self_ref). Equivalence is the format_for_agent rendering at neutral rank/total plus picture keys, bucketed under the qualified chunk id, so another database's copy or a different expansion of the same anchor formats in full. Search state now commits only after formatting and image construction succeed: a raising image build no longer leaves results citable that the model never saw, notes evidence for them, or suppresses a later sibling.
This commit is contained in:
parent
ddac328d05
commit
d9f489dcc8
7 changed files with 300 additions and 30 deletions
|
|
@ -7,6 +7,9 @@
|
||||||
- `qa.max_searches` counts search units: searches a model emits in one
|
- `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
|
response share a unit, up to 3 per unit; sequential searches pay one unit
|
||||||
each.
|
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
|
## [0.81.0] - 2026-09-01
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,6 +14,7 @@ from pydantic_ai import (
|
||||||
)
|
)
|
||||||
from pydantic_ai.capabilities import AbstractCapability
|
from pydantic_ai.capabilities import AbstractCapability
|
||||||
from pydantic_ai.messages import (
|
from pydantic_ai.messages import (
|
||||||
|
BinaryContent,
|
||||||
InstructionPart,
|
InstructionPart,
|
||||||
ModelMessage,
|
ModelMessage,
|
||||||
ModelRequest,
|
ModelRequest,
|
||||||
|
|
@ -29,6 +30,7 @@ from pydantic_ai.toolsets import AgentToolset
|
||||||
|
|
||||||
from haiku.rag.capabilities._tools import (
|
from haiku.rag.capabilities._tools import (
|
||||||
CodeExecutionEntry,
|
CodeExecutionEntry,
|
||||||
|
EvidenceKey,
|
||||||
merge_results,
|
merge_results,
|
||||||
search_corpus,
|
search_corpus,
|
||||||
)
|
)
|
||||||
|
|
@ -43,7 +45,7 @@ from haiku.rag.store.models.citation import (
|
||||||
ambiguous_citation,
|
ambiguous_citation,
|
||||||
resolve_citations,
|
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
|
CITATION_GRACE_REQUESTS = 2
|
||||||
"""Requests calling this capability's tools that its cite tool outlives the rest by.
|
"""Requests calling this capability's tools that its cite tool outlives the rest by.
|
||||||
|
|
@ -190,6 +192,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
||||||
"""The run_step whose searches are being priced and deduplicated."""
|
"""The run_step whose searches are being priced and deduplicated."""
|
||||||
step_searches: int = field(default=0, repr=False)
|
step_searches: int = field(default=0, repr=False)
|
||||||
step_rejected: bool = field(default=False, 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)
|
request_count: int = field(default=0, repr=False)
|
||||||
grace_requests_used: int = field(default=0, repr=False)
|
grace_requests_used: int = field(default=0, repr=False)
|
||||||
epoch: int = field(default=0, repr=False)
|
epoch: int = field(default=0, repr=False)
|
||||||
|
|
@ -242,6 +246,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
||||||
search_step=0,
|
search_step=0,
|
||||||
step_searches=0,
|
step_searches=0,
|
||||||
step_rejected=False,
|
step_rejected=False,
|
||||||
|
step_shown=set(),
|
||||||
|
step_pictures=set(),
|
||||||
request_count=0,
|
request_count=0,
|
||||||
grace_requests_used=0,
|
grace_requests_used=0,
|
||||||
epoch=0,
|
epoch=0,
|
||||||
|
|
@ -517,6 +523,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
||||||
self.search_step = run_step
|
self.search_step = run_step
|
||||||
self.step_searches = 0
|
self.step_searches = 0
|
||||||
self.step_rejected = False
|
self.step_rejected = False
|
||||||
|
self.step_shown = set()
|
||||||
|
self.step_pictures = set()
|
||||||
self.step_searches += 1
|
self.step_searches += 1
|
||||||
if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0:
|
if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0:
|
||||||
self.search_count += 1
|
self.search_count += 1
|
||||||
|
|
@ -527,23 +535,34 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
||||||
"the results you already have."
|
"the results you already have."
|
||||||
)
|
)
|
||||||
async with self.rag_lock:
|
async with self.rag_lock:
|
||||||
formatted, results, include_collection = await search_corpus(
|
formatted, results, rendered, include_collection = await search_corpus(
|
||||||
await self._ensure_rag(),
|
await self._ensure_rag(),
|
||||||
query,
|
query,
|
||||||
limit=limit,
|
limit=limit,
|
||||||
document_filter=self.state.document_filter,
|
document_filter=self.state.document_filter,
|
||||||
sources=self.state.sources,
|
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
|
state = self.state
|
||||||
# A model can search the same query twice with different limits, and the
|
# 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.
|
# narrower return must not drop what the wider one already showed it.
|
||||||
merge_results(state.searches.setdefault(query, []), results)
|
merge_results(state.searches.setdefault(query, []), results)
|
||||||
self._note_evidence()
|
self._note_evidence()
|
||||||
if self.vision and (
|
self.step_shown |= rendered
|
||||||
parts := build_image_content_from_results(
|
self.step_pictures |= emitted
|
||||||
results, include_collection=include_collection
|
if parts:
|
||||||
)
|
|
||||||
):
|
|
||||||
return ToolReturn(return_value=formatted, content=parts)
|
return ToolReturn(return_value=formatted, content=parts)
|
||||||
return formatted
|
return formatted
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,11 @@
|
||||||
from collections.abc import Iterable
|
from collections.abc import Iterable
|
||||||
|
from collections.abc import Set as AbstractSet
|
||||||
|
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG
|
from haiku.rag.client import HaikuRAG
|
||||||
from haiku.rag.store.models.chunk import SearchResult, qualified_id
|
from haiku.rag.store.models.chunk import SearchResult, qualified_id
|
||||||
|
from haiku.rag.tools.search import picture_keys
|
||||||
|
|
||||||
|
|
||||||
class CodeExecutionEntry(BaseModel):
|
class CodeExecutionEntry(BaseModel):
|
||||||
|
|
@ -13,14 +15,48 @@ class CodeExecutionEntry(BaseModel):
|
||||||
success: bool = True
|
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(
|
async def search_corpus(
|
||||||
rag: HaikuRAG,
|
rag: HaikuRAG,
|
||||||
query: str,
|
query: str,
|
||||||
limit: int | None = None,
|
limit: int | None = None,
|
||||||
document_filter: str | None = None,
|
document_filter: str | None = None,
|
||||||
sources: list[str] | None = None,
|
sources: list[str] | None = None,
|
||||||
) -> tuple[str, list[SearchResult], bool]:
|
shown: AbstractSet[EvidenceKey] = frozenset(),
|
||||||
"""Search and context-expand results, and whether they name their collection."""
|
) -> 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(
|
results = await rag.search(
|
||||||
query, limit=limit, filter=document_filter, sources=sources
|
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.
|
# two collections names them even when everything came back from one.
|
||||||
selected = rag.source_names if sources is None else sources
|
selected = rag.source_names if sources is None else sources
|
||||||
include_collection = len(set(selected)) > 1
|
include_collection = len(set(selected)) > 1
|
||||||
formatted = "\n\n---\n\n".join(
|
rendered: set[EvidenceKey] = set()
|
||||||
result.format_for_agent(
|
parts: list[str] = []
|
||||||
rank=index + 1, total=len(results), include_collection=include_collection
|
total = len(results)
|
||||||
)
|
for index, result in enumerate(results):
|
||||||
for index, result in enumerate(results)
|
key = evidence_key(result, include_collection)
|
||||||
)
|
if key in shown or key in rendered:
|
||||||
return formatted or "No results found.", list(results), include_collection
|
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(
|
def merge_results(
|
||||||
|
|
@ -56,6 +104,9 @@ def merge_results(
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"CodeExecutionEntry",
|
"CodeExecutionEntry",
|
||||||
|
"EvidenceKey",
|
||||||
|
"evidence_key",
|
||||||
|
"evidence_signature",
|
||||||
"merge_results",
|
"merge_results",
|
||||||
"search_corpus",
|
"search_corpus",
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,6 @@
|
||||||
import base64
|
import base64
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
|
from collections.abc import Set as AbstractSet
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
|
|
||||||
from PIL import Image
|
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:
|
def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
|
||||||
"""Wrap picture bytes for the wire, or return nothing if they will not decode.
|
"""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(
|
def build_image_content_from_results(
|
||||||
results: list[SearchResult],
|
results: list[SearchResult],
|
||||||
include_collection: bool = False,
|
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.
|
"""Decode and validate picture bytes attached to search results, labelled.
|
||||||
|
|
||||||
Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in
|
Returns the labelled content and the ``PictureKey`` of every picture it
|
||||||
different chunks is sent once, and a copy in another collection is its own. Pictures that fail
|
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
|
``PIL.Image.verify()`` are skipped — the model adapter renders one
|
||||||
vision placeholder per ``BinaryContent``, so emitting one for an
|
vision placeholder per ``BinaryContent``, so emitting one for an
|
||||||
image the server can't decode leaves the processor with 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.
|
to the vision API.
|
||||||
"""
|
"""
|
||||||
collected: list[tuple[str | None, str | None, str, BinaryContent]] = []
|
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:
|
for result in results:
|
||||||
if not result.image_data:
|
if not result.image_data:
|
||||||
continue
|
continue
|
||||||
|
|
@ -72,6 +93,7 @@ def build_image_content_from_results(
|
||||||
continue
|
continue
|
||||||
collected.append((result.source, result.chunk_id, self_ref, picture))
|
collected.append((result.source, result.chunk_id, self_ref, picture))
|
||||||
seen.add(key)
|
seen.add(key)
|
||||||
|
emitted.add(key)
|
||||||
|
|
||||||
content: list[str | BinaryContent] = []
|
content: list[str | BinaryContent] = []
|
||||||
total = len(collected)
|
total = len(collected)
|
||||||
|
|
@ -83,7 +105,7 @@ def build_image_content_from_results(
|
||||||
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
||||||
)
|
)
|
||||||
content.append(picture)
|
content.append(picture)
|
||||||
return content
|
return content, emitted
|
||||||
|
|
||||||
|
|
||||||
def create_search_toolset(
|
def create_search_toolset(
|
||||||
|
|
@ -174,7 +196,7 @@ def create_search_toolset(
|
||||||
if not config.qa.model.vision:
|
if not config.qa.model.vision:
|
||||||
return text
|
return text
|
||||||
|
|
||||||
image_content = build_image_content_from_results(
|
image_content, _ = build_image_content_from_results(
|
||||||
results_list, include_collection=include_collection
|
results_list, include_collection=include_collection
|
||||||
)
|
)
|
||||||
if image_content:
|
if image_content:
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,14 @@
|
||||||
|
import base64
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
|
from io import BytesIO
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from PIL import Image as PILImage
|
||||||
from pydantic_ai import Agent
|
from pydantic_ai import Agent
|
||||||
from pydantic_ai.messages import (
|
from pydantic_ai.messages import (
|
||||||
|
BinaryContent,
|
||||||
ModelResponse,
|
ModelResponse,
|
||||||
TextPart,
|
TextPart,
|
||||||
ToolCallPart,
|
ToolCallPart,
|
||||||
|
|
@ -12,8 +17,11 @@ from pydantic_ai.messages import (
|
||||||
from pydantic_ai.models.function import FunctionModel
|
from pydantic_ai.models.function import FunctionModel
|
||||||
from pydantic_ai.run import AgentRunResult
|
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.capabilities.rag import create_capability as create_rag
|
||||||
from haiku.rag.config.models import AppConfig
|
from haiku.rag.config.models import AppConfig
|
||||||
|
from haiku.rag.store.models.chunk import SearchResult
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
@ -154,3 +162,168 @@ async def test_unit_tracking_resets_between_runs(rag_db):
|
||||||
|
|
||||||
assert outcomes(first) == ["ok", "ok", "ok"]
|
assert outcomes(first) == ["ok", "ok", "ok"]
|
||||||
assert outcomes(second)[-3:] == ["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
|
||||||
|
|
|
||||||
|
|
@ -305,8 +305,10 @@ class TestWhenTheModelIsToldTheCollection:
|
||||||
async with HaikuRAG(config=config) as rag:
|
async with HaikuRAG(config=config) as rag:
|
||||||
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
|
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
|
||||||
|
|
||||||
spanning, _, spans = await search_corpus(rag, "cats")
|
spanning, _, _, spans = await search_corpus(rag, "cats")
|
||||||
narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"])
|
narrowed, _, _, narrows = await search_corpus(
|
||||||
|
rag, "cats", sources=["alpha"]
|
||||||
|
)
|
||||||
|
|
||||||
assert "Collection: alpha" in spanning
|
assert "Collection: alpha" in spanning
|
||||||
assert "Collection" not in narrowed
|
assert "Collection" not in narrowed
|
||||||
|
|
|
||||||
|
|
@ -311,7 +311,7 @@ class TestBuildImageContentFromResults:
|
||||||
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
|
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):
|
def test_duplicate_document_and_ref_is_attached_once(self):
|
||||||
from pydantic_ai.messages import BinaryContent
|
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)]
|
images = [item for item in content if isinstance(item, BinaryContent)]
|
||||||
assert len(images) == 1
|
assert len(images) == 1
|
||||||
|
|
@ -369,7 +369,7 @@ class TestBuildImageContentFromResults:
|
||||||
|
|
||||||
from haiku.rag.tools.search import build_image_content_from_results
|
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()
|
self._one_picture_in_two_collections()
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -381,7 +381,7 @@ class TestBuildImageContentFromResults:
|
||||||
reference."""
|
reference."""
|
||||||
from haiku.rag.tools.search import build_image_content_from_results
|
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
|
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):
|
def test_an_unasked_for_collection_is_not_named_on_an_image(self):
|
||||||
from haiku.rag.tools.search import build_image_content_from_results
|
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()
|
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.
|
# 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 [type(item) is str for item in content] == [True, False, True, False]
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue