Merge pull request #537 from lawrenceakka/help-speedup

Stop eagerly importing Store from haiku.rag.store

CHANGELOG.md

@@ -4,6 +4,8 @@
 ### Fixed
 
 - `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1.
+- haiku-rag CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`.
+- `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`.
 
 ## [0.73.0] - 2026-08-06
 

haiku_rag_slim/haiku/rag/cli.py

@@ -4,7 +4,7 @@ import sys
 import warnings
 from importlib.metadata import version
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import typer
 from dotenv import find_dotenv, load_dotenv
@@ -13,7 +13,6 @@ from dotenv import find_dotenv, load_dotenv
 # Env loading needs to be before config import; usecwd=True searches from cwd, not this .py file's location
 load_dotenv(find_dotenv(usecwd=True))
 
-from haiku.rag.app import HaikuRAGApp  # noqa: E402
 from haiku.rag.config import (  # noqa: E402
     AppConfig,
     find_config_file,
@@ -29,6 +28,9 @@ from haiku.rag.store.exceptions import (  # noqa: E402
 from haiku.rag.store.models.chunk import SearchType  # noqa: E402
 from haiku.rag.utils import is_up_to_date  # noqa: E402
 
+if TYPE_CHECKING:
+    from haiku.rag.app import HaikuRAGApp
+
 _cli = typer.Typer(
     context_settings={"help_option_names": ["-h", "--help"]},
     no_args_is_help=True,
@@ -48,7 +50,7 @@ def cli():
 _read_only: bool = False
 
 
-def create_app(db: Path | None = None) -> HaikuRAGApp:  # pragma: no cover
+def create_app(db: Path | None = None) -> "HaikuRAGApp":  # pragma: no cover
     """Create HaikuRAGApp with loaded config and resolved database path.
 
     Args:
@@ -57,6 +59,8 @@ def create_app(db: Path | None = None) -> HaikuRAGApp:  # pragma: no cover
     Returns:
         HaikuRAGApp instance with proper config and db path.
     """
+    from haiku.rag.app import 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)
@@ -402,6 +406,8 @@ def analyze(  # pragma: no cover
 
 @_cli.command("settings", help="Display current configuration settings")
 def settings():  # pragma: no cover
+    from haiku.rag.app import HaikuRAGApp
+
     config = get_config()
     app = HaikuRAGApp(db_path=Path(), config=config, read_only=True)
     app.show_settings()
@@ -719,6 +725,8 @@ def tag_restore(  # pragma: no cover
 
 @_cli.command("download-models", help="Download Docling and Ollama models per config")
 def download_models_cmd():  # pragma: no cover
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=Path(), config=get_config(), read_only=True)
     try:
         asyncio.run(app.download_models())

haiku_rag_slim/haiku/rag/ingester/cli.py

@@ -4,6 +4,7 @@ from collections.abc import Iterator
 from contextlib import contextmanager
 from datetime import UTC, datetime
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import typer
 import yaml
@@ -28,11 +29,6 @@ from haiku.rag.config import (  # noqa: E402
     load_yaml_config,
     set_config,
 )
-from haiku.rag.ingester.app import (  # noqa: E402
-    BatchProgress,
-    BatchProgressCallback,
-    IngesterApp,
-)
 from haiku.rag.ingester.batch import BatchManifest  # noqa: E402
 from haiku.rag.ingester.queue.migrations import open_queue  # noqa: E402
 from haiku.rag.logging import configure_cli_logging  # noqa: E402
@@ -41,6 +37,9 @@ from haiku.rag.store.exceptions import (  # noqa: E402
     ReadOnlyError,
 )
 
+if TYPE_CHECKING:
+    from haiku.rag.ingester.app import BatchProgress, BatchProgressCallback
+
 _cli = typer.Typer(
     name="haiku-ingester",
     no_args_is_help=True,
@@ -167,7 +166,7 @@ def _write_manifest(manifest: BatchManifest, path: Path) -> None:
 @contextmanager
 def _batch_progress(
     description: str,
-) -> Iterator[BatchProgressCallback | None]:  # pragma: no cover
+) -> "Iterator[BatchProgressCallback | None]":  # pragma: no cover
     console = Console(file=sys.stdout)
     if not console.is_terminal:
         yield None
@@ -184,7 +183,7 @@ def _batch_progress(
     )
     task_id = None
 
-    def _update(snapshot: BatchProgress) -> None:
+    def _update(snapshot: "BatchProgress") -> None:
         nonlocal task_id
         task_description = (
             f"{description} ({snapshot.succeeded} ok, {snapshot.dead} dead)"
@@ -244,6 +243,8 @@ def serve(
 ) -> None:
     """Run the production ingester: pollers + workers (and the HTTP API
     unless --no-api is set). Blocks until SIGINT/SIGTERM."""
+    from haiku.rag.ingester.app import IngesterApp
+
     app_config = get_config()
     if host is not None:
         app_config.ingester.api.host = host
@@ -309,6 +310,8 @@ async def _run_batch(
     output: Path | None = None,
     manifest_path: Path | None = None,
 ) -> None:
+    from haiku.rag.ingester.app import IngesterApp
+
     db = _resolve_db_path(app_config, db_path)
     app = IngesterApp(config=app_config, db_path=db)
     if dry_run:

haiku_rag_slim/haiku/rag/store/__init__.py

@@ -1,5 +1,4 @@
-from .engine import Store
 from .exceptions import MigrationRequiredError, ReadOnlyError
 from .models import Chunk, Document
 
-__all__ = ["Store", "Chunk", "Document", "MigrationRequiredError", "ReadOnlyError"]
+__all__ = ["Chunk", "Document", "MigrationRequiredError", "ReadOnlyError"]

tests/ingester/test_cli.py

@@ -1,6 +1,8 @@
 """haiku-ingester CLI: exercises every subcommand via CliRunner with
 IngesterApp / open_queue patched out so no real ingestion runs."""
 
+import subprocess
+import sys
 from contextlib import contextmanager
 from datetime import UTC, datetime
 from unittest.mock import AsyncMock, MagicMock
@@ -24,6 +26,22 @@ from haiku.rag.ingester.queue.models import JobOp
 runner = CliRunner()
 
 
+def test_importing_ingester_cli_does_not_load_lancedb():
+    """Test that lancedb is not imported automatically by the cli. Doing so is
+     expensive. Must be run in a subprocess because lancedb might be imported by
+    other tests in the same session."""
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import haiku.rag.ingester.cli, sys; assert 'lancedb' not in sys.modules",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0, result.stderr
+
+
 # --- helpers ---
 
 
@@ -41,7 +59,7 @@ def _config_with_queue(queue: QueueConfig):
 def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
     fake = AsyncMock()
     fake.run_batch.return_value = report
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
     return fake
 
 
@@ -78,14 +96,14 @@ def _manifest() -> BatchManifest:
 def _fake_dry_run_app(report: BatchDryRunReport, monkeypatch) -> AsyncMock:
     fake = AsyncMock()
     fake.run_batch_dry_run.return_value = report
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
     return fake
 
 
 def _fake_manifest_app(report: BatchReport, monkeypatch) -> AsyncMock:
     fake = AsyncMock()
     fake.run_batch_from_manifest.return_value = report
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
     return fake
 
 
@@ -111,7 +129,7 @@ def test_run_batch_passes_progress_callback_when_enabled(monkeypatch):
         return BatchReport(succeeded=1, dead=0)
 
     fake.run_batch.side_effect = _run_batch
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
     monkeypatch.setattr(
         "haiku.rag.ingester.cli._batch_progress",
         lambda _: _progress_context(lambda snapshot: None),
@@ -227,7 +245,7 @@ def test_run_batch_manifest_passes_progress_callback(monkeypatch, tmp_path):
         return BatchReport(succeeded=1, dead=0)
 
     fake.run_batch_from_manifest.side_effect = _run_manifest
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
     monkeypatch.setattr(
         "haiku.rag.ingester.cli._batch_progress",
         lambda _: _progress_context(lambda snapshot: None),
@@ -256,7 +274,7 @@ def test_run_batch_manifest_reports_validation_error(monkeypatch, tmp_path):
     _write_manifest(manifest_path)
     fake = AsyncMock()
     fake.run_batch_from_manifest.side_effect = ValueError("bad manifest")
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
 
     result = runner.invoke(cli, ["run-batch", "--manifest", str(manifest_path)])
 
@@ -322,7 +340,7 @@ def test_resolve_queue_config_keeps_dburi_when_path_override_present(tmp_path):
 
 def test_serve_invokes_app(monkeypatch):
     fake = AsyncMock()
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", lambda **_: fake)
 
     result = runner.invoke(cli, ["serve", "--db", "x.lancedb", "--no-api"])
 
@@ -338,7 +356,7 @@ def test_serve_passes_host_and_port(monkeypatch):
         captured.update(kwargs)
         return fake
 
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", _capture)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", _capture)
 
     result = runner.invoke(
         cli, ["serve", "--db", "x.lancedb", "--host", "0.0.0.0", "--port", "9999"]
@@ -357,7 +375,7 @@ def test_serve_passes_root_path(monkeypatch):
         captured.update(kwargs)
         return fake
 
-    monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", _capture)
+    monkeypatch.setattr("haiku.rag.ingester.app.IngesterApp", _capture)
 
     result = runner.invoke(
         cli, ["serve", "--db", "x.lancedb", "--root-path", "/ingester/"]

tests/store/test_engine.py

@@ -1,7 +1,6 @@
 import pytest
 
-from haiku.rag.store import Store
-from haiku.rag.store.engine import get_database_stats
+from haiku.rag.store.engine import Store, get_database_stats
 
 
 class TestGetDatabaseStats:

tests/store/test_migrations.py

@@ -2,7 +2,7 @@ from importlib import metadata
 
 import pytest
 
-from haiku.rag.store import Store
+from haiku.rag.store.engine import Store
 from haiku.rag.store.exceptions import MigrationRequiredError
 
 

tests/store/test_read_only.py

@@ -3,7 +3,8 @@ from pathlib import Path
 import pytest
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.store import ReadOnlyError, Store
+from haiku.rag.store import ReadOnlyError
+from haiku.rag.store.engine import Store
 from haiku.rag.store.models import Chunk, Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository

tests/store/test_restore.py

@@ -3,8 +3,8 @@ import re
 import pytest
 from lancedb.table import AsyncTable, AsyncTags
 
-from haiku.rag.store import ReadOnlyError, Store
-from haiku.rag.store.engine import RESTORE_TABLE_ORDER
+from haiku.rag.store import ReadOnlyError
+from haiku.rag.store.engine import RESTORE_TABLE_ORDER, Store
 from haiku.rag.store.models import Document
 from haiku.rag.store.repositories.document import DocumentRepository
 

tests/store/test_tags.py

@@ -3,8 +3,8 @@ import asyncio
 import pytest
 from lancedb.table import AsyncTags
 
-from haiku.rag.store import ReadOnlyError, Store
-from haiku.rag.store.engine import REQUIRED_TABLES
+from haiku.rag.store import ReadOnlyError
+from haiku.rag.store.engine import REQUIRED_TABLES, Store
 from haiku.rag.store.models import Document
 from haiku.rag.store.repositories.document import DocumentRepository
 

tests/store/test_v0_40_0_migration.py

@@ -8,9 +8,8 @@ introduced by later migrations (``picture_data`` in v0.45.0,
 
 import pytest
 
-from haiku.rag.store import Store
 from haiku.rag.store.compression import compress_docling_split
-from haiku.rag.store.engine import DocumentRecord
+from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.upgrades.v0_40_0 import _apply_populate_document_items
 
 

tests/store/test_v0_45_0_migration.py

@@ -12,9 +12,8 @@ import json
 import pyarrow as pa
 import pytest
 
-from haiku.rag.store import Store
 from haiku.rag.store.compression import compress_json, decompress_json
-from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord
+from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord, Store
 from haiku.rag.store.upgrades.v0_45_0 import _apply_extract_picture_bytes
 
 PNG_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes"

tests/test_cli.py

@@ -1,3 +1,5 @@
+import subprocess
+import sys
 from unittest.mock import patch
 
 import pytest
@@ -12,6 +14,22 @@ from haiku.rag.store.exceptions import MigrationRequiredError
 runner = CliRunner()
 
 
+def test_importing_cli_does_not_load_lancedb():
+    """Test that lancedb is not imported automatically by the cli. Doing so is
+     expensive. Must be run in a subprocess because lancedb might be imported by
+    other tests in the same session."""
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import haiku.rag.cli, sys; assert 'lancedb' not in sys.modules",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0, result.stderr
+
+
 class TestParseMetaOptions:
     def test_empty_input(self):
         assert _parse_meta_options(None) == {}