Lazy import when using CLI to speed up startup

This commit is contained in:
Yiorgis Gozadinos 2025-09-20 10:42:35 +03:00
parent 52febd01f9
commit f582fb1be0
No known key found for this signature in database
2 changed files with 56 additions and 28 deletions

View file

@ -3,29 +3,16 @@ import warnings
from importlib.metadata import version
from pathlib import Path
import logfire
import typer
from rich.console import Console
from haiku.rag.app import HaikuRAGApp
from haiku.rag.config import Config
from haiku.rag.logging import configure_cli_logging
from haiku.rag.migration import migrate_sqlite_to_lancedb
from haiku.rag.utils import is_up_to_date
if Config.ENV == "development":
logfire.configure(send_to_logfire="if-token-present")
logfire.instrument_pydantic_ai()
else:
configure_cli_logging()
warnings.filterwarnings("ignore")
cli = typer.Typer(
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
)
console = Console()
def complete_document_ids(ctx: typer.Context, incomplete: str):
"""Autocomplete document IDs from the selected DB."""
@ -90,16 +77,16 @@ async def check_version():
"""Check if haiku.rag is up to date and show warning if not."""
up_to_date, current_version, latest_version = await is_up_to_date()
if not up_to_date:
console.print(
f"[yellow]Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}[/yellow]"
typer.echo(
f"Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}",
)
console.print("[yellow]Please update.[/yellow]")
typer.echo("Please update.")
def version_callback(value: bool):
if value:
v = version("haiku.rag")
console.print(f"haiku.rag version {v}")
typer.echo(f"haiku.rag version {v}")
raise typer.Exit()
@ -114,11 +101,26 @@ def main(
),
):
"""haiku.rag CLI - Vector database RAG system"""
if Config.ENV != "development":
# Ensure only haiku.rag logs are emitted in CLI context
# Configure logging minimally for CLI context
if Config.ENV == "development":
# Lazy import logfire only in development
try:
import logfire # type: ignore
logfire.configure(send_to_logfire="if-token-present")
logfire.instrument_pydantic_ai()
except Exception:
pass
else:
configure_cli_logging()
warnings.filterwarnings("ignore")
# Run version check before any command
asyncio.run(check_version())
try:
asyncio.run(check_version())
except Exception:
# Do not block CLI on version check issues
pass
@cli.command("list", help="List all stored documents")
@ -129,6 +131,8 @@ def list_documents(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.list_documents())
@ -144,6 +148,8 @@ def add_document_text(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.add_document_from_text(text=text))
@ -160,6 +166,8 @@ def add_document_src(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.add_document_from_source(source=source))
@ -176,6 +184,8 @@ def get_document(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.get_document(doc_id=doc_id))
@ -192,6 +202,8 @@ def delete_document(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.delete_document(doc_id=doc_id))
@ -217,6 +229,8 @@ def search(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.search(query=query, limit=limit))
@ -237,6 +251,8 @@ def ask(
help="Include citations in the response",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.ask(question=question, cite=cite))
@ -273,6 +289,8 @@ def research(
help="Show verbose progress output",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(
app.research(
@ -287,6 +305,8 @@ def research(
@cli.command("settings", help="Display current configuration settings")
def settings():
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings
app.show_settings()
@ -302,6 +322,8 @@ def rebuild(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.rebuild())
@ -314,6 +336,8 @@ def vacuum(
help="Path to the LanceDB database file",
),
):
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.vacuum())
@ -339,9 +363,7 @@ def serve(
),
) -> None:
"""Start the MCP server."""
if stdio and sse:
console.print("[red]Error: Cannot use both --stdio and --http options[/red]")
raise typer.Exit(1)
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
@ -363,6 +385,9 @@ def migrate(
# Generate LanceDB path in same parent directory
lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
# Lazy import to avoid heavy deps on simple invocations
from haiku.rag.migration import migrate_sqlite_to_lancedb
success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path))
if not success:

View file

@ -9,10 +9,6 @@ from io import BytesIO
from pathlib import Path
from types import ModuleType
import httpx
from docling.document_converter import DocumentConverter
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.io import DocumentStream
from packaging.version import Version, parse
@ -82,6 +78,9 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
the running version and the latest version.
"""
# Lazy import to avoid pulling httpx (and its deps) on module import
import httpx
async with httpx.AsyncClient() as client:
running_version = parse(metadata.version("haiku.rag"))
try:
@ -94,7 +93,7 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
return running_version >= pypi_version, running_version, pypi_version
def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocument:
def text_to_docling_document(text: str, name: str = "content.md"):
"""Convert text content to a DoclingDocument.
Args:
@ -104,6 +103,10 @@ def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocu
Returns:
A DoclingDocument created from the text content.
"""
# Lazy import docling deps to keep import-time light
from docling.document_converter import DocumentConverter # type: ignore
from docling_core.types.io import DocumentStream # type: ignore
bytes_io = BytesIO(text.encode("utf-8"))
doc_stream = DocumentStream(name=name, stream=bytes_io)
converter = DocumentConverter()