diff --git a/CHANGELOG.md b/CHANGELOG.md index 4efc8e7e..8a15ffa4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed - Added `cachetools` as an explicit dependency (was only available transitively, causing `ModuleNotFoundError` for some installations) +- **download-models**: Show actionable error message when Ollama is not running instead of cryptic "All connection attempts failed" (#277) ## [0.30.1] - 2026-02-17 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 74482041..b1159876 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -1846,35 +1846,41 @@ class HaikuRAG: base_url = self._config.providers.ollama.base_url - async with httpx.AsyncClient(timeout=None) as client: - for model in sorted(required_models): - yield DownloadProgress(model=model, status="pulling") + 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", "") + 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 + 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") + 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.""" diff --git a/tests/test_download_models.py b/tests/test_download_models.py new file mode 100644 index 00000000..265d5b41 --- /dev/null +++ b/tests/test_download_models.py @@ -0,0 +1,119 @@ +from contextlib import asynccontextmanager +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from haiku.rag.client import HaikuRAG + + +@pytest.fixture(scope="module") +def vcr_cassette_dir(): + return str(Path(__file__).parent / "cassettes" / "test_download_models") + + +@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): + yield + + +@asynccontextmanager +async def _mock_httpx_client(stream_fn): + """Create a mock httpx.AsyncClient context manager with a given stream function.""" + mock_client = AsyncMock() + mock_client.stream = stream_fn + yield mock_client + + +async def test_download_models_ollama_connect_error(temp_db_path, 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 + + 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 + + assert "ollama serve" in str(exc_info.value) + + +async def test_download_models_ollama_pulls_models(temp_db_path, 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"}', + ] + + @asynccontextmanager + async def mock_stream(method, url, **kwargs): + mock_resp = AsyncMock() + + async def aiter_lines(): + for line in stream_lines: + yield line + + 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) + + # 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 + + 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): + """When no Ollama models are configured, no Ollama pull events are yielded.""" + from haiku.rag.config import AppConfig + + config = AppConfig() + config.embeddings.model.provider = "openai" + 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) + + models = {e.model for e in events} + assert "qwen3-embedding:4b" not in models + assert "gpt-oss" not in models