before parameter in client, app, inspector, cli
This commit is contained in:
parent
1e5eebfbf0
commit
77f9a2a1b9
4 changed files with 98 additions and 20 deletions
|
|
@ -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 = []
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue