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.
This commit is contained in:
Yiorgis Gozadinos
parent
a9e66b001b
commit
bddb32b469
5 changed files with 226 additions and 22 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Reference in a new issue