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

Resolve the databases a command works on, once

Browse source 
The CLI decides only what it knows — that --db and --database are the same
thing said twice, and whether a command reads more than one — and hands the
resolved scope down. Nothing rewrites the configuration, so a named database
keeps the name results and citations carry, and a remote one opens the URI
it was configured with rather than the local path standing in for it.

HaikuRAGApp, ChatApp and InspectorApp take that scope and nothing else.
Selection reaches the client through a private constructor, so the public
signature still takes a path or names.
This commit is contained in:
Yiorgis Gozadinos 2026-08-25 15:48:47 +03:00
parent 8db447e095
commit d367b1eb5a
No known key found for this signature in database
21 changed files with 641 additions and 391 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

204
haiku_rag_slim/haiku/rag/app.py
View file

@ -1,4 +1,5 @@
import logging
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING
@ -22,6 +23,7 @@ from haiku.rag.store.models.chunk import SearchType
from haiku.rag.store.models.document import Document
if TYPE_CHECKING:
from haiku.rag.client.scope import DatabaseRef, DatabaseScope
from haiku.rag.store.engine import Store
from haiku.rag.store.models import SearchResult
from haiku.rag.config import redact_secrets
@ -33,37 +35,82 @@ logger = logging.getLogger(__name__)
class HaikuRAGApp:
def __init__(
self,
db_path: Path,
scope: "DatabaseScope",
config: AppConfig | None = None,
read_only: bool = False,
federated: bool = False,
):
self.db_path = db_path
"""The databases this command works on, resolved by whoever built it.
One selector, not three: a command that took a path and a name and a
scope would have to decide between them, which is the deciding this
layer exists to have already done.
"""
self.scope = scope
self.config = config if config is not None else get_config()
self.read_only = read_only
self.federated = federated
self.console = Console()
from haiku.rag.store.engine import ConnectionMode
@property
def _one(self) -> "DatabaseRef":
"""The one database this command works on.
self._is_local = ConnectionMode.from_config(self.config) == ConnectionMode.LOCAL
self._display_path = self.db_path if self._is_local else self.config.lancedb.uri
Commands that cover a set never reach here: they read through the client
instead.
"""
[ref] = self.scope.databases
return ref
@cached_property
def _connection(self) -> "tuple[AppConfig, Path]":
"""How to open the one database this command works on, directly.
The configuration a command was given names a *set*; the one it opens
needs its own, or a named database behind a URI is opened as the local
path that stands in for it.
"""
from haiku.rag.client.session import default_db_path
config, db_path = self._one.connection(self.config)
return config, db_path or default_db_path(config)
@property
def _client_db_path(self) -> "Path | None":
"""None where the command covers `lancedb.databases`, so the client
opens the configured set rather than one path."""
return None if self.federated else self.db_path
def _store_config(self) -> AppConfig:
"""The configuration for opening the one database directly."""
return self._connection[0]
@property
def _is_local(self) -> bool:
"""Whether the database is a local path rather than a URI.
Read from the database this command resolved to, not from the
configuration: a database named in `lancedb.databases` can sit behind a
URI while the configuration's own `uri` is empty.
"""
return self._one.db_path is not None
@property
def _path(self) -> Path:
"""The path of the one database this command works on.
A database behind a URI has none of its own, and the default stands in:
the URI in `_store_config` is what decides where it connects.
"""
return self._connection[1]
@property
def _display_path(self) -> "Path | str":
"""What a one-database command calls the database it opened."""
return self._one.db_path or self._one.uri
async def init(self):
"""Initialize a new database."""
if self._is_local and self.db_path.exists():
if self._is_local and self._path.exists():
self.console.print(
f"[yellow]Database already exists at {self.db_path}[/yellow]"
f"[yellow]Database already exists at {self._path}[/yellow]"
)
return
async with HaikuRAG(db_path=self.db_path, config=self.config, create=True):
async with HaikuRAG._covering(self.scope, self.config, create=True):
pass
self.console.print(
f"[bold green]Database initialized at {self._display_path}[/bold green]"
@ -80,11 +127,11 @@ class HaikuRAGApp:
f" [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
)
if self._is_local and not self.db_path.exists():
if self._is_local and not self._path.exists():
self.console.print("[red]Database path does not exist.[/red]")
return
info = await gather_database_info(self.config, self.db_path)
info = await gather_database_info(self._store_config, self._path)
if not info.exists:
self.console.print(
@ -211,7 +258,7 @@ class HaikuRAGApp:
f" [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
)
if self._is_local and not self.db_path.exists():
if self._is_local and not self._path.exists():
self.console.print("[red]Database path does not exist.[/red]")
return True
@ -226,8 +273,8 @@ class HaikuRAGApp:
cm = status if status is not None else nullcontext()
with cm:
report = await run_doctor(
self.config,
self.db_path,
self._store_config,
self._path,
dict(os.environ),
duplicates_out=duplicates_out,
on_progress=on_progress,
@ -279,13 +326,13 @@ class HaikuRAGApp:
"""
from haiku.rag.store.engine import Store
if self._is_local and not self.db_path.exists():
if self._is_local and not self._path.exists():
self.console.print("[red]Database path does not exist.[/red]")
return
async with Store(
self.db_path,
config=self.config,
self._path,
config=self._store_config,
skip_validation=True,
read_only=True,
skip_migration_check=True,
@ -359,15 +406,15 @@ class HaikuRAGApp:
"""
from haiku.rag.store.engine import Store
return Store(self.db_path, config=self.config, read_only=self.read_only)
return Store(self._path, config=self._store_config, read_only=self.read_only)
def _tag_read_store(self) -> "Store":
"""Read-only store for tag inspection; works on old or drifted DBs."""
from haiku.rag.store.engine import Store
return Store(
self.db_path,
config=self.config,
self._path,
config=self._store_config,
skip_validation=True,
skip_migration_check=True,
read_only=True,
@ -375,16 +422,16 @@ class HaikuRAGApp:
async def create_tag(self, name: str):
"""Tag the current version of every table."""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
if self._is_local and not self._path.exists():
raise ValueError(f"Database path does not exist: {self._path}")
async with self._tag_write_store() as store:
await store.create_tag(name)
self.console.print(f"[green]Created tag '{escape(name)}'[/green]")
async def list_tags(self):
"""List database tags, flagging partial ones."""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
if self._is_local and not self._path.exists():
raise ValueError(f"Database path does not exist: {self._path}")
async with self._tag_read_store() as store:
tags = await store.list_tags()
@ -404,8 +451,8 @@ class HaikuRAGApp:
async def delete_tag(self, name: str):
"""Delete a tag from every table that has it."""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
if self._is_local and not self._path.exists():
raise ValueError(f"Database path does not exist: {self._path}")
async with self._tag_write_store() as store:
await store.delete_tag(name)
self.console.print(f"[green]Deleted tag '{escape(name)}'[/green]")
@ -419,8 +466,8 @@ class HaikuRAGApp:
Raises:
ValueError: If the database path does not exist.
"""
if self._is_local and not self.db_path.exists():
raise ValueError(f"Database path does not exist: {self.db_path}")
if self._is_local and not self._path.exists():
raise ValueError(f"Database path does not exist: {self._path}")
async with self._tag_write_store() as store:
safety_tag = await store.restore_tag(name)
self.console.print(f"[green]Restored database to tag '{escape(name)}'.[/green]")
@ -434,11 +481,8 @@ class HaikuRAGApp:
)
async def list_documents(self, filter: str | None = None):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=True,
skip_validation=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=True, skip_validation=True
) as self.client:
documents = await self.client.list_documents(filter=filter)
for doc in documents:
@ -447,10 +491,8 @@ class HaikuRAGApp:
async def add_document_from_text(
self, text: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
async with HaikuRAG._covering(
self.scope, self.config, read_only=self.read_only
) as self.client:
doc = await self.client.create_document(
text, title=title, metadata=metadata
@ -463,10 +505,8 @@ class HaikuRAGApp:
async def add_document_from_source(
self, source: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
async with HaikuRAG._covering(
self.scope, self.config, read_only=self.read_only
) as self.client:
result = await self.client.create_document_from_source(
source, title=title, metadata=metadata
@ -484,11 +524,8 @@ class HaikuRAGApp:
)
async def get_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=True,
skip_validation=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=True, skip_validation=True
) as self.client:
doc = await self.client.get_document_by_id(doc_id)
if doc is None:
@ -497,11 +534,8 @@ class HaikuRAGApp:
self._rich_print_document(doc, truncate=False)
async def delete_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
skip_validation=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=self.read_only, skip_validation=True
) as self.client:
deleted = await self.client.delete_document(doc_id)
if deleted:
@ -541,10 +575,8 @@ class HaikuRAGApp:
assert query is not None
search_input = query
async with HaikuRAG(
db_path=self._client_db_path,
config=self.config,
read_only=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=True
) as self.client:
results = await self.client.search(
search_input,
@ -562,11 +594,8 @@ class HaikuRAGApp:
"""Display visual grounding images for a chunk."""
from textual_image.renderable import Image as RichImage
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=True,
skip_validation=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=True, skip_validation=True
) as self.client:
chunk = await self.client.get_chunk_by_id(chunk_id)
if not chunk:
@ -608,10 +637,8 @@ class HaikuRAGApp:
filter: SQL WHERE clause to filter documents
images: Paths of images to attach to the question
"""
async with HaikuRAG(
db_path=self._client_db_path,
config=self.config,
read_only=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=True
) as self.client:
answer, citations = await self.client.ask(
question,
@ -641,10 +668,8 @@ class HaikuRAGApp:
filter: SQL WHERE clause to filter documents
images: Paths of images to attach to the question
"""
async with HaikuRAG(
db_path=self._client_db_path,
config=self.config,
read_only=True,
async with HaikuRAG._covering(
self.scope, self.config, read_only=True
) as self.client:
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
@ -667,11 +692,8 @@ class HaikuRAGApp:
self.console.print(renderable)
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
skip_validation=True,
read_only=self.read_only,
async with HaikuRAG._covering(
self.scope, self.config, skip_validation=True, read_only=self.read_only
) as client:
if mode == RebuildMode.SET_EMBEDDER:
async for _ in client.rebuild_database(mode=mode):
@ -710,11 +732,8 @@ class HaikuRAGApp:
async def vacuum(self):
"""Run database maintenance: optimize and cleanup table history."""
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
skip_validation=True,
read_only=self.read_only,
async with HaikuRAG._covering(
self.scope, self.config, skip_validation=True, read_only=self.read_only
) as client:
await client.vacuum()
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
@ -728,8 +747,8 @@ class HaikuRAGApp:
from haiku.rag.store.engine import Store
async with Store(
self.db_path,
config=self.config,
self._path,
config=self._store_config,
skip_validation=True,
skip_migration_check=True,
read_only=self.read_only,
@ -738,11 +757,8 @@ class HaikuRAGApp:
async def create_index(self):
"""Create vector index on the chunks table."""
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
skip_validation=True,
read_only=self.read_only,
async with HaikuRAG._covering(
self.scope, self.config, skip_validation=True, read_only=self.read_only
) as client:
row_count = await client.store.chunks_table.count_rows()
self.console.print(f"Chunks in database: {row_count}")
@ -899,13 +915,11 @@ class HaikuRAGApp:
port: int = 8001,
):
"""Run the MCP server until interrupted."""
async with HaikuRAG(
self.db_path,
config=self.config,
read_only=self.read_only,
async with HaikuRAG._covering(
self.scope, self.config, read_only=self.read_only
):
server = create_mcp_server(
self.db_path, config=self.config, read_only=self.read_only
self._path, config=self._store_config, read_only=self.read_only
)
try:
if transport == "stdio":

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

@ -1,4 +1,8 @@
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from haiku.rag.client.scope import DatabaseScope
def run_chat(
@ -6,11 +10,13 @@ def run_chat(
read_only: bool = False,
model: str | None = None,
capabilities: list[str] | None = None,
scope: "DatabaseScope | None" = None,
) -> None:
"""Run the chat TUI.
Args:
db_path: Path to the LanceDB database. If None, uses default from config.
db_path: Path to the LanceDB database, when no scope is given.
scope: The databases to cover, resolved by the caller.
read_only: Whether to open the database in read-only mode.
model: Model to use for the chat.
capabilities: Capabilities to enable ("rag", "analysis"). Defaults to ["rag"].
@ -26,14 +32,24 @@ def run_chat(
from haiku.rag.utils import get_model, parse_model_option
config = get_config()
if db_path is None and not config.lancedb.databases:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
if scope is None:
from haiku.rag.client.scope import DatabaseScope
scope = DatabaseScope.resolve(config, database_path=db_path)
if model:
model_config = parse_model_option(model)
config.qa.model = model_config
config.analysis.model = model_config
# The capabilities read the databases the scope covers, not whatever the
# configuration happens to name: `--db PATH` would otherwise leave them on
# the default database, and `--database NAME` on the whole set.
if scope.covers_multiple:
capability_config, capability_db_path = config, None
else:
capability_config, capability_db_path = scope.databases[0].connection(config)
enabled = capabilities or ["rag"]
capability_list = []
defer_loading = len(enabled) > 1
@ -52,8 +68,8 @@ def run_chat(
capability_list.append(
create_capability(
db_path=db_path,
config=config,
db_path=capability_db_path,
config=capability_config,
defer_loading=defer_loading,
vision=driving_model.vision,
)
@ -64,17 +80,17 @@ def run_chat(
capability_list.append(
create_capability(
db_path=db_path,
config=config,
db_path=capability_db_path,
config=capability_config,
defer_loading=defer_loading,
vision=driving_model.vision,
)
)
app = ChatApp(
db_path,
capabilities=capability_list,
read_only=read_only,
model=model or get_model(driving_model, config),
scope=scope,
)
app.run()

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

@ -3,7 +3,6 @@ import uuid
from collections.abc import Iterable, Sequence
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
import textual_image.widget # noqa: F401 - import early for renderer detection
@ -44,6 +43,8 @@ configure_telemetry(service_name="haiku-rag")
if TYPE_CHECKING:
from textual.app import ComposeResult
from haiku.rag.client.scope import DatabaseScope
RAG_STATE_NAMESPACE = "rag"
ANALYSIS_STATE_NAMESPACE = "analysis"
@ -86,13 +87,13 @@ class ChatApp(App):
def __init__(
self,
db_path: Path | None,
capabilities: Sequence[RAGCapabilityBase[Any]],
scope: "DatabaseScope",
read_only: bool = False,
model: str | None = None,
) -> None:
super().__init__()
self.db_path = db_path
self.scope = scope
self._capabilities = capabilities
self.read_only = read_only
self._model = model
@ -141,15 +142,16 @@ class ChatApp(App):
async def on_mount(self) -> None:
"""Initialize the app when mounted."""
client = HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
)
client = HaikuRAG._covering(self.scope, self.config, read_only=self.read_only)
# Assign only after a successful open: on_unmount must not tear down
# a client whose __aenter__ failed.
await client.__aenter__()
self.client = client
# The capabilities read through this one, rather than each opening its
# own: it is already the databases they were built for, and lending it
# means one connection per database instead of one per capability.
for capability in self._capabilities:
capability.borrowed_rag = client
self._agent = Agent(
self._model,
@ -387,7 +389,7 @@ class ChatApp(App):
from haiku.rag.inspector.widgets.info_modal import InfoModal
await self.push_screen(InfoModal(self.client, self.db_path))
await self.push_screen(InfoModal(self.client, None))
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
"""Handle citation selection."""

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

@ -33,6 +33,7 @@ from haiku.rag.utils import is_up_to_date # noqa: E402
if TYPE_CHECKING:
from haiku.rag.app import HaikuRAGApp
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config.models import AppConfig
_cli = typer.Typer(
@ -59,95 +60,53 @@ def cli():
# Module-level flags set by callback
_read_only: bool = False
_database: str | None = None
_database_path: Path | None = None
def create_app(db: Path | None = None, *, federated: bool = False) -> "HaikuRAGApp":
"""Create HaikuRAGApp with loaded config and resolved database path.
def create_app(db: Path | None = None, *, covers_set: bool = False) -> "HaikuRAGApp":
"""The application for a command, on the database(s) it works on.
Args:
db: Optional database path. If None, uses `--database`, then the path
from config.
federated: Whether this command works across `lancedb.databases`.
Returns:
HaikuRAGApp instance with proper config and db path.
`covers_set` is the command declaring that it can read several: `search`,
`ask`, `analyze` and `chat` can, and everything else names one.
Raises:
AmbiguousDatabaseError: Several databases are configured and this
AmbiguousDatabaseError: several databases are configured and this
command works on one, without `--db` or `--database` naming which.
"""
from haiku.rag.app import HaikuRAGApp
db_path = resolve_db_path(db, federated=federated)
return HaikuRAGApp(
db_path=db_path,
config=get_config(),
read_only=_read_only,
federated=federated and db is None and _database is None,
scope=resolve_scope(db, covers_set=covers_set),
)
def resolve_db_path(db: Path | None = None, *, federated: bool = False) -> Path:
"""The database a command works on, from `--db`, `--database`, or config."""
def resolve_scope(
db: Path | None = None, *, covers_set: bool = False
) -> "DatabaseScope":
"""The databases a command works on, resolved once.
The CLI decides only what it alone knows: that `--db` and `--database` are
the same thing said twice, and whether this command can read more than one.
Everything else — an unknown name, a legacy `uri`, the default location — is
`DatabaseScope.resolve`'s to answer, so there is one table and not two.
"""
from haiku.rag.client.scope import DatabaseScope
if db is not None and _database is not None:
raise AmbiguousDatabaseError(
"pass --db or --database, not both: they name the same thing"
)
require_one_database(get_config(), db, federated=federated)
if db is not None:
return db
if _database_path is not None:
return _database_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:
"""Refuse a one-database command that cannot tell which one to use."""
databases = config.lancedb.databases
if federated or db is not None or not databases:
return
raise AmbiguousDatabaseError(
f"lancedb.databases names {', '.join(sorted(databases))}; this command "
"works on a single database: pass --database NAME, or --db PATH."
scope = DatabaseScope.resolve(
get_config(), database_name=_database, database_path=db
)
def select_database(name: str) -> Path | None:
"""Point the configuration at one database from `lancedb.databases`.
Returns its local path, or None where it lives behind a URI. Rewriting the
configuration is what lets every command, the TUIs included, work on the
selected database without knowing the set exists.
"""
from haiku.rag.utils import locate_database
config = get_config()
databases = config.lancedb.databases
if name not in databases:
if scope.covers_multiple and not covers_set:
raise AmbiguousDatabaseError(
f"unknown database {name!r}; lancedb.databases names "
f"{', '.join(sorted(databases)) or 'nothing'}"
f"lancedb.databases names {', '.join(sorted(scope.names))}; this "
"command works on a single database: pass --database NAME, or "
"--db PATH."
)
uri, db_path = locate_database(databases[name])
selected = config.model_copy(deep=True)
selected.lancedb.databases = {}
selected.lancedb.uri = uri
set_config(selected)
return db_path
return scope
async def check_version():
@ -193,10 +152,9 @@ def main(
),
):
"""haiku.rag CLI - Vector database RAG system"""
global _read_only, _database, _database_path
global _read_only, _database
_read_only = read_only
_database = database
_database_path = None
# Load config from --config, local folder, or default directory
config_path = find_config_file(cli_path=config)
if config_path:
@ -205,9 +163,6 @@ def main(
else:
set_config(AppConfig())
if database is not None:
_database_path = select_database(database)
configure_cli_logging()
from haiku.rag.telemetry import configure as configure_telemetry
@ -402,7 +357,7 @@ def search(
help="Path to the LanceDB database file",
),
):
app = create_app(db, federated=True)
app = create_app(db, covers_set=True)
asyncio.run(
app.search(
query=query,
@ -456,7 +411,7 @@ def ask(
help="Path to an image to attach to the question (repeatable; requires a vision-capable model)",
),
):
app = create_app(db, federated=True)
app = create_app(db, covers_set=True)
asyncio.run(
app.ask(
question=question,
@ -488,7 +443,7 @@ def analyze(
help="Path to an image to attach to the question (repeatable; requires a vision-capable model)",
),
):
app = create_app(db, federated=True)
app = create_app(db, covers_set=True)
asyncio.run(
app.analyze(
question=question,
@ -503,7 +458,10 @@ def settings():
from haiku.rag.app import HaikuRAGApp
config = get_config()
app = HaikuRAGApp(db_path=Path(), config=config, read_only=True)
# Neither of these opens a database; the scope is whatever is configured.
app = HaikuRAGApp(
scope=resolve_scope(covers_set=True), config=config, read_only=True
)
app.show_settings()
@ -798,11 +756,11 @@ def tag_restore(
),
):
app = create_app(db)
if app._is_local and not app.db_path.exists():
typer.echo(f"Error: Database path does not exist: {app.db_path}", err=True)
if app._is_local and not app._path.exists():
typer.echo(f"Error: Database path does not exist: {app._path}", err=True)
raise typer.Exit(1)
if not yes:
typer.echo(f"Database: {app.db_path}")
typer.echo(f"Database: {app._display_path}")
typer.echo(f"Tag: {name}")
typer.echo("This changes the live database state across all tables.")
typer.echo("Stop all ingestion and other writers before continuing.")
@ -821,7 +779,9 @@ def tag_restore(
def download_models_cmd():
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=Path(), config=get_config(), read_only=True)
app = HaikuRAGApp(
scope=resolve_scope(covers_set=True), config=get_config(), read_only=True
)
try:
asyncio.run(app.download_models())
except Exception as e:
@ -844,7 +804,7 @@ def inspect(
typer.echo(f"Error: {e}", err=True)
raise typer.Exit(1) from e
run_inspector(resolve_db_path(db), read_only=True)
run_inspector(read_only=True, scope=resolve_scope(db))
@_cli.command("chat", help="Launch interactive chat TUI for conversational RAG")
@ -869,15 +829,15 @@ def chat(
"""Launch the chat TUI for conversational RAG."""
from haiku.rag.chat import run_chat
db_path = resolve_database_set(db)
scope = resolve_scope(db, covers_set=True)
capabilities = capability if capability else ["rag"]
try:
run_chat(
db_path,
read_only=True,
model=model,
capabilities=capabilities,
scope=scope,
)
except ImportError as e:
typer.echo(f"Error: {e}", err=True)

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

@ -21,6 +21,7 @@ from haiku.rag.client.session import (
FederatedSession,
SingleDatabaseSession,
aclose_quietly,
default_db_path,
)
from haiku.rag.config import AppConfig, get_config
from haiku.rag.converters import get_converter
@ -143,6 +144,7 @@ class HaikuRAG:
self._read_only = read_only
self._requested_sources = sources
self._clients: dict[str, HaikuRAG] = {}
self._scope: DatabaseScope | None = None
self._session: SingleDatabaseSession | FederatedSession | None = None
self._owns_session = True
@ -260,31 +262,21 @@ class HaikuRAG:
"""
return get_reranker(config=self._config)
def _selected(self) -> dict[str, str]:
"""The configured databases this client covers, name to location.
def _resolve_scope(self) -> DatabaseScope:
"""The databases this client covers.
Empty when the caller named a database itself: an explicit `db_path` says
which one to open, so it is not overridden by a configured set.
Resolved once, here or by whoever handed one in. An explicit `db_path`
says which database to open, so a configured set does not override it.
"""
declared = self._configured.lancedb.databases
if not declared or self._db_path_given:
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 self._scope is not None:
return self._scope
scope = DatabaseScope.resolve(
self._configured,
database_path=self._db_path if self._db_path_given else None,
)
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}
if self._requested_sources is not None and not self._db_path_given:
scope = scope.select(self._requested_sources)
return scope
async def __aenter__(self):
"""Async context manager entry — initializes store and repositories.
@ -294,34 +286,27 @@ class HaikuRAG:
repositories stay unset in that case, since they have no unambiguous
meaning across a set.
"""
selected = self._selected()
if len(selected) > 1:
scope = self._resolve_scope()
if scope.covers_multiple:
if self._create:
raise AmbiguousDatabaseError(
"create=True creates one database, and this client covers "
f"{', '.join(sorted(selected))}; name the one to create with "
"sources=[name]"
f"{', '.join(sorted(scope.names))}; name the one to create "
"with sources=[name]"
)
self._session = FederatedSession(
DatabaseScope(
tuple(
DatabaseRef.configured(name, location)
for name, location in selected.items()
)
),
self._config,
scope,
self._configured,
skip_validation=self._skip_validation,
read_only=self._read_only,
)
return self
source: str | None = None
if selected:
[(source, location)] = selected.items()
self._config, db_path = DatabaseRef.configured(source, location).connection(
self._configured
)
if db_path is not None:
self._db_path = db_path
[ref] = scope.databases
self._config, db_path = ref.connection(self._configured)
self._db_path = (
db_path if db_path is not None else default_db_path(self._config)
)
self._session = await SingleDatabaseSession(
self._db_path,
@ -329,7 +314,7 @@ class HaikuRAG:
skip_validation=self._skip_validation,
create=self._create,
read_only=self._read_only,
source=source,
source=ref.name,
).open()
return self
@ -361,6 +346,30 @@ class HaikuRAG:
self._clients[name] = facade
return facade
@classmethod
def _covering(
cls,
scope: DatabaseScope,
config: AppConfig | None = None,
*,
read_only: bool = False,
create: bool = False,
skip_validation: bool = False,
) -> "HaikuRAG":
"""A client over databases someone already resolved.
Internal: the public constructor takes a path or names, and resolving
those is its own job. This is for callers that did the resolving.
"""
client = cls(
config=config,
read_only=read_only,
create=create,
skip_validation=skip_validation,
)
client._scope = scope
return client
@classmethod
def _from_session(cls, session: SingleDatabaseSession) -> "HaikuRAG":
"""A client over a database another session opened and will close."""

6
haiku_rag_slim/haiku/rag/client/agents.py
View file

@ -63,8 +63,9 @@ async def ask(
)
from haiku.rag.utils import get_model
# No `db_path`: the lent client is what the capability reads through, and it
# already knows which databases that is.
capability = create_capability(
db_path=None if client.covers_multiple else client.store.db_path,
config=client._config,
rag=client,
defer_loading=False,
@ -123,8 +124,9 @@ async def analyze(
from haiku.rag.sandbox import AnalysisResult
from haiku.rag.utils import get_model
# No `db_path`: the lent client is what the capability reads through, and it
# already knows which databases that is.
capability = create_capability(
db_path=None if client.covers_multiple else client.store.db_path,
config=client._config,
rag=client,
defer_loading=False,

19
haiku_rag_slim/haiku/rag/client/scope.py
View file

@ -127,6 +127,25 @@ class DatabaseScope:
return cls((DatabaseRef.at(config.storage.data_dir / "haiku.rag.lancedb"),))
def select(self, names: list[str]) -> "DatabaseScope":
"""The databases in this scope named by `names`, in the order given.
Repeats collapse: a database named twice would be searched twice and
fused as two rank lists, which counts it double.
"""
if not names:
raise ValueError(
"sources=[] selects no database; pass None for all of them"
)
by_name = {ref.name: ref for ref in self.databases if ref.name is not None}
missing = [name for name in names if name not in by_name]
if missing:
raise KeyError(
f"unknown database(s) {', '.join(sorted(missing))}; "
f"configured: {', '.join(sorted(by_name))}"
)
return DatabaseScope(tuple(by_name[name] for name in dict.fromkeys(names)))
@property
def covers_multiple(self) -> bool:
"""Whether this scope covers more than one database."""

27
haiku_rag_slim/haiku/rag/inspector/app.py
View file

@ -16,6 +16,8 @@ from haiku.rag.inspector.widgets.search_modal import SearchModal
if TYPE_CHECKING:
from textual.app import ComposeResult
from haiku.rag.client.scope import DatabaseScope
class InspectorApp(App):
"""Textual TUI for inspecting LanceDB data."""
@ -66,9 +68,13 @@ class InspectorApp(App):
Binding("c", "show_context", "Context", show=True),
]
def __init__(self, db_path: Path, read_only: bool = False):
def __init__(
self,
scope: "DatabaseScope",
read_only: bool = False,
):
super().__init__()
self.db_path = db_path
self.scope = scope
self.read_only = read_only
self.client: HaikuRAG | None = None
@ -83,11 +89,7 @@ class InspectorApp(App):
async def on_mount(self) -> None:
"""Initialize the app when mounted."""
config = get_config()
client = HaikuRAG(
db_path=self.db_path,
config=config,
read_only=self.read_only,
)
client = HaikuRAG._covering(self.scope, config, read_only=self.read_only)
# Assign only after a successful open: on_unmount must not tear down
# a client whose __aenter__ failed.
await client.__aenter__()
@ -131,7 +133,7 @@ class InspectorApp(App):
if self.client:
from haiku.rag.inspector.widgets.info_modal import InfoModal
await self._switch_modal(InfoModal(self.client, self.db_path))
await self._switch_modal(InfoModal(self.client, None))
async def on_search_modal_chunk_selected(
self, message: SearchModal.ChunkSelected
@ -228,6 +230,7 @@ class InspectorApp(App):
def run_inspector(
db_path: Path | None = None,
read_only: bool = False,
scope: "DatabaseScope | None" = None,
) -> None:
"""Run the inspector TUI.
@ -236,8 +239,10 @@ def run_inspector(
read_only: Whether to open the database in read-only mode.
"""
config = get_config()
if db_path is None:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
if scope is None:
from haiku.rag.client.scope import DatabaseScope
app = InspectorApp(db_path, read_only=read_only)
scope = DatabaseScope.resolve(config, database_path=db_path)
app = InspectorApp(scope, read_only=read_only)
app.run()

113
tests/chat/test_chat_app.py
View file

@ -7,6 +7,7 @@ from typer.testing import CliRunner
from haiku.rag.capabilities.rag import RAGState, create_capability
from haiku.rag.cli import _cli as cli
from tests.conftest import _covering_returns, for_path
runner = CliRunner()
@ -58,7 +59,41 @@ def test_run_chat_covers_a_configured_set(tmp_path, monkeypatch):
run_chat(db_path=None)
assert app.call_args.args[0] is None
assert app.call_args.kwargs["scope"].names == ("a", "b")
def test_chat_capabilities_read_the_named_database(tmp_path, monkeypatch):
"""`--db PATH` and `--database NAME` have to reach the capabilities too, or
they answer from the default database or the whole set."""
import haiku.rag.config as config_module
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config import set_config
from haiku.rag.config.models import AppConfig, LanceDBConfig
monkeypatch.setattr(config_module, "_config", None)
config = AppConfig(
lancedb=LanceDBConfig(
databases={
"a": str(tmp_path / "a.lancedb"),
"b": str(tmp_path / "b.lancedb"),
}
)
)
set_config(config)
with patch("haiku.rag.chat.app.ChatApp") as chat_app:
from haiku.rag.chat import run_chat
run_chat(scope=DatabaseScope.resolve(config, database_name="b"))
[named] = chat_app.call_args.kwargs["capabilities"]
run_chat(scope=DatabaseScope.resolve(config))
[covering] = chat_app.call_args.kwargs["capabilities"]
assert named.db_path == tmp_path / "b.lancedb"
assert named.config.lancedb.databases == {}
assert covering.db_path is None
assert set(covering.config.lancedb.databases) == {"a", "b"}
def test_run_chat_defers_multiple_capabilities(temp_db_path: Path):
@ -134,7 +169,7 @@ def _make_app(db_path: Path, mock_client: AsyncMock | None = None):
mock_client = _make_mock_client()
return ChatApp(
db_path=db_path,
scope=for_path(db_path),
capabilities=[create_capability(db_path=db_path)],
read_only=True,
), mock_client
@ -148,7 +183,7 @@ def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None):
mock_client = _make_mock_client()
return ChatApp(
db_path=db_path,
scope=for_path(db_path),
capabilities=[create_capability(db_path=db_path)],
read_only=True,
), mock_client
@ -161,7 +196,10 @@ async def test_chat_app_has_required_widgets(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
chat_history = app.query_one(ChatHistory)
assert chat_history is not None
@ -177,7 +215,10 @@ async def test_chat_app_quit_binding(temp_db_path: Path):
"""Test that pressing ctrl+q quits the app."""
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test() as pilot:
assert app.is_running
await pilot.press("ctrl+q")
@ -191,7 +232,10 @@ async def test_chat_history_can_add_message(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
chat_history = app.query_one(ChatHistory)
@ -210,7 +254,10 @@ async def test_chat_history_can_add_tool_calls(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
chat_history = app.query_one(ChatHistory)
@ -232,7 +279,10 @@ async def test_chat_history_can_add_citations(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
chat_history = app.query_one(ChatHistory)
@ -272,7 +322,10 @@ async def test_chat_history_thinking_indicator(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test() as pilot:
chat_history = app.query_one(ChatHistory)
@ -293,7 +346,10 @@ async def test_clear_chat_resets_state(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test() as pilot:
chat_history = app.query_one(ChatHistory)
@ -317,7 +373,10 @@ async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
app, mock_client = _make_app(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test() as pilot:
chat_history = app.query_one(ChatHistory)
@ -356,7 +415,10 @@ async def test_show_citations_renders_from_flat_state(temp_db_path: Path):
app, mock_client = _make_app_with_state(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test() as pilot:
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
@ -391,7 +453,10 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path):
app, mock_client = _make_app_with_state(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
# The selection is document ids, so a repeated title cannot widen it.
selected = [
@ -421,7 +486,10 @@ async def test_document_filter_cleared_when_empty(temp_db_path: Path):
app, mock_client = _make_app_with_state(temp_db_path)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
# First set a filter
app.on_document_filter_modal_filter_changed(
@ -445,7 +513,7 @@ async def test_chat_app_open_failure_surfaces_real_error(tmp_path: Path):
AttributeError from tearing down a client that never opened."""
from haiku.rag.chat.app import ChatApp
app = ChatApp(db_path=tmp_path / "missing.lancedb", capabilities=[])
app = ChatApp(scope=for_path(tmp_path / "missing.lancedb"), capabilities=[])
with pytest.raises(FileNotFoundError):
async with app.run_test():
pass
@ -485,7 +553,10 @@ async def test_a_cancelled_run_does_not_advance_persisted_state(temp_db_path: Pa
async def __anext__(self):
raise asyncio.CancelledError
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, mock_client),
):
async with app.run_test():
app._state = {
"rag": {"evidence": {"question": 0, "latest_evidence_epoch": 0}}
@ -522,7 +593,10 @@ async def test_visual_grounding_uses_the_database_holding_the_citation(tmp_path)
covering.reader_for = AsyncMock(return_value=owner)
app, _ = _make_app(tmp_path / "unused.lancedb", covering)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=covering):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, covering),
):
async with app.run_test():
history = app.query_one(ChatHistory)
await history.add_citations(
@ -570,7 +644,10 @@ class TestDocumentSelectionIdentity:
modal = DocumentFilterModal(client=client)
app, _ = _make_app(temp_db_path, client)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=client):
with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
_covering_returns(_stub_rag, client),
):
async with app.run_test() as pilot:
await app.push_screen(modal)
await pilot.pause()

5
tests/chat/test_image_input.py
View file

@ -16,6 +16,7 @@ from haiku.rag.chat.widgets.prompt import (
PostableTextArea,
build_user_prompt,
)
from tests.conftest import for_path
def make_image_bytes(fmt: str = "PNG") -> bytes:
@ -131,7 +132,7 @@ class TestChatAppImageAttach:
async with HaikuRAG(temp_db_path, create=True):
pass
app = ChatApp(db_path=temp_db_path, capabilities=[])
app = ChatApp(scope=for_path(temp_db_path), capabilities=[])
async with app.run_test() as pilot:
data = make_image_bytes()
app.post_message(ImageAdded(Path("img.png"), data))
@ -150,7 +151,7 @@ class TestChatAppLayout:
async with HaikuRAG(temp_db_path, create=True):
pass
app = ChatApp(db_path=temp_db_path, capabilities=[])
app = ChatApp(scope=for_path(temp_db_path), capabilities=[])
async with app.run_test() as pilot:
await pilot.pause()
prompt = app.query_one(FlexibleInput)

30
tests/conftest.py
View file

@ -31,7 +31,9 @@ if TYPE_CHECKING:
from vcr import VCR
from haiku.rag.client import HaikuRAG
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.client.session import SingleDatabaseSession
from haiku.rag.config.models import AppConfig
setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
@ -230,3 +232,31 @@ def writing(client: "HaikuRAG") -> "SingleDatabaseSession":
assert isinstance(client._session, SingleDatabaseSession)
return client._session
def for_path(
db_path: "Path | str | None" = None, config: "AppConfig | None" = None
) -> "DatabaseScope":
"""A scope covering one database at `db_path`.
The application layer takes the databases it works on, already resolved.
Tests that hold a path rather than a scope go through here.
"""
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config import get_config
return DatabaseScope.resolve(
config if config is not None else get_config(), database_path=db_path
)
@contextmanager
def _covering_returns(stub, client):
"""Make a patched `HaikuRAG` hand back `client` however it is constructed.
The TUIs build their client through `HaikuRAG._covering`, so patching the
constructor alone leaves `_covering` answering with a fresh Mock.
"""
stub.return_value = client
stub._covering.return_value = client
yield stub

11
tests/store/test_read_only.py
View file

@ -9,6 +9,7 @@ from haiku.rag.store.models import Chunk, Document
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.settings import SettingsRepository
from tests.conftest import for_path
@pytest.fixture(scope="module")
@ -321,7 +322,7 @@ class TestAppReadVerbsDoNotWrite:
drift.embeddings.model.name = "different-model"
# list is a read verb — must open read-only and not raise on drift
app = HaikuRAGApp(db_path=temp_db_path, config=drift)
app = HaikuRAGApp(scope=for_path(temp_db_path, drift), config=drift)
await app.list_documents()
async with Store(temp_db_path, skip_validation=True, read_only=True) as store:
@ -347,7 +348,7 @@ class TestAppReadVerbsDoNotWrite:
drift.embeddings.model.name = "different-model"
drift.embeddings.model.vector_dim = 4096
app = HaikuRAGApp(db_path=temp_db_path, config=drift)
app = HaikuRAGApp(scope=for_path(temp_db_path, drift), config=drift)
await app.list_documents()
async with Store(temp_db_path, skip_validation=True, read_only=True) as store:
@ -368,7 +369,7 @@ class TestAppReadVerbsDoNotWrite:
drift = AppConfig()
drift.embeddings.model.vector_dim = 4096
app = HaikuRAGApp(db_path=temp_db_path, config=drift)
app = HaikuRAGApp(scope=for_path(temp_db_path, drift), config=drift)
assert doc.id is not None
await app.get_document(doc.id)
@ -386,7 +387,7 @@ class TestAppReadVerbsDoNotWrite:
drift = AppConfig()
drift.embeddings.model.vector_dim = 4096
app = HaikuRAGApp(db_path=temp_db_path, config=drift)
app = HaikuRAGApp(scope=for_path(temp_db_path, drift), config=drift)
await app.visualize_chunk("missing-chunk-id")
@pytest.mark.asyncio
@ -404,7 +405,7 @@ class TestAppReadVerbsDoNotWrite:
drift = AppConfig()
drift.embeddings.model.name = "different-model"
app = HaikuRAGApp(db_path=temp_db_path, config=drift)
app = HaikuRAGApp(scope=for_path(temp_db_path, drift), config=drift)
with pytest.raises(ConfigMismatchError):
await app.add_document_from_text("hello")

19
tests/test_app.py
View file

@ -14,6 +14,7 @@ from haiku.rag.client import RebuildMode
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document import Document
from tests.conftest import for_path
@pytest.fixture
@ -28,6 +29,10 @@ def app(tmp_path, client, monkeypatch):
def __init__(self, *args, **kwargs):
pass
@classmethod
def _covering(cls, *args, **kwargs):
return cls()
async def __aenter__(self):
return client
@ -37,7 +42,7 @@ def app(tmp_path, client, monkeypatch):
monkeypatch.setattr("haiku.rag.app.HaikuRAG", StubHaikuRAG)
db = tmp_path / "db.lancedb"
db.mkdir()
application = HaikuRAGApp(db_path=db, config=AppConfig())
application = HaikuRAGApp(scope=for_path(db), config=AppConfig())
application.console = Console(record=True, width=200)
return application
@ -317,7 +322,7 @@ async def test_create_index_rebuilds_an_existing_one(app, client):
def test_show_settings_hides_secrets(tmp_path):
config = AppConfig(lancedb=LanceDBConfig(uri="db://x", api_key="secret-value"))
app = HaikuRAGApp(db_path=tmp_path / "db", config=config)
app = HaikuRAGApp(scope=for_path(tmp_path / "db", config), config=config)
app.console = Console(record=True, width=200)
app.show_settings()
@ -329,7 +334,7 @@ def test_show_settings_hides_secrets(tmp_path):
def test_remote_uri_is_the_display_path(tmp_path):
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
app = HaikuRAGApp(db_path=tmp_path / "db", config=config)
app = HaikuRAGApp(scope=for_path(None, config), config=config)
assert app._display_path == "s3://bucket/path"
assert app._is_local is False
@ -364,7 +369,7 @@ async def test_init_reports_an_existing_database(app):
async def test_info_reports_a_missing_path(tmp_path):
application = HaikuRAGApp(db_path=tmp_path / "gone", config=AppConfig())
application = HaikuRAGApp(scope=for_path(tmp_path / "gone"), config=AppConfig())
application.console = Console(record=True, width=200)
await application.info()
@ -375,7 +380,7 @@ async def test_info_reports_a_missing_path(tmp_path):
async def test_history_reports_a_missing_path(tmp_path):
application = HaikuRAGApp(db_path=tmp_path / "gone", config=AppConfig())
application = HaikuRAGApp(scope=for_path(tmp_path / "gone"), config=AppConfig())
application.console = Console(record=True, width=200)
await application.history()
@ -459,7 +464,7 @@ async def test_restore_tag_reports_the_safety_tag(app, monkeypatch):
],
)
async def test_tag_operations_require_the_database(tmp_path, method, args):
application = HaikuRAGApp(db_path=tmp_path / "gone", config=AppConfig())
application = HaikuRAGApp(scope=for_path(tmp_path / "gone"), config=AppConfig())
with pytest.raises(ValueError, match="does not exist"):
await getattr(application, method)(*args)
@ -601,7 +606,7 @@ async def test_doctor_reports_the_duplicates_export(app, monkeypatch, tmp_path):
async def test_doctor_reports_a_missing_database(tmp_path):
application = HaikuRAGApp(db_path=tmp_path / "gone", config=AppConfig())
application = HaikuRAGApp(scope=for_path(tmp_path / "gone"), config=AppConfig())
application.console = Console(record=True, width=200)
assert await application.doctor() is True

3
tests/test_cascade_delete.py
View file

@ -3,6 +3,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import parent_uri_filter
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.document import Document
from tests.conftest import for_path
def test_parent_uri_filter_simple():
@ -138,7 +139,7 @@ async def test_delete_succeeds_with_embedding_dim_mismatch(temp_db_path):
mismatched = AppConfig()
mismatched.embeddings.model.vector_dim = 9999
app = HaikuRAGApp(db_path=temp_db_path, config=mismatched)
app = HaikuRAGApp(scope=for_path(temp_db_path, mismatched), config=mismatched)
await app.delete_document(doc.id)
async with HaikuRAG(

177
tests/test_cli.py
View file

@ -8,19 +8,15 @@ from click.exceptions import BadParameter
from typer.testing import CliRunner
from haiku.rag.cli import _cli as cli
from haiku.rag.cli import (
_parse_meta_options,
require_one_database,
resolve_db_path,
select_database,
)
from haiku.rag.cli import _parse_meta_options, resolve_scope
from haiku.rag.cli import cli as cli_wrapper
from haiku.rag.config import get_config, set_config
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.config.models import AppConfig, LanceDBConfig, StorageConfig
from haiku.rag.store.exceptions import (
AmbiguousDatabaseError,
MigrationRequiredError,
)
from tests.conftest import for_path
runner = CliRunner()
@ -109,47 +105,62 @@ class TestOneDatabaseCommands:
"""`lancedb.databases` names a set; most commands work on one database."""
@staticmethod
def _config(**databases):
return AppConfig(lancedb=LanceDBConfig(databases=databases))
def _install(monkeypatch, **databases):
import haiku.rag.config as config_module
monkeypatch.setattr(config_module, "_config", None)
monkeypatch.setattr("haiku.rag.cli._database", None)
set_config(AppConfig(lancedb=LanceDBConfig(databases=databases)))
def test_a_configured_set_refuses_a_one_database_command(self, monkeypatch):
self._install(monkeypatch, alpha="/db/a.lancedb", beta="/db/b.lancedb")
def test_a_configured_set_refuses_a_one_database_command(self):
with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
require_one_database(
self._config(alpha="/db/a.lancedb", beta="/db/b.lancedb"),
None,
federated=False,
)
resolve_scope(None)
def test_the_refusal_names_the_databases_and_not_their_locations(self):
def test_the_refusal_names_the_databases_and_not_their_locations(self, monkeypatch):
"""A location in an error message travels into logs and terminals; the
names exist so it does not have to."""
self._install(
monkeypatch,
alpha="s3://bucket/prefix/a.lancedb",
beta="s3://bucket/prefix/b.lancedb",
)
with pytest.raises(AmbiguousDatabaseError) as raised:
require_one_database(
self._config(alpha="s3://bucket/prefix/a.lancedb"),
None,
federated=False,
)
resolve_scope(None)
assert "alpha" in str(raised.value)
assert "s3://bucket/prefix/a.lancedb" not in str(raised.value)
assert "bucket" not in str(raised.value)
def test_a_command_covering_the_set_is_allowed(self):
require_one_database(
self._config(alpha="/db/a.lancedb", beta="/db/b.lancedb"),
None,
federated=True,
)
def test_a_configured_set_of_one_needs_no_choosing(self, monkeypatch):
"""Nothing is ambiguous about a set with one database in it, and it keeps
its name rather than being refused."""
self._install(monkeypatch, alpha="/db/a.lancedb")
def test_naming_a_path_is_allowed(self):
require_one_database(
self._config(alpha="/db/a.lancedb"),
Path("/db/other.lancedb"),
federated=False,
)
assert resolve_scope(None).names == ("alpha",)
def test_no_configured_databases_is_allowed(self):
require_one_database(AppConfig(), None, federated=False)
def test_a_command_covering_the_set_is_allowed(self, monkeypatch):
self._install(monkeypatch, alpha="/db/a.lancedb", beta="/db/b.lancedb")
assert resolve_scope(None, covers_set=True).names == ("alpha", "beta")
def test_naming_a_path_is_allowed(self, monkeypatch):
self._install(monkeypatch, alpha="/db/a.lancedb")
[ref] = resolve_scope(Path("/db/other.lancedb")).databases
assert ref.db_path == Path("/db/other.lancedb")
def test_no_configured_databases_is_allowed(self, monkeypatch, tmp_path):
import haiku.rag.config as config_module
monkeypatch.setattr(config_module, "_config", None)
monkeypatch.setattr("haiku.rag.cli._database", None)
set_config(AppConfig(storage=StorageConfig(data_dir=tmp_path)))
[ref] = resolve_scope(None).databases
assert ref.db_path == tmp_path / "haiku.rag.lancedb"
def test_the_refusal_exits_with_an_error(self):
with patch("haiku.rag.cli._cli") as mock_cli:
@ -171,46 +182,49 @@ class TestSelectingADatabaseByName:
monkeypatch.setattr(config_module, "_config", None)
set_config(AppConfig(lancedb=LanceDBConfig(databases=databases)))
def test_a_uri_location_becomes_the_configured_uri(self, monkeypatch):
def test_a_named_database_is_passed_on_by_name(self, monkeypatch):
"""Not resolved to a path: the name is what results and citations carry,
and rewriting the configuration is what used to lose it."""
self._install(monkeypatch, medic="s3://bucket/prefix/medic.lancedb")
monkeypatch.setattr("haiku.rag.cli._database", "medic")
db_path = select_database("medic")
assert resolve_scope(None).names == ("medic",)
def test_naming_a_database_leaves_the_configuration_alone(self, monkeypatch):
self._install(monkeypatch, st="/data/st.lancedb", other="/data/o.lancedb")
monkeypatch.setattr("haiku.rag.cli._database", "st")
resolve_scope(None)
assert db_path is None
config = get_config()
assert config.lancedb.uri == "s3://bucket/prefix/medic.lancedb"
assert config.lancedb.databases == {}
def test_a_local_location_becomes_the_database_path(self, monkeypatch):
self._install(monkeypatch, st="/data/st.lancedb")
db_path = select_database("st")
assert db_path == Path("/data/st.lancedb")
assert get_config().lancedb.uri == ""
assert config.lancedb.databases == {
"st": "/data/st.lancedb",
"other": "/data/o.lancedb",
}
assert config.lancedb.uri == ""
def test_an_unknown_name_names_the_configured_ones(self, monkeypatch):
self._install(monkeypatch, alpha="/data/a.lancedb", beta="/data/b.lancedb")
monkeypatch.setattr("haiku.rag.cli._database", "gamma")
with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
select_database("gamma")
resolve_scope(None)
def test_an_unknown_name_does_not_leak_locations(self, monkeypatch):
self._install(monkeypatch, medic="s3://bucket/prefix/medic.lancedb")
monkeypatch.setattr("haiku.rag.cli._database", "gamma")
with pytest.raises(AmbiguousDatabaseError) as raised:
select_database("gamma")
resolve_scope(None)
assert "bucket" not in str(raised.value)
def test_selecting_nothing_reports_an_empty_mapping(self, monkeypatch):
import haiku.rag.config as config_module
monkeypatch.setattr(config_module, "_config", None)
set_config(AppConfig())
self._install(monkeypatch)
monkeypatch.setattr("haiku.rag.cli._database", "medic")
with pytest.raises(AmbiguousDatabaseError, match="nothing"):
select_database("medic")
resolve_scope(None)
def test_the_callback_selects_before_a_command_runs(self, tmp_path, monkeypatch):
"""`--database` is resolved once the config is loaded, so every command
@ -237,7 +251,6 @@ class TestSelectingADatabaseByName:
monkeypatch.setattr(config_module, "_config", None)
monkeypatch.setattr(cli_module, "_database", None)
monkeypatch.setattr(cli_module, "_database_path", None)
config_file = tmp_path / "haiku.rag.yaml"
selected = tmp_path / "alpha.lancedb"
config_file.write_text(f"lancedb:\n databases:\n alpha: {selected}\n")
@ -245,24 +258,22 @@ class TestSelectingADatabaseByName:
runner.invoke(
cli, ["--config", str(config_file), "--database", "alpha", "info"]
)
assert cli_module._database_path == selected
assert cli_module._database == "alpha"
runner.invoke(cli, ["--config", str(config_file), "settings"])
assert cli_module._database_path is None
assert cli_module._database is None
def test_a_selection_does_not_outlive_its_invocation_in_process(
self, tmp_path, monkeypatch
):
"""Selecting rewrites the configuration, so a second invocation has to
start from a freshly loaded one rather than inherit the rewrite."""
"""The selection is per invocation, so a second one starts from a
freshly loaded configuration rather than inheriting the first."""
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)
config_file = tmp_path / "haiku.rag.yaml"
config_file.write_text(
f"lancedb:\n databases:\n alpha: {tmp_path / 'alpha.lancedb'}\n"
@ -272,24 +283,25 @@ class TestSelectingADatabaseByName:
runner.invoke(
cli, ["--config", str(config_file), "--database", "alpha", "settings"]
)
# Naming one leaves the configuration naming both.
assert get_config().lancedb.uri == ""
assert get_config().lancedb.databases == {}
assert set(get_config().lancedb.databases) == {"alpha", "beta"}
runner.invoke(cli, ["--config", str(config_file), "settings"])
assert set(get_config().lancedb.databases) == {"alpha", "beta"}
assert cli_module._database is None
def test_a_selection_does_not_outlive_an_invocation_without_a_config_file(
self, tmp_path, monkeypatch
):
"""No config file is still a load: the previous invocation's selected URI
"""No config file is still a load: the previous invocation's database
must not be what the next one talks to."""
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)
monkeypatch.delenv("HAIKU_RAG_CONFIG_PATH", raising=False)
monkeypatch.chdir(tmp_path)
config_file = tmp_path / "selected.yaml"
@ -300,32 +312,26 @@ class TestSelectingADatabaseByName:
runner.invoke(
cli, ["--config", str(config_file), "--database", "medic", "settings"]
)
assert get_config().lancedb.uri == "s3://bucket/medic.lancedb"
assert get_config().lancedb.databases == {"medic": "s3://bucket/medic.lancedb"}
runner.invoke(cli, ["settings"])
assert get_config().lancedb.uri == ""
assert get_config().lancedb.databases == {}
assert cli_module._database is None
class TestResolvingTheDatabasePath:
def test_a_path_wins_when_nothing_is_selected(self, monkeypatch):
monkeypatch.setattr("haiku.rag.cli._database", None)
monkeypatch.setattr("haiku.rag.cli._database_path", None)
assert resolve_db_path(Path("/data/one.lancedb")) == Path("/data/one.lancedb")
def test_the_selected_database_is_used_without_a_path(self, monkeypatch):
monkeypatch.setattr("haiku.rag.cli._database", "st")
monkeypatch.setattr("haiku.rag.cli._database_path", Path("/data/st.lancedb"))
assert resolve_db_path(None) == Path("/data/st.lancedb")
[ref] = resolve_scope(Path("/data/one.lancedb")).databases
assert ref.db_path == Path("/data/one.lancedb")
def test_naming_a_database_twice_is_refused(self, monkeypatch):
monkeypatch.setattr("haiku.rag.cli._database", "st")
monkeypatch.setattr("haiku.rag.cli._database_path", Path("/data/st.lancedb"))
with pytest.raises(AmbiguousDatabaseError, match="not both"):
resolve_db_path(Path("/data/other.lancedb"))
resolve_scope(Path("/data/other.lancedb"))
class TestCliConfigMismatchError:
@ -595,7 +601,7 @@ class TestAskAnalyzeImageOption:
with patch.object(
HaikuRAG, "ask", new_callable=AsyncMock, return_value=("answer", [])
) as mock_ask:
app = HaikuRAGApp(db_path=temp_db_path)
app = HaikuRAGApp(scope=for_path(temp_db_path))
await app.ask("q", images=[img_path])
assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()]
@ -622,7 +628,6 @@ class TestChatCoversTheSet:
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(
@ -630,8 +635,8 @@ class TestChatCoversTheSet:
)
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
# The scope covers both, which is what makes chat read the set.
assert run_chat.call_args.kwargs["scope"].names == ("arxiv", "wiki")
def test_naming_one_database_opens_that_one(self, tmp_path, monkeypatch):
import haiku.rag.cli as cli_module
@ -639,7 +644,6 @@ class TestChatCoversTheSet:
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(
@ -654,7 +658,8 @@ class TestChatCoversTheSet:
)
assert result.exit_code == 0, result.output
assert run_chat.call_args.args[0] == tmp_path / "w.lancedb"
# Passed on by name, so the database it opens keeps its identity.
assert run_chat.call_args.kwargs["scope"].names == ("wiki",)
def test_a_single_database_setup_is_unchanged(self, tmp_path, monkeypatch):
"""Without a configured set, chat opens the path it always did."""
@ -663,13 +668,13 @@ class TestChatCoversTheSet:
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"
[ref] = run_chat.call_args.kwargs["scope"].databases
assert ref.db_path == tmp_path / "one.lancedb"
class TestRenderingTheDatabase:
@ -681,7 +686,7 @@ class TestRenderingTheDatabase:
from haiku.rag.app import HaikuRAGApp
return HaikuRAGApp(
db_path=tmp_path / "unused",
scope=for_path(tmp_path / "unused"),
config=AppConfig(lancedb=LanceDBConfig(databases=databases)),
)
@ -730,7 +735,7 @@ def app_stub(monkeypatch, tmp_path):
stub = AsyncMock()
monkeypatch.setattr(
"haiku.rag.cli.create_app",
lambda db=None, *, federated=False: stub,
lambda db=None, *, covers_set=False: stub,
)
return stub

12
tests/test_database_scope.py
View file

@ -74,6 +74,18 @@ class TestResolution:
assert scope.databases == (DatabaseRef(None, "s3://bucket/one.lancedb", None),)
def test_a_path_selects_the_database_over_a_configured_uri(self):
"""`--db` exists to override what is configured, and the configuration
derived from the ref is what makes the connection follow it."""
config = _config(uri="s3://bucket/one.lancedb")
scope = DatabaseScope.resolve(config, database_path=Path("/data/local"))
[ref] = scope.databases
assert ref.db_path == Path("/data/local")
one, _ = ref.connection(config)
assert one.lancedb.uri == ""
def test_nothing_configured_falls_back_to_the_data_directory(self, tmp_path):
config = AppConfig(storage=StorageConfig(data_dir=tmp_path))

82
tests/test_info.py
View file

@ -6,6 +6,7 @@ import pytest
from haiku.rag.app import HaikuRAGApp
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.schema import DocumentItemRecord
from tests.conftest import for_path
@pytest.mark.asyncio
@ -61,7 +62,7 @@ async def test_app_info_outputs(temp_db_path, capsys):
[ChunkRecord(id="c1", document_id="doc-1", content="c", vector=[0.1, 0.2, 0.3])]
)
app = HaikuRAGApp(db_path=temp_db_path)
app = HaikuRAGApp(scope=for_path(temp_db_path))
await app.info()
out = capsys.readouterr().out
@ -148,7 +149,7 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
# Create vector index
await chunks_tbl.create_index("vector", config=IvfPq(distance_type="cosine"))
app = HaikuRAGApp(db_path=temp_db_path)
app = HaikuRAGApp(scope=for_path(temp_db_path))
await app.info()
out = capsys.readouterr().out
@ -164,16 +165,61 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
@pytest.mark.asyncio
async def test_app_info_opens_a_named_remote_database(tmp_path):
"""A database named in `lancedb.databases` can sit behind a URI while the
configuration's own `uri` is empty. Passing that configuration on would open
the local path that only stands in for it."""
from haiku.rag.client.scope import DatabaseScope
config = AppConfig(
lancedb=LanceDBConfig(databases={"medic": "s3://bucket/medic.lancedb"})
)
scope = DatabaseScope.resolve(config, database_name="medic")
app = HaikuRAGApp(scope=scope, config=config)
assert app._is_local is False
assert app._store_config.lancedb.uri == "s3://bucket/medic.lancedb"
assert app._store_config.lancedb.databases == {}
with patch(
"haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock
) as mock_connect:
mock_db = mock_connect.return_value
mock_list_result = MagicMock()
mock_list_result.tables = []
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
await app.info()
opened = mock_connect.call_args.args[0]
assert opened.lancedb.uri == "s3://bucket/medic.lancedb"
async def test_app_doctor_opens_a_named_remote_database():
"""`run_doctor` connects with the configuration it is handed, so it needs the
one derived for the database rather than the one naming the set."""
from haiku.rag.client.scope import DatabaseScope
config = AppConfig(
lancedb=LanceDBConfig(databases={"medic": "s3://bucket/medic.lancedb"})
)
app = HaikuRAGApp(scope=DatabaseScope.resolve(config, database_name="medic"))
with patch("haiku.rag.doctor.run_doctor", new_callable=AsyncMock) as run:
run.return_value = MagicMock(checks=[], ok=True, duplicates=None)
await app.doctor()
assert run.call_args.args[0].lancedb.uri == "s3://bucket/medic.lancedb"
async def test_app_info_uses_connect_lancedb_for_remote(tmp_path):
"""info() should use connect_lancedb() instead of direct lancedb.connect() for remote URIs."""
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
app = HaikuRAGApp(scope=for_path(None, config), config=config)
with patch(
"haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock
@ -185,7 +231,10 @@ async def test_app_info_uses_connect_lancedb_for_remote(tmp_path):
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
await app.info()
mock_connect.assert_called_once_with(config, nonexistent)
# The uri decides where it connects; the path argument is not read, so the
# claim here is the remote branch, not which placeholder path it carried.
mock_connect.assert_called_once()
assert mock_connect.call_args.args[0].lancedb.uri == "s3://bucket/path"
@pytest.mark.asyncio
@ -241,7 +290,7 @@ async def test_app_info_with_missing_document_items_table(temp_db_path, capsys):
[ChunkRecord(id="c1", document_id="doc-1", content="c", vector=[0.1, 0.2, 0.3])]
)
app = HaikuRAGApp(db_path=temp_db_path)
app = HaikuRAGApp(scope=for_path(temp_db_path))
await app.info()
out = capsys.readouterr().out
@ -311,7 +360,7 @@ async def test_app_info_reports_up_to_date(temp_db_path, capsys):
]
)
app = HaikuRAGApp(db_path=temp_db_path)
app = HaikuRAGApp(scope=for_path(temp_db_path))
await app.info()
out = capsys.readouterr().out
@ -322,35 +371,34 @@ async def test_app_info_reports_up_to_date(temp_db_path, capsys):
@pytest.mark.asyncio
async def test_app_init_skips_exists_check_for_remote(tmp_path):
"""init() should not check db_path.exists() for remote URIs."""
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
app = HaikuRAGApp(scope=for_path(None, config), config=config)
with patch("haiku.rag.app.HaikuRAG") as mock_client_cls:
mock_client = AsyncMock()
mock_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_client)
mock_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
covering = mock_client_cls._covering.return_value
covering.__aenter__ = AsyncMock(return_value=mock_client)
covering.__aexit__ = AsyncMock(return_value=False)
await app.init()
# Should have called HaikuRAG to create, not returned early
mock_client_cls.assert_called_once()
# Opened to create, rather than returning early on a missing local path.
mock_client_cls._covering.assert_called_once()
@pytest.mark.asyncio
async def test_app_history_skips_exists_check_for_remote(tmp_path):
"""history() should not check db_path.exists() for remote URIs."""
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
app = HaikuRAGApp(scope=for_path(None, config), config=config)
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
mock_store = AsyncMock()
@ -377,7 +425,7 @@ async def test_app_tag_rendering_escapes_markup(tmp_path):
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=tmp_path / "db.lancedb", config=config)
app = HaikuRAGApp(scope=for_path(None, config), config=config)
app.console = Console(record=True, width=200)
hostile = "[red]release[/red]"
@ -411,7 +459,7 @@ async def test_app_history_survives_tag_annotation_failure(tmp_path):
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=tmp_path / "db.lancedb", config=config)
app = HaikuRAGApp(scope=for_path(None, config), config=config)
app.console = Console(record=True, width=200)
with patch("haiku.rag.store.engine.Store") as mock_store_cls:

3
tests/test_inspector.py
View file

@ -9,6 +9,7 @@ from typer.testing import CliRunner
from haiku.rag.cli import _cli as cli
from haiku.rag.store.models import Chunk, Document, SearchResult
from tests.conftest import for_path
runner = CliRunner()
@ -326,7 +327,7 @@ async def test_inspector_open_failure_surfaces_real_error(tmp_path):
AttributeError from tearing down a client that never opened."""
from haiku.rag.inspector.app import InspectorApp
app = InspectorApp(db_path=tmp_path / "missing.lancedb", read_only=True)
app = InspectorApp(scope=for_path(tmp_path / "missing.lancedb"), read_only=True)
with pytest.raises(FileNotFoundError):
async with app.run_test():
pass

41
tests/test_multi_db.py
View file

@ -6,6 +6,7 @@ from docling_core.types.doc.labels import DocItemLabel
from pydantic import ValidationError
from haiku.rag.client import HaikuRAG
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.client.session import FederatedSession
from haiku.rag.config import get_config
from haiku.rag.config.models import AppConfig, LanceDBConfig
@ -379,6 +380,46 @@ class TestClosingASet:
assert sorted(drained) == ["alpha", "beta"]
class TestNamingOneOfTheSetOnTheCommandLine:
"""`--database NAME` reaches the application layer as a name, and every
client it opens has to honour it — one that ignores it covers the set and
quietly answers from the wrong database."""
@pytest.mark.asyncio
async def test_a_named_database_is_the_one_read(self, tmp_path, capsys):
from haiku.rag.app import HaikuRAGApp
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
scope = DatabaseScope.resolve(config).select(["beta"])
app = HaikuRAGApp(scope=scope, config=config, read_only=True)
await app.list_documents()
# Rich wraps long lines, so match the unwrapped part of the URI.
printed = capsys.readouterr().out
assert "test://beta/" in printed
assert "test://alpha/" not in printed
@pytest.mark.asyncio
async def test_naming_none_of_them_covers_the_set(self, tmp_path, capsys):
from haiku.rag.app import HaikuRAGApp
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
app = HaikuRAGApp(
scope=DatabaseScope.resolve(config), config=config, read_only=True
)
await app.list_documents()
printed = capsys.readouterr().out
assert "test://alpha/" in printed
assert "test://beta/" in printed
class TestPlacingADatabase:
"""What a client says about the databases it covers, so nothing outside has
to read its private state to find out."""

4
tests/test_rebuild.py
View file

@ -7,7 +7,7 @@ import pytest
from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import get_config
from tests.conftest import capture_logs, writing
from tests.conftest import capture_logs, for_path, writing
class ChunkData(TypedDict):
@ -1061,7 +1061,7 @@ async def test_rebuild_set_embedder_works_on_empty_database(temp_db_path):
drift = AppConfig()
drift.embeddings.model.name = "different-model"
app = HaikuRAGApp(db_path=temp_db_path, config=drift)
app = HaikuRAGApp(scope=for_path(temp_db_path, drift), config=drift)
await app.rebuild(mode=RebuildMode.SET_EMBEDDER)
async with HaikuRAG(temp_db_path, config=drift, skip_validation=True) as client:

5
tests/test_s3_integration.py
View file

@ -13,6 +13,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.s3 import make_s3_store
from haiku.rag.store.engine import Store
from tests.conftest import for_path
from tests.services import reachable
S3_ENDPOINT = "http://localhost:8333"
@ -137,7 +138,7 @@ async def test_app_info(tmp_path, capsys, config):
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
await rag.create_document("Info test document.", uri="test://info")
app = HaikuRAGApp(db_path=tmp_path / "unused", config=config)
app = HaikuRAGApp(scope=for_path(tmp_path / "unused", config), config=config)
await app.info()
out = capsys.readouterr().out
@ -148,7 +149,7 @@ async def test_app_info(tmp_path, capsys, config):
@pytest.mark.asyncio
async def test_app_info_empty_db(tmp_path, capsys, config):
app = HaikuRAGApp(db_path=tmp_path / "unused", config=config)
app = HaikuRAGApp(scope=for_path(tmp_path / "unused", config), config=config)
await app.info()
out = capsys.readouterr().out