Merge pull request #73 from ggozad/feat/prefetch-models

CLI command to prefetch models (Ollama & Docling)

docs/cli.md

@@ -163,6 +163,18 @@ when want to switch embeddings provider or model:
 haiku-rag rebuild
 ```
 
+### Download Models
+
+Download required runtime models:
+
+```bash
+haiku-rag download-models
+```
+
+This command:
+- Downloads Docling OCR/conversion models (no-op if already present).
+- Pulls Ollama models referenced in your configuration (embeddings, QA, research, rerank).
+
 ## Migration
 
 ### Migrate from SQLite to LanceDB

docs/installation.md

@@ -72,3 +72,13 @@ VLLM_RERANK_BASE_URL="http://localhost:8001"
 - Python 3.10+
 - Ollama (for default embeddings)
 - vLLM server (for vLLM provider)
+
+## Pre-download Models (Optional)
+
+You can prefetch all required runtime models before first use:
+
+```bash
+haiku-rag download-models
+```
+
+This will download Docling models and pull any Ollama models referenced by your current configuration.

src/haiku/rag/cli.py

@@ -361,6 +361,18 @@ def info(
     asyncio.run(app.info())
 
 
+@cli.command("download-models", help="Download Docling and Ollama models per config")
+def download_models_cmd():
+    from haiku.rag.utils import prefetch_models
+
+    try:
+        prefetch_models()
+        typer.echo("Models downloaded successfully.")
+    except Exception as e:
+        typer.echo(f"Error downloading models: {e}")
+        raise typer.Exit(1)
+
+
 @cli.command(
     "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)"
 )

src/haiku/rag/utils.py

@@ -163,3 +163,37 @@ def load_callable(path: str):
             f"Attribute '{func_name}' in module '{module_part}' is not callable"
         )
     return func
+
+
+def prefetch_models():
+    """Prefetch runtime models (Docling + Ollama as configured)."""
+    import httpx
+    from docling.utils.model_downloader import download_models
+
+    from haiku.rag.config import Config
+
+    download_models()
+
+    # Collect Ollama models from config
+    required_models: set[str] = set()
+    if Config.EMBEDDINGS_PROVIDER == "ollama":
+        required_models.add(Config.EMBEDDINGS_MODEL)
+    if Config.QA_PROVIDER == "ollama":
+        required_models.add(Config.QA_MODEL)
+    if Config.RESEARCH_PROVIDER == "ollama":
+        required_models.add(Config.RESEARCH_MODEL)
+    if Config.RERANK_PROVIDER == "ollama":
+        required_models.add(Config.RERANK_MODEL)
+
+    if not required_models:
+        return
+
+    base_url = Config.OLLAMA_BASE_URL
+
+    with httpx.Client(timeout=None) as client:
+        for model in sorted(required_models):
+            with client.stream(
+                "POST", f"{base_url}/api/pull", json={"model": model}
+            ) as r:
+                for _ in r.iter_lines():
+                    pass