diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 4a71ac9b..7e7a2cc7 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -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": diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index 593b4f33..9dd2c2b4 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index c4996e9b..c358e316 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index d9ed271e..db1d580c 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 7bbbbbb5..cdcf4fd9 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index f1836994..c64976a8 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/client/scope.py b/haiku_rag_slim/haiku/rag/client/scope.py index f8453e23..03fbaebc 100644 --- a/haiku_rag_slim/haiku/rag/client/scope.py +++ b/haiku_rag_slim/haiku/rag/client/scope.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index 61ef83a8..61502fe6 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -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() diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index b6ee7f03..1991ccbb 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -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() diff --git a/tests/chat/test_image_input.py b/tests/chat/test_image_input.py index 4a3f6255..ac0e9197 100644 --- a/tests/chat/test_image_input.py +++ b/tests/chat/test_image_input.py @@ -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) diff --git a/tests/conftest.py b/tests/conftest.py index 15f4654c..bad34436 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/store/test_read_only.py b/tests/store/test_read_only.py index 26e74b1d..f465c143 100644 --- a/tests/store/test_read_only.py +++ b/tests/store/test_read_only.py @@ -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") diff --git a/tests/test_app.py b/tests/test_app.py index bca18636..e1371f4c 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -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 diff --git a/tests/test_cascade_delete.py b/tests/test_cascade_delete.py index 5d68e375..300b14e9 100644 --- a/tests/test_cascade_delete.py +++ b/tests/test_cascade_delete.py @@ -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( diff --git a/tests/test_cli.py b/tests/test_cli.py index 8d9f4aaf..55cf1960 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -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 diff --git a/tests/test_database_scope.py b/tests/test_database_scope.py index c1f3eba1..2c131b3a 100644 --- a/tests/test_database_scope.py +++ b/tests/test_database_scope.py @@ -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)) diff --git a/tests/test_info.py b/tests/test_info.py index d6a48e48..71418787 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -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: diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 4e751861..95aebb0e 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -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 diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index 5e148983..a856b3e6 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -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.""" diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 96c860da..1cc7f985 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -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: diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index f916d07a..e2c61df1 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -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