71e4e4a40a
`resolve_db_path` manufactured the default path whenever `lancedb.databases` was empty, and `covers_several_databases` read coverage back out of the configuration, so a capability built without a client opened `storage.data_dir/haiku.rag.lancedb` instead of what `lancedb.uri` placed. The entry point resolves a `DatabaseScope` instead: instructions ask it what it covers and `_ensure_rag` opens it through `HaikuRAG._covering`, so coverage is decided once rather than encoded in a path and re-derived. `Sandbox._covering` takes the scope the capability already resolved, beside the public constructor that takes a path. The factory signatures are unchanged.
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.scope.names == ("alpha", "beta")
|
|
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"])
|