diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 2d0f16ac..ac60393b 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -127,7 +127,7 @@ def covers_several_databases( covers. Instructions follow coverage, not configuration. """ if rag is not None: - return bool(rag._federated) + return rag.covers_multiple return db_path is None and len(config.lancedb.databases) > 1 @@ -543,7 +543,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): owner, chunk = found if not chunk.document_id: continue - key = (owner._source, chunk.document_id) + key = (owner.source, chunk.document_id) if key not in documents: documents[key] = await owner.get_document_by_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_meta = document.metadata if document else {} result = SearchResult.from_chunk(chunk, score=1.0) - result.source = owner._source + result.source = owner.source synthetic.append(result) citations.extend(resolve_citations(missing, synthetic)) diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index a0d78513..c4996e9b 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -357,12 +357,10 @@ class ChatApp(App): citation = selected_widgets[0].citation # 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. - client = self.client - if client._federated: - if citation.source is None: - return - (client,) = await client.clients_for([citation.source]) + # cited chunk, which a client covering a set has to be asked for. + client = await self.client.reader_for(citation.source) + if client is None: + return chunk_ids = citation.chunk_ids or [citation.chunk_id] chunks = [] for cid in chunk_ids: diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index eb0128c7..7bbbbbb5 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -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_item import DocumentItemRepository 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: 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 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 if db_path is None: db_path = self._config.storage.data_dir / "haiku.rag.lancedb" @@ -139,20 +143,47 @@ class HaikuRAG: self._read_only = read_only self._requested_sources = sources self._clients: dict[str, HaikuRAG] = {} - self._source: str | None = None self._session: SingleDatabaseSession | FederatedSession | None = None self._owns_session = True @property - def _federated(self) -> dict[str, str]: - """The databases covered, name to location, empty while covering one. + def covers_multiple(self) -> bool: + """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 - "how many databases?" is what left every call site deciding for itself. + @property + 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): - return self._session.locations - return {} + return self._session.names + 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: """The one database this operation works on. @@ -163,7 +194,7 @@ class HaikuRAG: """ if isinstance(self._session, SingleDatabaseSession): return self._session - covered = ", ".join(sorted(self._federated)) + covered = ", ".join(sorted(self.source_names)) raise AmbiguousDatabaseError( f"{operation} works on one database, and this client covers " 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. 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 self.store.embedder @@ -235,7 +266,7 @@ class HaikuRAG: 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. """ - declared = self._config.lancedb.databases + declared = self._configured.lancedb.databases if not declared or self._db_path_given: return {} if self._requested_sources is not None and not self._requested_sources: @@ -283,12 +314,12 @@ class HaikuRAG: read_only=self._read_only, ) return self + source: str | None = None if selected: - [(self._source, location)] = selected.items() - uri, db_path = locate_database(location) - self._config = self._config.model_copy(deep=True) - self._config.lancedb.databases = {} - self._config.lancedb.uri = uri + [(source, location)] = selected.items() + self._config, db_path = DatabaseRef.configured(source, location).connection( + self._configured + ) if db_path is not None: self._db_path = db_path @@ -298,7 +329,7 @@ class HaikuRAG: skip_validation=self._skip_validation, create=self._create, read_only=self._read_only, - source=self._source, + source=source, ).open() return self @@ -337,7 +368,6 @@ class HaikuRAG: session.db_path, config=session.config, read_only=session.read_only ) client._session = session - client._source = session.source client._owns_session = False return client @@ -354,7 +384,7 @@ class HaikuRAG: `SettingsRepository` reports on open. """ recorded = [ - (client._source, client.store.stored_embedding) + (client.source, client.store.stored_embedding) for client in clients if client.store.stored_embedding is not None ] @@ -569,7 +599,7 @@ class HaikuRAG: Returns: The Document instance if found, None otherwise. """ - if self._federated: + if self.covers_multiple: return await self._from_any_covered( lambda owner: owner.get_document_by_id(document_id) ) @@ -586,7 +616,7 @@ class HaikuRAG: Returns: The Chunk instance if found, None otherwise. """ - if self._federated: + if self.covers_multiple: return await self._from_any_covered( lambda owner: owner.get_chunk_by_id(chunk_id) ) @@ -605,7 +635,7 @@ class HaikuRAG: Returns: 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( document_id, self_ref ) @@ -627,7 +657,7 @@ class HaikuRAG: Returns: The Document instance if found, None otherwise. """ - if self._federated: + if self.covers_multiple: return await self._from_any_covered( lambda owner: owner.get_document_by_uri(uri) ) @@ -685,7 +715,7 @@ class HaikuRAG: Returns: 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 # the window is applied to the merged listing: a limit means that # many documents in total, not that many per database. @@ -718,7 +748,7 @@ class HaikuRAG: Returns: Number of documents matching the criteria. """ - if self._federated: + if self.covers_multiple: counts = await asyncio.gather( *( owner.count_documents(filter=filter) @@ -745,19 +775,19 @@ class HaikuRAG: `None` for all of them. Every read honouring `sources` decides through this, so the rule cannot differ between one operation and another. """ - if self._federated: + if self.covers_multiple: 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: return [self] sources = _without_repeats(sources) if not sources: return [] - if sources != [self._source]: + if sources != [self.source]: raise KeyError( 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] @@ -772,7 +802,7 @@ class HaikuRAG: ) -> list[SearchResult]: from haiku.rag.client.search import search, search_sources - if self._federated: + if self.covers_multiple: return await search_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 # this client covers. Only a legacy single `uri` leaves source unset. for result in results: - result.source = self._source + result.source = self.source return results async def expand_context( diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index d68714fe..f1836994 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -64,7 +64,7 @@ async def ask( from haiku.rag.utils import get_model 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, rag=client, defer_loading=False, @@ -124,7 +124,7 @@ async def analyze( from haiku.rag.utils import get_model 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, rag=client, defer_loading=False, diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 3be0d96a..6362c72b 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -82,7 +82,7 @@ async def search_sources( if limit is None: 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: return [] selected = await client.clients_for(names) @@ -107,7 +107,7 @@ async def search_sources( results: list[SearchResult] = [] for owner, chunk, score in ranked: result = SearchResult.from_chunk(chunk, score) - result.source = owner._source + result.source = owner.source results.append(result) results = _dedup_picture_chunks(results) @@ -116,10 +116,11 @@ async def search_sources( for result in results: if result.source: by_owner.setdefault(result.source, []).append(result) + owners = await client.clients_for(list(by_owner)) await asyncio.gather( *( - _populate_image_data(client._clients[name], owned) - for name, owned in by_owner.items() + _populate_image_data(owner, by_owner[name]) + 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 # through the database it came from. - if client._federated: + if client.covers_multiple: by_source: dict[str, list[SearchResult]] = {} unsourced: list[SearchResult] = [] for result in search_results: @@ -441,9 +442,9 @@ async def expand_context( owners = await client.clients_for(list(by_source)) expanded_groups = await asyncio.gather( *( - expand_context(owner, by_source[owner._source]) + expand_context(owner, by_source[owner.source]) for owner in owners - if owner._source + if owner.source ) ) merged = unsourced + [r for group in expanded_groups for r in group] diff --git a/haiku_rag_slim/haiku/rag/client/session.py b/haiku_rag_slim/haiku/rag/client/session.py index 7ef34747..623e5b4b 100644 --- a/haiku_rag_slim/haiku/rag/client/session.py +++ b/haiku_rag_slim/haiku/rag/client/session.py @@ -280,9 +280,9 @@ class FederatedSession: self._lock = asyncio.Lock() @property - def locations(self) -> dict[str, str]: - """The databases covered, name to location as configured.""" - return {name: ref.uri or str(ref.db_path) for name, ref in self._refs.items()} + def names(self) -> tuple[str, ...]: + """The databases covered, in configured order.""" + return tuple(self._refs) async def sessions_for(self, names: list[str]) -> list[SingleDatabaseSession]: """The sessions for these databases, opening any not yet open. diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py index e816b22a..c18d309b 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -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 ended up covering says which of the two this is. """ - if client._federated: + if client.covers_multiple: return None 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 # belongs in the configuration. 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: lines.extend(block) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index d6703bf1..10bb7f65 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -214,7 +214,7 @@ class Sandbox: a database the question excluded cannot be mounted. """ async with self._connection() as rag: - if not rag._federated: + if not rag.covers_multiple: if not await rag.clients_covering(self._context.sources): return [], {} docs = await rag.list_documents(filter=self._context.filter) @@ -251,10 +251,10 @@ class Sandbox: if doc.id in holders: raise ValueError( 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 - held_by[doc.id] = owner._source + held_by[doc.id] = owner.source return holders def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any: @@ -354,7 +354,7 @@ class Sandbox: "title": d.title, "uri": d.uri, "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 ] diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index f62650cd..6c7fdcfd 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -394,7 +394,7 @@ async def format_citations_rich( idx = c.index if c.index is not None else (i + 1) 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) pages = _citation_pages(c) if pages: diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index de42dae3..746df6cc 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -144,13 +144,15 @@ def test_domain_preamble_is_added_to_capability_instructions(temp_db_path): def _single_database_client() -> AsyncMock: """A stand-in for a client covering one unnamed database. - A bare AsyncMock answers every attribute with a truthy Mock, so `_federated` - would read as a set of databases, `_source` would reach a validated field, and - `clients_covering` would return a Mock where the code iterates clients. + A bare AsyncMock answers every attribute with a truthy Mock, so + `covers_multiple` would read as a set of databases, `source` would reach a + validated field, and `clients_covering` would return a Mock where the code + iterates clients. """ client = AsyncMock() - client._federated = {} - client._source = None + client.covers_multiple = False + client.source_names = () + client.source = None client.clients_covering.return_value = [client] return client @@ -1537,7 +1539,8 @@ class TestSeveralDatabasesInstructions: @staticmethod def _client(federated): client = AsyncMock() - client._federated = federated + client.covers_multiple = len(federated) > 1 + client.source_names = tuple(federated) return client def test_one_database_is_instructed_as_before(self): diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 503a56d5..b6ee7f03 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -118,10 +118,11 @@ def _make_mock_client(): mock_client = AsyncMock() mock_client.__aenter__ = AsyncMock(return_value=mock_client) 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_client._federated = {} - mock_client._source = None + mock_client.covers_multiple = False + mock_client.source_names = () + mock_client.source = None 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 owner = _make_mock_client() - owner._source = "beta" + owner.source = "beta" owner.get_chunk_by_id.return_value = Chunk( id="c1", document_id="d1", content="cited body" ) covering = _make_mock_client() - covering._federated = {"alpha": "/a.lancedb", "beta": "/b.lancedb"} - covering.clients_for = AsyncMock(return_value=[owner]) + covering.covers_multiple = True + covering.source_names = ("alpha", "beta") + covering.reader_for = AsyncMock(return_value=owner) app, _ = _make_app(tmp_path / "unused.lancedb", 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: 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") assert push.await_args is not None assert push.await_args.args[0].client is owner @@ -558,7 +560,8 @@ class TestDocumentSelectionIdentity: from haiku.rag.store.models.document import Document client = AsyncMock() - client._federated = {"arxiv": "a", "wiki": "b"} + client.covers_multiple = True + client.source_names = ("arxiv", "wiki") client.list_documents.return_value = [ Document(id="id-one", content="", title="Capital region", source="arxiv"), Document(id="id-two", content="", title="Capital region", source="wiki"), diff --git a/tests/sandbox/test_sandbox_multi_db.py b/tests/sandbox/test_sandbox_multi_db.py index 8a8f8cd6..ba8dfa18 100644 --- a/tests/sandbox/test_sandbox_multi_db.py +++ b/tests/sandbox/test_sandbox_multi_db.py @@ -33,7 +33,7 @@ class TestDocumentsAcrossDatabases: "test://alpha/alpha 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 async def test_selected_databases_bound_the_corpus(self, tmp_path): @@ -46,7 +46,7 @@ class TestDocumentsAcrossDatabases: _, docs, owners = await _mounted(rag, sources=["alpha"]) 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 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: content = await owner.document_repository.get_content(doc.id) assert content is not None - assert owners[doc.id]._source is not None - assert owners[doc.id]._source in content + assert owners[doc.id].source is not None + assert owners[doc.id].source in content class TestExecutingAcrossDatabases: diff --git a/tests/test_inspector.py b/tests/test_inspector.py index d129b9bb..4e751861 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -339,7 +339,8 @@ class TestReportedDatabase: @staticmethod def _client(federated, store_path=None): client = MagicMock() - client._federated = federated + client.covers_multiple = len(federated) > 1 + client.source_names = tuple(federated) client.store.db_path = store_path return client diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index 00c26401..5e148983 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -126,8 +126,8 @@ class TestNamingADatabaseDirectly: await _seed(config, "alpha", ["alpha document about cats"]) async with HaikuRAG(temp_db_path, config=config, create=True) as rag: - assert rag._federated == {} - assert rag._source is None + assert not rag.covers_multiple + assert rag.source is None assert rag.store.db_path == temp_db_path @pytest.mark.asyncio @@ -137,8 +137,8 @@ class TestNamingADatabaseDirectly: await _seed(config, "alpha", ["alpha document about cats"]) async with HaikuRAG(config=config) as rag: - assert rag._federated == {} - assert rag._source == "alpha" + assert not rag.covers_multiple + assert rag.source == "alpha" results = await rag.search("cats", search_type="fts", limit=10) assert [r.source for r in results] == ["alpha"] @@ -168,7 +168,7 @@ class TestOpeningDatabases: rag._session._open = gated 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 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: 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 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: 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: @@ -379,6 +379,84 @@ class TestClosingASet: 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: """A client for one of a set wraps a database the set opened.""" diff --git a/tests/test_multi_db_ask.py b/tests/test_multi_db_ask.py index 2265f6da..673ae51c 100644 --- a/tests/test_multi_db_ask.py +++ b/tests/test_multi_db_ask.py @@ -293,7 +293,7 @@ class TestStandaloneCapabilities: await run._close() 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 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) rag = await capability._ensure_rag() try: - assert rag._source == "alpha" + assert rag.source == "alpha" finally: await capability._close() diff --git a/tests/test_utils.py b/tests/test_utils.py index 8d80acb8..300e7ad4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -785,7 +785,8 @@ async def test_format_citations_rich_names_the_database_when_federating(): source="medic", ) 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)) @@ -808,7 +809,8 @@ async def test_format_citations_rich_omits_the_database_for_one_database(): source="medic", ) client = AsyncMock() - client._federated = {} + client.covers_multiple = False + client.source_names = () output = _render_rich(await format_citations_rich([citation], client))