From f582fb1be0f2f5b78ceaa7c3590e267b13559f37 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Sat, 20 Sep 2025 10:42:35 +0300
Subject: [PATCH] Lazy import when using CLI to speed up startup

---
 src/haiku/rag/cli.py   | 71 ++++++++++++++++++++++++++++--------------
 src/haiku/rag/utils.py | 13 +++++---
 2 files changed, 56 insertions(+), 28 deletions(-)

diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py
index 210f3ea2..c23916c9 100644
--- a/src/haiku/rag/cli.py
+++ b/src/haiku/rag/cli.py
@@ -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:
diff --git a/src/haiku/rag/utils.py b/src/haiku/rag/utils.py
index 334276b1..2373798e 100644
--- a/src/haiku/rag/utils.py
+++ b/src/haiku/rag/utils.py
@@ -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()