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:
parent
746b663a2f
commit
0d7810c78a
7 changed files with 123 additions and 18 deletions
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Reference in a new issue