Render collection identity only for multi-collection searches

`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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-27 12:41:05 +03:00
parent 746b663a2f
commit 0d7810c78a
No known key found for this signature in database
7 changed files with 123 additions and 18 deletions

View file

@ -214,8 +214,13 @@ lancedb:
A location can be a URI or local path. `databases` and `uri` are mutually
exclusive.
Results, documents, citations, and model context use the configured name as
`source`. Commands such as `info` and path-related errors still show locations.
Results, documents, and citations use the configured name as `source`. Commands
such as `info` and path-related errors still show locations.
Searches spanning multiple databases identify each result with a model-facing
`Collection:` line. Searches over one database omit it. Structured `source`
fields on results, documents, citations, and analysis dictionaries are
unchanged.
Embedding compatibility is checked against two different things.

View file

@ -25,8 +25,14 @@ async def search_corpus(
query, limit=limit, filter=document_filter, sources=sources
)
results = await rag.expand_context(results)
# Named from the selection, not the hits: a search that could have drawn on
# 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))
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)

View file

@ -197,7 +197,11 @@ class SearchResult(BaseModel):
)
def format_for_agent(
self, rank: int | None = None, total: int | None = None
self,
rank: int | None = None,
total: int | None = None,
*,
include_collection: bool = False,
) -> str:
"""Format this search result for inclusion in agent context.
@ -209,8 +213,9 @@ class SearchResult(BaseModel):
the source and nature of the content. When rank is provided, shows
position instead of raw score to avoid confusing LLMs with low RRF scores.
The database is named only where one is named at all, so a single
unnamed database renders exactly as before.
`include_collection` is the caller's decision, not this result's: a
search spanning one collection has nothing to distinguish, whether or
not that collection is named.
"""
if rank is not None and total is not None:
parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"]
@ -219,8 +224,8 @@ class SearchResult(BaseModel):
else:
parts = [f"[{self.chunk_id}] (score: {self.score:.2f})"]
if self.source:
parts.append(f"Database: {self.source}")
if include_collection and self.source:
parts.append(f"Collection: {self.source}")
# Document source info
source_parts = []

View file

@ -160,8 +160,11 @@ def create_search_toolset(
return "No results found."
total = len(results_list)
include_collection = client.covers_multiple
formatted = [
r.format_for_agent(rank=i + 1, total=total)
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)

View file

@ -1,7 +1,10 @@
"""Asking and analyzing across the databases a question covers."""
from unittest.mock import AsyncMock
import pytest
from haiku.rag.capabilities._tools import search_corpus
from haiku.rag.capabilities.rag import RAGState, create_capability
from haiku.rag.client import HaikuRAG
from haiku.rag.sandbox import AnalysisContext, Sandbox
@ -140,23 +143,35 @@ class TestAnalyzeAcrossDatabases:
assert sandbox._context.sources == ["alpha"]
class TestDatabaseIdentityForTheModel:
def test_a_result_names_its_database(self):
"""The model has to attribute and compare evidence by database while it
composes the answer, not only afterwards through the citations."""
class TestCollectionIdentityForTheModel:
"""A collection is named to the model only when the search spans more than
one, and the caller decides that: a result cannot tell from its own fields
whether anything else was searched."""
def test_a_result_names_its_collection_when_asked(self):
"""The model has to attribute and compare evidence by collection while
it composes the answer, not only afterwards through the citations."""
result = SearchResult(content="body", score=0.9, source="alpha", chunk_id="c1")
assert "Database: alpha" in result.format_for_agent()
assert "Collection: alpha" in result.format_for_agent(include_collection=True)
def test_an_unnamed_database_is_not_mentioned(self):
"""A single unnamed database renders as it always has."""
def test_a_named_collection_is_silent_unless_asked(self):
"""One collection has nothing to distinguish, named or not."""
result = SearchResult(content="body", score=0.9, source="alpha", chunk_id="c1")
assert "Collection" not in result.format_for_agent()
def test_an_unnamed_collection_is_never_mentioned(self):
"""Nothing to name, whatever the caller asked for."""
result = SearchResult(content="body", score=0.9, chunk_id="c1")
assert "Database" not in result.format_for_agent()
assert "Collection" not in result.format_for_agent(include_collection=True)
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_in_code_search_names_the_database(self, tmp_path):
async def test_in_code_search_names_the_collection(self, tmp_path):
"""The dictionaries analysis code reads carry `source` whatever the
formatted output renders, since grouping by it is computation."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
@ -184,6 +199,34 @@ class TestDatabaseIdentityForTheModel:
assert result.stdout.count("['alpha', 'beta']") == 2
class TestWhenTheModelIsToldTheCollection:
"""The line is decided by what the search spans, not by whether a name
exists: one collection has nothing to distinguish."""
@pytest.mark.asyncio
async def test_a_search_spanning_a_set_names_every_result(
self, tmp_path, monkeypatch
):
"""Named from the selection, so a result is named even when every hit
came back from one collection: the search could have drawn on both."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
only_alpha = [
SearchResult(content="body", score=0.9, source="alpha", chunk_id="c1")
]
async with HaikuRAG(config=config) as rag:
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
spanning, _ = await search_corpus(rag, "cats")
narrowed, _ = await search_corpus(rag, "cats", sources=["alpha"])
assert "Collection: alpha" in spanning
assert "Collection" not in narrowed
class TestActionableFailures:
@pytest.mark.asyncio
async def test_a_migration_error_survives_being_named(self, tmp_path, temp_db_path):

View file

@ -37,6 +37,49 @@ class TestSearchToolset:
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):
from unittest.mock import AsyncMock
results = [
SearchResult(
content="body",
score=0.9,
source=source,
chunk_id="c1",
document_id="d1",
document_title="Report",
)
]
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
@pytest.mark.vcr()
class TestSearchToolExecution:
"""Tests for search tool execution."""