before parameter in client, app, inspector, cli

This commit is contained in:
Yiorgis Gozadinos 2025-12-19 11:41:10 +02:00
parent 1e5eebfbf0
commit 77f9a2a1b9
No known key found for this signature in database
4 changed files with 98 additions and 20 deletions

View file

@ -1,6 +1,7 @@
import asyncio
import json
import logging
from datetime import datetime
from importlib.metadata import version as pkg_version
from pathlib import Path
from typing import TYPE_CHECKING
@ -36,11 +37,16 @@ logger = logging.getLogger(__name__)
class HaikuRAGApp:
def __init__(
self, db_path: Path, config: AppConfig = Config, read_only: bool = False
self,
db_path: Path,
config: AppConfig = Config,
read_only: bool = False,
before: datetime | None = None,
):
self.db_path = db_path
self.config = config
self.read_only = read_only
self.before = before
self.console = Console()
async def init(self):
@ -217,7 +223,10 @@ class HaikuRAGApp:
async def list_documents(self, filter: str | None = None):
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
documents = await self.client.list_documents(filter=filter)
for doc in documents:
@ -225,7 +234,10 @@ class HaikuRAGApp:
async def add_document_from_text(self, text: str, metadata: dict | None = None):
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
doc = await self.client.create_document(text, metadata=metadata)
self._rich_print_document(doc, truncate=True)
@ -237,7 +249,10 @@ class HaikuRAGApp:
self, source: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
result = await self.client.create_document_from_source(
source, title=title, metadata=metadata
@ -256,7 +271,10 @@ class HaikuRAGApp:
async def get_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
doc = await self.client.get_document_by_id(doc_id)
if doc is None:
@ -266,7 +284,10 @@ class HaikuRAGApp:
async def delete_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
deleted = await self.client.delete_document(doc_id)
if deleted:
@ -282,7 +303,10 @@ class HaikuRAGApp:
self, query: str, limit: int | None = None, filter: str | None = None
):
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
results = await self.client.search(query, limit=limit, filter=filter)
if not results:
@ -296,7 +320,10 @@ class HaikuRAGApp:
from textual_image.renderable import Image as RichImage
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
chunk = await self.client.chunk_repository.get_by_id(chunk_id)
if not chunk:
@ -343,7 +370,10 @@ class HaikuRAGApp:
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
try:
citations = []
@ -414,7 +444,10 @@ class HaikuRAGApp:
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
db_path=self.db_path, config=self.config, read_only=self.read_only
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as client:
try:
self.console.print("[bold cyan]Starting research[/bold cyan]")
@ -510,6 +543,7 @@ class HaikuRAGApp:
config=self.config,
skip_validation=True,
read_only=self.read_only,
before=self.before,
) as client:
try:
documents = await client.list_documents()
@ -549,6 +583,7 @@ class HaikuRAGApp:
config=self.config,
skip_validation=True,
read_only=self.read_only,
before=self.before,
) as client:
await client.vacuum()
self.console.print(
@ -565,6 +600,7 @@ class HaikuRAGApp:
config=self.config,
skip_validation=True,
read_only=self.read_only,
before=self.before,
) as client:
row_count = client.store.chunks_table.count_rows()
self.console.print(f"Chunks in database: {row_count}")
@ -735,7 +771,10 @@ class HaikuRAGApp:
):
"""Start the server with selected services."""
async with HaikuRAG(
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,
) as client:
tasks = []

View file

@ -1,6 +1,7 @@
import asyncio
import json
import warnings
from datetime import datetime
from importlib.metadata import version
from pathlib import Path
from typing import Any
@ -26,8 +27,9 @@ cli = typer.Typer(
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
)
# Module-level read-only flag set by callback
# Module-level flags set by callback
_read_only: bool = False
_before: datetime | None = None
def create_app(db: Path | None = None) -> HaikuRAGApp:
@ -41,7 +43,9 @@ def create_app(db: Path | None = None) -> HaikuRAGApp:
"""
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)
return HaikuRAGApp(
db_path=db_path, config=config, read_only=_read_only, before=_before
)
async def check_version():
@ -80,10 +84,28 @@ def main(
"--read-only",
help="Open database in read-only mode",
),
before: str | None = typer.Option(
None,
"--before",
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)",
),
):
"""haiku.rag CLI - Vector database RAG system"""
global _read_only
global _read_only, _before
_read_only = read_only
# Parse and store before datetime
if before is not None:
from haiku.rag.utils import parse_datetime, to_utc
try:
_before = to_utc(parse_datetime(before))
except ValueError as e:
typer.echo(f"Error: {e}")
raise typer.Exit(1)
else:
_before = None
# Load config from --config, local folder, or default directory
config_path = find_config_file(cli_path=config)
if config_path:
@ -363,7 +385,9 @@ def research(
from haiku.rag.cli_chat import interactive_research
from haiku.rag.client import HaikuRAG
client = HaikuRAG(db_path=app.db_path, config=app.config, read_only=_read_only)
client = HaikuRAG(
db_path=app.db_path, config=app.config, read_only=_read_only, before=_before
)
try:
interactive_research(
client=client,
@ -531,7 +555,7 @@ def inspect(
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=_read_only)
run_inspector(db_path, read_only=_read_only, before=_before)
@cli.command(

View file

@ -61,6 +61,7 @@ class HaikuRAG:
skip_validation: bool = False,
create: bool = False,
read_only: bool = False,
before: datetime | None = None,
):
"""Initialize the RAG client with a database path.
@ -70,16 +71,20 @@ class HaikuRAG:
skip_validation: Whether to skip configuration validation on database load.
create: Whether to create the database if it doesn't exist.
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.
"""
self._config = config
if db_path is None:
db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
self.store = Store(
db_path,
config=self._config,
skip_validation=skip_validation,
create=create,
read_only=read_only,
before=before,
)
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)

View file

@ -1,4 +1,5 @@
# pyright: reportPossiblyUnboundVariable=false
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING
@ -74,10 +75,13 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
Binding("c", "show_context", "Context", show=True),
]
def __init__(self, db_path: Path, read_only: bool = False):
def __init__(
self, db_path: Path, read_only: bool = False, before: datetime | None = None
):
super().__init__()
self.db_path = db_path
self.read_only = read_only
self.before = before
self.client: HaikuRAG | None = None
def compose(self) -> "ComposeResult":
@ -92,7 +96,10 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
"""Initialize the app when mounted."""
config = get_config()
self.client = HaikuRAG(
db_path=self.db_path, config=config, read_only=self.read_only
db_path=self.db_path,
config=config,
read_only=self.read_only,
before=self.before,
)
await self.client.__aenter__()
@ -233,17 +240,20 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
def run_inspector(
db_path: Path | None = None, read_only: bool = False
db_path: Path | None = None,
read_only: bool = False,
before: datetime | None = None,
) -> None: # pragma: no cover
"""Run the inspector TUI.
Args:
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.
"""
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)
app = InspectorApp(db_path, read_only=read_only, before=before)
app.run()