momentum-system/youtube_claims/recap/llm.py

"""A small pool over the owner's several LAN Ollama endpoints: round-robin across the endpoints
that actually have the requested model, fail over on connection errors / timeouts / 5xx, and
park a dead endpoint for a few minutes. Endpoints come from RECAP_OLLAMA_HOSTS:

  RECAP_OLLAMA_HOSTS="http://127.0.0.1:11434, http://192.168.1.178:11434=qwen3:30b-a3b-instruct-2507-q4_K_M"

An entry may pin its own model with `=model` (and `#ctx=N` / `#think=V` options), otherwise it
serves RECAP_OLLAMA_MODEL with the global options. The transport (`call`) is injected so the
rotation/failover logic is unit-tested without a network."""

from __future__ import annotations

import json
import socket
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Callable

_THINK = {"false": False, "true": True, "none": None}


@dataclass
class Endpoint:
    host: str
    model: str | None = None          # None -> the pool's default model
    num_ctx: int | None = None        # None -> the pool's default
    think: object = "default"         # "default" -> the pool's default
    down_until: float = 0.0
    models: set | None = None
    models_at: float = 0.0


def parse_hosts(spec: str) -> list[Endpoint]:
    """'h1, h2=model#ctx=8192#think=false' -> endpoints (blank entries ignored)."""
    out = []
    for raw in (spec or "").split(","):
        raw = raw.strip()
        if not raw:
            continue
        head, *opts = raw.split("#")
        host, _, model = head.partition("=")
        ep = Endpoint(host=host.strip().rstrip("/"), model=model.strip() or None)
        for o in opts:
            k, _, v = o.partition("=")
            k, v = k.strip().lower(), v.strip()
            if k == "ctx" and v.isdigit():
                ep.num_ctx = int(v) or None
            elif k == "think":
                ep.think = _THINK.get(v.lower(), v)
        out.append(ep)
    return out


def _tags(host: str) -> set[str]:
    with urllib.request.urlopen(host + "/api/tags", timeout=5) as r:
        return {m["name"] for m in json.load(r).get("models", [])}


@dataclass
class OllamaPool:
    endpoints: list[Endpoint]
    default_model: str
    default_ctx: int | None = None
    default_think: object = None
    call: Callable | None = None                       # (system, user, model=, host=, ...) -> (str, stats)
    list_models: Callable[[str], set] = _tags
    down_seconds: float = 300.0
    tags_ttl: float = 300.0
    _rr: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock)

    def _serves(self, ep: Endpoint) -> str | None:
        """The model this endpoint would run, if it has it (tags cached); None otherwise."""
        want = ep.model or self.default_model
        now = time.monotonic()
        if ep.models is None or now - ep.models_at > self.tags_ttl:
            try:
                ep.models, ep.models_at = self.list_models(ep.host), now
            except Exception:  # noqa: BLE001 -- unreachable: park it
                ep.down_until = now + self.down_seconds
                ep.models, ep.models_at = None, 0.0
                return None
        names = ep.models or set()
        return want if (want in names or f"{want}:latest" in names) else None

    def candidates(self) -> list[tuple[Endpoint, str]]:
        now = time.monotonic()
        with self._lock:
            start = self._rr
            self._rr += 1
        n = len(self.endpoints)
        order = [self.endpoints[(start + i) % n] for i in range(n)] if n else []
        out = []
        for ep in order:
            if ep.down_until > now:
                continue
            m = self._serves(ep)
            if m:
                out.append((ep, m))
        return out

    def chat(self, system: str, user: str, **kw) -> tuple[str, dict]:
        """One chat call on the next healthy endpoint; fails over to the others."""
        cands = self.candidates()
        if not cands:
            raise RuntimeError(f"no reachable Ollama endpoint serves {self.default_model!r} "
                               f"(tried {[e.host for e in self.endpoints]})")
        last: Exception | None = None
        for ep, model in cands:
            think = self.default_think if ep.think == "default" else ep.think
            ctx = ep.num_ctx if ep.num_ctx is not None else self.default_ctx
            try:
                content, stats = self.call(system, user, model=model, host=ep.host, num_ctx=ctx, think=think, **kw)
                stats = {**stats, "host": ep.host, "model": model}
                return content, stats
            except (urllib.error.URLError, ConnectionError, TimeoutError, socket.timeout) as e:
                last = e
            except RuntimeError as e:          # ollama_chat's HTTP errors: only 5xx is the host's fault
                if "ollama 5" not in str(e):
                    raise
                last = e
            ep.down_until = time.monotonic() + self.down_seconds
        raise RuntimeError(f"all Ollama endpoints failed; last error: {last}")


def from_settings(call=None) -> OllamaPool:
    from youtube_claims.recap import settings
    from youtube_claims.recap.extract import ollama_chat
    return OllamaPool(parse_hosts(settings.OLLAMA_HOSTS), settings.OLLAMA_MODEL, settings.OLLAMA_NUM_CTX,
                      settings.OLLAMA_THINK, call=call or ollama_chat)