121 lines
4.3 KiB
Python
121 lines
4.3 KiB
Python
"""`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}
|