diff --git a/backend/ollama_admin.py b/backend/ollama_admin.py index 02a0fbd..d97460b 100644 --- a/backend/ollama_admin.py +++ b/backend/ollama_admin.py @@ -10,9 +10,11 @@ from urllib.parse import urlparse import httpx from .app_settings import get_embed_model_preference, get_ollama_api_url, normalize_embed_model +from .whisper_admin import DEFAULT_WHISPER_MODEL, ensure_whisper_model_downloaded, inspect_whisper_model LOCAL_OLLAMA_HOSTS = {"127.0.0.1", "localhost", "::1"} +_OLLAMA_PULL_LOCK = asyncio.Lock() def _ollama_binary() -> Optional[str]: @@ -56,6 +58,7 @@ async def inspect_ollama_startup() -> Dict[str, Any]: embed_model = get_embed_model_preference() ollama_bin = _ollama_binary() is_local = _is_local_ollama_url(ollama_url) + whisper_status = inspect_whisper_model(DEFAULT_WHISPER_MODEL) available_models: List[str] = [] error = "" running = False @@ -75,6 +78,9 @@ async def inspect_ollama_startup() -> Dict[str, Any]: "selected_embed_model": embed_model, "embedding_model_available": available, "available_models": available_models, + "whisper_model": whisper_status["model"], + "whisper_model_available": bool(whisper_status["available"]), + "whisper_error": whisper_status["error"], "error": error, } @@ -109,32 +115,78 @@ async def start_local_ollama() -> Dict[str, Any]: async def pull_local_model(model: Optional[str] = None) -> Dict[str, Any]: + async with _OLLAMA_PULL_LOCK: + status = await inspect_ollama_startup() + if not status["can_manage_locally"]: + raise RuntimeError("Heimgeist can only pull models automatically when the configured Ollama URL points to this machine.") + if not status["ollama_running"]: + raise RuntimeError("Ollama must be running before Heimgeist can pull a model.") + + ollama_bin = _ollama_binary() + if not ollama_bin: + raise FileNotFoundError("Could not find the 'ollama' executable in PATH.") + + model_name = normalize_embed_model(model or status["selected_embed_model"]) + if bool(set(status["available_models"]) & _model_aliases(model_name)): + return { + "model": model_name, + "downloaded": False, + "status": status, + } + + process = await asyncio.create_subprocess_exec( + ollama_bin, + "pull", + model_name, + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.PIPE, + ) + _stdout, stderr = await process.communicate() + if process.returncode != 0: + detail = (stderr or b"").decode("utf-8", errors="ignore").strip() + raise RuntimeError(detail or f"'ollama pull {model_name}' failed with exit code {process.returncode}.") + + status = await inspect_ollama_startup() + return { + "model": model_name, + "downloaded": True, + "status": status, + } + + +async def prepare_startup_models() -> Dict[str, Any]: status = await inspect_ollama_startup() - if not status["can_manage_locally"]: - raise RuntimeError("Heimgeist can only pull models automatically when the configured Ollama URL points to this machine.") - if not status["ollama_running"]: - raise RuntimeError("Ollama must be running before Heimgeist can pull a model.") - - ollama_bin = _ollama_binary() - if not ollama_bin: - raise FileNotFoundError("Could not find the 'ollama' executable in PATH.") - - model_name = normalize_embed_model(model or status["selected_embed_model"]) - process = await asyncio.create_subprocess_exec( - ollama_bin, - "pull", - model_name, - stdin=asyncio.subprocess.DEVNULL, - stdout=asyncio.subprocess.DEVNULL, - stderr=asyncio.subprocess.PIPE, - ) - _stdout, stderr = await process.communicate() - if process.returncode != 0: - detail = (stderr or b"").decode("utf-8", errors="ignore").strip() - raise RuntimeError(detail or f"'ollama pull {model_name}' failed with exit code {process.returncode}.") - + whisper_result = await asyncio.to_thread(ensure_whisper_model_downloaded, status["whisper_model"]) status = await inspect_ollama_startup() - return { - "model": model_name, - "status": status, + + embedding_result: Dict[str, Any] = { + "model": status["selected_embed_model"], + "available": bool(status["embedding_model_available"]), + "downloaded": False, + "skipped": False, + "reason": "", + } + + if not status["ollama_running"]: + embedding_result["skipped"] = True + embedding_result["reason"] = "Ollama is not running." + elif not status["can_manage_locally"]: + embedding_result["skipped"] = True + embedding_result["reason"] = "Automatic model pulls are only available for local Ollama." + elif not status["embedding_model_available"]: + pulled = await pull_local_model(status["selected_embed_model"]) + status = pulled["status"] + embedding_result = { + "model": pulled["model"], + "available": bool(status["embedding_model_available"]), + "downloaded": bool(pulled.get("downloaded")), + "skipped": False, + "reason": "", + } + + return { + "ollama": status, + "whisper": whisper_result, + "embedding_model": embedding_result, }