"""`tht ollama` -- embeddings preflight (ensure Ollama up + model warm). The system REQUIRES embeddings: any condition that makes them unavailable is a hard error (the caller refuses the session). "Load" = warm the already-installed model. """ from __future__ import annotations import subprocess import time # --- low-level ops (real implementations; injected as fakes in tests) ---------- def _probe(base_url: str, timeout: float = 2.0) -> bool: import requests try: return requests.get(f"{base_url.rstrip('/')}/api/tags", timeout=timeout).status_code == 200 except requests.RequestException: return False def _installed_models(base_url: str, timeout: float = 5.0) -> set[str]: import requests resp = requests.get(f"{base_url.rstrip('/')}/api/tags", timeout=timeout) resp.raise_for_status() return {m.get("name", "") for m in resp.json().get("models", [])} def _start(start_cmd: list[str]) -> None: # Detached so the server outlives this short-lived CLI process. subprocess.Popen( # noqa: S603 start_cmd, start_new_session=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) def _warm(cfg) -> None: from tht.vectorstore.embeddings import OllamaEmbeddings OllamaEmbeddings(cfg.embeddings).embed_query("ping") def _model_present(installed: set[str], model: str) -> bool: """Match the configured model against installed names, allowing the implicit ':latest'.""" if model in installed: return True base = model.split(":")[0] return any(name == base or name.split(":")[0] == base for name in installed) # --- orchestration (pure: returns a result dict, never raises for control flow) ---- def ensure_ollama( cfg, *, timeout: int, no_start: bool, probe=_probe, installed_models=_installed_models, start=_start, warm=_warm, sleep=time.sleep, clock=time.monotonic, ) -> dict: if cfg.embeddings is None: return {"ok": False, "stage": "config", "error": "il sistema richiede embeddings ma il workspace non li configura"} emb = cfg.embeddings base_url = emb.base_url start_cmd = emb.start_cmd if emb.start_cmd is not None else [emb.bin, "serve"] server_state = "up" try: reachable = probe(base_url) except Exception as e: # noqa: BLE001 - a probe failure is a hard server-stage error return {"ok": False, "stage": "server", "error": f"probe di Ollama fallito su {base_url}: {e}"} if not reachable: if no_start or start_cmd == []: return {"ok": False, "stage": "server", "error": f"Ollama non raggiungibile su {base_url} e avvio disabilitato"} try: start(start_cmd) except Exception as e: # noqa: BLE001 - failing to launch is a hard server-stage error return {"ok": False, "stage": "server", "error": f"avvio di Ollama fallito ({start_cmd}): {e}"} server_state = "started" deadline = clock() + timeout up = False while clock() < deadline: sleep(1.0) try: if probe(base_url): up = True break except Exception: # noqa: BLE001 - transient during startup; keep polling continue if not up: return {"ok": False, "stage": "server", "error": f"Ollama non raggiungibile su {base_url} entro {timeout}s"} try: installed = installed_models(base_url) except Exception as e: # noqa: BLE001 - any read failure is a hard error return {"ok": False, "stage": "server", "error": f"impossibile leggere i modelli da {base_url}: {e}"} if not _model_present(installed, emb.model): return {"ok": False, "stage": "model", "error": f"modello '{emb.model}' non installato in Ollama: " f"esegui `ollama pull {emb.model}` o importalo"} try: warm(cfg) except Exception as e: # noqa: BLE001 - warm failure is a hard error return {"ok": False, "stage": "warm", "error": f"warm del modello '{emb.model}' fallito: {e}"} return {"ok": True, "server": server_state, "model": "warmed", "model_name": emb.model}