Extract download_models into client/downloads.py
This commit is contained in:
Yiorgis Gozadinos
parent
43f2774c14
commit
e814459e1f
4 changed files with 218 additions and 217 deletions
|
|
@ -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]"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
155
haiku_rag_slim/haiku/rag/client/downloads.py
Normal file
155
haiku_rag_slim/haiku/rag/client/downloads.py
Normal file
|
|
@ -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'."
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue