From bddb32b469251b5e9286788554c7cc21c34c7726 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 24 Aug 2026 14:54:56 +0300 Subject: [PATCH] Build the embedder for a set from configuration An embedder is a function of configuration, not of a database, and the databases in a selection share one, so a client covering a set builds it on first use and closes it on teardown. Operations that need one database say so instead of surfacing a missing store. --- CHANGELOG.md | 1 + docs/configuration/storage.md | 6 + haiku_rag_slim/haiku/rag/client/__init__.py | 95 ++++++++++--- haiku_rag_slim/haiku/rag/store/exceptions.py | 6 +- tests/test_multi_db.py | 140 ++++++++++++++++++- 5 files changed, 226 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f725eaa2..5193a05d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ ### Fixed +- `client.chunk()` and `client.embedder` work on a client covering `lancedb.databases`: an embedder is a function of configuration, so the client builds one on first use and closes it on teardown. Operations that need one database (`create_document`, `import_document(s)`, `create_document_from_source`, `update_document`, `delete_document`, `rebuild_database`, `vacuum`, `visualize_chunk`, `close`) raise `AmbiguousDatabaseError` naming the databases covered, instead of `AttributeError`. - The chat TUI's document filter selects documents by id and names each document's database, instead of matching the displayed title or URI as a substring across every database. - `doctor`'s docling-serve probe sends `X-Api-Key`, so an instance requiring a key is reported reachable rather than unreachable. - The picture-description request to the public OpenAI endpoint sends `OPENAI_API_KEY`; it carried no authorization header. diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 2aba421e..3f82f894 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -232,6 +232,12 @@ the same queries: retrieval MAP 0.9914 with a reranker against 0.9918 for the same corpus in a single database, and 0.6044 without one against 0.9798. The cost is that a reranker scores candidates in proportion to the number of databases. +Converting, chunking and title generation are functions of the configuration +rather than of a database, so they work on a client covering the set. Writing, +rebuilding and vacuuming name one database: asking a set-covering client raises +`AmbiguousDatabaseError`, and `client.clients_for(["name"])` returns a client for +one of them, writable when the covering client is. + A database that cannot be opened fails the whole query and is named in the error. A result set silently missing one of the databases asked for cannot be told apart from a complete one. diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index fe34e255..d29aa1a8 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -18,9 +18,11 @@ import httpx from haiku.rag.client.documents import DocumentImport from haiku.rag.config import AppConfig, get_config from haiku.rag.converters import get_converter +from haiku.rag.embeddings import get_embedder from haiku.rag.reranking import get_reranker from haiku.rag.store.engine import Store from haiku.rag.store.exceptions import ( + AmbiguousDatabaseError, ConfigMismatchError, MigrationRequiredError, ReadOnlyError, @@ -78,6 +80,18 @@ async def first_found( return None +async def _aclose_quietly(closeable: Any, what: str) -> None: + """Close, reporting failure to the log rather than raising. + + Teardown can run while an exception unwinds, so a raising close must + neither mask that exception nor stop a sibling from being closed. + """ + try: + await closeable.aclose() + except Exception: + logger.debug("Closing the %s failed on teardown", what, exc_info=True) + + def _spell(embedding: tuple[str | None, str | None, int | None]) -> str: """An embedder identity, for an error message.""" provider, name, vector_dim = embedding @@ -158,9 +172,17 @@ class HaikuRAG: """ return self._read_only - @property + @cached_property def embedder(self) -> "EmbedderWrapper": - """The embedder owned by the Store, reused across all operations.""" + """The embedder for the databases this client covers. + + An embedder is a function of configuration rather than of a database, + and the databases in a selection are required to share one, so a client + 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: + return get_embedder(config=self._config) return self.store.embedder @cached_property @@ -357,29 +379,30 @@ class HaikuRAG: # store either. if self._federated: await self._close_clients() - # The set shares one reranker, this client's, so this is the only - # place it is closed — and only if a text query ever built it. - reranker = self.__dict__.get("reranker") - if reranker is not None: - try: - await reranker.aclose() - except Exception: - logger.debug("Closing the reranker failed", exc_info=True) + # The set shares this client's embedder and reranker, so this is the + # only place they are closed — and only if anything built them. + await self._aclose_cached("embedder") + await self._aclose_cached("reranker") return False await self._await_vacuum_tasks() - # Best-effort: __aexit__ may run during exception unwinding, and a - # raising close must not mask the original exception. The reranker is - # a cached_property — close it only if it was materialized. - try: - await self.embedder.aclose() - reranker = self.__dict__.get("reranker") - if reranker is not None: - await reranker.aclose() - except Exception: - logger.debug("Closing embedder/reranker failed on teardown", exc_info=True) + # Accessed so the store's embedder is closed even where nothing used it; + # `cached_property` stores it, which is what `_aclose_cached` discards. + _ = self.embedder + await self._aclose_cached("embedder") + await self._aclose_cached("reranker") self.close() return False + async def _aclose_cached(self, name: str) -> None: + """Close a cached_property this client materialized, and discard it. + + Discarded rather than left in place so that re-entering the client + builds a fresh one instead of reusing something already closed. + """ + cached = self.__dict__.pop(name, None) + if cached is not None: + await _aclose_quietly(cached, name) + async def _await_vacuum_tasks(self) -> None: """Drain background vacuum work and run a final collapse before teardown. @@ -481,6 +504,8 @@ class HaikuRAG: ) -> Document: from haiku.rag.client.documents import create_document + self._require_one_database("create_document") + return await create_document(self, content, uri, title, metadata, format) async def import_document( @@ -493,6 +518,8 @@ class HaikuRAG: ) -> Document: from haiku.rag.client.documents import import_document + self._require_one_database("import_document") + return await import_document( self, docling_document, chunks, uri, title, metadata ) @@ -503,6 +530,8 @@ class HaikuRAG: ) -> list[Document]: from haiku.rag.client.documents import import_documents + self._require_one_database("import_documents") + return await import_documents(self, imports) async def create_document_from_source( @@ -518,6 +547,8 @@ class HaikuRAG: ) -> Document | list[Document]: from haiku.rag.client.documents import create_document_from_source + self._require_one_database("create_document_from_source") + return await create_document_from_source( self, source, @@ -542,6 +573,8 @@ class HaikuRAG: ) -> Document: from haiku.rag.client.documents import update_document + self._require_one_database("update_document") + return await update_document( self, document_id, @@ -659,6 +692,8 @@ class HaikuRAG: """ from haiku.rag.client.documents import parent_uri_filter + self._require_one_database("delete_document") + async with self.store.write_transaction(): # Resolve existence and collect the subtree under the lock so two # concurrent deletes of the same id can't both proceed, and children @@ -754,6 +789,20 @@ class HaikuRAG: return sum(counts) return await self.document_repository.count(filter=filter) + def _require_one_database(self, operation: str) -> None: + """Refuse an operation that has no meaning across a set of databases. + + Writing, rebuilding and vacuuming all have to name a database. Raised as + a domain error rather than surfacing the missing repository, so a caller + can tell an unsupported selection from a bug. + """ + if self._federated: + raise AmbiguousDatabaseError( + f"{operation} works on one database, and this client covers " + f"{', '.join(sorted(self._federated))}; select one with " + "clients_for([name])" + ) + def _name(self, document: Document | None) -> Document | None: """`document`, told which configured database it came from. @@ -858,6 +907,8 @@ class HaikuRAG: ) -> list: from haiku.rag.client.search import visualize_chunk + self._require_one_database("visualize_chunk") + return await visualize_chunk(self, chunk, refs, expand) async def rebuild_database( @@ -865,13 +916,17 @@ class HaikuRAG: ) -> AsyncGenerator[str, None]: from haiku.rag.client.rebuild import rebuild_database + self._require_one_database("rebuild_database") + async for doc_id in rebuild_database(self, mode): yield doc_id async def vacuum(self) -> None: """Optimize and clean up old versions across all tables.""" + self._require_one_database("vacuum") await self.store.vacuum() def close(self): """Close the underlying store connection.""" + self._require_one_database("close") self.store.close() diff --git a/haiku_rag_slim/haiku/rag/store/exceptions.py b/haiku_rag_slim/haiku/rag/store/exceptions.py index 3ad60070..26b256f7 100644 --- a/haiku_rag_slim/haiku/rag/store/exceptions.py +++ b/haiku_rag_slim/haiku/rag/store/exceptions.py @@ -17,7 +17,11 @@ class MigrationRequiredError(Exception): class AmbiguousDatabaseError(Exception): - """A command that works on one database was run against a configured set.""" + """An operation that works on one database was asked of a configured set. + + Raised by the CLI for a command that cannot tell which database to use, and + by the client for a method that has no meaning across several. + """ class SourceUnavailableError(Exception): diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index da2e3c7f..9298fa4d 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -8,7 +8,11 @@ from pydantic import ValidationError from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config from haiku.rag.config.models import AppConfig, LanceDBConfig -from haiku.rag.store.exceptions import ConfigMismatchError, SourceUnavailableError +from haiku.rag.store.exceptions import ( + AmbiguousDatabaseError, + ConfigMismatchError, + SourceUnavailableError, +) from haiku.rag.store.models import Chunk, DocumentItem from haiku.rag.utils import locate_database @@ -338,6 +342,140 @@ class TestLookupByIdentifier: assert await rag.get_document_by_uri("test://nowhere") is None +class TestDatabaseIndependentWork: + """Converting, chunking and titling are functions of the configuration, not + of a database, so covering a set does not stop them.""" + + @pytest.mark.asyncio + async def test_chunking_opens_no_database(self, tmp_path, monkeypatch): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + opened: list[str] = [] + + async def refuse(self, name, location): + opened.append(name) + raise AssertionError("opened a database to chunk a document") + + monkeypatch.setattr(HaikuRAG, "_open_client", refuse) + + doc = DoclingDocument(name="note") + doc.add_text( + label=DocItemLabel.TEXT, text="Boltzmann machines are energy based." + ) + + async with HaikuRAG(config=config, read_only=True) as rag: + chunks = await rag.chunk(doc) + + assert opened == [] + assert [c.content for c in chunks] + + @pytest.mark.asyncio + async def test_the_embedder_is_built_once_and_closed_once( + self, tmp_path, monkeypatch + ): + """The parent owns the embedder it built, so leaving the context closes + it, once.""" + from haiku.rag.embeddings import EmbedderWrapper + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + + closed: list[object] = [] + original = EmbedderWrapper.aclose + + async def counting(self): + closed.append(self) + return await original(self) + + monkeypatch.setattr(EmbedderWrapper, "aclose", counting) + + rag = HaikuRAG(config=config, read_only=True) + async with rag: + built = rag.embedder + assert rag.embedder is built + + assert closed == [built] + + @pytest.mark.asyncio + async def test_re_entering_a_set_builds_a_fresh_embedder(self, tmp_path): + """Teardown closes the embedder, so keeping it would hand the next + context one that is already closed.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + + rag = HaikuRAG(config=config, read_only=True) + async with rag: + first = rag.embedder + async with rag: + assert rag.embedder is not first + + @pytest.mark.asyncio + async def test_re_entering_one_database_builds_a_fresh_embedder(self, temp_db_path): + """One database opens a new store on re-entry, and the embedder is that + store's.""" + rag = HaikuRAG(temp_db_path, create=True) + async with rag: + first = rag.embedder + async with rag: + assert rag.embedder is rag.store.embedder + assert rag.embedder is not first + + @pytest.mark.asyncio + async def test_a_set_nobody_asked_anything_of_builds_no_embedder(self, tmp_path): + """Built on first use, so a client that answered nothing holds nothing.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + + async with HaikuRAG(config=config, read_only=True) as rag: + assert "embedder" not in rag.__dict__ + + @pytest.mark.asyncio + async def test_one_database_still_uses_its_store_s_embedder(self, temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as rag: + assert rag.embedder is rag.store.embedder + + +class TestOperationsThatNeedOneDatabase: + @pytest.mark.asyncio + async def test_writing_names_the_databases_it_covers(self, tmp_path): + """A domain error, so a caller can tell an unsupported selection from a + missing attribute.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"): + await rag.create_document("orphan") + with pytest.raises(AmbiguousDatabaseError, match="clients_for"): + await rag.vacuum() + + @pytest.mark.asyncio + async def test_a_selected_database_is_still_writable(self, tmp_path): + """Naming one of the set is how a write picks its database.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + dim = get_config().embeddings.model.vector_dim + written = DoclingDocument(name="written") + written.add_text(label=DocItemLabel.TEXT, text="written") + + async with HaikuRAG(config=config) as rag: + alpha = (await rag.clients_for(["alpha"]))[0] + assert alpha.is_read_only is False + document = await alpha.import_document( + written, + [Chunk(content="written", embedding=[0.1] * dim, order=0)], + uri="test://alpha/written", + ) + assert await alpha.count_documents() == 2 + + assert document.id is not None + + class TestOneQueryVector: @pytest.mark.asyncio async def test_a_search_embeds_the_query_once_for_the_whole_set(