diff --git a/CHANGELOG.md b/CHANGELOG.md index 996d04d1..98d8e70a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - `api_key` on model and embedding-model config, overriding the provider's environment variable. Honored on the `openai` and `ollama` providers, `vllm` embedders and rerankers, the picture-description VLM endpoint, and `doctor`'s endpoint probes; other providers raise. - `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. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask` and `analyze` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`. +- `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. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask`, `analyze` and `chat` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`. - `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another. - `client.analyze(..., sources=[...])` analyzes across the selected databases: the sandbox mounts their documents under one flat `/documents/{id}/` namespace, resolving each id to the database holding it, and in-code `search()` covers the same selection. diff --git a/docs/cli.md b/docs/cli.md index 808e1d9c..8bea7bdb 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -24,8 +24,8 @@ The `haiku-rag` CLI provides complete document management functionality. haiku-rag add -h ``` - With `lancedb.databases` configured, `search`, `ask` and `analyze` cover - every database in it. Every other command works on one, named with + With `lancedb.databases` configured, `search`, `ask`, `analyze` and `chat` + cover every database in it. Every other command works on one, named with `--database` or `--db`. See [Several Databases](configuration/storage.md#several-databases). diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index dc149e5d..6212ccd5 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -219,8 +219,9 @@ from a complete one. ### Commands that work on one database -`haiku-rag search`, `ask` and `analyze` cover the configured set. Every other -command works on a single database, named with the global `--database` option: +`haiku-rag search`, `ask`, `analyze` and `chat` cover the configured set. Every +other command works on a single database, named with the global `--database` +option: ```bash haiku-rag search "query" # every configured database diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index a149d084..593b4f33 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -26,7 +26,7 @@ def run_chat( from haiku.rag.utils import get_model, parse_model_option config = get_config() - if db_path is None: + if db_path is None and not config.lancedb.databases: db_path = config.storage.data_dir / "haiku.rag.lancedb" if model: diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index ed438149..77650495 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -86,7 +86,7 @@ class ChatApp(App): def __init__( self, - db_path: Path, + db_path: Path | None, capabilities: Sequence[RAGCapabilityBase[Any]], read_only: bool = False, model: str | None = None, @@ -356,10 +356,17 @@ class ChatApp(App): return 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]) chunk_ids = citation.chunk_ids or [citation.chunk_id] chunks = [] for cid in chunk_ids: - chunk = await self.client.get_chunk_by_id(cid) + chunk = await client.get_chunk_by_id(cid) if chunk: chunks.append(chunk) if not chunks: @@ -370,7 +377,7 @@ class ChatApp(App): await self.push_screen( VisualGroundingModal( chunk=chunks, - client=self.client, + client=client, refs=citation.doc_item_refs or None, ) ) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index c7ab75cd..d9ed271e 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -102,6 +102,17 @@ def resolve_db_path(db: Path | None = None, *, federated: bool = False) -> Path: return get_config().storage.data_dir / "haiku.rag.lancedb" +def resolve_database_set(db: Path | None = None) -> Path | None: + """The database a set-covering command opens, or None to cover the set. + + None where `lancedb.databases` names the databases and the caller named none + of them, so the client resolves the set itself. + """ + if db is None and _database is None and get_config().lancedb.databases: + return None + return resolve_db_path(db, federated=True) + + def require_one_database( config: "AppConfig", db: Path | None, *, federated: bool ) -> None: @@ -858,7 +869,7 @@ def chat( """Launch the chat TUI for conversational RAG.""" from haiku.rag.chat import run_chat - db_path = resolve_db_path(db) + db_path = resolve_database_set(db) capabilities = capability if capability else ["rag"] try: diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 8003381f..68765514 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -638,6 +638,22 @@ class HaikuRAG: Returns: List of Document instances matching the criteria. """ + if self._federated: + # 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. + wanted = None if limit is None else limit + (offset or 0) + groups = await asyncio.gather( + *( + owner.list_documents( + limit=wanted, filter=filter, include_content=include_content + ) + for owner in await self.clients_covering() + ) + ) + merged = [doc for group in groups for doc in group] + start = offset or 0 + return merged[start:] if limit is None else merged[start : start + limit] return await self.document_repository.list_all( limit=limit, offset=offset, filter=filter, include_content=include_content ) @@ -651,6 +667,14 @@ class HaikuRAG: Returns: Number of documents matching the criteria. """ + if self._federated: + counts = await asyncio.gather( + *( + owner.count_documents(filter=filter) + for owner in await self.clients_covering() + ) + ) + return sum(counts) return await self.document_repository.count(filter=filter) async def clients_covering( 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 b5f8fd12..d1bda19f 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -1,3 +1,4 @@ +import asyncio import json from pathlib import Path from typing import TYPE_CHECKING @@ -14,6 +15,123 @@ if TYPE_CHECKING: from haiku.rag.client import HaikuRAG +def reported_database(client: "HaikuRAG", db_path: "Path | None") -> "Path | None": + """The database whose statistics to report, or None where a set is covered. + + A caller passing no path leaves the choice to the client, and a client given + 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: + return None + return db_path if db_path is not None else client.store.db_path + + +async def database_lines(client: "HaikuRAG", db_path: Path) -> list[str]: + """What one database reports about itself, without naming its location. + + A failure is reported as a line rather than raised, so one unreachable + database does not cost the report on the others. + """ + from haiku.rag.store.engine import ConnectionMode, connect_lancedb + from haiku.rag.store.info import get_database_stats + + lines: list[str] = [] + config = client.store._config + + if client.store._connection_mode == ConnectionMode.LOCAL and not db_path.exists(): + return ["[red]Database path does not exist.[/red]"] + + try: + db = await connect_lancedb(config, db_path) + stats = await get_database_stats(db) + except Exception as e: + return [f"[red]Failed to open database: {e}[/red]"] + + stored_version = "unknown" + embed_provider: str | None = None + embed_model: str | None = None + vector_dim: int | None = None + + if stats["settings"]["exists"]: + settings_tbl = await db.open_table("settings") + arrow = await settings_tbl.query().where("id = 'settings'").limit(1).to_arrow() + rows = arrow.to_pylist() if arrow is not None else [] + if rows: + raw = rows[0].get("settings") or "{}" + data = json.loads(raw) if isinstance(raw, str) else (raw or {}) + stored_version = str(data.get("version", stored_version)) + embeddings = data.get("embeddings", {}) + embed_model_obj = embeddings.get("model", {}) + embed_provider = embed_model_obj.get("provider") + embed_model = embed_model_obj.get("name") + vector_dim = embed_model_obj.get("vector_dim") + + num_docs = stats["documents"].get("num_rows", 0) + num_chunks = stats["chunks"].get("num_rows", 0) + has_vector_index = stats["chunks"].get("has_vector_index", False) + num_unindexed_rows = stats["chunks"].get("num_unindexed_rows", 0) + + lines.append( + f"[bold $accent]haiku.rag version (db)[/bold $accent]: {stored_version}" + ) + + if embed_provider or embed_model or vector_dim: + provider_part = embed_provider or "unknown" + model_part = embed_model or "unknown" + dim_part = f"{vector_dim}" if vector_dim is not None else "unknown" + lines.append( + f"[bold $accent]embeddings[/bold $accent]: " + f"{provider_part}/{model_part} (dim: {dim_part})" + ) + else: + lines.append("[bold $accent]embeddings[/bold $accent]: unknown") + + lines.append( + f"[bold $accent]documents[/bold $accent]: {num_docs} " + f"({format_bytes(stats['documents'].get('total_bytes', 0))})" + ) + lines.append( + f"[bold $accent]document_meta[/bold $accent]: " + f"{stats['document_meta'].get('num_rows', 0)} " + f"({format_bytes(stats['document_meta'].get('total_bytes', 0))})" + ) + lines.append( + f"[bold $accent]chunks[/bold $accent]: {num_chunks} " + f"({format_bytes(stats['chunks'].get('total_bytes', 0))})" + ) + + if has_vector_index: + lines.append("[bold $accent]vector index[/bold $accent]: ✓ exists") + lines.append( + f"[bold $accent]indexed chunks[/bold $accent]: " + f"{stats['chunks'].get('num_indexed_rows', 0)}" + ) + colour = "[yellow]" if num_unindexed_rows > 0 else "" + close = "[/yellow]" if num_unindexed_rows > 0 else "" + lines.append( + f"[bold $accent]unindexed chunks[/bold $accent]: " + f"{colour}{num_unindexed_rows}{close}" + ) + elif num_chunks >= 256: + lines.append( + "[bold $accent]vector index[/bold $accent]: [yellow]✗ not created[/yellow]" + ) + else: + lines.append( + f"[bold $accent]vector index[/bold $accent]: ✗ not created " + f"(need {256 - num_chunks} more chunks)" + ) + + for table in ("documents", "document_meta", "chunks"): + lines.append( + f"[bold $accent]versions ({table})[/bold $accent]: " + f"{stats[table].get('num_versions', 0)}" + ) + lines.append("") + return lines + + class InfoModal(ModalScreen): """Modal screen for displaying database information.""" @@ -50,7 +168,7 @@ class InfoModal(ModalScreen): } """ - def __init__(self, client: "HaikuRAG", db_path: Path): + def __init__(self, client: "HaikuRAG", db_path: Path | None): super().__init__() self.client = client self.db_path = db_path @@ -64,130 +182,25 @@ class InfoModal(ModalScreen): async def on_mount(self) -> None: """Load and display database info.""" - from haiku.rag.store.engine import ConnectionMode, connect_lancedb - from haiku.rag.store.info import get_database_stats - lines: list[str] = [] - lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}") - - is_local = self.client.store._connection_mode == ConnectionMode.LOCAL - if is_local and not self.db_path.exists(): - lines.append("[red]Database path does not exist.[/red]") - self._content_widget.update("\n".join(lines)) - return - - # Connect to get table info - config = self.client.store._config - try: - db = await connect_lancedb(config, self.db_path) - stats = await get_database_stats(db) - except Exception as e: - lines.append(f"[red]Failed to open database: {e}[/red]") - self._content_widget.update("\n".join(lines)) - return - - versions = get_package_versions() - - stored_version = "unknown" - embed_provider: str | None = None - embed_model: str | None = None - vector_dim: int | None = None - - if stats["settings"]["exists"]: - settings_tbl = await db.open_table("settings") - arrow = await ( - settings_tbl.query().where("id = 'settings'").limit(1).to_arrow() - ) - rows = arrow.to_pylist() if arrow is not None else [] - if rows: - raw = rows[0].get("settings") or "{}" - data = json.loads(raw) if isinstance(raw, str) else (raw or {}) - stored_version = str(data.get("version", stored_version)) - embeddings = data.get("embeddings", {}) - embed_model_obj = embeddings.get("model", {}) - embed_provider = embed_model_obj.get("provider") - embed_model = embed_model_obj.get("name") - vector_dim = embed_model_obj.get("vector_dim") - - num_docs = stats["documents"].get("num_rows", 0) - doc_bytes = stats["documents"].get("total_bytes", 0) - - num_meta = stats["document_meta"].get("num_rows", 0) - meta_bytes = stats["document_meta"].get("total_bytes", 0) - meta_versions = stats["document_meta"].get("num_versions", 0) - - num_chunks = stats["chunks"].get("num_rows", 0) - chunk_bytes = stats["chunks"].get("total_bytes", 0) - - has_vector_index = stats["chunks"].get("has_vector_index", False) - num_indexed_rows = stats["chunks"].get("num_indexed_rows", 0) - num_unindexed_rows = stats["chunks"].get("num_unindexed_rows", 0) - - doc_versions = stats["documents"].get("num_versions", 0) - chunk_versions = stats["chunks"].get("num_versions", 0) - - lines.append( - f"[bold $accent]haiku.rag version (db)[/bold $accent]: {stored_version}" - ) - - if embed_provider or embed_model or vector_dim: - provider_part = embed_provider or "unknown" - model_part = embed_model or "unknown" - dim_part = f"{vector_dim}" if vector_dim is not None else "unknown" - lines.append( - f"[bold $accent]embeddings[/bold $accent]: " - f"{provider_part}/{model_part} (dim: {dim_part})" + db_path = reported_database(self.client, self.db_path) + if db_path is None: + # Covering a set: report each database under its configured name, and + # each on its own, so one that cannot be opened costs its own block + # 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)) ) + for block in blocks: + lines.extend(block) else: - lines.append("[bold $accent]embeddings[/bold $accent]: unknown") + lines.append(f"[bold $accent]path[/bold $accent]: {db_path}") + lines.extend(await database_lines(self.client, db_path)) - lines.append( - f"[bold $accent]documents[/bold $accent]: {num_docs} ({format_bytes(doc_bytes)})" - ) - lines.append( - f"[bold $accent]document_meta[/bold $accent]: {num_meta} ({format_bytes(meta_bytes)})" - ) - lines.append( - f"[bold $accent]chunks[/bold $accent]: {num_chunks} ({format_bytes(chunk_bytes)})" - ) - - if has_vector_index: - lines.append("[bold $accent]vector index[/bold $accent]: ✓ exists") - lines.append( - f"[bold $accent]indexed chunks[/bold $accent]: {num_indexed_rows}" - ) - if num_unindexed_rows > 0: - lines.append( - f"[bold $accent]unindexed chunks[/bold $accent]: [yellow]{num_unindexed_rows}[/yellow]" - ) - else: - lines.append( - f"[bold $accent]unindexed chunks[/bold $accent]: {num_unindexed_rows}" - ) - else: - if num_chunks >= 256: - lines.append( - "[bold $accent]vector index[/bold $accent]: [yellow]✗ not created[/yellow]" - ) - else: - lines.append( - f"[bold $accent]vector index[/bold $accent]: ✗ not created " - f"(need {256 - num_chunks} more chunks)" - ) - - lines.append( - f"[bold $accent]versions (documents)[/bold $accent]: {doc_versions}" - ) - lines.append( - f"[bold $accent]versions (document_meta)[/bold $accent]: {meta_versions}" - ) - lines.append( - f"[bold $accent]versions (chunks)[/bold $accent]: {chunk_versions}" - ) - - lines.append("") lines.append("[bold]Versions[/bold]") + versions = get_package_versions() lines.append(f"[bold $accent]haiku.rag[/bold $accent]: {versions['haiku_rag']}") lines.append(f"[bold $accent]lancedb[/bold $accent]: {versions['lancedb']}") lines.append(f"[bold $accent]docling[/bold $accent]: {versions['docling']}") @@ -195,10 +208,22 @@ class InfoModal(ModalScreen): f"[bold $accent]pydantic-ai[/bold $accent]: {versions['pydantic_ai']}" ) lines.append( - f"[bold $accent]docling-document schema[/bold $accent]: {versions['docling_document_schema']}" + f"[bold $accent]docling-document schema[/bold $accent]: " + f"{versions['docling_document_schema']}" ) self._content_widget.update("\n".join(lines)) + async def _report(self, name: str) -> list[str]: + """One database's block, including its own failure to open.""" + lines = [f"[bold]{name}[/bold]"] + try: + (owner,) = await self.client.clients_for([name]) + except Exception as e: + # The client names a configured database by name and never by + # location, so its message is safe to show. + return [*lines, f"[red]{e}[/red]", ""] + return [*lines, *await database_lines(owner, owner.store.db_path)] + async def action_dismiss(self, result=None) -> None: self.app.pop_screen() diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index b1fdd345..ea3f9b33 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -35,6 +35,32 @@ def test_run_chat_creates_app_and_runs(temp_db_path: Path): assert attached[0].defer_loading is False +def test_run_chat_covers_a_configured_set(tmp_path, monkeypatch): + """A configured set has no single path to fall back to: the app is handed + None so the client opens the set.""" + import haiku.rag.config as config_module + from haiku.rag.config import AppConfig, set_config + from haiku.rag.config.models import LanceDBConfig + + monkeypatch.setattr(config_module, "_config", None) + set_config( + AppConfig( + lancedb=LanceDBConfig( + databases={ + "a": str(tmp_path / "a.lancedb"), + "b": str(tmp_path / "b.lancedb"), + } + ) + ) + ) + with patch("haiku.rag.chat.app.ChatApp") as app: + from haiku.rag.chat import run_chat + + run_chat(db_path=None) + + assert app.call_args.args[0] is None + + def test_run_chat_defers_multiple_capabilities(temp_db_path: Path): """Test chat only defers capabilities when routing between multiple choices.""" with patch("haiku.rag.chat.app.ChatApp") as mock_app: @@ -92,6 +118,10 @@ 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 + # Mock, which would send every read down the covering-a-set branch. + mock_client._federated = {} + mock_client._source = None return mock_client @@ -465,3 +495,47 @@ async def test_a_cancelled_run_does_not_advance_persisted_state(temp_db_path: Pa assert app._state == before assert app._messages == [] + + +async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path): + """Chunks, pages and boxes come from one database. Covering a set, the + citation's source says which, and a covering client has no repositories.""" + from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget + from haiku.rag.store.models import Chunk + from haiku.rag.store.models.citation import Citation + + owner = _make_mock_client() + 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]) + + app, _ = _make_app(tmp_path / "unused.lancedb", covering) + with patch("haiku.rag.chat.app.HaikuRAG", return_value=covering): + async with app.run_test(): + history = app.query_one(ChatHistory) + await history.add_citations( + [ + Citation( + chunk_id="c1", + document_id="d1", + document_uri="test://beta/one", + content="cited body", + source="beta", + ) + ] + ) + widget = next(iter(app.query(CitationWidget))) + widget.add_class("selected") + + with patch.object(app, "push_screen", new=AsyncMock()) as push: + await app.action_show_visual() + + covering.clients_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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 597fe2a4..8d9f4aaf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -601,6 +601,77 @@ class TestAskAnalyzeImageOption: assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()] +class TestChatCoversTheSet: + """Chat is a read verb: it answers with the same capabilities `ask` uses, so + it covers the configured set rather than demanding one database.""" + + @staticmethod + def _config_file(tmp_path): + config_file = tmp_path / "haiku.rag.yaml" + config_file.write_text( + f"lancedb:\n databases:\n arxiv: {tmp_path / 'a.lancedb'}\n" + f" wiki: {tmp_path / 'w.lancedb'}\n" + ) + return config_file + + def test_a_configured_set_is_covered_rather_than_refused( + self, tmp_path, monkeypatch + ): + import haiku.rag.cli as cli_module + import haiku.rag.config as config_module + + monkeypatch.setattr(config_module, "_config", None) + monkeypatch.setattr(cli_module, "_database", None) + monkeypatch.setattr(cli_module, "_database_path", None) + + with patch("haiku.rag.chat.run_chat") as run_chat: + result = runner.invoke( + cli, ["--config", str(self._config_file(tmp_path)), "chat"] + ) + + assert result.exit_code == 0, result.output + # None is what makes the client resolve the set for itself. + assert run_chat.call_args.args[0] is None + + def test_naming_one_database_opens_that_one(self, tmp_path, monkeypatch): + import haiku.rag.cli as cli_module + import haiku.rag.config as config_module + + monkeypatch.setattr(config_module, "_config", None) + monkeypatch.setattr(cli_module, "_database", None) + monkeypatch.setattr(cli_module, "_database_path", None) + + with patch("haiku.rag.chat.run_chat") as run_chat: + result = runner.invoke( + cli, + [ + "--config", + str(self._config_file(tmp_path)), + "--database", + "wiki", + "chat", + ], + ) + + assert result.exit_code == 0, result.output + assert run_chat.call_args.args[0] == tmp_path / "w.lancedb" + + def test_a_single_database_setup_is_unchanged(self, tmp_path, monkeypatch): + """Without a configured set, chat opens the path it always did.""" + import haiku.rag.cli as cli_module + import haiku.rag.config as config_module + + monkeypatch.setattr(config_module, "_config", None) + monkeypatch.setattr(cli_module, "_database", None) + monkeypatch.setattr(cli_module, "_database_path", None) + + with patch("haiku.rag.chat.run_chat") as run_chat: + result = runner.invoke(cli, ["chat", "--db", str(tmp_path / "one.lancedb")]) + + assert result.exit_code == 0, result.output + assert run_chat.call_args.args[0] == tmp_path / "one.lancedb" + + class TestRenderingTheDatabase: """Across databases a result has to say which one it came from. One database needs no such label, so single-database output is unchanged.""" diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 52c69c8d..35d95d5b 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -1,5 +1,6 @@ import base64 from io import BytesIO +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -329,3 +330,61 @@ async def test_inspector_open_failure_surfaces_real_error(tmp_path): with pytest.raises(FileNotFoundError): async with app.run_test(): pass + + +class TestReportedDatabase: + """`db_path=None` means the client chose, and a client handed one named + database opens it rather than covering a set.""" + + @staticmethod + def _client(federated, store_path=None): + client = MagicMock() + client._federated = federated + client.store.db_path = store_path + return client + + def test_a_covered_set_reports_no_single_database(self): + from haiku.rag.inspector.widgets.info_modal import reported_database + + client = self._client({"a": "/a.lancedb", "b": "/b.lancedb"}) + + assert reported_database(client, None) is None + + def test_one_named_database_reports_the_path_it_opened(self): + """The path is None because the client resolved the name, not because + there is a set to cover.""" + from haiku.rag.inspector.widgets.info_modal import reported_database + + client = self._client({}, store_path=Path("/data/alpha.lancedb")) + + assert reported_database(client, None) == Path("/data/alpha.lancedb") + + def test_an_explicit_path_is_reported_as_given(self): + from haiku.rag.inspector.widgets.info_modal import reported_database + + client = self._client({}, store_path=Path("/data/other.lancedb")) + + assert reported_database(client, Path("/data/given.lancedb")) == Path( + "/data/given.lancedb" + ) + + +class TestReportingEachDatabase: + @pytest.mark.asyncio + async def test_a_database_that_cannot_be_opened_reports_itself(self): + """One unreachable database must not cost the report on the others.""" + from haiku.rag.inspector.widgets.info_modal import InfoModal + from haiku.rag.store.exceptions import SourceUnavailableError + + modal = InfoModal.__new__(InfoModal) + client = AsyncMock() + client.clients_for.side_effect = SourceUnavailableError( + "database 'beta' could not be opened: OSError" + ) + modal.client = client + modal.db_path = None + + lines = await modal._report("beta") + + assert lines[0] == "[bold]beta[/bold]" + assert "could not be opened" in lines[1] diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index 188f27ed..ca3d2aaa 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -157,6 +157,58 @@ class TestOpeningDatabases: assert [client._source for client in covering] == ["alpha"] +class TestListingAcrossDatabases: + """The chat TUI's document filter lists documents through the client, and a + client covering a set has no repositories of its own.""" + + @pytest.mark.asyncio + async def test_listing_covers_every_database(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one", "alpha two"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + docs = await rag.list_documents() + + assert {d.uri for d in docs} == { + "test://alpha/alpha one", + "test://alpha/alpha two", + "test://beta/beta one", + } + + @pytest.mark.asyncio + async def test_counting_covers_every_database(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one", "alpha two"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + assert await rag.count_documents() == 3 + + @pytest.mark.asyncio + async def test_a_limit_bounds_the_merged_listing(self, tmp_path): + """A limit is that many documents in total, not that many per database.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one", "alpha two"]) + await _seed(config, "beta", ["beta one", "beta two"]) + + async with HaikuRAG(config=config) as rag: + assert len(await rag.list_documents(limit=3)) == 3 + assert len(await rag.list_documents(limit=2, offset=2)) == 2 + assert len(await rag.list_documents(offset=3)) == 1 + + @pytest.mark.asyncio + async def test_a_filter_reaches_every_database(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) as rag: + docs = await rag.list_documents(filter="uri LIKE 'test://beta/%'") + + assert [d.uri for d in docs] == ["test://beta/beta one"] + + class TestFederatedSearch: @pytest.mark.asyncio async def test_results_carry_their_source(self, tmp_path):