ThothII/harness/tht/cli/ollama_cmd.py

"""`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 json
import subprocess
import time
from pathlib import Path

import typer

from tht.cli.config_cmd import CONFIG_OPT
from tht.cli.schema_cmd import _load_config_or_exit

ollama_app = typer.Typer(help="Ollama (embeddings) -- preflight.")


# --- 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}


# --- CLI command ----------


@ollama_app.command("ensure")
def ensure_cmd(
    timeout: int = typer.Option(60, "--timeout", help="Secondi di attesa per l'avvio di Ollama."),
    no_start: bool = typer.Option(False, "--no-start", help="Non avviare Ollama (solo verifica)."),
    json_out: bool = typer.Option(False, "--json", help="Emetti JSON puro su stdout."),
    config: Path = CONFIG_OPT,
) -> None:
    """Assicura Ollama attivo + modello di embedding caricato; errore se non possibile."""
    cfg = _load_config_or_exit(config)
    result = ensure_ollama(cfg, timeout=timeout, no_start=no_start)
    if json_out:
        typer.echo(json.dumps(result, ensure_ascii=False))
    elif result["ok"]:
        typer.secho(
            f"OK: Ollama {result['server']}, modello {result['model_name']} {result['model']}.",
            fg=typer.colors.GREEN,
        )
    else:
        typer.secho(f"ERRORE [{result['stage']}]: {result['error']}", fg=typer.colors.RED, err=True)
    if not result["ok"]:
        raise typer.Exit(code=1)