danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cover the configured set in the chat TUI

Browse source 
Chat answers with the same capabilities `ask` does, so it federates as
naturally as `ask` and `analyze` — but it went through the one-database
guard and refused a configured set outright, which left no way to chat
across several databases.

The guard was the visible half. `run_chat` also defaulted `db_path` to the
single default path whenever it was None, so lifting the refusal alone
would still have opened one database. It now leaves the path unresolved
when `lancedb.databases` names the set, and the client resolves it.

Listing and counting documents fan out over the set, which is what the
document filter reads, and visual grounding resolves the database holding
the cited chunk through the citation's source: chunks, pages and bounding
boxes all come from that one database. A limit on a listing means that
many documents in total, not that many per database.

The info modal reports every database it covers, each under its
configured name and without its location, since names are the only
identity that leaves the configuration. `database_lines` is what one
database reports about itself, shared by both paths, and it reports a
failure as a line so one unreachable database does not cost the report on
the others.

`inspect` stays a one-database command. It browses one database's
documents and chunks, so a set has nothing to show it.
This commit is contained in:
Yiorgis Gozadinos 2026-08-21 13:14:31 +03:00
parent 5c9639bbbb
commit bb92ccf67e
No known key found for this signature in database
12 changed files with 454 additions and 129 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

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

4
docs/cli.md
View file

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

5
docs/configuration/storage.md
View file

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

2
haiku_rag_slim/haiku/rag/chat/__init__.py
View file

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

13
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -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,
)
)

13
haiku_rag_slim/haiku/rag/cli.py
View file

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

24
haiku_rag_slim/haiku/rag/client/__init__.py
View file

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

265
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
View file

@ -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()

74
tests/chat/test_chat_app.py
View file

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

71
tests/test_cli.py
View file

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

59
tests/test_inspector.py
View file

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

52
tests/test_multi_db.py
View file

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