d9ac221ca0
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.
814 lines
31 KiB
Python
814 lines
31 KiB
Python
import pytest
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
from docling_core.types.doc.labels import DocItemLabel
|
|
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 Citation, resolve_citations
|
|
from tests.test_multi_db import _config, _seed
|
|
|
|
|
|
async def _seed_expandable(config, name, sentences):
|
|
"""One document whose chunk covers a single item, so expansion has
|
|
neighbours to pull in and rebuilds the result rather than passing it
|
|
through."""
|
|
dim = get_config().embeddings.model.vector_dim
|
|
doc = DoclingDocument(name=name)
|
|
for sentence in sentences:
|
|
doc.add_text(label=DocItemLabel.TEXT, text=sentence)
|
|
async with HaikuRAG(config=config, create=True, sources=[name]) as rag:
|
|
await rag.import_document(
|
|
doc,
|
|
[
|
|
Chunk(
|
|
content=sentences[0],
|
|
embedding=[0.1] * dim,
|
|
order=0,
|
|
metadata={"doc_item_refs": ["#/texts/0"]},
|
|
)
|
|
],
|
|
uri=f"test://{name}/expandable",
|
|
)
|
|
|
|
|
|
class TestExpansionRouting:
|
|
@pytest.mark.asyncio
|
|
async def test_expansion_routes_each_result_to_its_database(self, tmp_path):
|
|
"""A federating client has no repositories of its own, so expansion has
|
|
to go through the database each result came from."""
|
|
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:
|
|
results = await rag.search("cats", search_type="fts", limit=10)
|
|
expanded = await rag.expand_context(results)
|
|
|
|
assert {r.source for r in expanded} == {"alpha", "beta"}
|
|
for r in expanded:
|
|
assert r.source is not None
|
|
assert r.source in r.content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_an_expanded_result_keeps_its_source(self, tmp_path):
|
|
"""Expansion rebuilds the result, and the rebuilt one has to name the
|
|
database it was expanded through."""
|
|
config = _config(tmp_path, ["alpha"])
|
|
await _seed_expandable(
|
|
config, "alpha", ["cats sleep often", "cats also hunt", "cats purr"]
|
|
)
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
results = await rag.search("cats", search_type="fts", limit=10)
|
|
expanded = await rag.expand_context(results)
|
|
|
|
assert len(expanded) == 1
|
|
assert "cats also hunt" in expanded[0].content, "expansion did not run"
|
|
assert expanded[0].source == "alpha"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_federated_result_is_expanded_by_its_own_database(self, tmp_path):
|
|
"""Routing is not enough: each result has to come back carrying the
|
|
neighbours of the database it was expanded through, and only those."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed_expandable(
|
|
config, "alpha", ["cats sleep often", "alpha follows on"]
|
|
)
|
|
await _seed_expandable(config, "beta", ["cats also hunt", "beta follows on"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
results = await rag.search("cats", search_type="fts", limit=10)
|
|
expanded = await rag.expand_context(results)
|
|
|
|
content = {r.source: r.content for r in expanded}
|
|
assert "alpha follows on" in content["alpha"]
|
|
assert "beta follows on" not in content["alpha"]
|
|
assert "beta follows on" in content["beta"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expansion_keeps_tied_results_in_fused_order(self, tmp_path):
|
|
"""Fused scores tie often, so grouping by database must not reorder
|
|
them: the tiebreak is the order they arrived in."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha one about cats", "alpha two about cats"])
|
|
await _seed(config, "beta", ["beta one about cats"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
found = await rag.search("cats", search_type="fts", limit=10)
|
|
by_source: dict[str, list[SearchResult]] = {}
|
|
for result in found:
|
|
by_source.setdefault(result.source or "", []).append(result)
|
|
# Interleaved, so grouping by database is visible as a reordering.
|
|
fused = [by_source["alpha"][0], by_source["beta"][0], by_source["alpha"][1]]
|
|
for result in fused:
|
|
result.score = 0.5
|
|
|
|
expanded = await rag.expand_context(fused)
|
|
|
|
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(
|
|
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 TestAskAcrossDatabases:
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_the_capability_searches_the_selected_databases(self, tmp_path):
|
|
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:
|
|
capability = create_capability(config=config, rag=rag, defer_loading=False)
|
|
capability.state = RAGState(sources=["alpha"])
|
|
|
|
formatted = await capability._search("cats", limit=10)
|
|
|
|
assert isinstance(formatted, str)
|
|
assert "alpha" in formatted
|
|
assert "beta document" not in formatted
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_searching_all_databases_reaches_both(self, tmp_path):
|
|
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:
|
|
capability = create_capability(config=config, rag=rag, defer_loading=False)
|
|
capability.state = RAGState()
|
|
|
|
formatted = await capability._search("cats", limit=10)
|
|
|
|
assert isinstance(formatted, str)
|
|
assert "alpha document" in formatted
|
|
assert "beta document" in formatted
|
|
|
|
|
|
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])
|
|
|
|
|
|
class TestStandaloneCapabilities:
|
|
"""A capability nobody hands a client opens its own. It has to reach the
|
|
configured set, or a host that only registers capabilities gets one
|
|
database while the configuration names several."""
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_a_rag_capability_opens_the_configured_set(self, tmp_path):
|
|
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"])
|
|
|
|
capability = create_capability(config=config, defer_loading=False)
|
|
assert capability.db_path is None
|
|
run = await capability.for_run(make_context(Deps()))
|
|
try:
|
|
formatted = await run._search("cats", limit=10)
|
|
finally:
|
|
await run._close()
|
|
|
|
assert isinstance(formatted, str)
|
|
assert "alpha document" in formatted
|
|
assert "beta document" in formatted
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_an_analysis_capability_mounts_the_configured_set(self, tmp_path):
|
|
from haiku.rag.capabilities.analysis import (
|
|
create_capability as create_analysis,
|
|
)
|
|
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"])
|
|
|
|
capability = create_analysis(config=config, defer_loading=False)
|
|
run = await capability.for_run(make_context(Deps()))
|
|
try:
|
|
sandbox = await run._ensure_sandbox()
|
|
docs, owners = await sandbox._documents()
|
|
finally:
|
|
await run._close()
|
|
|
|
assert len(docs) == 2
|
|
assert {owner.source for owner in owners.values()} == {"alpha", "beta"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_single_configured_database_is_still_opened(self, tmp_path):
|
|
"""One named database is a set of one, not a path to guess."""
|
|
config = _config(tmp_path, ["alpha"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
|
|
capability = create_capability(config=config, defer_loading=False)
|
|
rag = await capability._ensure_rag()
|
|
try:
|
|
assert rag.source == "alpha"
|
|
finally:
|
|
await capability._close()
|
|
|
|
|
|
class TestAnalyzeAcrossDatabases:
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_the_capability_searches_the_selected_databases(self, tmp_path):
|
|
"""`analysis_search` is the same tool as the RAG one, and the sandbox is
|
|
scoped by the same selection."""
|
|
from haiku.rag.capabilities.analysis import AnalysisState
|
|
from haiku.rag.capabilities.analysis import (
|
|
create_capability as create_analysis,
|
|
)
|
|
|
|
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:
|
|
capability = create_analysis(config=config, rag=rag, defer_loading=False)
|
|
capability.state = AnalysisState(sources=["alpha"])
|
|
|
|
formatted = await capability._search("cats", limit=10)
|
|
sandbox = await capability._ensure_sandbox()
|
|
await capability._close()
|
|
|
|
assert isinstance(formatted, str)
|
|
assert "alpha document" in formatted
|
|
assert "beta document" not in formatted
|
|
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."""
|
|
result = SearchResult(content="body", score=0.9, source="alpha", chunk_id="c1")
|
|
|
|
assert "Database: alpha" in result.format_for_agent()
|
|
|
|
def test_an_unnamed_database_is_not_mentioned(self):
|
|
"""A single unnamed database renders as it always has."""
|
|
result = SearchResult(content="body", score=0.9, chunk_id="c1")
|
|
|
|
assert "Database" not in result.format_for_agent()
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.vcr()
|
|
async def test_in_code_search_names_the_database(self, tmp_path):
|
|
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
|
|
|
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:
|
|
sandbox = Sandbox(
|
|
db_path=None,
|
|
config=config,
|
|
context=AnalysisContext(),
|
|
rag=rag,
|
|
)
|
|
try:
|
|
result = await sandbox.execute(
|
|
"rows = await search('cats', limit=10)\n"
|
|
"print(sorted(r['source'] for r in rows))\n"
|
|
"docs = await list_documents()\n"
|
|
"print(sorted(d['source'] for d in docs))"
|
|
)
|
|
finally:
|
|
await sandbox.close()
|
|
|
|
assert result.success, result.stderr
|
|
assert "['alpha', 'beta']" in result.stdout
|
|
assert result.stdout.count("['alpha', 'beta']") == 2
|
|
|
|
|
|
class TestActionableFailures:
|
|
@pytest.mark.asyncio
|
|
async def test_a_migration_error_survives_being_named(self, tmp_path, temp_db_path):
|
|
"""The remedy is the whole value of the message, and it names no location,
|
|
so it is not replaced by the database's name."""
|
|
from haiku.rag.store.exceptions import MigrationRequiredError
|
|
|
|
config = _config(tmp_path, ["alpha"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
|
|
async with HaikuRAG(config=config, sources=["alpha"]) as rag:
|
|
await rag.store.set_haiku_version("0.20.0")
|
|
|
|
with pytest.raises(MigrationRequiredError) as raised:
|
|
async with HaikuRAG(config=config, sources=["alpha"]):
|
|
pass
|
|
|
|
# Both halves: which database failed, and what to run about it.
|
|
assert "haiku-rag migrate" in str(raised.value)
|
|
assert "alpha" in str(raised.value)
|
|
assert str(tmp_path) not in str(raised.value)
|
|
|
|
|
|
class TestPictureDeduplication:
|
|
"""One picture yields two chunks — a text-embedded one and an image-embedded
|
|
one — that collapse to the best. Two databases holding the same picture are
|
|
two results, not a duplicate."""
|
|
|
|
@staticmethod
|
|
def _picture(source, score):
|
|
return SearchResult(
|
|
content="a figure",
|
|
score=score,
|
|
source=source,
|
|
chunk_id=f"{source}-c",
|
|
document_id="doc-1",
|
|
doc_item_refs=["#/pictures/0"],
|
|
)
|
|
|
|
def test_the_same_picture_in_two_databases_survives(self):
|
|
from haiku.rag.client.search import _dedup_picture_chunks
|
|
|
|
kept = _dedup_picture_chunks(
|
|
[self._picture("alpha", 0.9), self._picture("clone", 0.5)]
|
|
)
|
|
|
|
assert [r.source for r in kept] == ["alpha", "clone"]
|
|
|
|
def test_duplicates_within_one_database_still_collapse(self):
|
|
from haiku.rag.client.search import _dedup_picture_chunks
|
|
|
|
lower = self._picture("alpha", 0.5)
|
|
higher = self._picture("alpha", 0.9)
|
|
|
|
kept = _dedup_picture_chunks([lower, higher])
|
|
|
|
assert kept == [higher]
|
|
|
|
|
|
class TestPictureRouting:
|
|
@pytest.mark.asyncio
|
|
async def test_a_picture_is_fetched_from_the_database_that_holds_it(self, tmp_path):
|
|
"""A `self_ref` repeats across databases, so the citation's source is
|
|
what decides where the bytes come from."""
|
|
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:
|
|
beta = (await rag.clients_for(["beta"]))[0]
|
|
[document] = await beta.document_repository.list_all(limit=1)
|
|
assert document.id is not None
|
|
await beta.document_item_repository.create_all(
|
|
[
|
|
DocumentItem(
|
|
document_id=document.id,
|
|
self_ref="#/pictures/0",
|
|
position=99,
|
|
label="picture",
|
|
text="",
|
|
picture_data=b"beta-picture",
|
|
)
|
|
]
|
|
)
|
|
|
|
assert (
|
|
await rag.get_picture_bytes(document.id, "#/pictures/0", "beta")
|
|
== b"beta-picture"
|
|
)
|
|
assert (
|
|
await rag.get_picture_bytes(document.id, "#/pictures/0", "alpha")
|
|
is None
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_single_database_needs_no_source(self, temp_db_path):
|
|
"""One database is where the picture is, named or not."""
|
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
document = await rag.document_repository.create(
|
|
Document(content="body", uri="test://one")
|
|
)
|
|
assert document.id is not None
|
|
await rag.document_item_repository.create_all(
|
|
[
|
|
DocumentItem(
|
|
document_id=document.id,
|
|
self_ref="#/pictures/0",
|
|
position=0,
|
|
label="picture",
|
|
text="",
|
|
picture_data=b"the-picture",
|
|
)
|
|
]
|
|
)
|
|
|
|
assert (
|
|
await rag.get_picture_bytes(document.id, "#/pictures/0")
|
|
== b"the-picture"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_picture_lookup_without_a_source_is_refused(self, tmp_path):
|
|
"""Federating, nothing can say which database holds an unqualified
|
|
reference."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
with pytest.raises(ValueError, match="source"):
|
|
await rag.get_picture_bytes("doc-1", "#/pictures/0")
|
|
|
|
|
|
class TestFederatedEdges:
|
|
@pytest.mark.asyncio
|
|
async def test_expansion_passes_through_results_without_a_source(self, tmp_path):
|
|
"""A caller can hand `expand_context` results it built itself. Those name
|
|
no database, so there is nowhere to expand them from."""
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
await _seed(config, "beta", ["beta document about cats"])
|
|
|
|
handmade = SearchResult(content="handmade", score=0.4, doc_item_refs=[])
|
|
|
|
async with HaikuRAG(config=config) as rag:
|
|
found = await rag.search("cats", search_type="fts", limit=10)
|
|
expanded = await rag.expand_context([*found, handmade])
|
|
|
|
assert "handmade" in [r.content for r in expanded]
|
|
scores = [r.score for r in expanded]
|
|
assert scores == sorted(scores, reverse=True), "merged in score order"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_chunk_without_a_document_is_not_cited(self, tmp_path):
|
|
"""`Chunk.document_id` is optional, and a citation without a document has
|
|
nothing to point at."""
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from haiku.rag.capabilities.rag import RAGCapability
|
|
from haiku.rag.store.models import Chunk
|
|
from tests.capabilities.test_capabilities import (
|
|
Deps,
|
|
_single_database_client,
|
|
make_context,
|
|
)
|
|
|
|
config = _config(tmp_path, ["alpha", "beta"])
|
|
await _seed(config, "alpha", ["alpha document about cats"])
|
|
|
|
orphan = _single_database_client()
|
|
orphan.get_chunk_by_id.return_value = Chunk(
|
|
id="orphan", document_id=None, content="no document"
|
|
)
|
|
|
|
capability = create_capability(config=config, defer_loading=False)
|
|
run = await capability.for_run(make_context(Deps()))
|
|
with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=orphan)):
|
|
with pytest.raises(ModelRetry):
|
|
await run._cite(["orphan"])
|