ca2e28559e
Two files of 1,405 and 814 lines become seven: scope resolution, lifecycle, search, documents, expansion, citations and capabilities. `_config`, `_seed` and the rest move to `helpers.py`, importable by the sandbox tests that share them, and the package points VCR back at `tests/cassettes/multi_db/`.
365 lines
13 KiB
Python
365 lines
13 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."""
|
|
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(
|
|
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])
|