From 0d7810c78a05d576f2518a98e31e533c6d598085 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 27 Aug 2026 12:41:05 +0300 Subject: [PATCH] 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. --- docs/configuration/storage.md | 9 ++- .../haiku/rag/capabilities/_tools.py | 8 ++- .../haiku/rag/store/models/chunk.py | 15 +++-- haiku_rag_slim/haiku/rag/tools/search.py | 5 +- ..._in_code_search_names_the_collection.yaml} | 0 tests/multi_db/test_capabilities.py | 61 ++++++++++++++++--- tests/tools/test_search.py | 43 +++++++++++++ 7 files changed, 123 insertions(+), 18 deletions(-) rename tests/cassettes/multi_db/test_capabilities/{TestDatabaseIdentityForTheModel.test_in_code_search_names_the_database.yaml => TestCollectionIdentityForTheModel.test_in_code_search_names_the_collection.yaml} (100%) diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index d4b895d9..87d63e4c 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_tools.py b/haiku_rag_slim/haiku/rag/capabilities/_tools.py index 314e73db..3d636988 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_tools.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_tools.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 9ae77f65..e04fbb55 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -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 = [] diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 4f192a74..ccf51c82 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -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) diff --git a/tests/cassettes/multi_db/test_capabilities/TestDatabaseIdentityForTheModel.test_in_code_search_names_the_database.yaml b/tests/cassettes/multi_db/test_capabilities/TestCollectionIdentityForTheModel.test_in_code_search_names_the_collection.yaml similarity index 100% rename from tests/cassettes/multi_db/test_capabilities/TestDatabaseIdentityForTheModel.test_in_code_search_names_the_database.yaml rename to tests/cassettes/multi_db/test_capabilities/TestCollectionIdentityForTheModel.test_in_code_search_names_the_collection.yaml diff --git a/tests/multi_db/test_capabilities.py b/tests/multi_db/test_capabilities.py index aefa2274..ca9c375d 100644 --- a/tests/multi_db/test_capabilities.py +++ b/tests/multi_db/test_capabilities.py @@ -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): diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index eb3efa99..38bb3fe4 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -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."""