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

@@ -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.

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.

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
+            # The set shares this client's embedder and reranker, so this is the
-            # place it is closed — and only if a text query ever built it.
+            # only place they are closed — and only if anything built them.
-            reranker = self.__dict__.get("reranker")
+            await self._aclose_cached("embedder")
-            if reranker is not None:
+            await self._aclose_cached("reranker")
-                try:
-                    await reranker.aclose()
-                except Exception:
-                    logger.debug("Closing the reranker failed", exc_info=True)
             return False
         await self._await_vacuum_tasks()
-        # Best-effort: __aexit__ may run during exception unwinding, and a
+        # Accessed so the store's embedder is closed even where nothing used it;
-        # raising close must not mask the original exception. The reranker is
+        # `cached_property` stores it, which is what `_aclose_cached` discards.
-        # a cached_property — close it only if it was materialized.
+        _ = self.embedder
-        try:
+        await self._aclose_cached("embedder")
-            await self.embedder.aclose()
+        await self._aclose_cached("reranker")
-            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)
         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()

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):

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(