diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 9166aa1c..768a6e5c 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/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 - - client = HaikuRAG(db_path=None, config=self.config) + from haiku.rag.client.downloads import download_models 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]" diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 244d430c..7d6aec10 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/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() diff --git a/haiku_rag_slim/haiku/rag/client/downloads.py b/haiku_rag_slim/haiku/rag/client/downloads.py new file mode 100644 index 00000000..9578cfbb --- /dev/null +++ b/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'." + ) diff --git a/tests/test_download_models.py b/tests/test_download_models.py index 448da403..92b20f04 100644 --- a/tests/test_download_models.py +++ b/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 - async def failing_stream(method, url, **kwargs): - raise httpx.ConnectError("All connection attempts failed") - yield # unreachable, but needed for generator syntax + @asynccontextmanager + async def failing_stream(method, url, **kwargs): + raise httpx.ConnectError("All connection attempts failed") + yield # unreachable, but needed for generator syntax - with patch( - "haiku.rag.client.httpx.AsyncClient", - return_value=_mock_httpx_client(failing_stream), - ): - with pytest.raises( - ConnectionError, match="Cannot connect to Ollama" - ) as exc_info: - async for _ in client.download_models(): - pass + with patch( + "haiku.rag.client.downloads.httpx.AsyncClient", + return_value=_mock_httpx_client(failing_stream), + ): + with pytest.raises( + ConnectionError, match="Cannot connect to Ollama" + ) as exc_info: + async for _ in download_models(Config): + 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 = [ - '{"status": "pulling manifest"}', - "", - '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}', - '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}', - "not valid json", - '{"status": "verifying sha256 digest"}', - '{"status": "writing manifest"}', - '{"status": "success"}', - ] + stream_lines = [ + '{"status": "pulling manifest"}', + "", + '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}', + '{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}', + "not valid json", + '{"status": "verifying sha256 digest"}', + '{"status": "writing manifest"}', + '{"status": "success"}', + ] - @asynccontextmanager - async def mock_stream(method, url, **kwargs): - mock_resp = AsyncMock() + @asynccontextmanager + async def mock_stream(method, url, **kwargs): + mock_resp = AsyncMock() - async def aiter_lines(): - for line in stream_lines: - yield line + async def aiter_lines(): + for line in stream_lines: + yield line - mock_resp.aiter_lines = aiter_lines - yield mock_resp + mock_resp.aiter_lines = aiter_lines + yield mock_resp - with patch( - "haiku.rag.client.httpx.AsyncClient", - return_value=_mock_httpx_client(mock_stream), - ): - events = [] - async for progress in client.download_models(): - events.append(progress) + with patch( + "haiku.rag.client.downloads.httpx.AsyncClient", + return_value=_mock_httpx_client(mock_stream), + ): + events = [] + async for progress in download_models(Config): + events.append(progress) - # Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss - ollama_models = {"gpt-oss", "qwen3-embedding:4b"} - 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"] - done_events = [e for e in ollama_events if e.status == "done"] - download_events = [e for e in ollama_events if e.status == "downloading"] + # Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss + ollama_models = {"gpt-oss", "qwen3-embedding:4b"} + 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"] + done_events = [e for e in ollama_events if e.status == "done"] + download_events = [e for e in ollama_events if e.status == "downloading"] - assert len(pulling_events) == 2 - assert len(done_events) == 2 - assert len(download_events) > 0 + assert len(pulling_events) == 2 + assert len(done_events) == 2 + assert len(download_events) > 0 - for de in download_events: - assert de.digest == "sha256:abc" - assert de.total == 1000 - assert de.completed > 0 + for de in download_events: + assert de.digest == "sha256:abc" + assert de.total == 1000 + 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 = [] - async for progress in client.download_models(): - events.append(progress) + events = [] + async for progress in download_models(config): + events.append(progress) models = {e.model for e in events} assert "qwen3-embedding:4b" not in models