Add tag CLI commands and history tag annotations
This commit is contained in:
parent
0cbde6b7a2
commit
0813a1c980
20 changed files with 693 additions and 33 deletions
|
|
@ -1,6 +1,14 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
### 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.
|
- `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.
|
- Unknown `reranking.model.provider` raises `ValueError` instead of silently disabling reranking.
|
||||||
- `search.max_context_chars` default lowered from 10000 to 5000.
|
- `search.max_context_chars` default lowered from 10000 to 5000.
|
||||||
- `lancedb` bumped to 0.34.0.
|
|
||||||
|
|
||||||
### Removed
|
### Removed
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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.)
|
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
|
||||||
- **Visual grounding** — View chunks highlighted on original page images
|
- **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).
|
- **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
|
- **Inspector** — TUI for browsing documents, chunks, and search results
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
|
||||||
34
docs/cli.md
34
docs/cli.md
|
|
@ -8,6 +8,7 @@ The `haiku-rag` CLI provides complete document management functionality.
|
||||||
- `--config` - Specify custom configuration file
|
- `--config` - Specify custom configuration file
|
||||||
- `--read-only` - Open database in read-only mode (blocks writes, skips upgrades)
|
- `--read-only` - Open database in read-only mode (blocks writes, skips upgrades)
|
||||||
- `--before` - Query database as it existed before a datetime (implies `--read-only`)
|
- `--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
|
- `--version` / `-v` - Show version and exit
|
||||||
|
|
||||||
Per-command options:
|
Per-command options:
|
||||||
|
|
@ -572,6 +573,31 @@ Supported datetime formats:
|
||||||
!!! note
|
!!! note
|
||||||
Time travel mode automatically enables read-only mode. You cannot modify the database while viewing historical state.
|
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
|
### Version History
|
||||||
|
|
||||||
View version history for database tables:
|
View version history for database tables:
|
||||||
|
|
@ -587,20 +613,20 @@ haiku-rag history --table documents
|
||||||
haiku-rag history --limit 10
|
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
|
Version History
|
||||||
|
|
||||||
documents
|
documents
|
||||||
v5: 2025-01-15 14:30:00
|
v5: 2025-01-15 14:30:00 <- release-1
|
||||||
v4: 2025-01-14 10:00:00
|
v4: 2025-01-14 10:00:00
|
||||||
v3: 2025-01-13 09:15:00
|
v3: 2025-01-13 09:15:00
|
||||||
|
|
||||||
chunks
|
chunks
|
||||||
v8: 2025-01-15 14:30:00
|
v8: 2025-01-15 14:30:00 <- release-1
|
||||||
v7: 2025-01-14 10:00:00
|
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`.
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
from rich.markdown import Markdown
|
from rich.markdown import Markdown
|
||||||
|
from rich.markup import escape
|
||||||
from rich.progress import (
|
from rich.progress import (
|
||||||
BarColumn,
|
BarColumn,
|
||||||
DownloadColumn,
|
DownloadColumn,
|
||||||
|
|
@ -22,6 +23,7 @@ from haiku.rag.store.models.chunk import SearchType
|
||||||
from haiku.rag.store.models.document import Document
|
from haiku.rag.store.models.document import Document
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
from haiku.rag.store.models import SearchResult
|
from haiku.rag.store.models import SearchResult
|
||||||
from haiku.rag.utils import format_bytes, format_citations_rich
|
from haiku.rag.utils import format_bytes, format_citations_rich
|
||||||
|
|
||||||
|
|
@ -35,11 +37,13 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config: AppConfig = Config,
|
config: AppConfig = Config,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
before: datetime | None = None,
|
before: datetime | None = None,
|
||||||
|
at_tag: str | None = None,
|
||||||
):
|
):
|
||||||
self.db_path = db_path
|
self.db_path = db_path
|
||||||
self.config = config
|
self.config = config
|
||||||
self.read_only = read_only
|
self.read_only = read_only
|
||||||
self.before = before
|
self.before = before
|
||||||
|
self.at_tag = at_tag
|
||||||
self.console = Console()
|
self.console = Console()
|
||||||
|
|
||||||
from haiku.rag.store.engine import ConnectionMode
|
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
|
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(
|
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
|
# Basic: show path/URI
|
||||||
|
|
@ -289,6 +293,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
read_only=True,
|
read_only=True,
|
||||||
skip_migration_check=True,
|
skip_migration_check=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as store:
|
) as store:
|
||||||
tables = [
|
tables = [
|
||||||
"documents",
|
"documents",
|
||||||
|
|
@ -307,6 +312,8 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
|
|
||||||
self.console.print("[bold]Version History[/bold]")
|
self.console.print("[bold]Version History[/bold]")
|
||||||
|
|
||||||
|
tags = await store.list_tags()
|
||||||
|
|
||||||
for table_name in tables:
|
for table_name in tables:
|
||||||
versions = await store.list_table_versions(table_name)
|
versions = await store.list_table_versions(table_name)
|
||||||
|
|
||||||
|
|
@ -316,6 +323,12 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
if limit:
|
if limit:
|
||||||
versions = versions[: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]")
|
self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]")
|
||||||
|
|
||||||
if not versions:
|
if not versions:
|
||||||
|
|
@ -325,16 +338,91 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
for v in versions:
|
for v in versions:
|
||||||
version_num = v["version"]
|
version_num = v["version"]
|
||||||
timestamp = v["timestamp"]
|
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(
|
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 def list_documents(self, filter: str | None = None):
|
||||||
async with HaikuRAG(
|
async with HaikuRAG(
|
||||||
db_path=self.db_path,
|
db_path=self.db_path,
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
documents = await self.client.list_documents(filter=filter)
|
documents = await self.client.list_documents(filter=filter)
|
||||||
for doc in documents:
|
for doc in documents:
|
||||||
|
|
@ -348,6 +436,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
doc = await self.client.create_document(
|
doc = await self.client.create_document(
|
||||||
text, title=title, metadata=metadata
|
text, title=title, metadata=metadata
|
||||||
|
|
@ -365,6 +454,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
result = await self.client.create_document_from_source(
|
result = await self.client.create_document_from_source(
|
||||||
source, title=title, metadata=metadata
|
source, title=title, metadata=metadata
|
||||||
|
|
@ -387,6 +477,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
doc = await self.client.get_document_by_id(doc_id)
|
doc = await self.client.get_document_by_id(doc_id)
|
||||||
if doc is None:
|
if doc is None:
|
||||||
|
|
@ -400,6 +491,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
deleted = await self.client.delete_document(doc_id)
|
deleted = await self.client.delete_document(doc_id)
|
||||||
if deleted:
|
if deleted:
|
||||||
|
|
@ -444,6 +536,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
results = await self.client.search(
|
results = await self.client.search(
|
||||||
search_input,
|
search_input,
|
||||||
|
|
@ -466,6 +559,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
chunk = await self.client.get_chunk_by_id(chunk_id)
|
chunk = await self.client.get_chunk_by_id(chunk_id)
|
||||||
if not chunk:
|
if not chunk:
|
||||||
|
|
@ -510,6 +604,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
answer, citations = await self.client.ask(question, filter=filter)
|
answer, citations = await self.client.ask(question, filter=filter)
|
||||||
|
|
||||||
|
|
@ -538,6 +633,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
@ -562,6 +658,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
skip_validation=True,
|
skip_validation=True,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as client:
|
) as client:
|
||||||
if mode == RebuildMode.SET_EMBEDDER:
|
if mode == RebuildMode.SET_EMBEDDER:
|
||||||
async for _ in client.rebuild_database(mode=mode):
|
async for _ in client.rebuild_database(mode=mode):
|
||||||
|
|
@ -606,6 +703,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
skip_validation=True,
|
skip_validation=True,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as client:
|
) as client:
|
||||||
await client.vacuum()
|
await client.vacuum()
|
||||||
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
|
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
|
||||||
|
|
@ -635,6 +733,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
skip_validation=True,
|
skip_validation=True,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
) as client:
|
) as client:
|
||||||
row_count = await client.store.chunks_table.count_rows()
|
row_count = await client.store.chunks_table.count_rows()
|
||||||
self.console.print(f"Chunks in database: {row_count}")
|
self.console.print(f"Chunks in database: {row_count}")
|
||||||
|
|
@ -803,9 +902,14 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
):
|
):
|
||||||
server = create_mcp_server(
|
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:
|
try:
|
||||||
if transport == "stdio":
|
if transport == "stdio":
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ def run_chat(
|
||||||
db_path: Path | None = None,
|
db_path: Path | None = None,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
before: datetime | None = None,
|
before: datetime | None = None,
|
||||||
|
at_tag: str | None = None,
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
skills: list[str] | None = None,
|
skills: list[str] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
|
|
@ -15,6 +16,7 @@ def run_chat(
|
||||||
db_path: Path to the LanceDB database. If None, uses default from config.
|
db_path: Path to the LanceDB database. If None, uses default from config.
|
||||||
read_only: Whether to open the database in read-only mode.
|
read_only: Whether to open the database in read-only mode.
|
||||||
before: Query database as it existed before this datetime.
|
before: Query database as it existed before this datetime.
|
||||||
|
at_tag: Query database at this tag.
|
||||||
model: Model to use for the chat.
|
model: Model to use for the chat.
|
||||||
skills: Skills to enable ("rag", "analysis"). Defaults to ["rag"].
|
skills: Skills to enable ("rag", "analysis"). Defaults to ["rag"].
|
||||||
"""
|
"""
|
||||||
|
|
@ -56,6 +58,7 @@ def run_chat(
|
||||||
skills=skill_list,
|
skills=skill_list,
|
||||||
read_only=read_only,
|
read_only=read_only,
|
||||||
before=before,
|
before=before,
|
||||||
|
at_tag=at_tag,
|
||||||
model=model or get_model(config.qa.model, config),
|
model=model or get_model(config.qa.model, config),
|
||||||
)
|
)
|
||||||
app.run()
|
app.run()
|
||||||
|
|
|
||||||
|
|
@ -86,6 +86,7 @@ class ChatApp(App):
|
||||||
skills: list[Skill],
|
skills: list[Skill],
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
before: datetime | None = None,
|
before: datetime | None = None,
|
||||||
|
at_tag: str | None = None,
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
@ -93,6 +94,7 @@ class ChatApp(App):
|
||||||
self._skills = skills
|
self._skills = skills
|
||||||
self.read_only = read_only
|
self.read_only = read_only
|
||||||
self.before = before
|
self.before = before
|
||||||
|
self.at_tag = at_tag
|
||||||
self._model = model
|
self._model = model
|
||||||
self.client: HaikuRAG | None = None
|
self.client: HaikuRAG | None = None
|
||||||
self.config = get_config()
|
self.config = get_config()
|
||||||
|
|
@ -151,6 +153,7 @@ class ChatApp(App):
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
)
|
)
|
||||||
await self.client.__aenter__()
|
await self.client.__aenter__()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -48,6 +48,7 @@ def cli():
|
||||||
# Module-level flags set by callback
|
# Module-level flags set by callback
|
||||||
_read_only: bool = False
|
_read_only: bool = False
|
||||||
_before: datetime | None = None
|
_before: datetime | None = None
|
||||||
|
_at_tag: str | None = None
|
||||||
|
|
||||||
|
|
||||||
def create_app(db: Path | None = None) -> HaikuRAGApp: # pragma: no cover
|
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()
|
config = get_config()
|
||||||
db_path = db if db else config.storage.data_dir / "haiku.rag.lancedb"
|
db_path = db if db else config.storage.data_dir / "haiku.rag.lancedb"
|
||||||
return HaikuRAGApp(
|
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). "
|
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)",
|
"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"""
|
"""haiku.rag CLI - Vector database RAG system"""
|
||||||
global _read_only, _before
|
global _read_only, _before, _at_tag
|
||||||
_read_only = read_only
|
_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
|
# Parse and store before datetime
|
||||||
if before is not None: # pragma: no cover
|
if before is not None: # pragma: no cover
|
||||||
from haiku.rag.utils import parse_datetime, to_utc
|
from haiku.rag.utils import parse_datetime, to_utc
|
||||||
|
|
@ -124,6 +139,7 @@ def main(
|
||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
else:
|
else:
|
||||||
_before = None
|
_before = None
|
||||||
|
_at_tag = at
|
||||||
# Load config from --config, local folder, or default directory
|
# Load config from --config, local folder, or default directory
|
||||||
config_path = find_config_file(cli_path=config)
|
config_path = find_config_file(cli_path=config)
|
||||||
if config_path:
|
if config_path:
|
||||||
|
|
@ -635,6 +651,70 @@ def history( # pragma: no cover
|
||||||
asyncio.run(app.history(table=table, limit=limit))
|
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")
|
@_cli.command("download-models", help="Download Docling and Ollama models per config")
|
||||||
def download_models_cmd(): # pragma: no cover
|
def download_models_cmd(): # pragma: no cover
|
||||||
app = HaikuRAGApp(db_path=Path(), config=get_config(), read_only=True)
|
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
|
raise typer.Exit(1) from e
|
||||||
|
|
||||||
db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
|
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")
|
@_cli.command("chat", help="Launch interactive chat TUI for conversational RAG")
|
||||||
|
|
@ -693,6 +773,7 @@ def chat( # pragma: no cover
|
||||||
db_path,
|
db_path,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
before=_before,
|
before=_before,
|
||||||
|
at_tag=_at_tag,
|
||||||
model=model,
|
model=model,
|
||||||
skills=skills,
|
skills=skills,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -73,6 +73,7 @@ class HaikuRAG:
|
||||||
create: bool = False,
|
create: bool = False,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
before: datetime | None = None,
|
before: datetime | None = None,
|
||||||
|
at_tag: str | None = None,
|
||||||
):
|
):
|
||||||
"""Initialize the RAG client with a database path.
|
"""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.
|
read_only: Whether to open the database in read-only mode.
|
||||||
before: Query the database as it existed at this datetime.
|
before: Query the database as it existed at this datetime.
|
||||||
Implies read_only=True.
|
Implies read_only=True.
|
||||||
|
at_tag: Query the database at this tag. Implies read_only=True;
|
||||||
|
mutually exclusive with before.
|
||||||
"""
|
"""
|
||||||
self._config = config
|
self._config = config
|
||||||
if db_path is None:
|
if db_path is None:
|
||||||
|
|
@ -94,6 +97,7 @@ class HaikuRAG:
|
||||||
self._create = create
|
self._create = create
|
||||||
self._read_only = read_only
|
self._read_only = read_only
|
||||||
self._before = before
|
self._before = before
|
||||||
|
self._at_tag = at_tag
|
||||||
self._vacuum_tasks: set[asyncio.Task] = set()
|
self._vacuum_tasks: set[asyncio.Task] = set()
|
||||||
self._last_vacuum_at: float | None = None
|
self._last_vacuum_at: float | None = None
|
||||||
self._vacuum_dirty = False
|
self._vacuum_dirty = False
|
||||||
|
|
@ -126,6 +130,7 @@ class HaikuRAG:
|
||||||
create=self._create,
|
create=self._create,
|
||||||
read_only=self._read_only,
|
read_only=self._read_only,
|
||||||
before=self._before,
|
before=self._before,
|
||||||
|
at_tag=self._at_tag,
|
||||||
)
|
)
|
||||||
# If _initialize fails mid-way (e.g. migration check raises after
|
# If _initialize fails mid-way (e.g. migration check raises after
|
||||||
# connect), close the store so we don't leak the LanceDB connection —
|
# connect), close the store so we don't leak the LanceDB connection —
|
||||||
|
|
|
||||||
|
|
@ -71,15 +71,30 @@ async def rebuild_database(
|
||||||
"""Rebuild the database with the specified mode.
|
"""Rebuild the database with the specified mode.
|
||||||
|
|
||||||
Yields the ID of each document as it is processed.
|
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
|
from haiku.rag.client import RebuildMode
|
||||||
|
|
||||||
if mode is None:
|
if mode is None:
|
||||||
mode = RebuildMode.FULL
|
mode = RebuildMode.FULL
|
||||||
|
|
||||||
if mode == RebuildMode.SET_EMBEDDER:
|
async with client.store._rebuild_lock:
|
||||||
await _set_embedder(client)
|
if mode == RebuildMode.SET_EMBEDDER:
|
||||||
return
|
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
|
# Resolve any leftover staging/marker tables from a previously
|
||||||
# interrupted rebuild. Returns True only when phase 1 was already
|
# interrupted rebuild. Returns True only when phase 1 was already
|
||||||
|
|
|
||||||
|
|
@ -68,12 +68,17 @@ class InspectorApp(App):
|
||||||
]
|
]
|
||||||
|
|
||||||
def __init__(
|
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__()
|
super().__init__()
|
||||||
self.db_path = db_path
|
self.db_path = db_path
|
||||||
self.read_only = read_only
|
self.read_only = read_only
|
||||||
self.before = before
|
self.before = before
|
||||||
|
self.at_tag = at_tag
|
||||||
self.client: HaikuRAG | None = None
|
self.client: HaikuRAG | None = None
|
||||||
|
|
||||||
def compose(self) -> "ComposeResult":
|
def compose(self) -> "ComposeResult":
|
||||||
|
|
@ -92,6 +97,7 @@ class InspectorApp(App):
|
||||||
config=config,
|
config=config,
|
||||||
read_only=self.read_only,
|
read_only=self.read_only,
|
||||||
before=self.before,
|
before=self.before,
|
||||||
|
at_tag=self.at_tag,
|
||||||
)
|
)
|
||||||
await self.client.__aenter__()
|
await self.client.__aenter__()
|
||||||
|
|
||||||
|
|
@ -235,6 +241,7 @@ def run_inspector(
|
||||||
db_path: Path | None = None,
|
db_path: Path | None = None,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
before: datetime | None = None,
|
before: datetime | None = None,
|
||||||
|
at_tag: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Run the inspector TUI.
|
"""Run the inspector TUI.
|
||||||
|
|
||||||
|
|
@ -242,10 +249,11 @@ def run_inspector(
|
||||||
db_path: Path to the LanceDB database. If None, uses default from config.
|
db_path: Path to the LanceDB database. If None, uses default from config.
|
||||||
read_only: Whether to open the database in read-only mode.
|
read_only: Whether to open the database in read-only mode.
|
||||||
before: Query database as it existed before this datetime.
|
before: Query database as it existed before this datetime.
|
||||||
|
at_tag: Query database at this tag.
|
||||||
"""
|
"""
|
||||||
config = get_config()
|
config = get_config()
|
||||||
if db_path is None:
|
if db_path is None:
|
||||||
db_path = config.storage.data_dir / "haiku.rag.lancedb"
|
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()
|
app.run()
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,4 @@
|
||||||
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
|
@ -11,7 +12,11 @@ from haiku.rag.utils import format_citations
|
||||||
|
|
||||||
|
|
||||||
def create_mcp_server(
|
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:
|
) -> FastMCP:
|
||||||
"""Create an MCP server with the specified database path.
|
"""Create an MCP server with the specified database path.
|
||||||
|
|
||||||
|
|
@ -19,7 +24,10 @@ def create_mcp_server(
|
||||||
db_path: Path to the database file.
|
db_path: Path to the database file.
|
||||||
config: Configuration to use.
|
config: Configuration to use.
|
||||||
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
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")
|
mcp = FastMCP("haiku-rag")
|
||||||
|
|
||||||
# Write tools - only registered when not in read-only mode
|
# 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).
|
response (smaller JSON payload for plain-text consumers).
|
||||||
"""
|
"""
|
||||||
try:
|
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(
|
return await rag.search(
|
||||||
query, limit=limit, include_images=include_images
|
query, limit=limit, include_images=include_images
|
||||||
)
|
)
|
||||||
|
|
@ -135,7 +149,13 @@ def create_mcp_server(
|
||||||
except Exception:
|
except Exception:
|
||||||
return []
|
return []
|
||||||
try:
|
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(
|
return await rag.search(
|
||||||
raw, limit=limit, include_images=include_images
|
raw, limit=limit, include_images=include_images
|
||||||
)
|
)
|
||||||
|
|
@ -146,7 +166,13 @@ def create_mcp_server(
|
||||||
async def get_document(document_id: str) -> Document | None:
|
async def get_document(document_id: str) -> Document | None:
|
||||||
"""Get a document by its ID."""
|
"""Get a document by its ID."""
|
||||||
try:
|
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)
|
return await rag.get_document_by_id(document_id)
|
||||||
except Exception:
|
except Exception:
|
||||||
return None
|
return None
|
||||||
|
|
@ -165,7 +191,13 @@ def create_mcp_server(
|
||||||
filter: Optional SQL WHERE clause to filter documents.
|
filter: Optional SQL WHERE clause to filter documents.
|
||||||
"""
|
"""
|
||||||
try:
|
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)
|
documents = await rag.list_documents(limit, offset, filter)
|
||||||
|
|
||||||
return [
|
return [
|
||||||
|
|
@ -195,7 +227,13 @@ def create_mcp_server(
|
||||||
The answer as a string.
|
The answer as a string.
|
||||||
"""
|
"""
|
||||||
try:
|
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)
|
answer, citations = await rag.ask(question)
|
||||||
if cite and citations:
|
if cite and citations:
|
||||||
answer += "\n\n" + format_citations(citations)
|
answer += "\n\n" + format_citations(citations)
|
||||||
|
|
@ -222,7 +260,13 @@ def create_mcp_server(
|
||||||
The answer as a string.
|
The answer as a string.
|
||||||
"""
|
"""
|
||||||
try:
|
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)
|
result = await rag.analyze(question, filter=filter)
|
||||||
return result.answer
|
return result.answer
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|
|
||||||
|
|
@ -371,18 +371,25 @@ class Store:
|
||||||
create: bool = False,
|
create: bool = False,
|
||||||
read_only: bool = False,
|
read_only: bool = False,
|
||||||
before: datetime | None = None,
|
before: datetime | None = None,
|
||||||
|
at_tag: str | None = None,
|
||||||
skip_migration_check: bool = False,
|
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.db_path: Path = db_path
|
||||||
self._config = config
|
self._config = config
|
||||||
self._before = before
|
self._before = before
|
||||||
|
self._at_tag = at_tag
|
||||||
# Time-travel mode is always read-only
|
# 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._create = create
|
||||||
self._skip_validation = skip_validation
|
self._skip_validation = skip_validation
|
||||||
self._skip_migration_check = skip_migration_check
|
self._skip_migration_check = skip_migration_check
|
||||||
self._vacuum_lock = asyncio.Lock()
|
self._vacuum_lock = asyncio.Lock()
|
||||||
self._write_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
|
self._is_new_db = False
|
||||||
|
|
||||||
# Check if database exists (for local filesystem only)
|
# Check if database exists (for local filesystem only)
|
||||||
|
|
@ -432,9 +439,11 @@ class Store:
|
||||||
# pending, before creating any newly-introduced table.
|
# pending, before creating any newly-introduced table.
|
||||||
await self._init_tables(is_new_db)
|
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:
|
if self._before is not None:
|
||||||
await self._checkout_tables_before(self._before)
|
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.
|
# Set version for new databases.
|
||||||
if is_new_db and not self._read_only:
|
if is_new_db and not self._read_only:
|
||||||
|
|
@ -497,6 +506,13 @@ class Store:
|
||||||
if self._read_only:
|
if self._read_only:
|
||||||
raise ReadOnlyError("Cannot modify database in read-only mode")
|
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:
|
async def vacuum(self, retention_seconds: int | None = None) -> None:
|
||||||
"""Optimize and clean up old versions across all tables to reduce disk usage.
|
"""Optimize and clean up old versions across all tables to reduce disk usage.
|
||||||
|
|
||||||
|
|
@ -871,11 +887,12 @@ class Store:
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ReadOnlyError: If the store is in read-only mode.
|
ReadOnlyError: If the store is in read-only mode.
|
||||||
ValueError: If the tag already exists on any table. A partial tag
|
ValueError: If a rebuild is in progress, or if the tag already
|
||||||
(present on some tables only) must be deleted before the name
|
exists on any table. A partial tag (present on some tables
|
||||||
can be reused.
|
only) must be deleted before the name can be reused.
|
||||||
"""
|
"""
|
||||||
self._assert_writable()
|
self._assert_writable()
|
||||||
|
self._assert_not_rebuilding()
|
||||||
tables = self._tables()
|
tables = self._tables()
|
||||||
|
|
||||||
async with self._write_lock:
|
async with self._write_lock:
|
||||||
|
|
@ -928,9 +945,10 @@ class Store:
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
ReadOnlyError: If the store is in read-only mode.
|
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_writable()
|
||||||
|
self._assert_not_rebuilding()
|
||||||
async with self._write_lock:
|
async with self._write_lock:
|
||||||
found = False
|
found = False
|
||||||
for table in self._tables().values():
|
for table in self._tables().values():
|
||||||
|
|
@ -994,6 +1012,17 @@ class Store:
|
||||||
# Checkout to the found version
|
# Checkout to the found version
|
||||||
await table.checkout(best_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]]:
|
async def list_table_versions(self, table_name: str) -> list[dict[str, Any]]:
|
||||||
"""List version history for a table.
|
"""List version history for a table.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -11,14 +11,18 @@ runner = CliRunner()
|
||||||
|
|
||||||
|
|
||||||
def test_chat_command():
|
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:
|
with patch("haiku.rag.chat.run_chat") as mock_chat:
|
||||||
mock_chat.return_value = None
|
mock_chat.return_value = None
|
||||||
|
|
||||||
result = runner.invoke(cli, ["chat"])
|
result = runner.invoke(cli, ["--at", "release-1", "chat"])
|
||||||
|
|
||||||
assert result.exit_code == 0
|
assert result.exit_code == 0
|
||||||
mock_chat.assert_called_once()
|
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):
|
def test_run_chat_creates_app_and_runs(temp_db_path: Path):
|
||||||
|
|
|
||||||
|
|
@ -175,6 +175,24 @@ async def test_vacuum_cleans_untagged_versions_and_keeps_tagged(temp_db_path):
|
||||||
assert rows == 2
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_vacuum_waits_for_write_lock(temp_db_path):
|
async def test_vacuum_waits_for_write_lock(temp_db_path):
|
||||||
"""Vacuum serializes with writers and tag operations so a tag cannot be
|
"""Vacuum serializes with writers and tag operations so a tag cannot be
|
||||||
|
|
|
||||||
|
|
@ -73,6 +73,67 @@ class TestStoreTimeTravel:
|
||||||
pass
|
pass
|
||||||
assert "No data exists before" in str(exc_info.value)
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_current_table_versions_returns_versions(self, temp_db_path):
|
async def test_current_table_versions_returns_versions(self, temp_db_path):
|
||||||
"""current_table_versions returns dict of table versions."""
|
"""current_table_versions returns dict of table versions."""
|
||||||
|
|
|
||||||
|
|
@ -84,3 +84,98 @@ class TestCliMigrationError:
|
||||||
with pytest.raises(SystemExit) as exc_info:
|
with pytest.raises(SystemExit) as exc_info:
|
||||||
cli_wrapper()
|
cli_wrapper()
|
||||||
assert exc_info.value.code == 1
|
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
|
||||||
|
|
|
||||||
|
|
@ -2264,3 +2264,32 @@ async def test_rebuild_rechunk_with_url_prefixed_stored_content(
|
||||||
assert doc_after is not None
|
assert doc_after is not None
|
||||||
assert "example.com" in doc_after.content
|
assert "example.com" in doc_after.content
|
||||||
assert "Stored" 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"
|
||||||
|
|
|
||||||
|
|
@ -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:
|
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
|
||||||
mock_store = AsyncMock()
|
mock_store = AsyncMock()
|
||||||
mock_store.list_table_versions = AsyncMock(return_value=[])
|
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.__aenter__ = AsyncMock(return_value=mock_store)
|
||||||
mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||||
await app.history()
|
await app.history()
|
||||||
mock_store_cls.assert_called_once()
|
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
|
||||||
|
|
|
||||||
|
|
@ -267,3 +267,47 @@ class TestMCPImageQuery:
|
||||||
# in search_documents_by_image rejects it.
|
# in search_documents_by_image rejects it.
|
||||||
results = await search_by_image(image_base64="!!! not base64 !!!")
|
results = await search_by_image(image_base64="!!! not base64 !!!")
|
||||||
assert results == []
|
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
|
||||||
|
|
|
||||||
|
|
@ -1049,3 +1049,49 @@ async def test_rebuild_set_embedder_raises_on_vector_dim_mismatch(temp_db_path):
|
||||||
with pytest.raises(ConfigMismatchError):
|
with pytest.raises(ConfigMismatchError):
|
||||||
async for _ in client.rebuild_database(mode=RebuildMode.SET_EMBEDDER):
|
async for _ in client.rebuild_database(mode=RebuildMode.SET_EMBEDDER):
|
||||||
pass
|
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"}
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue