Ask the client what it covers

covers_multiple, source_names, source and reader_for replace the private
state seven modules were reading to work out how many databases they had.
The configured selection is kept intact, so entering a client twice derives
the same database rather than the last derivation.
This commit is contained in:
Yiorgis Gozadinos 2026-08-25 14:16:22 +03:00
parent 16b7319c48
commit 8db447e095
No known key found for this signature in database
16 changed files with 206 additions and 90 deletions

View file

@ -127,7 +127,7 @@ def covers_several_databases(
covers. Instructions follow coverage, not configuration. covers. Instructions follow coverage, not configuration.
""" """
if rag is not None: if rag is not None:
return bool(rag._federated) return rag.covers_multiple
return db_path is None and len(config.lancedb.databases) > 1 return db_path is None and len(config.lancedb.databases) > 1
@ -543,7 +543,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
owner, chunk = found owner, chunk = found
if not chunk.document_id: if not chunk.document_id:
continue continue
key = (owner._source, chunk.document_id) key = (owner.source, chunk.document_id)
if key not in documents: if key not in documents:
documents[key] = await owner.get_document_by_id( documents[key] = await owner.get_document_by_id(
chunk.document_id chunk.document_id
@ -553,7 +553,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
chunk.document_title = document.title if document else None chunk.document_title = document.title if document else None
chunk.document_meta = document.metadata if document else {} chunk.document_meta = document.metadata if document else {}
result = SearchResult.from_chunk(chunk, score=1.0) result = SearchResult.from_chunk(chunk, score=1.0)
result.source = owner._source result.source = owner.source
synthetic.append(result) synthetic.append(result)
citations.extend(resolve_citations(missing, synthetic)) citations.extend(resolve_citations(missing, synthetic))

View file

@ -357,12 +357,10 @@ class ChatApp(App):
citation = selected_widgets[0].citation citation = selected_widgets[0].citation
# Chunks, pages and bounding boxes all come from the database holding the # Chunks, pages and bounding boxes all come from the database holding the
# cited chunk. A client covering a set has no repositories of its own. # cited chunk, which a client covering a set has to be asked for.
client = self.client client = await self.client.reader_for(citation.source)
if client._federated: if client is None:
if citation.source is None: return
return
(client,) = await client.clients_for([citation.source])
chunk_ids = citation.chunk_ids or [citation.chunk_id] chunk_ids = citation.chunk_ids or [citation.chunk_id]
chunks = [] chunks = []
for cid in chunk_ids: for cid in chunk_ids:

View file

@ -41,7 +41,7 @@ from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.document_item import DocumentItemRepository from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository from haiku.rag.store.repositories.settings import SettingsRepository
from haiku.rag.utils import escape_sql_string, locate_database from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
@ -128,7 +128,11 @@ class HaikuRAG:
None means all of them. Ignored when a single ``uri`` or an None means all of them. Ignored when a single ``uri`` or an
explicit ``db_path`` is given. explicit ``db_path`` is given.
""" """
self._config = config if config is not None else get_config() self._configured = config if config is not None else get_config()
# What the caller configured, kept intact: entering derives a
# single-database configuration from it, and asking again has to see the
# same set rather than the answer from last time.
self._config = self._configured
self._db_path_given = db_path is not None self._db_path_given = db_path is not None
if db_path is None: if db_path is None:
db_path = self._config.storage.data_dir / "haiku.rag.lancedb" db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
@ -139,20 +143,47 @@ class HaikuRAG:
self._read_only = read_only self._read_only = read_only
self._requested_sources = sources self._requested_sources = sources
self._clients: dict[str, HaikuRAG] = {} self._clients: dict[str, HaikuRAG] = {}
self._source: str | None = None
self._session: SingleDatabaseSession | FederatedSession | None = None self._session: SingleDatabaseSession | FederatedSession | None = None
self._owns_session = True self._owns_session = True
@property @property
def _federated(self) -> dict[str, str]: def covers_multiple(self) -> bool:
"""The databases covered, name to location, empty while covering one. """Whether this client reads from more than one database."""
return isinstance(self._session, FederatedSession)
Derived from the session rather than stored beside it: two answers to @property
"how many databases?" is what left every call site deciding for itself. def source_names(self) -> tuple[str, ...]:
"""The configured databases this client covers, in configured order.
A single database contributes its own name, or nothing where the
configuration named none.
""" """
if isinstance(self._session, FederatedSession): if isinstance(self._session, FederatedSession):
return self._session.locations return self._session.names
return {} return () if self.source is None else (self.source,)
@property
def source(self) -> str | None:
"""The configured database this client reads, or None while covering a
set or reading a database the configuration did not name."""
if isinstance(self._session, SingleDatabaseSession):
return self._session.source
return None
async def reader_for(self, source: str | None) -> "HaikuRAG | None":
"""The client that can read `source` — itself, where it reads one
database.
None where the database cannot be placed: a client covering a set needs
the name, and evidence recorded before databases could be named carries
none.
"""
if not self.covers_multiple:
return self
if source is None:
return None
(owner,) = await self.clients_for([source])
return owner
def _single_session(self, operation: str) -> SingleDatabaseSession: def _single_session(self, operation: str) -> SingleDatabaseSession:
"""The one database this operation works on. """The one database this operation works on.
@ -163,7 +194,7 @@ class HaikuRAG:
""" """
if isinstance(self._session, SingleDatabaseSession): if isinstance(self._session, SingleDatabaseSession):
return self._session return self._session
covered = ", ".join(sorted(self._federated)) covered = ", ".join(sorted(self.source_names))
raise AmbiguousDatabaseError( raise AmbiguousDatabaseError(
f"{operation} works on one database, and this client covers " f"{operation} works on one database, and this client covers "
f"{covered}; select one with clients_for([name])" f"{covered}; select one with clients_for([name])"
@ -216,7 +247,7 @@ class HaikuRAG:
covering a set has an unambiguous embedder without opening any of them. covering a set has an unambiguous embedder without opening any of them.
Built on first use and owned by this client, which closes it. Built on first use and owned by this client, which closes it.
""" """
if self._federated: if self.covers_multiple:
return get_embedder(config=self._config) return get_embedder(config=self._config)
return self.store.embedder return self.store.embedder
@ -235,7 +266,7 @@ class HaikuRAG:
Empty when the caller named a database itself: an explicit `db_path` says Empty when the caller named a database itself: an explicit `db_path` says
which one to open, so it is not overridden by a configured set. which one to open, so it is not overridden by a configured set.
""" """
declared = self._config.lancedb.databases declared = self._configured.lancedb.databases
if not declared or self._db_path_given: if not declared or self._db_path_given:
return {} return {}
if self._requested_sources is not None and not self._requested_sources: if self._requested_sources is not None and not self._requested_sources:
@ -283,12 +314,12 @@ class HaikuRAG:
read_only=self._read_only, read_only=self._read_only,
) )
return self return self
source: str | None = None
if selected: if selected:
[(self._source, location)] = selected.items() [(source, location)] = selected.items()
uri, db_path = locate_database(location) self._config, db_path = DatabaseRef.configured(source, location).connection(
self._config = self._config.model_copy(deep=True) self._configured
self._config.lancedb.databases = {} )
self._config.lancedb.uri = uri
if db_path is not None: if db_path is not None:
self._db_path = db_path self._db_path = db_path
@ -298,7 +329,7 @@ class HaikuRAG:
skip_validation=self._skip_validation, skip_validation=self._skip_validation,
create=self._create, create=self._create,
read_only=self._read_only, read_only=self._read_only,
source=self._source, source=source,
).open() ).open()
return self return self
@ -337,7 +368,6 @@ class HaikuRAG:
session.db_path, config=session.config, read_only=session.read_only session.db_path, config=session.config, read_only=session.read_only
) )
client._session = session client._session = session
client._source = session.source
client._owns_session = False client._owns_session = False
return client return client
@ -354,7 +384,7 @@ class HaikuRAG:
`SettingsRepository` reports on open. `SettingsRepository` reports on open.
""" """
recorded = [ recorded = [
(client._source, client.store.stored_embedding) (client.source, client.store.stored_embedding)
for client in clients for client in clients
if client.store.stored_embedding is not None if client.store.stored_embedding is not None
] ]
@ -569,7 +599,7 @@ class HaikuRAG:
Returns: Returns:
The Document instance if found, None otherwise. The Document instance if found, None otherwise.
""" """
if self._federated: if self.covers_multiple:
return await self._from_any_covered( return await self._from_any_covered(
lambda owner: owner.get_document_by_id(document_id) lambda owner: owner.get_document_by_id(document_id)
) )
@ -586,7 +616,7 @@ class HaikuRAG:
Returns: Returns:
The Chunk instance if found, None otherwise. The Chunk instance if found, None otherwise.
""" """
if self._federated: if self.covers_multiple:
return await self._from_any_covered( return await self._from_any_covered(
lambda owner: owner.get_chunk_by_id(chunk_id) lambda owner: owner.get_chunk_by_id(chunk_id)
) )
@ -605,7 +635,7 @@ class HaikuRAG:
Returns: Returns:
The picture bytes if found, None otherwise. The picture bytes if found, None otherwise.
""" """
if not self._federated: if not self.covers_multiple:
return await self.document_item_repository.get_picture_bytes( return await self.document_item_repository.get_picture_bytes(
document_id, self_ref document_id, self_ref
) )
@ -627,7 +657,7 @@ class HaikuRAG:
Returns: Returns:
The Document instance if found, None otherwise. The Document instance if found, None otherwise.
""" """
if self._federated: if self.covers_multiple:
return await self._from_any_covered( return await self._from_any_covered(
lambda owner: owner.get_document_by_uri(uri) lambda owner: owner.get_document_by_uri(uri)
) )
@ -685,7 +715,7 @@ class HaikuRAG:
Returns: Returns:
List of Document instances matching the criteria. List of Document instances matching the criteria.
""" """
if self._federated: if self.covers_multiple:
# Each database is asked for enough rows to satisfy the window, and # Each database is asked for enough rows to satisfy the window, and
# the window is applied to the merged listing: a limit means that # the window is applied to the merged listing: a limit means that
# many documents in total, not that many per database. # many documents in total, not that many per database.
@ -718,7 +748,7 @@ class HaikuRAG:
Returns: Returns:
Number of documents matching the criteria. Number of documents matching the criteria.
""" """
if self._federated: if self.covers_multiple:
counts = await asyncio.gather( counts = await asyncio.gather(
*( *(
owner.count_documents(filter=filter) owner.count_documents(filter=filter)
@ -745,19 +775,19 @@ class HaikuRAG:
`None` for all of them. Every read honouring `sources` decides through `None` for all of them. Every read honouring `sources` decides through
this, so the rule cannot differ between one operation and another. this, so the rule cannot differ between one operation and another.
""" """
if self._federated: if self.covers_multiple:
return await self.clients_for( return await self.clients_for(
list(self._federated) if sources is None else sources list(self.source_names) if sources is None else sources
) )
if sources is None: if sources is None:
return [self] return [self]
sources = _without_repeats(sources) sources = _without_repeats(sources)
if not sources: if not sources:
return [] return []
if sources != [self._source]: if sources != [self.source]:
raise KeyError( raise KeyError(
f"unknown database(s) {', '.join(sources) or '(none)'}; this " f"unknown database(s) {', '.join(sources) or '(none)'}; this "
f"client covers {self._source or 'a single unnamed database'}" f"client covers {self.source or 'a single unnamed database'}"
) )
return [self] return [self]
@ -772,7 +802,7 @@ class HaikuRAG:
) -> list[SearchResult]: ) -> list[SearchResult]:
from haiku.rag.client.search import search, search_sources from haiku.rag.client.search import search, search_sources
if self._federated: if self.covers_multiple:
return await search_sources( return await search_sources(
self, query, limit, search_type, filter, include_images, sources self, query, limit, search_type, filter, include_images, sources
) )
@ -782,7 +812,7 @@ class HaikuRAG:
# A database named in config keeps its name even when it is the only one # A database named in config keeps its name even when it is the only one
# this client covers. Only a legacy single `uri` leaves source unset. # this client covers. Only a legacy single `uri` leaves source unset.
for result in results: for result in results:
result.source = self._source result.source = self.source
return results return results
async def expand_context( async def expand_context(

View file

@ -64,7 +64,7 @@ async def ask(
from haiku.rag.utils import get_model from haiku.rag.utils import get_model
capability = create_capability( capability = create_capability(
db_path=None if client._federated else client.store.db_path, db_path=None if client.covers_multiple else client.store.db_path,
config=client._config, config=client._config,
rag=client, rag=client,
defer_loading=False, defer_loading=False,
@ -124,7 +124,7 @@ async def analyze(
from haiku.rag.utils import get_model from haiku.rag.utils import get_model
capability = create_capability( capability = create_capability(
db_path=None if client._federated else client.store.db_path, db_path=None if client.covers_multiple else client.store.db_path,
config=client._config, config=client._config,
rag=client, rag=client,
defer_loading=False, defer_loading=False,

View file

@ -82,7 +82,7 @@ async def search_sources(
if limit is None: if limit is None:
limit = client._config.search.limit limit = client._config.search.limit
names = list(client._federated) if sources is None else list(sources) names = list(client.source_names) if sources is None else list(sources)
if not names: if not names:
return [] return []
selected = await client.clients_for(names) selected = await client.clients_for(names)
@ -107,7 +107,7 @@ async def search_sources(
results: list[SearchResult] = [] results: list[SearchResult] = []
for owner, chunk, score in ranked: for owner, chunk, score in ranked:
result = SearchResult.from_chunk(chunk, score) result = SearchResult.from_chunk(chunk, score)
result.source = owner._source result.source = owner.source
results.append(result) results.append(result)
results = _dedup_picture_chunks(results) results = _dedup_picture_chunks(results)
@ -116,10 +116,11 @@ async def search_sources(
for result in results: for result in results:
if result.source: if result.source:
by_owner.setdefault(result.source, []).append(result) by_owner.setdefault(result.source, []).append(result)
owners = await client.clients_for(list(by_owner))
await asyncio.gather( await asyncio.gather(
*( *(
_populate_image_data(client._clients[name], owned) _populate_image_data(owner, by_owner[name])
for name, owned in by_owner.items() for name, owner in zip(by_owner, owners, strict=True)
) )
) )
@ -430,7 +431,7 @@ async def expand_context(
""" """
# A federating client has no repositories of its own, so each result expands # A federating client has no repositories of its own, so each result expands
# through the database it came from. # through the database it came from.
if client._federated: if client.covers_multiple:
by_source: dict[str, list[SearchResult]] = {} by_source: dict[str, list[SearchResult]] = {}
unsourced: list[SearchResult] = [] unsourced: list[SearchResult] = []
for result in search_results: for result in search_results:
@ -441,9 +442,9 @@ async def expand_context(
owners = await client.clients_for(list(by_source)) owners = await client.clients_for(list(by_source))
expanded_groups = await asyncio.gather( expanded_groups = await asyncio.gather(
*( *(
expand_context(owner, by_source[owner._source]) expand_context(owner, by_source[owner.source])
for owner in owners for owner in owners
if owner._source if owner.source
) )
) )
merged = unsourced + [r for group in expanded_groups for r in group] merged = unsourced + [r for group in expanded_groups for r in group]

View file

@ -280,9 +280,9 @@ class FederatedSession:
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
@property @property
def locations(self) -> dict[str, str]: def names(self) -> tuple[str, ...]:
"""The databases covered, name to location as configured.""" """The databases covered, in configured order."""
return {name: ref.uri or str(ref.db_path) for name, ref in self._refs.items()} return tuple(self._refs)
async def sessions_for(self, names: list[str]) -> list[SingleDatabaseSession]: async def sessions_for(self, names: list[str]) -> list[SingleDatabaseSession]:
"""The sessions for these databases, opening any not yet open. """The sessions for these databases, opening any not yet open.

View file

@ -22,7 +22,7 @@ def reported_database(client: "HaikuRAG", db_path: "Path | None") -> "Path | Non
one named database opens it rather than covering a set. Only what the client one named database opens it rather than covering a set. Only what the client
ended up covering says which of the two this is. ended up covering says which of the two this is.
""" """
if client._federated: if client.covers_multiple:
return None return None
return db_path if db_path is not None else client.store.db_path return db_path if db_path is not None else client.store.db_path
@ -194,7 +194,7 @@ class InfoModal(ModalScreen):
# rather than the whole panel. Names only, no paths — a location # rather than the whole panel. Names only, no paths — a location
# belongs in the configuration. # belongs in the configuration.
blocks = await asyncio.gather( blocks = await asyncio.gather(
*(self._report(name) for name in sorted(self.client._federated)) *(self._report(name) for name in sorted(self.client.source_names))
) )
for block in blocks: for block in blocks:
lines.extend(block) lines.extend(block)

View file

@ -214,7 +214,7 @@ class Sandbox:
a database the question excluded cannot be mounted. a database the question excluded cannot be mounted.
""" """
async with self._connection() as rag: async with self._connection() as rag:
if not rag._federated: if not rag.covers_multiple:
if not await rag.clients_covering(self._context.sources): if not await rag.clients_covering(self._context.sources):
return [], {} return [], {}
docs = await rag.list_documents(filter=self._context.filter) docs = await rag.list_documents(filter=self._context.filter)
@ -251,10 +251,10 @@ class Sandbox:
if doc.id in holders: if doc.id in holders:
raise ValueError( raise ValueError(
f"document {doc.id} is in databases {held_by[doc.id]!r} and " f"document {doc.id} is in databases {held_by[doc.id]!r} and "
f"{owner._source!r}; analysis mounts one document per id" f"{owner.source!r}; analysis mounts one document per id"
) )
holders[doc.id] = owner holders[doc.id] = owner
held_by[doc.id] = owner._source held_by[doc.id] = owner.source
return holders return holders
def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any: def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
@ -354,7 +354,7 @@ class Sandbox:
"title": d.title, "title": d.title,
"uri": d.uri, "uri": d.uri,
"created_at": str(d.created_at), "created_at": str(d.created_at),
"source": owners[d.id]._source if d.id in owners else None, "source": owners[d.id].source if d.id in owners else None,
} }
for d in docs for d in docs
] ]

View file

@ -394,7 +394,7 @@ async def format_citations_rich(
idx = c.index if c.index is not None else (i + 1) idx = c.index if c.index is not None else (i + 1)
header_parts: list[str] = [f"[{idx}] {_citation_label(c)}"] header_parts: list[str] = [f"[{idx}] {_citation_label(c)}"]
if c.source and client is not None and client._federated: if c.source and client is not None and client.covers_multiple:
header_parts.append(c.source) header_parts.append(c.source)
pages = _citation_pages(c) pages = _citation_pages(c)
if pages: if pages:

View file

@ -144,13 +144,15 @@ def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
def _single_database_client() -> AsyncMock: def _single_database_client() -> AsyncMock:
"""A stand-in for a client covering one unnamed database. """A stand-in for a client covering one unnamed database.
A bare AsyncMock answers every attribute with a truthy Mock, so `_federated` A bare AsyncMock answers every attribute with a truthy Mock, so
would read as a set of databases, `_source` would reach a validated field, and `covers_multiple` would read as a set of databases, `source` would reach a
`clients_covering` would return a Mock where the code iterates clients. validated field, and `clients_covering` would return a Mock where the code
iterates clients.
""" """
client = AsyncMock() client = AsyncMock()
client._federated = {} client.covers_multiple = False
client._source = None client.source_names = ()
client.source = None
client.clients_covering.return_value = [client] client.clients_covering.return_value = [client]
return client return client
@ -1537,7 +1539,8 @@ class TestSeveralDatabasesInstructions:
@staticmethod @staticmethod
def _client(federated): def _client(federated):
client = AsyncMock() client = AsyncMock()
client._federated = federated client.covers_multiple = len(federated) > 1
client.source_names = tuple(federated)
return client return client
def test_one_database_is_instructed_as_before(self): def test_one_database_is_instructed_as_before(self):

View file

@ -118,10 +118,11 @@ def _make_mock_client():
mock_client = AsyncMock() mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client) mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None) mock_client.__aexit__ = AsyncMock(return_value=None)
# Covers one database: a bare AsyncMock answers `_federated` with a truthy # Covers one database: a bare AsyncMock answers `covers_multiple` with a truthy
# Mock, which would send every read down the covering-a-set branch. # Mock, which would send every read down the covering-a-set branch.
mock_client._federated = {} mock_client.covers_multiple = False
mock_client._source = None mock_client.source_names = ()
mock_client.source = None
return mock_client return mock_client
@ -510,14 +511,15 @@ async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path)
from haiku.rag.store.models.citation import Citation from haiku.rag.store.models.citation import Citation
owner = _make_mock_client() owner = _make_mock_client()
owner._source = "beta" owner.source = "beta"
owner.get_chunk_by_id.return_value = Chunk( owner.get_chunk_by_id.return_value = Chunk(
id="c1", document_id="d1", content="cited body" id="c1", document_id="d1", content="cited body"
) )
covering = _make_mock_client() covering = _make_mock_client()
covering._federated = {"alpha": "/a.lancedb", "beta": "/b.lancedb"} covering.covers_multiple = True
covering.clients_for = AsyncMock(return_value=[owner]) covering.source_names = ("alpha", "beta")
covering.reader_for = AsyncMock(return_value=owner)
app, _ = _make_app(tmp_path / "unused.lancedb", covering) app, _ = _make_app(tmp_path / "unused.lancedb", covering)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=covering): with patch("haiku.rag.chat.app.HaikuRAG", return_value=covering):
@ -540,7 +542,7 @@ async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path)
with patch.object(app, "push_screen", new=AsyncMock()) as push: with patch.object(app, "push_screen", new=AsyncMock()) as push:
await app.action_show_visual() await app.action_show_visual()
covering.clients_for.assert_awaited_once_with(["beta"]) covering.reader_for.assert_awaited_once_with("beta")
owner.get_chunk_by_id.assert_awaited_once_with("c1") owner.get_chunk_by_id.assert_awaited_once_with("c1")
assert push.await_args is not None assert push.await_args is not None
assert push.await_args.args[0].client is owner assert push.await_args.args[0].client is owner
@ -558,7 +560,8 @@ class TestDocumentSelectionIdentity:
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
client = AsyncMock() client = AsyncMock()
client._federated = {"arxiv": "a", "wiki": "b"} client.covers_multiple = True
client.source_names = ("arxiv", "wiki")
client.list_documents.return_value = [ client.list_documents.return_value = [
Document(id="id-one", content="", title="Capital region", source="arxiv"), Document(id="id-one", content="", title="Capital region", source="arxiv"),
Document(id="id-two", content="", title="Capital region", source="wiki"), Document(id="id-two", content="", title="Capital region", source="wiki"),

View file

@ -33,7 +33,7 @@ class TestDocumentsAcrossDatabases:
"test://alpha/alpha document about cats", "test://alpha/alpha document about cats",
"test://beta/beta document about cats", "test://beta/beta document about cats",
} }
assert {owner._source for owner in owners.values()} == {"alpha", "beta"} assert {owner.source for owner in owners.values()} == {"alpha", "beta"}
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_selected_databases_bound_the_corpus(self, tmp_path): async def test_selected_databases_bound_the_corpus(self, tmp_path):
@ -46,7 +46,7 @@ class TestDocumentsAcrossDatabases:
_, docs, owners = await _mounted(rag, sources=["alpha"]) _, docs, owners = await _mounted(rag, sources=["alpha"])
assert [d.uri for d in docs] == ["test://alpha/alpha document about cats"] assert [d.uri for d in docs] == ["test://alpha/alpha document about cats"]
assert {owner._source for owner in owners.values()} == {"alpha"} assert {owner.source for owner in owners.values()} == {"alpha"}
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_one_database_needs_no_owners(self, tmp_path, temp_db_path): async def test_one_database_needs_no_owners(self, tmp_path, temp_db_path):
@ -76,8 +76,8 @@ class TestDocumentsAcrossDatabases:
async with sandbox._connection(owners[doc.id]) as owner: async with sandbox._connection(owners[doc.id]) as owner:
content = await owner.document_repository.get_content(doc.id) content = await owner.document_repository.get_content(doc.id)
assert content is not None assert content is not None
assert owners[doc.id]._source is not None assert owners[doc.id].source is not None
assert owners[doc.id]._source in content assert owners[doc.id].source in content
class TestExecutingAcrossDatabases: class TestExecutingAcrossDatabases:

View file

@ -339,7 +339,8 @@ class TestReportedDatabase:
@staticmethod @staticmethod
def _client(federated, store_path=None): def _client(federated, store_path=None):
client = MagicMock() client = MagicMock()
client._federated = federated client.covers_multiple = len(federated) > 1
client.source_names = tuple(federated)
client.store.db_path = store_path client.store.db_path = store_path
return client return client

View file

@ -126,8 +126,8 @@ class TestNamingADatabaseDirectly:
await _seed(config, "alpha", ["alpha document about cats"]) await _seed(config, "alpha", ["alpha document about cats"])
async with HaikuRAG(temp_db_path, config=config, create=True) as rag: async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
assert rag._federated == {} assert not rag.covers_multiple
assert rag._source is None assert rag.source is None
assert rag.store.db_path == temp_db_path assert rag.store.db_path == temp_db_path
@pytest.mark.asyncio @pytest.mark.asyncio
@ -137,8 +137,8 @@ class TestNamingADatabaseDirectly:
await _seed(config, "alpha", ["alpha document about cats"]) await _seed(config, "alpha", ["alpha document about cats"])
async with HaikuRAG(config=config) as rag: async with HaikuRAG(config=config) as rag:
assert rag._federated == {} assert not rag.covers_multiple
assert rag._source == "alpha" assert rag.source == "alpha"
results = await rag.search("cats", search_type="fts", limit=10) results = await rag.search("cats", search_type="fts", limit=10)
assert [r.source for r in results] == ["alpha"] assert [r.source for r in results] == ["alpha"]
@ -168,7 +168,7 @@ class TestOpeningDatabases:
rag._session._open = gated rag._session._open = gated
clients = await asyncio.wait_for(rag.clients_for(names), timeout=15) clients = await asyncio.wait_for(rag.clients_for(names), timeout=15)
assert {client._source for client in clients} == set(names) assert {client.source for client in clients} == set(names)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_failed_open_does_not_leak_the_ones_that_worked(self, tmp_path): async def test_a_failed_open_does_not_leak_the_ones_that_worked(self, tmp_path):
@ -195,7 +195,7 @@ class TestOpeningDatabases:
async with HaikuRAG(config=config) as rag: async with HaikuRAG(config=config) as rag:
clients = await rag.clients_for(["alpha", "alpha", "beta"]) clients = await rag.clients_for(["alpha", "alpha", "beta"])
assert [client._source for client in clients] == ["alpha", "beta"] assert [client.source for client in clients] == ["alpha", "beta"]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_database_named_twice_returns_each_result_once(self, tmp_path): async def test_a_database_named_twice_returns_each_result_once(self, tmp_path):
@ -220,7 +220,7 @@ class TestOpeningDatabases:
async with HaikuRAG(config=config) as rag: async with HaikuRAG(config=config) as rag:
covering = await rag.clients_covering(["alpha", "alpha"]) covering = await rag.clients_covering(["alpha", "alpha"])
assert [client._source for client in covering] == ["alpha"] assert [client.source for client in covering] == ["alpha"]
class TestListingAcrossDatabases: class TestListingAcrossDatabases:
@ -379,6 +379,84 @@ class TestClosingASet:
assert sorted(drained) == ["alpha", "beta"] assert sorted(drained) == ["alpha", "beta"]
class TestPlacingADatabase:
"""What a client says about the databases it covers, so nothing outside has
to read its private state to find out."""
@pytest.mark.asyncio
async def test_a_set_names_every_database_it_covers(self, tmp_path):
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha one"])
await _seed(config, "beta", ["beta one"])
async with HaikuRAG(config=config, read_only=True) as rag:
assert rag.covers_multiple
assert rag.source_names == ("alpha", "beta")
assert rag.source is None
@pytest.mark.asyncio
async def test_one_named_database_names_itself(self, tmp_path):
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha one"])
async with HaikuRAG(config=config, read_only=True, sources=["alpha"]) as rag:
assert not rag.covers_multiple
assert rag.source_names == ("alpha",)
assert rag.source == "alpha"
@pytest.mark.asyncio
async def test_a_named_database_keeps_its_name_on_re_entry(self, tmp_path):
"""Entering derives a single-database configuration from what was
configured. Deriving it from the last derivation loses the name."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
rag = HaikuRAG(config=config, read_only=True, sources=["alpha"])
async with rag:
assert rag.source == "alpha"
async with rag:
assert rag.source == "alpha"
assert rag.source_names == ("alpha",)
results = await rag.search("cats", search_type="fts")
assert {r.source for r in results} == {"alpha"}
@pytest.mark.asyncio
async def test_an_unnamed_database_names_nothing(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
assert rag.source_names == ()
assert rag.source is None
@pytest.mark.asyncio
async def test_the_reader_for_a_database_is_the_client_holding_it(self, tmp_path):
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha one"])
await _seed(config, "beta", ["beta one"])
async with HaikuRAG(config=config, read_only=True) as rag:
reader = await rag.reader_for("beta")
assert reader is not None
assert reader.source == "beta"
# Asked twice, the same wrapper comes back.
assert await rag.reader_for("beta") is reader
@pytest.mark.asyncio
async def test_a_client_reading_one_database_is_its_own_reader(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
assert await rag.reader_for(None) is rag
assert await rag.reader_for("anything") is rag
@pytest.mark.asyncio
async def test_a_set_cannot_place_evidence_that_names_no_database(self, tmp_path):
"""Evidence recorded before databases could be named carries no source."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha one"])
async with HaikuRAG(config=config, read_only=True) as rag:
assert await rag.reader_for(None) is None
class TestBorrowedDatabases: class TestBorrowedDatabases:
"""A client for one of a set wraps a database the set opened.""" """A client for one of a set wraps a database the set opened."""

View file

@ -293,7 +293,7 @@ class TestStandaloneCapabilities:
await run._close() await run._close()
assert len(docs) == 2 assert len(docs) == 2
assert {owner._source for owner in owners.values()} == {"alpha", "beta"} assert {owner.source for owner in owners.values()} == {"alpha", "beta"}
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_single_configured_database_is_still_opened(self, tmp_path): async def test_a_single_configured_database_is_still_opened(self, tmp_path):
@ -304,7 +304,7 @@ class TestStandaloneCapabilities:
capability = create_capability(config=config, defer_loading=False) capability = create_capability(config=config, defer_loading=False)
rag = await capability._ensure_rag() rag = await capability._ensure_rag()
try: try:
assert rag._source == "alpha" assert rag.source == "alpha"
finally: finally:
await capability._close() await capability._close()

View file

@ -785,7 +785,8 @@ async def test_format_citations_rich_names_the_database_when_federating():
source="medic", source="medic",
) )
client = AsyncMock() client = AsyncMock()
client._federated = {"medic": "/data/medic.lancedb", "st": "/data/st.lancedb"} client.covers_multiple = True
client.source_names = ("medic", "st")
output = _render_rich(await format_citations_rich([citation], client)) output = _render_rich(await format_citations_rich([citation], client))
@ -808,7 +809,8 @@ async def test_format_citations_rich_omits_the_database_for_one_database():
source="medic", source="medic",
) )
client = AsyncMock() client = AsyncMock()
client._federated = {} client.covers_multiple = False
client.source_names = ()
output = _render_rich(await format_citations_rich([citation], client)) output = _render_rich(await format_citations_rich([citation], client))