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

View file

@ -9,10 +9,6 @@ from io import BytesIO
from pathlib import Path from pathlib import Path
from types import ModuleType 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 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. 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: async with httpx.AsyncClient() as client:
running_version = parse(metadata.version("haiku.rag")) running_version = parse(metadata.version("haiku.rag"))
try: try:
@ -94,7 +93,7 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
return running_version >= pypi_version, running_version, pypi_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. """Convert text content to a DoclingDocument.
Args: Args:
@ -104,6 +103,10 @@ def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocu
Returns: Returns:
A DoclingDocument created from the text content. 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")) bytes_io = BytesIO(text.encode("utf-8"))
doc_stream = DocumentStream(name=name, stream=bytes_io) doc_stream = DocumentStream(name=name, stream=bytes_io)
converter = DocumentConverter() converter = DocumentConverter()