87367759f9
The documented rule was existence in more than one selected database. The rule is narrower and does not need to be wider: a copy the search never returned grounded nothing, so the retrieved result is the citation and its database is recorded rather than guessed. The page now states the whole rule, including the fallback's own check for an id no search returned. Two retrieved copies are refused, and so is an id already cited from another database in the conversation. The case between them had no test.
403 lines
15 KiB
Python
403 lines
15 KiB
Python
"""Citing evidence drawn from several databases."""
|
|
|
|
import pytest
|
|
from pydantic_ai import ModelRetry
|
|
|
|
from haiku.rag.capabilities.rag import RAGState, create_capability
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.store.exceptions import (
|
|
AmbiguousCitationError,
|
|
)
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.store.models.citation import Citation, resolve_citations
|
|
from tests.multi_db.helpers import (
|
|
_config,
|
|
_seed,
|
|
)
|
|
|
|
|
|
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, and it resolves rather than raising. Which
|
|
copy supplies the content is `resolve_citations`' own rule, pinned in
|
|
`tests/store/test_citation.py`."""
|
|
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.chunk_id == "c1"
|
|
assert citation.source == "alpha"
|
|
|
|
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"])
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_retrieved_copy_of_a_shared_id_cites_where_it_came_from(
|
|
self, tmp_path
|
|
):
|
|
"""Rejection is about evidence the run holds, not about what the other
|
|
databases contain. A copy the search never returned did not ground the
|
|
answer, so the retrieved one is the citation and its database is not a
|
|
guess. A later question that does retrieve the twin is refused by the
|
|
citation index."""
|
|
from tests.capabilities.test_capabilities import Deps, make_context
|
|
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
capability = create_capability(config=config, defer_loading=False)
|
|
deps = Deps(state={"rag": RAGState().model_dump(mode="json")})
|
|
run = await capability.for_run(make_context(deps))
|
|
assert run.state is not None
|
|
# What a run holds when fusion returned one copy of a shared id and
|
|
# truncated the other: the run never saw the twin.
|
|
run.state.searches["cats"] = [
|
|
SearchResult(
|
|
content="alpha body",
|
|
score=0.9,
|
|
source="alpha",
|
|
chunk_id="c1",
|
|
document_id="d1",
|
|
document_uri="test://alpha/one",
|
|
)
|
|
]
|
|
|
|
await run._cite(["c1"])
|
|
|
|
[cited] = run.state.citation_index.values()
|
|
assert cited.source == "alpha"
|
|
assert cited.content == "alpha body"
|
|
|
|
|
|
class TestCitationSource:
|
|
def test_a_citation_carries_the_result_source(self):
|
|
result = SearchResult(
|
|
content="body",
|
|
score=0.9,
|
|
source="alpha",
|
|
chunk_id="c1",
|
|
document_id="d1",
|
|
document_uri="test://alpha/one",
|
|
)
|
|
|
|
[citation] = resolve_citations(["c1"], [result])
|
|
|
|
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",
|
|
score=0.9,
|
|
chunk_id="c1",
|
|
document_id="d1",
|
|
document_uri="test://one",
|
|
)
|
|
|
|
[citation] = resolve_citations(["c1"], [result])
|
|
|
|
assert citation.source is None
|
|
|
|
|
|
class TestCiteFallback:
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_an_id_from_a_selected_database_resolves_with_its_source(
|
|
self, tmp_path
|
|
):
|
|
"""The fallback exists for a real id this run's searches did not return.
|
|
Across databases it looks through the selected ones and records which
|
|
held it."""
|
|
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 cats"])
|
|
|
|
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)
|
|
deps = Deps(
|
|
state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
|
|
)
|
|
run = await capability.for_run(make_context(deps))
|
|
# The search returns the cats chunk, never the aardvark one.
|
|
await run._search("cats", limit=10)
|
|
|
|
await run._cite([aardvark.id])
|
|
|
|
assert run.state is not None
|
|
[citation] = list(run.state.citation_index.values())
|
|
assert citation.chunk_id == aardvark.id
|
|
assert citation.source == "alpha"
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_an_id_outside_the_selected_databases_does_not_resolve(
|
|
self, tmp_path
|
|
):
|
|
"""A question scoped to one database must not produce a citation from
|
|
another: the fallback looks only where the question looked."""
|
|
from tests.capabilities.test_capabilities import Deps, make_context
|
|
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
await _seed(config, "beta", ["beta document about dogs"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
beta = (await rag.clients_for(["beta"]))[0]
|
|
[outside] = await beta.chunk_repository.list_all(limit=1)
|
|
assert outside.id is not None
|
|
|
|
capability = create_capability(config=config, rag=rag, defer_loading=False)
|
|
deps = Deps(
|
|
state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
|
|
)
|
|
run = await capability.for_run(make_context(deps))
|
|
await run._search("cats", limit=10)
|
|
|
|
with pytest.raises(ModelRetry):
|
|
await run._cite([outside.id])
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_selecting_no_databases_cites_nothing(self, tmp_path):
|
|
"""`sources=[]` selected nothing, which is not the same as everything:
|
|
the fallback must not go looking where the question never looked."""
|
|
from tests.capabilities.test_capabilities import Deps, make_context
|
|
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
await _seed(config, "beta", ["beta document about cats"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
alpha = (await rag.clients_for(["alpha"]))[0]
|
|
[chunk] = await alpha.chunk_repository.list_all(limit=1)
|
|
assert chunk.id is not None
|
|
|
|
capability = create_capability(config=config, rag=rag, defer_loading=False)
|
|
deps = Deps(state={"rag": RAGState(sources=[]).model_dump(mode="json")})
|
|
run = await capability.for_run(make_context(deps))
|
|
|
|
with pytest.raises(ModelRetry):
|
|
await run._cite([chunk.id])
|