From 0e1bf3e00d0b547d291681e046bd56c1d35594c5 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 23 Sep 2025 16:30:56 +0300
Subject: [PATCH] CLI command to prefetch models (Ollama & Docling)

---
 docs/cli.md            | 12 ++++++++++++
 docs/installation.md   | 10 ++++++++++
 src/haiku/rag/cli.py   | 12 ++++++++++++
 src/haiku/rag/utils.py | 34 ++++++++++++++++++++++++++++++++++
 4 files changed, 68 insertions(+)

diff --git a/docs/cli.md b/docs/cli.md
index ae89da6e..373e231c 100644
--- a/docs/cli.md
+++ b/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
diff --git a/docs/installation.md b/docs/installation.md
index eb1e0750..40487998 100644
--- a/docs/installation.md
+++ b/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.
diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py
index 85327677..517d7c71 100644
--- a/src/haiku/rag/cli.py
+++ b/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)"
 )
diff --git a/src/haiku/rag/utils.py b/src/haiku/rag/utils.py
index 2373798e..433c6f1e 100644
--- a/src/haiku/rag/utils.py
+++ b/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