From 0813a1c98005ee84c59f2d3439033cc55fa8633e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 14 Jul 2026 16:15:28 +0300 Subject: [PATCH] Add tag CLI commands and history tag annotations --- CHANGELOG.md | 9 +- README.md | 2 +- docs/cli.md | 34 +++++- haiku_rag_slim/haiku/rag/app.py | 112 +++++++++++++++++++- haiku_rag_slim/haiku/rag/chat/__init__.py | 3 + haiku_rag_slim/haiku/rag/chat/app.py | 3 + haiku_rag_slim/haiku/rag/cli.py | 87 ++++++++++++++- haiku_rag_slim/haiku/rag/client/__init__.py | 5 + haiku_rag_slim/haiku/rag/client/rebuild.py | 21 +++- haiku_rag_slim/haiku/rag/inspector/app.py | 12 ++- haiku_rag_slim/haiku/rag/mcp.py | 58 ++++++++-- haiku_rag_slim/haiku/rag/store/engine.py | 41 +++++-- tests/chat/test_chat_app.py | 8 +- tests/store/test_tags.py | 18 ++++ tests/store/test_time_travel.py | 61 +++++++++++ tests/test_cli.py | 95 +++++++++++++++++ tests/test_client.py | 29 +++++ tests/test_info.py | 38 +++++++ tests/test_mcp.py | 44 ++++++++ tests/test_rebuild.py | 46 ++++++++ 20 files changed, 693 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3bc491c2..a67157bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- Database tags: `haiku-rag tag create/list/delete`, `--at TAG` time travel, tags shown in `history`. Vacuum retains versions back to the oldest tag. + +### Changed + +- `lancedb` bumped to 0.34.0. + ### Fixed - `docling-local` text conversion no longer misroutes markdown/HTML content whose first bytes collide with a binary magic signature (e.g. `BM`, `ID3`) to an image or audio backend. @@ -11,7 +19,6 @@ - Unknown `reranking.model.provider` raises `ValueError` instead of silently disabling reranking. - `search.max_context_chars` default lowered from 10000 to 5000. -- `lancedb` bumped to 0.34.0. ### Removed diff --git a/README.md b/README.md index 6e57596d..b0eae73f 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p - **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.) - **Visual grounding** — View chunks highlighted on original page images - **Production ingester** — Long-lived `haiku-ingester` service with persistent SQLite queue, async worker pool with retries and a dead-letter queue, FS / HTTP / S3 / WebDAV source adapters, FastAPI control plane, and a browser dashboard for operators. See [docs/ingester.md](docs/ingester.md). -- **Time travel** — Query the database at any historical point with `--before` +- **Time travel** — Query the database at any historical point with `--before`, or tag states and query them with `--at` - **Inspector** — TUI for browsing documents, chunks, and search results ## Installation diff --git a/docs/cli.md b/docs/cli.md index 1fe6d14d..cb058673 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -8,6 +8,7 @@ The `haiku-rag` CLI provides complete document management functionality. - `--config` - Specify custom configuration file - `--read-only` - Open database in read-only mode (blocks writes, skips upgrades) - `--before` - Query database as it existed before a datetime (implies `--read-only`) + - `--at` - Query database at a tag (implies `--read-only`, mutually exclusive with `--before`) - `--version` / `-v` - Show version and exit Per-command options: @@ -572,6 +573,31 @@ Supported datetime formats: !!! note Time travel mode automatically enables read-only mode. You cannot modify the database while viewing historical state. +### Tags + +Tags name the current database state so you can return to it without remembering timestamps. A tag covers every table in the database. Tagged versions survive `vacuum`; everything older than your oldest tag is retained until that tag is deleted, so remove tags you no longer need. + +```bash +# Tag the current state, e.g. at deploy time or after an ingestion run +haiku-rag tag create release-1 + +# List tags with the versions they point to +haiku-rag tag list + +# Delete a tag, releasing its versions for cleanup +haiku-rag tag delete release-1 +``` + +Query the database at a tag with `--at`: + +```bash +haiku-rag --at release-1 list +haiku-rag --at release-1 search "machine learning" +haiku-rag --at release-1 ask "What documents existed?" +``` + +`--at` implies read-only mode and is mutually exclusive with `--before`. + ### Version History View version history for database tables: @@ -587,20 +613,20 @@ haiku-rag history --table documents haiku-rag history --limit 10 ``` -Output shows version numbers and timestamps, sorted newest first: +Output shows version numbers and timestamps, sorted newest first, with tags marked: ``` Version History documents - v5: 2025-01-15 14:30:00 + v5: 2025-01-15 14:30:00 <- release-1 v4: 2025-01-14 10:00:00 v3: 2025-01-13 09:15:00 chunks - v8: 2025-01-15 14:30:00 + v8: 2025-01-15 14:30:00 <- release-1 v7: 2025-01-14 10:00:00 ... ``` -Use the timestamps from `history` to construct `--before` queries. +Use the timestamps from `history` to construct `--before` queries, or tag names with `--at`. diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index b3eba176..fc037eac 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING from rich.console import Console from rich.markdown import Markdown +from rich.markup import escape from rich.progress import ( BarColumn, DownloadColumn, @@ -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.store.engine import Store from haiku.rag.store.models import SearchResult from haiku.rag.utils import format_bytes, format_citations_rich @@ -35,11 +37,13 @@ class HaikuRAGApp: # pragma: no cover config: AppConfig = Config, read_only: bool = False, before: datetime | None = None, + at_tag: str | None = None, ): self.db_path = db_path self.config = config self.read_only = read_only self.before = before + self.at_tag = at_tag self.console = Console() from haiku.rag.store.engine import ConnectionMode @@ -67,9 +71,9 @@ class HaikuRAGApp: # pragma: no cover from haiku.rag.store.engine import gather_database_info - if self.before is not None: + if self.before is not None or self.at_tag is not None: self.console.print( - "[yellow]Note: --before is not supported by info; showing current state.[/yellow]" + "[yellow]Note: --before/--at is not supported by info; showing current state.[/yellow]" ) # Basic: show path/URI @@ -289,6 +293,7 @@ class HaikuRAGApp: # pragma: no cover read_only=True, skip_migration_check=True, before=self.before, + at_tag=self.at_tag, ) as store: tables = [ "documents", @@ -307,6 +312,8 @@ class HaikuRAGApp: # pragma: no cover self.console.print("[bold]Version History[/bold]") + tags = await store.list_tags() + for table_name in tables: versions = await store.list_table_versions(table_name) @@ -316,6 +323,12 @@ class HaikuRAGApp: # pragma: no cover if limit: versions = versions[:limit] + version_tags: dict[int, list[str]] = {} + for tag_name, info in tags.items(): + tagged_version = info.tables.get(table_name) + if tagged_version is not None: + version_tags.setdefault(tagged_version, []).append(tag_name) + self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]") if not versions: @@ -325,16 +338,91 @@ class HaikuRAGApp: # pragma: no cover for v in versions: version_num = v["version"] timestamp = v["timestamp"] + suffix = "" + if version_num in version_tags: + names = ", ".join( + escape(n) for n in sorted(version_tags[version_num]) + ) + suffix = f" [magenta]<- {names}[/magenta]" self.console.print( - f" [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}" + f" [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}{suffix}" ) + def _tag_write_store(self) -> "Store": + """Writable store for tag create/delete. + + Migration checks stay on: a coordinated tag is only reliable when the + database schema is current, and a writable open of a legacy database + would create missing tables as a side effect. + """ + from haiku.rag.store.engine import Store + + return Store( + self.db_path, + config=self.config, + skip_validation=True, + 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, + skip_validation=True, + skip_migration_check=True, + read_only=True, + ) + + async def create_tag(self, name: str): + """Tag the current version of every table.""" + if self._is_local and not self.db_path.exists(): + self.console.print("[red]Database path does not exist.[/red]") + return + 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(): + self.console.print("[red]Database path does not exist.[/red]") + return + async with self._tag_read_store() as store: + tags = await store.list_tags() + + if not tags: + self.console.print("No tags") + return + + self.console.print("[bold]Tags[/bold]") + for name in sorted(tags): + info = tags[name] + versions = " ".join(f"{t}=v{v}" for t, v in info.tables.items()) + line = f" [repr.attrib_name]{escape(name)}[/repr.attrib_name]: {versions}" + if not info.complete: + missing = ", ".join(info.missing_tables) + line += f" [yellow](partial - missing: {missing})[/yellow]" + self.console.print(line) + + 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(): + self.console.print("[red]Database path does not exist.[/red]") + return + async with self._tag_write_store() as store: + await store.delete_tag(name) + self.console.print(f"[green]Deleted tag '{escape(name)}'[/green]") + async def list_documents(self, filter: str | None = None): async with HaikuRAG( db_path=self.db_path, config=self.config, read_only=True, before=self.before, + at_tag=self.at_tag, ) as self.client: documents = await self.client.list_documents(filter=filter) for doc in documents: @@ -348,6 +436,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) as self.client: doc = await self.client.create_document( text, title=title, metadata=metadata @@ -365,6 +454,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) as self.client: result = await self.client.create_document_from_source( source, title=title, metadata=metadata @@ -387,6 +477,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=True, before=self.before, + at_tag=self.at_tag, ) as self.client: doc = await self.client.get_document_by_id(doc_id) if doc is None: @@ -400,6 +491,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) as self.client: deleted = await self.client.delete_document(doc_id) if deleted: @@ -444,6 +536,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=True, before=self.before, + at_tag=self.at_tag, ) as self.client: results = await self.client.search( search_input, @@ -466,6 +559,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=True, before=self.before, + at_tag=self.at_tag, ) as self.client: chunk = await self.client.get_chunk_by_id(chunk_id) if not chunk: @@ -510,6 +604,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=True, before=self.before, + at_tag=self.at_tag, ) as self.client: answer, citations = await self.client.ask(question, filter=filter) @@ -538,6 +633,7 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=True, before=self.before, + at_tag=self.at_tag, ) as self.client: self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() @@ -562,6 +658,7 @@ class HaikuRAGApp: # pragma: no cover skip_validation=True, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) as client: if mode == RebuildMode.SET_EMBEDDER: async for _ in client.rebuild_database(mode=mode): @@ -606,6 +703,7 @@ class HaikuRAGApp: # pragma: no cover skip_validation=True, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) as client: await client.vacuum() self.console.print("[bold green]Vacuum completed successfully.[/bold green]") @@ -635,6 +733,7 @@ class HaikuRAGApp: # pragma: no cover skip_validation=True, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) as client: row_count = await client.store.chunks_table.count_rows() self.console.print(f"Chunks in database: {row_count}") @@ -803,9 +902,14 @@ class HaikuRAGApp: # pragma: no cover config=self.config, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ): server = create_mcp_server( - self.db_path, config=self.config, read_only=self.read_only + self.db_path, + config=self.config, + read_only=self.read_only, + before=self.before, + at_tag=self.at_tag, ) 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 00674990..b7e4ad18 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -6,6 +6,7 @@ def run_chat( db_path: Path | None = None, read_only: bool = False, before: datetime | None = None, + at_tag: str | None = None, model: str | None = None, skills: list[str] | None = None, ) -> None: @@ -15,6 +16,7 @@ def run_chat( db_path: Path to the LanceDB database. If None, uses default from config. read_only: Whether to open the database in read-only mode. before: Query database as it existed before this datetime. + at_tag: Query database at this tag. model: Model to use for the chat. skills: Skills to enable ("rag", "analysis"). Defaults to ["rag"]. """ @@ -56,6 +58,7 @@ def run_chat( skills=skill_list, read_only=read_only, before=before, + at_tag=at_tag, model=model or get_model(config.qa.model, config), ) app.run() diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 70439286..b1dc7d8a 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -86,6 +86,7 @@ class ChatApp(App): skills: list[Skill], read_only: bool = False, before: datetime | None = None, + at_tag: str | None = None, model: str | None = None, ) -> None: super().__init__() @@ -93,6 +94,7 @@ class ChatApp(App): self._skills = skills self.read_only = read_only self.before = before + self.at_tag = at_tag self._model = model self.client: HaikuRAG | None = None self.config = get_config() @@ -151,6 +153,7 @@ class ChatApp(App): config=self.config, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) await self.client.__aenter__() diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index a14e5bfe..50796a7f 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -48,6 +48,7 @@ def cli(): # Module-level flags set by callback _read_only: bool = False _before: datetime | None = None +_at_tag: str | None = None def create_app(db: Path | None = None) -> HaikuRAGApp: # pragma: no cover @@ -62,7 +63,11 @@ def create_app(db: Path | None = None) -> HaikuRAGApp: # pragma: no cover config = get_config() db_path = db if db else config.storage.data_dir / "haiku.rag.lancedb" return HaikuRAGApp( - db_path=db_path, config=config, read_only=_read_only, before=_before + db_path=db_path, + config=config, + read_only=_read_only, + before=_before, + at_tag=_at_tag, ) @@ -108,11 +113,21 @@ def main( help="Query database as it existed before this datetime (implies --read-only). " "Accepts ISO 8601 format (e.g., 2025-01-15T14:30:00) or date (e.g., 2025-01-15)", ), + at: str | None = typer.Option( + None, + "--at", + help="Query database at this tag (implies --read-only). " + "Mutually exclusive with --before", + ), ): """haiku.rag CLI - Vector database RAG system""" - global _read_only, _before + global _read_only, _before, _at_tag _read_only = read_only + if before is not None and at is not None: # pragma: no cover + typer.echo("Error: --before and --at are mutually exclusive") + raise typer.Exit(1) + # Parse and store before datetime if before is not None: # pragma: no cover from haiku.rag.utils import parse_datetime, to_utc @@ -124,6 +139,7 @@ def main( raise typer.Exit(1) else: _before = None + _at_tag = at # Load config from --config, local folder, or default directory config_path = find_config_file(cli_path=config) if config_path: @@ -635,6 +651,70 @@ def history( # pragma: no cover asyncio.run(app.history(table=table, limit=limit)) +tag_cli = typer.Typer( + help="Manage database tags (named versions across all tables)", + no_args_is_help=True, +) +_cli.add_typer(tag_cli, name="tag") + + +def _reject_time_travel(operation: str) -> None: + """Writable tag operations act on the live database state; combining them + with a historical checkout would tag something other than what the user + sees.""" + if _before is not None or _at_tag is not None: + typer.echo(f"Error: --before/--at cannot be used with {operation}", err=True) + raise typer.Exit(1) + + +@tag_cli.command("create", help="Tag the current database state") +def tag_create( # pragma: no cover + name: str = typer.Argument(help="Name of the tag to create"), + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + _reject_time_travel("tag create") + app = create_app(db) + try: + asyncio.run(app.create_tag(name)) + except (ValueError, RuntimeError) as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) + + +@tag_cli.command("list", help="List database tags") +def tag_list( # pragma: no cover + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + app = create_app(db) + asyncio.run(app.list_tags()) + + +@tag_cli.command("delete", help="Delete a tag") +def tag_delete( # pragma: no cover + name: str = typer.Argument(help="Name of the tag to delete"), + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + _reject_time_travel("tag delete") + app = create_app(db) + try: + asyncio.run(app.delete_tag(name)) + except (ValueError, RuntimeError) as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) + + @_cli.command("download-models", help="Download Docling and Ollama models per config") def download_models_cmd(): # pragma: no cover app = HaikuRAGApp(db_path=Path(), config=get_config(), read_only=True) @@ -661,7 +741,7 @@ def inspect( # pragma: no cover raise typer.Exit(1) from e db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb" - run_inspector(db_path, read_only=True, before=_before) + run_inspector(db_path, read_only=True, before=_before, at_tag=_at_tag) @_cli.command("chat", help="Launch interactive chat TUI for conversational RAG") @@ -693,6 +773,7 @@ def chat( # pragma: no cover db_path, read_only=True, before=_before, + at_tag=_at_tag, model=model, skills=skills, ) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index c817180f..0f7dccf9 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -73,6 +73,7 @@ class HaikuRAG: create: bool = False, read_only: bool = False, before: datetime | None = None, + at_tag: str | None = None, ): """Initialize the RAG client with a database path. @@ -84,6 +85,8 @@ class HaikuRAG: read_only: Whether to open the database in read-only mode. before: Query the database as it existed at this datetime. Implies read_only=True. + at_tag: Query the database at this tag. Implies read_only=True; + mutually exclusive with before. """ self._config = config if db_path is None: @@ -94,6 +97,7 @@ class HaikuRAG: self._create = create self._read_only = read_only self._before = before + self._at_tag = at_tag self._vacuum_tasks: set[asyncio.Task] = set() self._last_vacuum_at: float | None = None self._vacuum_dirty = False @@ -126,6 +130,7 @@ class HaikuRAG: create=self._create, read_only=self._read_only, before=self._before, + at_tag=self._at_tag, ) # If _initialize fails mid-way (e.g. migration check raises after # connect), close the store so we don't leak the LanceDB connection — diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 38be804e..fb7f5b5a 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -71,15 +71,30 @@ async def rebuild_database( """Rebuild the database with the specified mode. Yields the ID of each document as it is processed. + + Holds the store's rebuild lock for the whole run so tag operations fail + fast instead of snapshotting a half-rebuilt database. The lock is held + across yields; an abandoned generator releases it when closed or + garbage-collected. """ from haiku.rag.client import RebuildMode if mode is None: mode = RebuildMode.FULL - if mode == RebuildMode.SET_EMBEDDER: - await _set_embedder(client) - return + async with client.store._rebuild_lock: + if mode == RebuildMode.SET_EMBEDDER: + await _set_embedder(client) + return + + async for doc_id in _rebuild_locked(client, mode): + yield doc_id + + +async def _rebuild_locked( + client: "HaikuRAG", mode: "RebuildMode" +) -> AsyncGenerator[str, None]: + from haiku.rag.client import RebuildMode # Resolve any leftover staging/marker tables from a previously # interrupted rebuild. Returns True only when phase 1 was already diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index b744ffe9..4752d572 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -68,12 +68,17 @@ class InspectorApp(App): ] def __init__( - self, db_path: Path, read_only: bool = False, before: datetime | None = None + self, + db_path: Path, + read_only: bool = False, + before: datetime | None = None, + at_tag: str | None = None, ): super().__init__() self.db_path = db_path self.read_only = read_only self.before = before + self.at_tag = at_tag self.client: HaikuRAG | None = None def compose(self) -> "ComposeResult": @@ -92,6 +97,7 @@ class InspectorApp(App): config=config, read_only=self.read_only, before=self.before, + at_tag=self.at_tag, ) await self.client.__aenter__() @@ -235,6 +241,7 @@ def run_inspector( db_path: Path | None = None, read_only: bool = False, before: datetime | None = None, + at_tag: str | None = None, ) -> None: """Run the inspector TUI. @@ -242,10 +249,11 @@ def run_inspector( db_path: Path to the LanceDB database. If None, uses default from config. read_only: Whether to open the database in read-only mode. before: Query database as it existed before this datetime. + at_tag: Query database at this tag. """ config = get_config() if db_path is None: db_path = config.storage.data_dir / "haiku.rag.lancedb" - app = InspectorApp(db_path, read_only=read_only, before=before) + app = InspectorApp(db_path, read_only=read_only, before=before, at_tag=at_tag) app.run() diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index e2e9d1bd..3f00ccc3 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,3 +1,4 @@ +from datetime import datetime from pathlib import Path from typing import Any @@ -11,7 +12,11 @@ from haiku.rag.utils import format_citations def create_mcp_server( - db_path: Path, config: AppConfig = Config, read_only: bool = False + db_path: Path, + config: AppConfig = Config, + read_only: bool = False, + before: datetime | None = None, + at_tag: str | None = None, ) -> FastMCP: """Create an MCP server with the specified database path. @@ -19,7 +24,10 @@ def create_mcp_server( db_path: Path to the database file. config: Configuration to use. read_only: If True, write tools (add_document_*, delete_document) are not registered. + before: Serve the database as it existed at this datetime. Implies read_only. + at_tag: Serve the database at this tag. Implies read_only. """ + read_only = read_only or before is not None or at_tag is not None mcp = FastMCP("haiku-rag") # Write tools - only registered when not in read-only mode @@ -100,7 +108,13 @@ def create_mcp_server( response (smaller JSON payload for plain-text consumers). """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: + async with HaikuRAG( + db_path, + config=config, + read_only=read_only, + before=before, + at_tag=at_tag, + ) as rag: return await rag.search( query, limit=limit, include_images=include_images ) @@ -135,7 +149,13 @@ def create_mcp_server( except Exception: return [] try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: + async with HaikuRAG( + db_path, + config=config, + read_only=read_only, + before=before, + at_tag=at_tag, + ) as rag: return await rag.search( raw, limit=limit, include_images=include_images ) @@ -146,7 +166,13 @@ def create_mcp_server( async def get_document(document_id: str) -> Document | None: """Get a document by its ID.""" try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: + async with HaikuRAG( + db_path, + config=config, + read_only=read_only, + before=before, + at_tag=at_tag, + ) as rag: return await rag.get_document_by_id(document_id) except Exception: return None @@ -165,7 +191,13 @@ def create_mcp_server( filter: Optional SQL WHERE clause to filter documents. """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: + async with HaikuRAG( + db_path, + config=config, + read_only=read_only, + before=before, + at_tag=at_tag, + ) as rag: documents = await rag.list_documents(limit, offset, filter) return [ @@ -195,7 +227,13 @@ def create_mcp_server( The answer as a string. """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: + async with HaikuRAG( + db_path, + config=config, + read_only=read_only, + before=before, + at_tag=at_tag, + ) as rag: answer, citations = await rag.ask(question) if cite and citations: answer += "\n\n" + format_citations(citations) @@ -222,7 +260,13 @@ def create_mcp_server( The answer as a string. """ try: - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: + async with HaikuRAG( + db_path, + config=config, + read_only=read_only, + before=before, + at_tag=at_tag, + ) as rag: result = await rag.analyze(question, filter=filter) return result.answer except Exception as e: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 35b9cffe..4fb4bf8a 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -371,18 +371,25 @@ class Store: create: bool = False, read_only: bool = False, before: datetime | None = None, + at_tag: str | None = None, skip_migration_check: bool = False, ): + if before is not None and at_tag is not None: + raise ValueError("before and at_tag are mutually exclusive") self.db_path: Path = db_path self._config = config self._before = before + self._at_tag = at_tag # Time-travel mode is always read-only - self._read_only = read_only or (before is not None) + self._read_only = read_only or before is not None or at_tag is not None self._create = create self._skip_validation = skip_validation self._skip_migration_check = skip_migration_check self._vacuum_lock = asyncio.Lock() self._write_lock = asyncio.Lock() + # Held by rebuild_database for its whole run; tag operations check it + # and fail fast instead of snapshotting a half-rebuilt database. + self._rebuild_lock = asyncio.Lock() self._is_new_db = False # Check if database exists (for local filesystem only) @@ -432,9 +439,11 @@ class Store: # pending, before creating any newly-introduced table. await self._init_tables(is_new_db) - # Checkout tables to historical state if before is specified + # Checkout tables to historical state if before or at_tag is specified if self._before is not None: await self._checkout_tables_before(self._before) + if self._at_tag is not None: + await self._checkout_tables_at_tag(self._at_tag) # Set version for new databases. if is_new_db and not self._read_only: @@ -497,6 +506,13 @@ class Store: if self._read_only: raise ReadOnlyError("Cannot modify database in read-only mode") + def _assert_not_rebuilding(self) -> None: + """Raise if a rebuild is in progress in this process.""" + if self._rebuild_lock.locked(): + raise ValueError( + "Rebuild in progress; tag operations are unavailable until it completes" + ) + async def vacuum(self, retention_seconds: int | None = None) -> None: """Optimize and clean up old versions across all tables to reduce disk usage. @@ -871,11 +887,12 @@ class Store: Raises: ReadOnlyError: If the store is in read-only mode. - ValueError: If the tag already exists on any table. A partial tag - (present on some tables only) must be deleted before the name - can be reused. + ValueError: If a rebuild is in progress, or if the tag already + exists on any table. A partial tag (present on some tables + only) must be deleted before the name can be reused. """ self._assert_writable() + self._assert_not_rebuilding() tables = self._tables() async with self._write_lock: @@ -928,9 +945,10 @@ class Store: Raises: ReadOnlyError: If the store is in read-only mode. - ValueError: If no table has the tag. + ValueError: If a rebuild is in progress or no table has the tag. """ self._assert_writable() + self._assert_not_rebuilding() async with self._write_lock: found = False for table in self._tables().values(): @@ -994,6 +1012,17 @@ class Store: # Checkout to the found version await table.checkout(best_version) + async def _checkout_tables_at_tag(self, name: str) -> None: + """Checkout all tables at the version the tag points to. + + Raises: + ValueError: If any table is missing the tag. + """ + for table_name, table in self._tables().items(): + if name not in await table.tags.list(): + raise ValueError(f"Tag '{name}' does not exist on table '{table_name}'") + await table.checkout(name) + async def list_table_versions(self, table_name: str) -> list[dict[str, Any]]: """List version history for a table. diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index e225cc54..baf8eaa9 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -11,14 +11,18 @@ runner = CliRunner() def test_chat_command(): - """Test chat command launches chat TUI.""" + """Test chat command launches chat TUI with the global time-travel flags.""" with patch("haiku.rag.chat.run_chat") as mock_chat: mock_chat.return_value = None - result = runner.invoke(cli, ["chat"]) + result = runner.invoke(cli, ["--at", "release-1", "chat"]) assert result.exit_code == 0 mock_chat.assert_called_once() + kwargs = mock_chat.call_args.kwargs + assert kwargs["read_only"] is True + assert kwargs["before"] is None + assert kwargs["at_tag"] == "release-1" def test_run_chat_creates_app_and_runs(temp_db_path: Path): diff --git a/tests/store/test_tags.py b/tests/store/test_tags.py index 70eaf03a..1f741814 100644 --- a/tests/store/test_tags.py +++ b/tests/store/test_tags.py @@ -175,6 +175,24 @@ async def test_vacuum_cleans_untagged_versions_and_keeps_tagged(temp_db_path): assert rows == 2 +@pytest.mark.asyncio +async def test_tag_operations_rejected_during_rebuild(temp_db_path): + """While a rebuild holds the rebuild lock, tag operations fail fast + instead of snapshotting a half-rebuilt database.""" + async with Store(temp_db_path, create=True) as store: + await store.create_tag("keep") + + async with store._rebuild_lock: + with pytest.raises(ValueError, match="[Rr]ebuild in progress"): + await store.create_tag("release-1") + with pytest.raises(ValueError, match="[Rr]ebuild in progress"): + await store.delete_tag("keep") + + await store.create_tag("release-1") + await store.delete_tag("keep") + assert set(await store.list_tags()) == {"release-1"} + + @pytest.mark.asyncio async def test_vacuum_waits_for_write_lock(temp_db_path): """Vacuum serializes with writers and tag operations so a tag cannot be diff --git a/tests/store/test_time_travel.py b/tests/store/test_time_travel.py index db0a7e85..35aeafab 100644 --- a/tests/store/test_time_travel.py +++ b/tests/store/test_time_travel.py @@ -73,6 +73,67 @@ class TestStoreTimeTravel: pass assert "No data exists before" in str(exc_info.value) + @pytest.mark.asyncio + async def test_store_with_at_tag_is_read_only(self, temp_db_path): + """Store with at_tag parameter is automatically read-only.""" + async with Store(temp_db_path, create=True) as store: + await store.create_tag("release-1") + + async with Store(temp_db_path, at_tag="release-1") as store: + assert store.is_read_only is True + with pytest.raises(ReadOnlyError): + store._assert_writable() + + @pytest.mark.asyncio + async def test_store_at_tag_checks_out_tagged_state(self, temp_db_path): + """Store with at_tag checks out every table at the tagged version.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentRepository(store) + await repo.create(Document(content="First document")) + await store.create_tag("release-1") + await repo.create(Document(content="Second document")) + + async with Store(temp_db_path, at_tag="release-1") as store: + docs = await DocumentRepository(store).list_all(include_content=True) + assert len(docs) == 1 + assert docs[0].content == "First document" + + async with Store(temp_db_path) as store: + docs = await DocumentRepository(store).list_all() + assert len(docs) == 2 + + @pytest.mark.asyncio + async def test_store_at_tag_unknown_raises(self, temp_db_path): + async with Store(temp_db_path, create=True): + pass + + with pytest.raises(ValueError, match="nope"): + async with Store(temp_db_path, at_tag="nope"): + pass + + @pytest.mark.asyncio + async def test_store_at_tag_partial_raises_naming_table(self, temp_db_path): + """A tag missing from some tables fails with the table named.""" + async with Store(temp_db_path, create=True) as store: + version = await store.chunks_table.version() + await store.chunks_table.tags.create("stale", version) + + with pytest.raises(ValueError, match="documents"): + async with Store(temp_db_path, at_tag="stale"): + pass + + @pytest.mark.asyncio + async def test_store_at_tag_and_before_mutually_exclusive(self, temp_db_path): + async with Store(temp_db_path, create=True) as store: + await store.create_tag("release-1") + + with pytest.raises(ValueError, match="mutually exclusive"): + Store( + temp_db_path, + at_tag="release-1", + before=datetime.now(UTC), + ) + @pytest.mark.asyncio async def test_current_table_versions_returns_versions(self, temp_db_path): """current_table_versions returns dict of table versions.""" diff --git a/tests/test_cli.py b/tests/test_cli.py index fbc1acfb..32a19d2b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -84,3 +84,98 @@ class TestCliMigrationError: with pytest.raises(SystemExit) as exc_info: cli_wrapper() assert exc_info.value.code == 1 + + +class TestTagCommands: + def test_tag_round_trip(self, temp_db_path): + db = str(temp_db_path) + + result = runner.invoke(cli, ["init", "--db", db]) + assert result.exit_code == 0 + + result = runner.invoke(cli, ["tag", "create", "release-1", "--db", db]) + assert result.exit_code == 0 + assert "release-1" in result.output + + result = runner.invoke(cli, ["tag", "list", "--db", db]) + assert result.exit_code == 0 + assert "release-1" in result.output + assert "partial" not in result.output + + result = runner.invoke(cli, ["history", "--db", db, "-t", "documents"]) + assert result.exit_code == 0 + assert "release-1" in result.output + + result = runner.invoke(cli, ["tag", "create", "release-1", "--db", db]) + assert result.exit_code == 1 + assert "already exists" in result.output + + result = runner.invoke(cli, ["tag", "delete", "release-1", "--db", db]) + assert result.exit_code == 0 + + result = runner.invoke(cli, ["tag", "list", "--db", db]) + assert result.exit_code == 0 + assert "No tags" in result.output + + result = runner.invoke(cli, ["tag", "delete", "release-1", "--db", db]) + assert result.exit_code == 1 + assert "does not exist" in result.output + + def test_tag_create_rejected_when_migrations_pending(self, temp_db_path): + """A writable tag operation must hit the migration gate and must not + mutate a legacy database (e.g. by creating missing tables).""" + import asyncio + + import lancedb + + from haiku.rag.store.engine import Store + + async def _prepare_legacy_db(): + async with Store(temp_db_path, create=True) as store: + await store.set_haiku_version("0.19.0") + db = await lancedb.connect_async(temp_db_path.absolute()) + await db.drop_table("document_meta") + db.close() + + asyncio.run(_prepare_legacy_db()) + + result = runner.invoke( + cli, ["tag", "create", "release-1", "--db", str(temp_db_path)] + ) + assert result.exit_code == 1 + assert isinstance(result.exception, MigrationRequiredError) + + async def _table_names() -> list[str]: + db = await lancedb.connect_async(temp_db_path.absolute()) + tables = (await db.list_tables()).tables + db.close() + return tables + + assert "document_meta" not in asyncio.run(_table_names()) + + def test_tag_write_commands_reject_time_travel(self, temp_db_path): + db = str(temp_db_path) + result = runner.invoke(cli, ["init", "--db", db]) + assert result.exit_code == 0 + + result = runner.invoke(cli, ["--at", "x", "tag", "create", "r1", "--db", db]) + assert result.exit_code == 1 + assert "--at" in result.output + + result = runner.invoke( + cli, ["--before", "2025-01-01", "tag", "delete", "r1", "--db", db] + ) + assert result.exit_code == 1 + assert "--before" in result.output + + def test_tag_create_invalid_name_fails_cleanly(self, temp_db_path): + """lance restricts ref names to alphanumeric, '.', '-', '_'; the CLI + surfaces that as a clean error instead of a traceback.""" + db = str(temp_db_path) + result = runner.invoke(cli, ["init", "--db", db]) + assert result.exit_code == 0 + + result = runner.invoke(cli, ["tag", "create", "[red]release[/red]", "--db", db]) + assert result.exit_code == 1 + assert "Error:" in result.output + assert "Ref characters" in result.output diff --git a/tests/test_client.py b/tests/test_client.py index 40884c0a..0fb7d7b1 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -2264,3 +2264,32 @@ async def test_rebuild_rechunk_with_url_prefixed_stored_content( assert doc_after is not None assert "example.com" in doc_after.content assert "Stored" in doc_after.content + + +async def test_client_at_tag_opens_tagged_state_read_only(temp_db_path): + """HaikuRAG(at_tag=...) opens the database read-only at the tagged state.""" + dim = Config.embeddings.model.vector_dim + + def _doc(name: str, text: str) -> DoclingDocument: + doc = DoclingDocument(name=name) + doc.add_text(label=DocItemLabel.TEXT, text=text) + return doc + + async with HaikuRAG(temp_db_path, create=True) as client: + await client.import_document( + _doc("first", "First document"), + [Chunk(content="First document", embedding=[0.1] * dim, order=0)], + uri="mem://first", + ) + await client.store.create_tag("release-1") + await client.import_document( + _doc("second", "Second document"), + [Chunk(content="Second document", embedding=[0.1] * dim, order=0)], + uri="mem://second", + ) + + async with HaikuRAG(temp_db_path, at_tag="release-1") as client: + assert client.is_read_only is True + docs = await client.list_documents() + assert len(docs) == 1 + assert docs[0].uri == "mem://first" diff --git a/tests/test_info.py b/tests/test_info.py index 42206fef..d88d777b 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -355,7 +355,45 @@ async def test_app_history_skips_exists_check_for_remote(tmp_path): with patch("haiku.rag.store.engine.Store") as mock_store_cls: mock_store = AsyncMock() mock_store.list_table_versions = AsyncMock(return_value=[]) + mock_store.list_tags = AsyncMock(return_value={}) mock_store_cls.return_value.__aenter__ = AsyncMock(return_value=mock_store) mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False) await app.history() mock_store_cls.assert_called_once() + + +@pytest.mark.asyncio +async def test_app_tag_rendering_escapes_markup(tmp_path): + """lance forbids markup characters in ref names, but externally created + tags are rendered defensively: markup-looking names must come out as + literal text in tag list and history, not be interpreted by Rich.""" + from rich.console import Console + + from haiku.rag.store.engine import TagInfo + + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", + storage_options={"endpoint": "http://localhost:9000"}, + ) + ) + app = HaikuRAGApp(db_path=tmp_path / "db.lancedb", config=config) + app.console = Console(record=True, width=200) + + hostile = "[red]release[/red]" + tags = {hostile: TagInfo(tables={"documents": 1}, missing_tables=[])} + + with patch("haiku.rag.store.engine.Store") as mock_store_cls: + mock_store = AsyncMock() + mock_store.list_tags = AsyncMock(return_value=tags) + mock_store.list_table_versions = AsyncMock( + return_value=[{"version": 1, "timestamp": "2026-07-14 10:00:00"}] + ) + mock_store_cls.return_value.__aenter__ = AsyncMock(return_value=mock_store) + mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False) + + await app.list_tags() + await app.history(table="documents") + + output = app.console.export_text() + assert output.count(hostile) == 2 diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 74bd866a..76fa571a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -267,3 +267,47 @@ class TestMCPImageQuery: # in search_documents_by_image rejects it. results = await search_by_image(image_base64="!!! not base64 !!!") assert results == [] + + +class TestMCPTimeTravel: + @pytest.mark.asyncio + async def test_at_tag_disables_write_tools(self, mcp_db): + async with HaikuRAG(mcp_db) as rag: + await rag.store.create_tag("release-1") + + mcp = create_mcp_server(mcp_db, at_tag="release-1") + tool_names = [t.name for t in await mcp.list_tools()] + assert "add_document_from_text" not in tool_names + assert "delete_document" not in tool_names + + @pytest.mark.asyncio + async def test_before_disables_write_tools(self, mcp_db): + from datetime import UTC, datetime + + mcp = create_mcp_server(mcp_db, before=datetime.now(UTC)) + tool_names = [t.name for t in await mcp.list_tools()] + assert "add_document_from_text" not in tool_names + assert "delete_document" not in tool_names + + @pytest.mark.asyncio + async def test_at_tag_serves_tagged_state(self, mcp_db): + """Read tools on a tagged server must see the tagged state, not the + live database.""" + async with HaikuRAG(mcp_db) as rag: + await rag.store.create_tag("release-1") + await rag.create_document( + "A document added after the tag.", + title="Post-tag Doc", + uri="test://post-tag", + ) + + mcp = create_mcp_server(mcp_db, at_tag="release-1") + list_docs = await _get_tool(mcp, "list_documents") + docs = await list_docs() + uris = {d.uri for d in docs} + assert "test://post-tag" not in uris + assert len(docs) == 2 + + mcp_live = create_mcp_server(mcp_db, read_only=True) + list_docs = await _get_tool(mcp_live, "list_documents") + assert len(await list_docs()) == 3 diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 9649a625..c371bae2 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -1049,3 +1049,49 @@ async def test_rebuild_set_embedder_raises_on_vector_dim_mismatch(temp_db_path): with pytest.raises(ConfigMismatchError): async for _ in client.rebuild_database(mode=RebuildMode.SET_EMBEDDER): pass + + +async def test_rebuild_blocks_tag_operations(temp_db_path, monkeypatch): + """rebuild_database holds the rebuild lock for its whole run: tag + operations fail mid-rebuild and work again once it completes.""" + import random + + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + from haiku.rag.embeddings import EmbedderWrapper + from haiku.rag.store.models.chunk import Chunk + + async def fake_embed_documents(self, texts): + result = [] + for t in texts: + random.seed(hash(t) % (2**32)) + result.append([random.random() for _ in range(2560)]) + return result + + monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents) + + dim = Config.embeddings.model.vector_dim + docling_doc = DoclingDocument(name="d") + docling_doc.add_text(label=DocItemLabel.TEXT, text="body") + + async with HaikuRAG(temp_db_path, create=True) as client: + await client.import_document( + docling_doc, + [Chunk(content="body", embedding=[0.1] * dim, order=0)], + uri="mem://rebuild", + ) + + rebuild = client.rebuild_database(mode=RebuildMode.EMBED_ONLY) + await anext(rebuild) + + assert client.store._rebuild_lock.locked() + with pytest.raises(ValueError, match="[Rr]ebuild in progress"): + await client.store.create_tag("mid-rebuild") + + async for _ in rebuild: + pass + + assert not client.store._rebuild_lock.locked() + await client.store.create_tag("post-rebuild") + assert set(await client.store.list_tags()) == {"post-rebuild"}