Extract download_models into client/downloads.py

haiku_rag_slim/haiku/rag/app.py

@@ -645,16 +645,14 @@ class HaikuRAGApp:  # pragma: no cover
 
     async def download_models(self):
         """Download Docling, HuggingFace tokenizer, and Ollama models per config."""
-        from haiku.rag.client import HaikuRAG
+        from haiku.rag.client.downloads import download_models
-
-        client = HaikuRAG(db_path=None, config=self.config)
 
         progress: Progress | None = None
         task_id: TaskID | None = None
         current_model = ""
         current_digest = ""
 
-        async for event in client.download_models():
+        async for event in download_models(self.config):
             if event.status == "start":
                 self.console.print(
                     f"[bold blue]Downloading {event.model}...[/bold blue]"

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -5,7 +5,6 @@ import logging
 import mimetypes
 import tempfile
 from collections.abc import AsyncGenerator
-from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
@@ -48,17 +47,6 @@ class RebuildMode(Enum):
     TITLE_ONLY = "title_only"  # Only generate titles for untitled documents
 
 
-@dataclass
-class DownloadProgress:
-    """Progress event for model downloads."""
-
-    model: str
-    status: str
-    completed: int = 0
-    total: int = 0
-    digest: str = ""
-
-
 class HaikuRAG:
     """High-level haiku-rag client."""
 
@@ -1718,144 +1706,6 @@ class HaikuRAG:
         """Optimize and clean up old versions across all tables."""
         await self.store.vacuum()
 
-    async def download_models(self) -> AsyncGenerator[DownloadProgress, None]:
-        """Download required models, yielding progress events.
-
-        Yields DownloadProgress events for:
-        - Docling models
-        - HuggingFace tokenizer
-        - Sentence-transformers embedder (if configured)
-        - HuggingFace reranker models (mxbai, jina-local)
-        - Ollama models
-        """
-        # Docling models
-        try:
-            from docling.utils.model_downloader import download_models
-
-            yield DownloadProgress(model="docling", status="start")
-            await asyncio.to_thread(download_models)
-            yield DownloadProgress(model="docling", status="done")
-        except ImportError:
-            pass
-
-        # HuggingFace tokenizer
-        from transformers import AutoTokenizer
-
-        tokenizer_name = self._config.processing.chunking_tokenizer
-        yield DownloadProgress(model=tokenizer_name, status="start")
-        await asyncio.to_thread(AutoTokenizer.from_pretrained, tokenizer_name)
-        yield DownloadProgress(model=tokenizer_name, status="done")
-
-        # Sentence-transformers embedder
-        if (
-            self._config.embeddings.model.provider == "sentence-transformers"
-        ):  # pragma: no cover
-            try:
-                from sentence_transformers import (  # type: ignore[import-not-found]  # ty: ignore[unresolved-import]
-                    SentenceTransformer,
-                )
-
-                model_name = self._config.embeddings.model.name
-                yield DownloadProgress(model=model_name, status="start")
-                await asyncio.to_thread(SentenceTransformer, model_name)
-                yield DownloadProgress(model=model_name, status="done")
-            except ImportError:
-                pass
-
-        # HuggingFace reranker models
-        if self._config.reranking.model:  # pragma: no cover
-            provider = self._config.reranking.model.provider
-            model_name = self._config.reranking.model.name
-
-            if provider == "mxbai":
-                try:
-                    from mxbai_rerank import MxbaiRerankV2
-
-                    yield DownloadProgress(model=model_name, status="start")
-                    await asyncio.to_thread(
-                        MxbaiRerankV2, model_name, disable_transformers_warnings=True
-                    )
-                    yield DownloadProgress(model=model_name, status="done")
-                except ImportError:
-                    pass
-
-            elif provider == "jina-local":
-                try:
-                    from transformers import AutoModel
-
-                    yield DownloadProgress(model=model_name, status="start")
-                    await asyncio.to_thread(
-                        AutoModel.from_pretrained,
-                        model_name,
-                        trust_remote_code=True,
-                    )
-                    yield DownloadProgress(model=model_name, status="done")
-                except ImportError:
-                    pass
-
-        # Collect Ollama models from config
-        required_models: set[str] = set()
-        if self._config.embeddings.model.provider == "ollama":
-            required_models.add(self._config.embeddings.model.name)
-        if self._config.qa.model.provider == "ollama":
-            required_models.add(self._config.qa.model.name)
-        if self._config.research.model.provider == "ollama":
-            required_models.add(self._config.research.model.name)
-        if (
-            self._config.reranking.model
-            and self._config.reranking.model.provider == "ollama"
-        ):
-            required_models.add(self._config.reranking.model.name)
-        pic_desc = self._config.processing.conversion_options.picture_description
-        if pic_desc.enabled and pic_desc.model.provider == "ollama":
-            required_models.add(pic_desc.model.name)
-        if (
-            self._config.processing.auto_title
-            and self._config.processing.title_model.provider == "ollama"
-        ):
-            required_models.add(self._config.processing.title_model.name)
-
-        if not required_models:
-            return
-
-        base_url = self._config.providers.ollama.base_url
-
-        try:
-            async with httpx.AsyncClient(timeout=None) as client:
-                for model in sorted(required_models):
-                    yield DownloadProgress(model=model, status="pulling")
-
-                    async with client.stream(
-                        "POST", f"{base_url}/api/pull", json={"model": model}
-                    ) as r:
-                        async for line in r.aiter_lines():
-                            if not line:
-                                continue
-                            try:
-                                data = json.loads(line)
-                                status = data.get("status", "")
-                                digest = data.get("digest", "")
-
-                                if digest and "total" in data:
-                                    yield DownloadProgress(
-                                        model=model,
-                                        status="downloading",
-                                        total=data.get("total", 0),
-                                        completed=data.get("completed", 0),
-                                        digest=digest,
-                                    )
-                                elif status:
-                                    yield DownloadProgress(model=model, status=status)
-                            except json.JSONDecodeError:
-                                pass
-
-                    yield DownloadProgress(model=model, status="done")
-        except httpx.ConnectError:
-            raise ConnectionError(
-                f"Cannot connect to Ollama at {base_url}. "
-                "Is Ollama running? Start it with 'ollama serve'."
-            )
-
     def close(self):
         """Close the underlying store connection."""
         self.store.close()

haiku_rag_slim/haiku/rag/client/downloads.py

@@ -0,0 +1,155 @@
+import asyncio
+import json
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+
+import httpx
+
+from haiku.rag.config import AppConfig
+
+
+@dataclass
+class DownloadProgress:
+    """Progress event for model downloads."""
+
+    model: str
+    status: str
+    completed: int = 0
+    total: int = 0
+    digest: str = ""
+
+
+async def download_models(
+    config: AppConfig,
+) -> AsyncGenerator[DownloadProgress, None]:
+    """Download required models per config, yielding progress events.
+
+    Yields DownloadProgress events for:
+    - Docling models
+    - HuggingFace tokenizer
+    - Sentence-transformers embedder (if configured)
+    - HuggingFace reranker models (mxbai, jina-local)
+    - Ollama models
+    """
+    # Docling models
+    try:
+        from docling.utils.model_downloader import download_models
+
+        yield DownloadProgress(model="docling", status="start")
+        await asyncio.to_thread(download_models)
+        yield DownloadProgress(model="docling", status="done")
+    except ImportError:
+        pass
+
+    # HuggingFace tokenizer
+    from transformers import AutoTokenizer
+
+    tokenizer_name = config.processing.chunking_tokenizer
+    yield DownloadProgress(model=tokenizer_name, status="start")
+    await asyncio.to_thread(AutoTokenizer.from_pretrained, tokenizer_name)
+    yield DownloadProgress(model=tokenizer_name, status="done")
+
+    # Sentence-transformers embedder
+    if config.embeddings.model.provider == "sentence-transformers":  # pragma: no cover
+        try:
+            from sentence_transformers import (  # type: ignore[import-not-found]  # ty: ignore[unresolved-import]
+                SentenceTransformer,
+            )
+
+            model_name = config.embeddings.model.name
+            yield DownloadProgress(model=model_name, status="start")
+            await asyncio.to_thread(SentenceTransformer, model_name)
+            yield DownloadProgress(model=model_name, status="done")
+        except ImportError:
+            pass
+
+    # HuggingFace reranker models
+    if config.reranking.model:  # pragma: no cover
+        provider = config.reranking.model.provider
+        model_name = config.reranking.model.name
+
+        if provider == "mxbai":
+            try:
+                from mxbai_rerank import MxbaiRerankV2
+
+                yield DownloadProgress(model=model_name, status="start")
+                await asyncio.to_thread(
+                    MxbaiRerankV2, model_name, disable_transformers_warnings=True
+                )
+                yield DownloadProgress(model=model_name, status="done")
+            except ImportError:
+                pass
+
+        elif provider == "jina-local":
+            try:
+                from transformers import AutoModel
+
+                yield DownloadProgress(model=model_name, status="start")
+                await asyncio.to_thread(
+                    AutoModel.from_pretrained,
+                    model_name,
+                    trust_remote_code=True,
+                )
+                yield DownloadProgress(model=model_name, status="done")
+            except ImportError:
+                pass
+
+    # Collect Ollama models from config
+    required_models: set[str] = set()
+    if config.embeddings.model.provider == "ollama":
+        required_models.add(config.embeddings.model.name)
+    if config.qa.model.provider == "ollama":
+        required_models.add(config.qa.model.name)
+    if config.research.model.provider == "ollama":
+        required_models.add(config.research.model.name)
+    if config.reranking.model and config.reranking.model.provider == "ollama":
+        required_models.add(config.reranking.model.name)
+    pic_desc = config.processing.conversion_options.picture_description
+    if pic_desc.enabled and pic_desc.model.provider == "ollama":
+        required_models.add(pic_desc.model.name)
+    if (
+        config.processing.auto_title
+        and config.processing.title_model.provider == "ollama"
+    ):
+        required_models.add(config.processing.title_model.name)
+
+    if not required_models:
+        return
+
+    base_url = config.providers.ollama.base_url
+
+    try:
+        async with httpx.AsyncClient(timeout=None) as client:
+            for model in sorted(required_models):
+                yield DownloadProgress(model=model, status="pulling")
+
+                async with client.stream(
+                    "POST", f"{base_url}/api/pull", json={"model": model}
+                ) as r:
+                    async for line in r.aiter_lines():
+                        if not line:
+                            continue
+                        try:
+                            data = json.loads(line)
+                            status = data.get("status", "")
+                            digest = data.get("digest", "")
+
+                            if digest and "total" in data:
+                                yield DownloadProgress(
+                                    model=model,
+                                    status="downloading",
+                                    total=data.get("total", 0),
+                                    completed=data.get("completed", 0),
+                                    digest=digest,
+                                )
+                            elif status:
+                                yield DownloadProgress(model=model, status=status)
+                        except json.JSONDecodeError:
+                            pass
+
+                yield DownloadProgress(model=model, status="done")
+    except httpx.ConnectError:
+        raise ConnectionError(
+            f"Cannot connect to Ollama at {base_url}. "
+            "Is Ollama running? Start it with 'ollama serve'."
+        )

tests/test_download_models.py

@@ -4,13 +4,14 @@ from unittest.mock import AsyncMock, patch
 import httpx
 import pytest
 
-from haiku.rag.client import HaikuRAG
+from haiku.rag.client.downloads import download_models
+from haiku.rag.config import Config
 
 
 @pytest.fixture
 def mock_to_thread():
     """Patch asyncio.to_thread to skip docling/tokenizer downloads."""
-    with patch("haiku.rag.client.asyncio.to_thread", new_callable=AsyncMock):
+    with patch("haiku.rag.client.downloads.asyncio.to_thread", new_callable=AsyncMock):
         yield
 
 
@@ -22,79 +23,77 @@ async def _mock_httpx_client(stream_fn):
     yield mock_client
 
 
-async def test_download_models_ollama_connect_error(temp_db_path, mock_to_thread):
+async def test_download_models_ollama_connect_error(mock_to_thread):
     """When Ollama is not running, download_models raises ConnectionError."""
-    async with HaikuRAG(temp_db_path, create=True) as client:
 
-        @asynccontextmanager
+    @asynccontextmanager
-        async def failing_stream(method, url, **kwargs):
+    async def failing_stream(method, url, **kwargs):
-            raise httpx.ConnectError("All connection attempts failed")
+        raise httpx.ConnectError("All connection attempts failed")
-            yield  # unreachable, but needed for generator syntax
+        yield  # unreachable, but needed for generator syntax
 
-        with patch(
+    with patch(
-            "haiku.rag.client.httpx.AsyncClient",
+        "haiku.rag.client.downloads.httpx.AsyncClient",
-            return_value=_mock_httpx_client(failing_stream),
+        return_value=_mock_httpx_client(failing_stream),
-        ):
+    ):
-            with pytest.raises(
+        with pytest.raises(
-                ConnectionError, match="Cannot connect to Ollama"
+            ConnectionError, match="Cannot connect to Ollama"
-            ) as exc_info:
+        ) as exc_info:
-                async for _ in client.download_models():
+            async for _ in download_models(Config):
-                    pass
+                pass
 
-            assert "ollama serve" in str(exc_info.value)
+        assert "ollama serve" in str(exc_info.value)
 
 
-async def test_download_models_ollama_pulls_models(temp_db_path, mock_to_thread):
+async def test_download_models_ollama_pulls_models(mock_to_thread):
     """download_models yields correct progress events for Ollama model pulls."""
-    async with HaikuRAG(temp_db_path, create=True) as client:
+    stream_lines = [
-        stream_lines = [
+        '{"status": "pulling manifest"}',
-            '{"status": "pulling manifest"}',
+        "",
-            "",
+        '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}',
-            '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}',
+        '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}',
-            '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}',
+        "not valid json",
-            "not valid json",
+        '{"status": "verifying sha256 digest"}',
-            '{"status": "verifying sha256 digest"}',
+        '{"status": "writing manifest"}',
-            '{"status": "writing manifest"}',
+        '{"status": "success"}',
-            '{"status": "success"}',
+    ]
-        ]
 
-        @asynccontextmanager
+    @asynccontextmanager
-        async def mock_stream(method, url, **kwargs):
+    async def mock_stream(method, url, **kwargs):
-            mock_resp = AsyncMock()
+        mock_resp = AsyncMock()
 
-            async def aiter_lines():
+        async def aiter_lines():
-                for line in stream_lines:
+            for line in stream_lines:
-                    yield line
+                yield line
 
-            mock_resp.aiter_lines = aiter_lines
+        mock_resp.aiter_lines = aiter_lines
-            yield mock_resp
+        yield mock_resp
 
-        with patch(
+    with patch(
-            "haiku.rag.client.httpx.AsyncClient",
+        "haiku.rag.client.downloads.httpx.AsyncClient",
-            return_value=_mock_httpx_client(mock_stream),
+        return_value=_mock_httpx_client(mock_stream),
-        ):
+    ):
-            events = []
+        events = []
-            async for progress in client.download_models():
+        async for progress in download_models(Config):
-                events.append(progress)
+            events.append(progress)
 
-        # Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss
+    # Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss
-        ollama_models = {"gpt-oss", "qwen3-embedding:4b"}
+    ollama_models = {"gpt-oss", "qwen3-embedding:4b"}
-        ollama_events = [e for e in events if e.model in ollama_models]
+    ollama_events = [e for e in events if e.model in ollama_models]
-        pulling_events = [e for e in ollama_events if e.status == "pulling"]
+    pulling_events = [e for e in ollama_events if e.status == "pulling"]
-        done_events = [e for e in ollama_events if e.status == "done"]
+    done_events = [e for e in ollama_events if e.status == "done"]
-        download_events = [e for e in ollama_events if e.status == "downloading"]
+    download_events = [e for e in ollama_events if e.status == "downloading"]
 
-        assert len(pulling_events) == 2
+    assert len(pulling_events) == 2
-        assert len(done_events) == 2
+    assert len(done_events) == 2
-        assert len(download_events) > 0
+    assert len(download_events) > 0
 
-        for de in download_events:
+    for de in download_events:
-            assert de.digest == "sha256:abc"
+        assert de.digest == "sha256:abc"
-            assert de.total == 1000
+        assert de.total == 1000
-            assert de.completed > 0
+        assert de.completed > 0
 
 
-async def test_download_models_no_ollama_models(temp_db_path, mock_to_thread):
+async def test_download_models_no_ollama_models(mock_to_thread):
     """When no Ollama models are configured, no Ollama pull events are yielded."""
     from haiku.rag.config import AppConfig
 
@@ -103,10 +102,9 @@ async def test_download_models_no_ollama_models(temp_db_path, mock_to_thread):
     config.qa.model.provider = "openai"
     config.research.model.provider = "openai"
 
-    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
+    events = []
-        events = []
+    async for progress in download_models(config):
-        async for progress in client.download_models():
+        events.append(progress)
-            events.append(progress)
 
     models = {e.model for e in events}
     assert "qwen3-embedding:4b" not in models