diff --git a/CHANGELOG.md b/CHANGELOG.md index 84a8f398..b165753a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,5 @@ # Changelog + ## [Unreleased] ### Added @@ -37,6 +38,7 @@ - `haiku.rag.capabilities.EvidenceState`: the state base `RAGState` and `AnalysisState` derive from, with `begin_invocation()` for the per-question reset. `RAGCapabilityBase.evidence_record()` and `citation_index()` expose what a capability recorded, so a host reads it without reaching into `capability.state`. - The `haiku.rag` package declares the `jina` extra, so `provider: jina-local` is supported by declaration rather than through `cross-encoder`'s transitive `transformers` and `torch`. Raises the full package's torch floor to 2.0. - `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout. +- `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured. ### Changed diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index e97dd695..2df90d77 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -19,6 +19,7 @@ from haiku.rag.config import AppConfig, get_config from haiku.rag.converters import get_converter from haiku.rag.reranking import get_reranker from haiku.rag.store.engine import Store +from haiku.rag.store.exceptions import SourceUnavailableError from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import extract_items @@ -71,6 +72,7 @@ class HaikuRAG: skip_validation: bool = False, create: bool = False, read_only: bool = False, + sources: list[str] | None = None, ): """Initialize the RAG client with a database path. @@ -81,6 +83,9 @@ class HaikuRAG: skip_validation: Whether to skip configuration validation on database load. create: Whether to create the database if it doesn't exist. read_only: Whether to open the database in read-only mode. + sources: Names from ``config.lancedb.databases`` this client covers. + None means all of them. Ignored when a single ``uri`` is + configured. """ self._config = config if config is not None else get_config() if db_path is None: @@ -93,6 +98,11 @@ class HaikuRAG: self._vacuum_tasks: set[asyncio.Task] = set() self._last_vacuum_at: float | None = None self._vacuum_dirty = False + self._requested_sources = sources + self._clients: dict[str, HaikuRAG] = {} + self._federated: dict[str, str] = {} + self._clients_lock = asyncio.Lock() + self._source: str | None = None @property def is_read_only(self) -> bool: @@ -113,30 +123,157 @@ class HaikuRAG: """ return get_reranker(config=self._config) - async def __aenter__(self): - """Async context manager entry — initializes store and repositories.""" - self.store = Store( - self._db_path, - config=self._config, - skip_validation=self._skip_validation, - create=self._create, - read_only=self._read_only, + @staticmethod + def _locate(location: str) -> tuple[str, "Path | None"]: + """Split a configured location into (uri, db_path). + + A value with a scheme is a `lancedb.uri`; anything else is a local path. + Routing a local path through `uri` would have `ConnectionMode` classify it + as object storage, which opens it without the existence check a local + database gets. + """ + if "://" in location: + return location, None + return "", Path(location) + + def _selected(self) -> dict[str, str]: + """The configured databases this client covers, name to location.""" + declared = self._config.lancedb.databases + if not declared: + return {} + if self._requested_sources is not None and not self._requested_sources: + raise ValueError( + "sources=[] selects no database; pass None for all of them" + ) + names = ( + list(declared) + if self._requested_sources is None + else list(self._requested_sources) ) - # If _initialize fails mid-way (e.g. migration check raises after - # connect), close the store so we don't leak the LanceDB connection — - # __aexit__ won't run because the `async with` never entered. + missing = [n for n in names if n not in declared] + if missing: + raise KeyError( + f"unknown database(s) {', '.join(sorted(missing))}; " + f"configured: {', '.join(sorted(declared))}" + ) + return {n: declared[n] for n in names} + + async def __aenter__(self): + """Async context manager entry — initializes store and repositories. + + A client covering several databases opens none of them here: which are + searched is a per-query choice, so they open on first use. `store` and the + repositories stay unset in that case, since they have no unambiguous + meaning across a set. + """ + selected = self._selected() + if len(selected) > 1: + self._federated = selected + return self + if selected: + [(self._source, location)] = selected.items() + uri, db_path = self._locate(location) + self._config = self._config.model_copy(deep=True) + self._config.lancedb.databases = {} + self._config.lancedb.uri = uri + if db_path is not None: + self._db_path = db_path + + failure: str | None = None try: - await self.store._initialize() - except BaseException: - self.store.close() - raise + self.store = Store( + self._db_path, + config=self._config, + skip_validation=self._skip_validation, + create=self._create, + read_only=self._read_only, + ) + # If _initialize fails mid-way (e.g. migration check raises after + # connect), close the store so we don't leak the LanceDB connection — + # __aexit__ won't run because the `async with` never entered. + try: + await self.store._initialize() + except BaseException: + self.store.close() + raise + except Exception as error: + # A legacy `uri` or `db_path` client has no name to report instead, so + # its error passes through as it always has. + if self._source is None: + raise + failure = type(error).__name__ + if failure is not None: + # Raised outside the except block on purpose. A database named in + # config is reported by name, and the original spells out the path or + # the bucket: `from None` would only stop it being *printed*, leaving + # it on `__context__` for anything that walks the chain. + raise SourceUnavailableError( + f"database {self._source!r} could not be opened: {failure}" + ) self.document_repository = DocumentRepository(self.store) self.chunk_repository = ChunkRepository(self.store) self.document_item_repository = DocumentItemRepository(self.store) return self + async def clients_for(self, names: list[str]) -> list["HaikuRAG"]: + """The clients for these databases, opening any not yet open. + + Opening is per query rather than at entry: a set of 25 configured + databases is typically queried a few at a time, and a database nobody + asked for must not be able to fail a query, or be opened for nothing. + """ + unknown = [n for n in names if n not in self._federated] + if unknown: + raise KeyError( + f"unknown database(s) {', '.join(sorted(unknown))}; configured: " + f"{', '.join(sorted(self._federated))}" + ) + async with self._clients_lock: + for name in names: + if name not in self._clients: + self._clients[name] = await self._open_client( + name, self._federated[name] + ) + return [self._clients[n] for n in names] + + async def _open_client(self, name: str, location: str) -> "HaikuRAG": + uri, db_path = self._locate(location) + config = self._config.model_copy(deep=True) + config.lancedb.databases = {} + config.lancedb.uri = uri + client = HaikuRAG( + db_path, + config=config, + skip_validation=self._skip_validation, + read_only=self._read_only, + ) + client._source = name + return await client.__aenter__() + + async def _close_clients(self) -> None: + for client in self._clients.values(): + try: + await client.__aexit__(None, None, None) + except Exception: + logger.debug("Closing a database failed on teardown", exc_info=True) + self._clients.clear() + async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 """Async context manager exit.""" + # Branch on what this client covers, not on what it happened to open: + # a federating client that answered no query has nothing open and no + # 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) + 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 @@ -463,10 +600,27 @@ class HaikuRAG: search_type: SearchType | None = None, filter: str | None = None, include_images: bool = True, + sources: list[str] | None = None, ) -> list[SearchResult]: - from haiku.rag.client.search import search + from haiku.rag.client.search import search, search_sources - return await search(self, query, limit, search_type, filter, include_images) + if self._federated: + return await search_sources( + self, query, limit, search_type, filter, include_images, sources + ) + if sources is not None and not sources: + return [] + if sources is not None and sources != [self._source]: + raise KeyError( + f"unknown database(s) {', '.join(sources) or '(none)'}; this " + f"client covers {self._source or 'a single unnamed database'}" + ) + results = await search(self, query, limit, search_type, filter, include_images) + # 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 + return results async def expand_context( self, diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index c47d9b98..d8336de1 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -1,3 +1,4 @@ +import asyncio import base64 from collections.abc import Sequence from typing import TYPE_CHECKING @@ -38,7 +39,9 @@ async def search( if limit is None: limit = client._config.search.limit - candidates = await _fetch(client, query, limit, search_type, filter) + candidates = await _fetch( + client, query, _fetch_limit(client, query, limit), search_type, filter + ) chunk_results = await _rank(client, query, candidates, limit) results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] @@ -50,10 +53,132 @@ async def search( return results +async def search_sources( + client: "HaikuRAG", + query: "str | bytes | PILImage.Image", + limit: int | None = None, + search_type: SearchType | None = None, + filter: str | None = None, + include_images: bool = True, + sources: list[str] | None = None, +) -> list[SearchResult]: + """Search several databases and fuse their results into one ranked list. + + Fetch, fuse, truncate, then enrich: enrichment runs on the survivors through + the database each came from, so it costs the same as a single-database search + rather than multiplying by the number searched. + """ + if limit is None: + limit = client._config.search.limit + + names = list(client._federated) if sources is None else list(sources) + if not names: + return [] + selected = await client.clients_for(names) + + # One over-fetch decision, and one reranker, for the whole set. + fetch_limit = _fetch_limit(client, query, limit) + per_source = await asyncio.gather( + *(_fetch(c, query, fetch_limit, search_type, filter) for c in selected) + ) + + ranked = await _fuse(client, selected, query, per_source, limit) + + results: list[SearchResult] = [] + for owner, chunk, score in ranked: + result = SearchResult.from_chunk(chunk, score) + result.source = owner._source + results.append(result) + results = _dedup_picture_chunks(results) + + if include_images: + by_owner: dict[str, list[SearchResult]] = {} + for result in results: + if result.source: + by_owner.setdefault(result.source, []).append(result) + await asyncio.gather( + *( + _populate_image_data(client._clients[name], owned) + for name, owned in by_owner.items() + ) + ) + + return results + + +async def _fuse( + federator: "HaikuRAG", + clients: list["HaikuRAG"], + query: "str | bytes | PILImage.Image", + per_source: list[list[tuple[Chunk, float]]], + limit: int, +) -> list[tuple["HaikuRAG", Chunk, float]]: + """One ranked list from several, keeping each candidate's owner. + + A configured reranker scores the union directly, which is what makes ranking + across databases tractable: it compares query against document and does not + care where a candidate came from. Without one, reciprocal rank fusion over the + per-database rankings, since scores from separate indexes are not comparable. + """ + owned = [ + (client, chunk, score) + for client, candidates in zip(clients, per_source, strict=True) + for chunk, score in candidates + ] + if not owned: + return [] + + # An image query has no text for a reranker to score against, and the check + # precedes `reranker`, which builds the reranker on first access. + if isinstance(query, str): + reranker = federator.reranker + if reranker is not None: + chunks = [chunk for _, chunk, _ in owned] + if federator._config.reranking.multimodal: + await asyncio.gather( + *( + _attach_picture_data( + c, [chunk for owner, chunk, _ in owned if owner is c] + ) + for c in clients + ) + ) + reranked = await reranker.rerank(query, chunks, top_n=limit) + owner_of = {id(chunk): client for client, chunk, _ in owned} + return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked] + + scored: list[tuple[float, HaikuRAG, Chunk]] = [] + for client, candidates in zip(clients, per_source, strict=True): + for rank, (chunk, _) in enumerate(candidates): + scored.append((1.0 / (_RRF_K + rank + 1), client, chunk)) + scored.sort(key=lambda item: item[0], reverse=True) + return [(client, chunk, score) for score, client, chunk in scored[:limit]] + + +# Reciprocal rank fusion's smoothing constant, the value the literature uses. +_RRF_K = 60 + + # Candidates per requested result when a reranker will re-order them. _RERANK_OVERFETCH = 10 +def _fetch_limit( + client: "HaikuRAG", + query: "str | bytes | PILImage.Image", + limit: int, +) -> int: + """How many candidates to fetch per database. + + Only a text query with a reranker over-fetches: an image query keeps its + vector ranking, and the type is checked before `reranker`, which loads model + weights for a local one on first access. + """ + if not isinstance(query, str): + return limit + return limit * _RERANK_OVERFETCH if client.reranker else limit + + async def _fetch( client: "HaikuRAG", query: "str | bytes | PILImage.Image", @@ -63,17 +188,14 @@ async def _fetch( ) -> list[tuple[Chunk, float]]: """Candidates from one database, ranked by that database. - Over-fetches when a reranker will re-order them. Separate from `_rank` so a - caller searching several databases can fuse their candidates before anything - is ranked or enriched. + `limit` is how many to fetch, already including any over-fetch the caller + wants. Deciding that here would have each database consult its own reranker, + and a local reranker loads model weights per instance. """ if isinstance(query, str): if search_type is None: search_type = "hybrid" - fetch_limit = limit * _RERANK_OVERFETCH if client.reranker else limit - return await client.chunk_repository.search( - query, fetch_limit, search_type, filter - ) + return await client.chunk_repository.search(query, limit, search_type, filter) embedder = client.embedder if not embedder.supports_images: diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 181f70b9..06074531 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -101,16 +101,30 @@ class LanceDBConfig(ConfigModel): never re-checks, so a long-lived reader never sees another process's writes. The cache sizes are per process, since the session is shared across connections. + + `databases` maps a name to a location, for searching several at once. The + name is what results and citations carry, so a location never leaves the + configuration. Mutually exclusive with `uri`. """ uri: str = "" api_key: str = "" region: str = "" storage_options: dict[str, str] = Field(default_factory=dict) + databases: dict[str, str] = Field(default_factory=dict) read_consistency_interval_seconds: float | None = Field(default=30, ge=0) index_cache_size_bytes: int | None = Field(default=None, ge=0) metadata_cache_size_bytes: int | None = Field(default=None, ge=0) + @model_validator(mode="after") + def _one_way_of_naming_databases(self) -> "LanceDBConfig": + if self.uri and self.databases: + raise ValueError( + "lancedb.uri and lancedb.databases are mutually exclusive: " + "use uri for one database, databases for several" + ) + return self + class EmbeddingsConfig(ConfigModel): model: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig) diff --git a/haiku_rag_slim/haiku/rag/store/__init__.py b/haiku_rag_slim/haiku/rag/store/__init__.py index fe1d3c1f..564e9a69 100644 --- a/haiku_rag_slim/haiku/rag/store/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/__init__.py @@ -1,4 +1,14 @@ -from .exceptions import MigrationRequiredError, ReadOnlyError +from .exceptions import ( + MigrationRequiredError, + ReadOnlyError, + SourceUnavailableError, +) from .models import Chunk, Document -__all__ = ["Chunk", "Document", "MigrationRequiredError", "ReadOnlyError"] +__all__ = [ + "Chunk", + "Document", + "MigrationRequiredError", + "ReadOnlyError", + "SourceUnavailableError", +] diff --git a/haiku_rag_slim/haiku/rag/store/exceptions.py b/haiku_rag_slim/haiku/rag/store/exceptions.py index a811b5c4..20f8d796 100644 --- a/haiku_rag_slim/haiku/rag/store/exceptions.py +++ b/haiku_rag_slim/haiku/rag/store/exceptions.py @@ -8,3 +8,12 @@ class MigrationRequiredError(Exception): """Database requires migration. Run 'haiku-rag migrate' to upgrade.""" pass + + +class SourceUnavailableError(Exception): + """A configured database could not be opened. + + Carries the configured name and never the location: a path or URI in an + error message travels into logs and into whatever a consumer renders, and + the point of naming databases is that locations stay in the configuration. + """ diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 2a75fd4a..f355d614 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -136,10 +136,15 @@ class SearchResult(BaseModel): ``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not include the metadata of any other chunks merged with it. Never part of ``format_for_agent`` output. + + ``source`` names the configured database a result came from, and is None when + only one is configured. It is the name from ``lancedb.databases``, never a + path or URI, so a location cannot travel in a result, a citation or a log. """ content: str score: float + source: str | None = None chunk_id: str | None = None chunk_ids: list[str] = [] chunk_meta: dict = {} diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py new file mode 100644 index 00000000..9668d14f --- /dev/null +++ b/tests/test_multi_db.py @@ -0,0 +1,460 @@ +import pytest +from docling_core.types.doc.document import DoclingDocument +from docling_core.types.doc.labels import DocItemLabel +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 SourceUnavailableError +from haiku.rag.store.models import Chunk, DocumentItem + + +class TestConfig: + def test_databases_and_uri_are_mutually_exclusive(self): + with pytest.raises(ValidationError, match="databases"): + LanceDBConfig( + uri="s3://b/one.lancedb", databases={"one": "s3://b/one.lancedb"} + ) + + def test_databases_alone_is_fine(self): + config = LanceDBConfig(databases={"one": "s3://b/one.lancedb"}) + assert config.databases == {"one": "s3://b/one.lancedb"} + + def test_uri_alone_is_fine(self): + assert LanceDBConfig(uri="s3://b/one.lancedb").databases == {} + + +def _config(tmp_path, names) -> AppConfig: + return AppConfig( + lancedb=LanceDBConfig( + databases={n: str(tmp_path / f"{n}.lancedb") for n in names} + ) + ) + + +async def _seed(config, name, contents): + """Precomputed embeddings and FTS queries keep the embedder out of the way: + these tests are about fusion, not retrieval quality.""" + dim = get_config().embeddings.model.vector_dim + async with HaikuRAG(config=config, create=True, sources=[name]) as rag: + for content in contents: + doc = DoclingDocument(name=content) + doc.add_text(label=DocItemLabel.TEXT, text=content) + await rag.import_document( + doc, + [Chunk(content=content, embedding=[0.1] * dim, order=0)], + uri=f"test://{name}/{content}", + ) + + +class TestFederatedSearch: + @pytest.mark.asyncio + async def test_results_carry_their_source(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + results = await rag.search("cats", limit=10, search_type="fts") + + assert {r.source for r in results} == {"alpha", "beta"} + for r in results: + assert r.source is not None + assert r.source in r.content + + @pytest.mark.asyncio + async def test_sources_selects_a_subset(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + results = await rag.search( + "cats", limit=10, search_type="fts", sources=["alpha"] + ) + + assert {r.source for r in results} == {"alpha"} + + @pytest.mark.asyncio + async def test_unknown_source_is_rejected(self, tmp_path): + config = _config(tmp_path, ["alpha"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + async with HaikuRAG(config=config) as rag: + with pytest.raises(KeyError, match="nope"): + await rag.search("cats", search_type="fts", sources=["nope"]) + + @pytest.mark.asyncio + async def test_an_unopenable_database_fails_the_query(self, tmp_path): + config = _config(tmp_path, ["alpha", "missing"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + with pytest.raises(SourceUnavailableError, match="missing"): + async with HaikuRAG(config=config) as rag: + await rag.search("cats", search_type="fts") + + +class TestSingleDatabaseUnchanged: + @pytest.mark.asyncio + async def test_source_is_unset_without_configured_databases(self, temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as rag: + doc = DoclingDocument(name="one") + doc.add_text(label=DocItemLabel.TEXT, text="a document about cats") + await rag.import_document( + doc, + [ + Chunk( + content="a document about cats", + embedding=[0.1] * get_config().embeddings.model.vector_dim, + order=0, + ) + ], + uri="test://one", + ) + results = await rag.search("cats", search_type="fts") + + assert results + assert all(r.source is None for r in results) + + +class TestLocate: + def test_a_scheme_is_a_uri(self): + assert HaikuRAG._locate("s3://bucket/one.lancedb") == ( + "s3://bucket/one.lancedb", + None, + ) + + def test_anything_else_is_a_local_path(self): + uri, db_path = HaikuRAG._locate("/data/one.lancedb") + assert uri == "" + assert db_path is not None and str(db_path) == "/data/one.lancedb" + + +class TestSelection: + @pytest.mark.asyncio + async def test_unknown_source_at_construction_is_rejected(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + + with pytest.raises(KeyError, match="nope"): + async with HaikuRAG(config=config, sources=["nope"]): + pass + + @pytest.mark.asyncio + async def test_unknown_source_across_several_databases_is_rejected(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + with pytest.raises(KeyError, match="nope"): + await rag.search("cats", search_type="fts", sources=["nope"]) + + @pytest.mark.asyncio + async def test_no_matches_anywhere_returns_nothing(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + assert await rag.search("aardvarks", search_type="fts") == [] + + +class StubReranker: + """Scores the union, reversing it so the ordering is unmistakably its own.""" + + def __init__(self): + self.seen: list[str] = [] + + async def rerank(self, query, chunks, top_n): + self.seen = [c.content for c in chunks] + # Whatever the caller attached before handing them over. + self.attached = { + c.content.split()[0]: c._picture_data + for c in chunks + if getattr(c, "_picture_data", None) + } + return [(c, 1.0 - i) for i, c in enumerate(reversed(chunks))][:top_n] + + +class TestRerankerFusion: + @pytest.mark.asyncio + async def test_the_reranker_scores_the_union_and_owners_survive( + self, tmp_path, monkeypatch + ): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + stub = StubReranker() + monkeypatch.setattr(HaikuRAG, "reranker", property(lambda self: stub)) + + async with HaikuRAG(config=config) as rag: + results = await rag.search("cats", limit=2, search_type="fts") + + # It saw both databases' candidates, not one database at a time. + assert len(stub.seen) == 2 + assert {c.split()[0] for c in stub.seen} == {"alpha", "beta"} + # Each result still knows which database it came from. + for r in results: + assert r.source is not None + assert r.content.startswith(r.source) + + @pytest.mark.asyncio + async def test_a_closing_failure_does_not_mask_the_exit(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + rag = HaikuRAG(config=config) + await rag.__aenter__() + await rag.clients_for(["alpha", "beta"]) + + async def boom(exc_type, exc_val, exc_tb): + raise RuntimeError("close failed") + + rag._clients["alpha"].__aexit__ = boom # ty: ignore[invalid-assignment] + + await rag.__aexit__(None, None, None) + + assert rag._clients == {} + + @pytest.mark.asyncio + async def test_multimodal_reranking_attaches_each_database_own_pictures( + self, tmp_path, monkeypatch + ): + """Picture self_refs repeat across databases exactly as they do across + documents, so the pre-rerank attach must stay per database.""" + config = _config(tmp_path, ["alpha", "beta"]) + config.reranking.multimodal = True + dim = get_config().embeddings.model.vector_dim + + for name in ("alpha", "beta"): + async with HaikuRAG(config=config, create=True, sources=[name]) as rag: + doc = DoclingDocument(name=name) + doc.add_text(label=DocItemLabel.TEXT, text=f"{name} figure of cats") + await rag.import_document( + doc, + [ + Chunk( + content=f"{name} figure of cats", + embedding=[0.1] * dim, + order=0, + metadata={ + "doc_item_refs": ["#/pictures/0"], + "labels": ["picture"], + }, + ) + ], + uri=f"test://{name}/figure", + ) + [document] = await rag.list_documents() + assert document.id is not None + await rag.document_item_repository.create_items( + document.id, + [ + DocumentItem( + document_id=document.id, + position=0, + self_ref="#/pictures/0", + label="picture", + text=f"caption {name}", + picture_data=f"bytes-{name}".encode(), + ) + ], + ) + + stub = StubReranker() + monkeypatch.setattr(HaikuRAG, "reranker", property(lambda self: stub)) + + async with HaikuRAG(config=config) as rag: + await rag.search("cats", limit=2, search_type="fts") + + assert stub.attached == {"alpha": b"bytes-alpha", "beta": b"bytes-beta"} + + +class TestLazyOpening: + @pytest.mark.asyncio + async def test_entering_opens_nothing(self, tmp_path): + """25 configured databases queried a few at a time must not all open.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + assert rag._clients == {} + + @pytest.mark.asyncio + async def test_only_the_selected_database_opens(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + await rag.search("cats", search_type="fts", sources=["alpha"]) + assert list(rag._clients) == ["alpha"] + + @pytest.mark.asyncio + async def test_an_unselected_broken_database_does_not_break_a_query(self, tmp_path): + """A database nobody asked for cannot fail a query.""" + config = _config(tmp_path, ["alpha", "missing"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + async with HaikuRAG(config=config) as rag: + results = await rag.search("cats", search_type="fts", sources=["alpha"]) + + assert [r.source for r in results] == ["alpha"] + + @pytest.mark.asyncio + async def test_a_selected_broken_database_fails_the_query(self, tmp_path): + config = _config(tmp_path, ["alpha", "missing"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + async with HaikuRAG(config=config) as rag: + with pytest.raises(SourceUnavailableError, match="missing"): + await rag.search("cats", search_type="fts") + + +class TestOneNamedDatabase: + @pytest.mark.asyncio + async def test_a_single_named_database_keeps_its_name(self, tmp_path): + """Named in config is named in results, even as the only entry.""" + config = _config(tmp_path, ["alpha"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + async with HaikuRAG(config=config) as rag: + results = await rag.search("cats", search_type="fts") + + assert results + assert all(r.source == "alpha" for r in results) + + @pytest.mark.asyncio + async def test_selecting_nothing_at_construction_is_rejected(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + + with pytest.raises(ValueError, match="selects no database"): + async with HaikuRAG(config=config, sources=[]): + pass + + @pytest.mark.asyncio + async def test_selecting_nothing_means_the_same_with_one_database(self, tmp_path): + """`sources=[]` selects nothing whether one database is configured or + several, rather than raising on one path and returning nothing on the + other.""" + config = _config(tmp_path, ["alpha"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + async with HaikuRAG(config=config) as rag: + assert await rag.search("cats", search_type="fts", sources=[]) == [] + + @pytest.mark.asyncio + async def test_selecting_nothing_per_query_returns_nothing(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + assert await rag.search("cats", search_type="fts", sources=[]) == [] + + +class TestOneReranker: + @pytest.mark.asyncio + async def test_the_set_builds_one_reranker_for_a_text_query( + self, tmp_path, monkeypatch + ): + """Local rerankers load model weights per instance, so a set of + databases must build one, not one each.""" + built = [] + monkeypatch.setattr( + "haiku.rag.client.get_reranker", + lambda config: built.append(config) or StubReranker(), + ) + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + built.clear() + + async with HaikuRAG(config=config) as rag: + await rag.search("cats", limit=2, search_type="fts") + + assert len(built) == 1, f"built {len(built)} rerankers" + + @pytest.mark.asyncio + async def test_an_image_query_builds_no_reranker(self, tmp_path, monkeypatch): + """Opening a database must not build one either: an image query has no + text to score against and never uses it.""" + built = [] + monkeypatch.setattr( + "haiku.rag.client.get_reranker", + lambda config: built.append(config) or StubReranker(), + ) + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + built.clear() + + async with HaikuRAG(config=config) as rag: + await rag.clients_for(["alpha", "beta"]) + + assert built == [] + + @pytest.mark.asyncio + async def test_the_reranker_is_closed_once(self, tmp_path, monkeypatch): + """Handing the same object to every database and letting each close it + would close it N times, and the federator not at all.""" + closes = [] + + class CountingReranker(StubReranker): + async def aclose(self): + closes.append(1) + + monkeypatch.setattr( + "haiku.rag.client.get_reranker", lambda config: CountingReranker() + ) + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + await _seed(config, "beta", ["beta document about cats"]) + + async with HaikuRAG(config=config) as rag: + await rag.search("cats", limit=2, search_type="fts") + + assert closes == [1], f"closed {len(closes)} times" + + +class TestFailureNaming: + @pytest.mark.asyncio + async def test_a_single_named_database_is_reported_by_name(self, tmp_path): + """One configured database is still a named one: it must not fall back to + the raw error, which spells out the path.""" + config = _config(tmp_path, ["alpha"]) + + with pytest.raises(SourceUnavailableError, match="alpha") as caught: + async with HaikuRAG(config=config): + pass + + assert str(tmp_path) not in str(caught.value) + assert caught.value.__cause__ is None + + @pytest.mark.asyncio + async def test_a_legacy_uri_client_keeps_its_error(self, tmp_path): + """Nothing named it, so there is no name to report instead.""" + with pytest.raises(FileNotFoundError): + async with HaikuRAG(tmp_path / "nope.lancedb"): + pass + + @pytest.mark.asyncio + async def test_the_location_is_absent_from_the_whole_chain(self, tmp_path): + config = _config(tmp_path, ["alpha", "missing"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + with pytest.raises(SourceUnavailableError) as caught: + async with HaikuRAG(config=config) as rag: + await rag.search("cats", search_type="fts") + + rendered = str(caught.value) + error = caught.value.__cause__ or caught.value.__context__ + assert "missing.lancedb" not in rendered + assert error is None, "the location-bearing cause is still attached"