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.
This commit is contained in:
parent
05c204071d
commit
d9ac221ca0
13 changed files with 434 additions and 32 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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__ = [
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from .exceptions import (
|
||||
AmbiguousCitationError,
|
||||
AmbiguousDatabaseError,
|
||||
ConfigMismatchError,
|
||||
MigrationRequiredError,
|
||||
|
|
@ -12,6 +13,7 @@ __all__ = [
|
|||
"Document",
|
||||
"MigrationRequiredError",
|
||||
"ReadOnlyError",
|
||||
"AmbiguousCitationError",
|
||||
"AmbiguousDatabaseError",
|
||||
"ConfigMismatchError",
|
||||
"SourceUnavailableError",
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Reference in a new issue