From d9ac221ca0a9142030d5aff108a4bf3176b33bd3 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 25 Aug 2026 16:27:34 +0300 Subject: [PATCH] Refuse a chunk id that names a chunk in two databases A chunk id is unique within a database and says nothing across them, so a database copied from another holds the same ids. `qualified_id` keys the two in-memory identity sites on the database and the id together: `merge_results` was dropping the second database's result when a query repeated, and the arrival map that breaks fused score ties was ranking one of the pair as the other. Everything serialized records the id alone, so there ambiguity is refused rather than qualified. `resolve_citations` raises `AmbiguousCitationError` for a cited id held by two of the databases searched, where it used to resolve to whichever result came last; `_register_citations` raises for one already cited from another database in an earlier question. `_cite` turns both into a `ModelRetry` asking for other evidence. The direct-id fallback asks every database the question covers instead of taking the first that answers, so an id no search returned is refused on the same terms. `all_found` collects them and `first_found` reads its first, which document reads keep doing on purpose. Also drop a duplicated 0.77.0 heading from the changelog. --- CHANGELOG.md | 4 +- docs/configuration/storage.md | 6 + .../haiku/rag/capabilities/_base.py | 60 ++++- .../haiku/rag/capabilities/_tools.py | 13 +- haiku_rag_slim/haiku/rag/client/__init__.py | 32 ++- haiku_rag_slim/haiku/rag/client/search.py | 13 +- haiku_rag_slim/haiku/rag/store/__init__.py | 2 + haiku_rag_slim/haiku/rag/store/exceptions.py | 9 + .../haiku/rag/store/models/chunk.py | 15 ++ .../haiku/rag/store/models/citation.py | 32 ++- tests/capabilities/test_capabilities.py | 24 ++ tests/test_multi_db.py | 22 ++ tests/test_multi_db_ask.py | 234 +++++++++++++++++- 13 files changed, 434 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5193a05d..09aa32c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - `api_key` on model and embedding-model config, overriding the provider's environment variable. Honored on the `openai` and `ollama` providers, `vllm` embedders and rerankers, the picture-description VLM endpoint, and `doctor`'s endpoint probes; other providers raise. - `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. `Document.source` names it on a document from a listing or a lookup, so a listing that spans databases says which one each came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured. Databases searched together must have been written with the same embedder; two that disagree raise `ConfigMismatchError`. The query is embedded once for the whole selection. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask`, `analyze` and `chat` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`. -- `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another. +- `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another. A chunk id held by two of the databases searched raises `AmbiguousCitationError`, which reaches the model as a retry; the fallback refuses it too, rather than answering from the first database that holds it. - `client.analyze(..., sources=[...])` analyzes across the selected databases: the sandbox mounts their documents under one flat `/documents/{id}/` namespace, resolving each id to the database holding it, and in-code `search()` covers the same selection. ### Fixed @@ -19,8 +19,6 @@ ## [0.77.0] - 2026-08-21 -## [0.77.0] - 2026-08-21 - ### Added - The four capabilities support Pydantic AI agent specs via `from_spec`, registered with diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 695932e1..f1021f85 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -223,6 +223,12 @@ citation. A document from `list_documents`, `get_document_by_id`, on the command line points the configuration at it, so commands that work on one database report no name. +Chunk ids are unique within a database and say nothing across them, so a database +copied from another holds the same ids. Results are told apart by the database +and the id together. A chunk id held by two of the databases searched cannot be +cited: `resolve_citations` raises `AmbiguousCitationError`, and the capability +asks the model for other evidence instead. + **Configure a reranker when searching several databases.** Reciprocal rank fusion compares ranks, not scores, so every database contributes its own best matches whether or not they are relevant to the question, and results from databases diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index ac60393b..76ba1bf8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -33,10 +33,15 @@ from haiku.rag.capabilities._tools import ( search_corpus, ) from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef -from haiku.rag.client import HaikuRAG, first_found +from haiku.rag.client import HaikuRAG, all_found from haiku.rag.config.models import AppConfig +from haiku.rag.store.exceptions import AmbiguousCitationError from haiku.rag.store.models.chunk import SearchResult -from haiku.rag.store.models.citation import Citation, resolve_citations +from haiku.rag.store.models.citation import ( + Citation, + ambiguous_citation, + resolve_citations, +) from haiku.rag.tools.search import build_image_content_from_results CITATION_GRACE_REQUESTS = 2 @@ -56,6 +61,18 @@ duplicating a character or a whole group stays above 0.75, so the gap is wide. """ +def _ambiguous_retry(error: AmbiguousCitationError) -> ModelRetry: + """The only way out of an id that names a chunk in two databases. + + The model cannot say which it meant, so it is asked for other evidence + rather than for the same id again. + """ + return ModelRetry( + f"{error}. Cite a chunk id that appears once across the databases " + "searched, or cite nothing." + ) + + def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str: """Recover a chunk id the model damaged while transcribing it. @@ -521,7 +538,10 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): all_results.extend(results) known_ids = [result.chunk_id for result in all_results if result.chunk_id] requested = [_nearest_known_id(cid.strip("[]"), known_ids) for cid in chunk_ids] - citations = resolve_citations(requested, all_results) + try: + citations = resolve_citations(requested, all_results) + except AmbiguousCitationError as error: + raise _ambiguous_retry(error) from error resolved = {citation.chunk_id for citation in citations} missing = [cid for cid in requested if cid not in resolved] @@ -535,12 +555,21 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): synthetic: list[SearchResult] = [] documents: dict[tuple[str | None, str], Any] = {} for chunk_id in missing: - found = await first_found( + # Every database that has it, not the first: taking the first + # would attribute the answer to one of them without ever + # seeing that another had it too. + holders = await all_found( lookups, lambda owner: owner.get_chunk_by_id(chunk_id) ) - if found is None: + if len(holders) > 1: + raise _ambiguous_retry( + ambiguous_citation( + chunk_id, [owner.source for owner, _ in holders] + ) + ) + if not holders: continue - owner, chunk = found + [(owner, chunk)] = holders if not chunk.document_id: continue key = (owner.source, chunk.document_id) @@ -562,7 +591,10 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): f"None of the supplied chunk_ids {list(chunk_ids)} could be resolved. " "Copy chunk_ids verbatim from search results." ) - self._register_citations(citations) + try: + self._register_citations(citations) + except AmbiguousCitationError as error: + raise _ambiguous_retry(error) from error self._declare(citations) resolved = {citation.chunk_id for citation in citations} unresolved = [cid for cid in missing if cid not in resolved] @@ -579,8 +611,22 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): return f"Registered {len(citations)} citation(s)." def _register_citations(self, citations: list[Citation]) -> None: + """Index the citations, numbering the ones not already registered. + + The index is keyed by chunk id and outlives the question, so an id + already registered from another database is refused here rather than + silently keeping the earlier one's content and database. + """ assert self.state is not None state = self.state + for citation in citations: + held = state.citation_index.get(citation.chunk_id) + if held is not None and held.source != citation.source: + raise AmbiguousCitationError( + f"chunk id {citation.chunk_id} was already cited from " + "another database in this conversation; a citation records " + "the id alone and cannot say which" + ) next_index = len(state.citation_index) + 1 for citation in citations: if citation.chunk_id not in state.citation_index: diff --git a/haiku_rag_slim/haiku/rag/capabilities/_tools.py b/haiku_rag_slim/haiku/rag/capabilities/_tools.py index a732157c..314e73db 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_tools.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_tools.py @@ -3,7 +3,7 @@ from collections.abc import Iterable from pydantic import BaseModel from haiku.rag.client import HaikuRAG -from haiku.rag.store.models.chunk import SearchResult +from haiku.rag.store.models.chunk import SearchResult, qualified_id class CodeExecutionEntry(BaseModel): @@ -37,14 +37,15 @@ def merge_results( ) -> None: """Add the results not already held. - Identity is the chunk id, which every stored chunk carries; results built by - hand without one cannot be told apart and collapse to the first. + Identity is the database and the chunk id: results built by hand carry + neither and cannot be told apart, so they collapse to the first. """ - seen = {result.chunk_id for result in existing} + seen = {qualified_id(result.source, result.chunk_id) for result in existing} for result in incoming: - if result.chunk_id not in seen: + key = qualified_id(result.source, result.chunk_id) + if key not in seen: existing.append(result) - seen.add(result.chunk_id) + seen.add(key) __all__ = [ diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 53ab292c..e5e2730b 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -58,22 +58,36 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +async def all_found( + clients: "list[HaikuRAG]", + lookup: "Callable[[HaikuRAG], Coroutine[Any, Any, Any]]", +) -> "list[tuple[HaikuRAG, Any]]": + """Every client for which `lookup` finds something, and what each found. + + An id or a URI says nothing about which database holds it, so every one is + asked at once. Asking in turn would cost a round trip per database for an + identifier that is missing or held by the last of them. + """ + found_by_client = await asyncio.gather(*(lookup(client) for client in clients)) + return [ + (client, found) + for client, found in zip(clients, found_by_client, strict=True) + if found is not None + ] + + async def first_found( clients: "list[HaikuRAG]", lookup: "Callable[[HaikuRAG], Coroutine[Any, Any, Any]]", ) -> "tuple[HaikuRAG, Any] | None": """The first of `clients` for which `lookup` finds something, and what it found. - An id or a URI says nothing about which database holds it, so every one is - asked at once and the first that has it, in the order given, answers. Asking - in turn would cost a round trip per database for an identifier that is - missing or held by the last of them. + For a lookup that has an answer wherever it is found, as a document read + does. Where holding the same identifier in two databases means something, + ask `all_found` and decide. """ - found_by_client = await asyncio.gather(*(lookup(client) for client in clients)) - for client, found in zip(clients, found_by_client, strict=True): - if found is not None: - return client, found - return None + found = await all_found(clients, lookup) + return found[0] if found else None def _spell(embedding: tuple[str | None, str | None, int | None]) -> str: diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 0ccabf52..34b45387 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -3,7 +3,12 @@ import base64 from collections.abc import Sequence from typing import TYPE_CHECKING -from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType +from haiku.rag.store.models.chunk import ( + Chunk, + SearchResult, + SearchType, + qualified_id, +) from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX if TYPE_CHECKING: @@ -429,7 +434,7 @@ async def expand_sources( # Grouping by database must not become the tiebreak: fused scores tie often, # so equal scores keep the order they were fused in. arrival = { - result.chunk_id: rank + qualified_id(result.source, result.chunk_id): rank for rank, result in enumerate(search_results) if result.chunk_id } @@ -437,9 +442,9 @@ async def expand_sources( def fused_rank(result: SearchResult) -> int: return min( ( - arrival[cid] + arrival[key] for cid in (result.chunk_id, *result.chunk_ids) - if cid in arrival + if (key := qualified_id(result.source, cid)) in arrival ), default=len(arrival), ) diff --git a/haiku_rag_slim/haiku/rag/store/__init__.py b/haiku_rag_slim/haiku/rag/store/__init__.py index 81521f90..e6126a10 100644 --- a/haiku_rag_slim/haiku/rag/store/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/__init__.py @@ -1,4 +1,5 @@ from .exceptions import ( + AmbiguousCitationError, AmbiguousDatabaseError, ConfigMismatchError, MigrationRequiredError, @@ -12,6 +13,7 @@ __all__ = [ "Document", "MigrationRequiredError", "ReadOnlyError", + "AmbiguousCitationError", "AmbiguousDatabaseError", "ConfigMismatchError", "SourceUnavailableError", diff --git a/haiku_rag_slim/haiku/rag/store/exceptions.py b/haiku_rag_slim/haiku/rag/store/exceptions.py index 26b256f7..81cc4ae2 100644 --- a/haiku_rag_slim/haiku/rag/store/exceptions.py +++ b/haiku_rag_slim/haiku/rag/store/exceptions.py @@ -24,6 +24,15 @@ class AmbiguousDatabaseError(Exception): """ +class AmbiguousCitationError(Exception): + """A cited chunk id names a chunk in more than one database. + + A citation records the id alone, so nothing downstream can say which + database it came from. Raised rather than resolved: picking one attributes + the answer to a database it may not have come from. + """ + + class SourceUnavailableError(Exception): """A configured database could not be opened. diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 8f37c53f..c363d64f 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -116,6 +116,21 @@ class Chunk(BaseModel): SearchType = Literal["vector", "fts", "hybrid"] +def qualified_id(source: str | None, id: str | None) -> tuple[str | None, str | None]: + """What tells one chunk from another. + + A chunk id is unique within a database and says nothing across them: a + database copied from another holds the same ids, so the database is part of + the identity. `id` is optional because a result built by hand carries none, + and those cannot be told apart. + + Only for structures held in memory. Everything serialized — citations, + evidence refs, compaction — records the id alone, so ambiguity there is + rejected rather than qualified. + """ + return (source, id) + + class SearchResult(BaseModel): """Search result with optional provenance information for citations. diff --git a/haiku_rag_slim/haiku/rag/store/models/citation.py b/haiku_rag_slim/haiku/rag/store/models/citation.py index 3c4c531e..10a834e3 100644 --- a/haiku_rag_slim/haiku/rag/store/models/citation.py +++ b/haiku_rag_slim/haiku/rag/store/models/citation.py @@ -1,7 +1,9 @@ +from collections.abc import Iterable from typing import TYPE_CHECKING from pydantic import BaseModel, Field +from haiku.rag.store.exceptions import AmbiguousCitationError from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX if TYPE_CHECKING: @@ -55,16 +57,42 @@ class Citation(BaseModel): picture_refs: list[str] = Field(default_factory=list) +def ambiguous_citation( + chunk_id: str, sources: Iterable[str | None] +) -> AmbiguousCitationError: + """The refusal for an id that names a chunk in more than one database.""" + named = ", ".join(sorted(s or "unnamed" for s in sources)) + return AmbiguousCitationError( + f"chunk id {chunk_id} names a chunk in more than one database " + f"({named}); a citation records the id alone and cannot say which" + ) + + def resolve_citations( cited_chunk_ids: list[str], search_results: "list[SearchResult]", ) -> list[Citation]: - """Resolve chunk IDs to full Citation objects with metadata.""" - by_id = {r.chunk_id: r for r in search_results if r.chunk_id} + """Resolve chunk IDs to full Citation objects with metadata. + + Raises ``AmbiguousCitationError`` when a cited id names a chunk in more than + one of the databases searched. An id held by two of them, as after copying a + database, resolves to whichever result came last, attributing the answer to a + database it may not have come from. + """ + by_id: dict[str, SearchResult] = {} + ambiguous: dict[str, set[str | None]] = {} + for r in search_results: + # A result built by hand carries no id and nothing can cite it. + if cid := r.chunk_id: + held = by_id.setdefault(cid, r) + if held.source != r.source: + ambiguous.setdefault(cid, {held.source}).add(r.source) citations = [] for raw_id in cited_chunk_ids: chunk_id = raw_id.strip("[]") + if chunk_id in ambiguous: + raise ambiguous_citation(chunk_id, ambiguous[chunk_id]) r = by_id.get(chunk_id) if not r: continue diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 746df6cc..e11c51f8 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -396,6 +396,30 @@ async def test_a_narrower_repeat_keeps_what_the_wider_search_returned(temp_db_pa ] +@pytest.mark.asyncio +async def test_two_databases_holding_one_chunk_id_both_survive(temp_db_path): + """A database copied from another holds the same chunk ids, so what tells + two results apart is the database and the id together.""" + capability = create_rag(db_path=temp_db_path, config=AppConfig()) + capability.state = RAGState() + capability.borrowed_rag = _stub_client( + [SearchResult(content="alpha", score=1.0, chunk_id="c1", source="alpha")], + [ + SearchResult(content="alpha", score=1.0, chunk_id="c1", source="alpha"), + SearchResult(content="beta", score=0.9, chunk_id="c1", source="beta"), + ], + ) + + await capability._search("cats", 20) + await capability._search("cats", None) + + stored = capability.state.searches["cats"] + assert [(r.source, r.chunk_id) for r in stored] == [ + ("alpha", "c1"), + ("beta", "c1"), + ] + + @pytest.mark.asyncio async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index a856b3e6..ae02c964 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -328,6 +328,28 @@ class TestLookupByIdentifier: assert found is not None and found.content == "beta one" + @pytest.mark.asyncio + async def test_a_document_held_by_two_databases_answers_from_the_first( + self, tmp_path + ): + """A database copied from another holds the same ids. A read has an + answer wherever it finds one, and which one it is has to be the + configured order rather than whichever replied first.""" + import shutil + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "beta.lancedb") + + async with HaikuRAG(config=config) as rag: + beta = (await rag.clients_for(["beta"]))[0] + [target] = await beta.document_repository.list_all(limit=1) + assert target.id is not None + + found = await rag.get_document_by_id(target.id) + + assert found is not None and found.source == "alpha" + @pytest.mark.asyncio async def test_an_unknown_identifier_is_absent_rather_than_an_error(self, tmp_path): config = _config(tmp_path, ["alpha", "beta"]) diff --git a/tests/test_multi_db_ask.py b/tests/test_multi_db_ask.py index d48b69b9..fec7aee3 100644 --- a/tests/test_multi_db_ask.py +++ b/tests/test_multi_db_ask.py @@ -6,8 +6,9 @@ from pydantic_ai import ModelRetry from haiku.rag.capabilities.rag import RAGState, create_capability from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config +from haiku.rag.store.exceptions import AmbiguousCitationError from haiku.rag.store.models import Chunk, Document, DocumentItem, SearchResult -from haiku.rag.store.models.citation import resolve_citations +from haiku.rag.store.models.citation import Citation, resolve_citations from tests.test_multi_db import _config, _seed @@ -111,6 +112,220 @@ class TestExpansionRouting: assert [r.chunk_id for r in expanded] == [r.chunk_id for r in fused] +class TestSharedChunkIds: + """A database copied from another holds the same chunk ids.""" + + @pytest.mark.asyncio + async def test_a_shared_id_does_not_confuse_the_fused_order(self, tmp_path): + """Arrival order breaks score ties, so it has to tell two databases' + identically-numbered chunks apart.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one about cats"]) + await _seed(config, "beta", ["beta one about cats"]) + fused = [ + SearchResult(content="a0", score=0.5, chunk_id="a0", source="alpha"), + SearchResult(content="beta", score=0.5, chunk_id="shared", source="beta"), + SearchResult(content="alpha", score=0.5, chunk_id="shared", source="alpha"), + ] + + async with HaikuRAG(config=config) as rag: + expanded = await rag.expand_context(fused) + + assert [(r.source, r.chunk_id) for r in expanded] == [ + (r.source, r.chunk_id) for r in fused + ] + + def test_a_shared_id_cannot_be_cited(self): + """A citation records the id alone, so resolving one held by two + databases would attribute the answer to whichever came last.""" + results = [ + SearchResult( + content="alpha body", + score=0.9, + source="alpha", + chunk_id="c1", + document_id="d1", + document_uri="test://alpha/one", + ), + SearchResult( + content="beta body", + score=0.8, + source="beta", + chunk_id="c1", + document_id="d1", + document_uri="test://beta/one", + ), + ] + + with pytest.raises(AmbiguousCitationError, match="c1"): + resolve_citations(["c1"], results) + + def test_a_repeated_id_from_one_database_still_collapses(self): + """One database cannot hold two chunks under one id, so seeing it twice + is the same chunk seen twice.""" + results = [ + SearchResult( + content="first", + score=0.9, + source="alpha", + chunk_id="c1", + document_id="d1", + document_uri="test://alpha/one", + ), + SearchResult( + content="second", + score=0.8, + source="alpha", + chunk_id="c1", + document_id="d1", + document_uri="test://alpha/one", + ), + ] + + [citation] = resolve_citations(["c1"], results) + + assert citation.content == "first" + + def test_only_a_cited_id_has_to_be_unambiguous(self): + """An id the answer never cites attributes nothing.""" + shared = [ + SearchResult( + content=f"{name} body", + score=0.9, + source=name, + chunk_id="c1", + document_id="d1", + document_uri=f"test://{name}/one", + ) + for name in ("alpha", "beta") + ] + own = SearchResult( + content="alpha only", + score=0.7, + source="alpha", + chunk_id="c2", + document_id="d2", + document_uri="test://alpha/two", + ) + + [citation] = resolve_citations(["c2"], [*shared, own]) + + assert citation.source == "alpha" + + @pytest.mark.asyncio + async def test_an_unsearched_shared_id_is_refused_by_the_fallback(self, tmp_path): + """The direct lookup is the only place a collision shows for an id no + search returned, so it has to ask every database rather than take the + first that answers.""" + import shutil + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed( + config, "alpha", ["alpha document about cats", "alpha on aardvarks"] + ) + shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "beta.lancedb") + + async with HaikuRAG(config=config) as rag: + alpha = (await rag.clients_for(["alpha"]))[0] + chunks = await alpha.chunk_repository.list_all() + [aardvark] = [c for c in chunks if "aardvark" in c.content] + assert aardvark.id is not None + + capability = create_capability(config=config, rag=rag, defer_loading=False) + capability.state = RAGState() + + # No search ran, so the id can only resolve through the fallback. + with pytest.raises(ModelRetry, match="more than one database"): + await capability._cite([aardvark.id]) + + @pytest.mark.asyncio + async def test_an_unsearched_id_in_one_database_still_resolves(self, tmp_path): + """The refusal is for a collision, not for looking through several + databases: an id only one of them holds still resolves.""" + from tests.capabilities.test_capabilities import Deps, make_context + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed( + config, "alpha", ["alpha document about cats", "alpha on aardvarks"] + ) + await _seed(config, "beta", ["beta document about dogs"]) + + async with HaikuRAG(config=config) as rag: + alpha = (await rag.clients_for(["alpha"]))[0] + chunks = await alpha.chunk_repository.list_all() + [aardvark] = [c for c in chunks if "aardvark" in c.content] + assert aardvark.id is not None + + capability = create_capability(config=config, rag=rag, defer_loading=False) + run = await capability.for_run(make_context(Deps())) + + await run._cite([aardvark.id]) + + assert run.state is not None + [citation] = list(run.state.citation_index.values()) + assert citation.source == "alpha" + + @pytest.mark.asyncio + async def test_cite_asks_for_other_evidence(self, tmp_path): + capability = create_capability( + config=_config(tmp_path, ["alpha", "beta"]), defer_loading=False + ) + capability.state = RAGState( + searches={ + "cats": [ + SearchResult( + content=f"{name} body", + score=0.9, + source=name, + chunk_id="c1", + document_id="d1", + document_uri=f"test://{name}/one", + ) + for name in ("alpha", "beta") + ] + } + ) + + with pytest.raises(ModelRetry, match="appears once"): + await capability._cite(["c1"]) + + @pytest.mark.asyncio + async def test_cite_refuses_an_id_already_cited_from_another_database( + self, tmp_path + ): + """The citation index outlives the question, so the collision can arrive + a turn later than the search that would have shown it.""" + capability = create_capability( + config=_config(tmp_path, ["alpha", "beta"]), defer_loading=False + ) + capability.state = RAGState( + citation_index={ + "c1": Citation( + document_id="d1", + source="alpha", + chunk_id="c1", + document_uri="test://alpha/one", + content="alpha body", + ) + }, + searches={ + "cats": [ + SearchResult( + content="beta body", + score=0.9, + source="beta", + chunk_id="c1", + document_id="d1", + document_uri="test://beta/one", + ) + ] + }, + ) + + with pytest.raises(ModelRetry, match="another database"): + await capability._cite(["c1"]) + + class TestCitationSource: def test_a_citation_carries_the_result_source(self): result = SearchResult( @@ -126,6 +341,23 @@ class TestCitationSource: assert citation.source == "alpha" + def test_a_result_without_an_id_is_skipped(self): + """A result built by hand carries no chunk id, so nothing can cite it + and it takes part in no collision.""" + handmade = SearchResult(content="loose text", score=0.5) + real = SearchResult( + content="body", + score=0.9, + source="alpha", + chunk_id="c1", + document_id="d1", + document_uri="test://alpha/one", + ) + + [citation] = resolve_citations(["c1"], [handmade, real]) + + assert citation.chunk_id == "c1" + def test_a_single_database_citation_has_no_source(self): result = SearchResult( content="body",