ARIA-AGENT/bridge/llm_cluster.py

"""LLM-Cluster: ein Eingang fuer alle lokalen Sprachmodelle aller KI-Boxen.

Clients waehlen ein MODELL, nicht Box oder Karte. Wo es laeuft, entscheidet
dieser Router:

  * Auswahl     — unter allen online Boxen, die das Modell GELADEN haben
                  (Karten-Belegung aus worker_hello/worker_ping), die mit den
                  wenigsten laufenden Anfragen; bei Gleichstand reihum.
                  Ist nur eine Karte da oder alles frei, landet alles dort —
                  "verteilen, wenn noetig, sonst nicht".
  * Ausfall     — verschwindet eine Box mitten in der Anfrage (kein Ping mehr)
                  oder antwortet nicht, geht die Anfrage EINMAL an eine andere
                  Box mit demselben Modell. Fehler des Modells selbst (z.B.
                  Kontext zu lang) werden nicht wiederholt.
  * Fehler      — kommen als FEHLER zurueck (ok=false / HTTP-Status), nie als
                  normale Antwort. Ein Client wuerde einen Fehlertext sonst
                  als Aussage der KI in seinen Verlauf schreiben — genau das
                  hat bei ARIA schon mal einen ganzen Chat vergiftet.

Eingang ist der RVS:

  llm_cluster_models_request → llm_cluster_models
  llm_cluster_request        → llm_cluster_response

Programme brauchen nur eine ausgehende Verbindung zum RVS, egal in welchem
Netz sie stehen. Werkzeuge, die nur "OpenAI-kompatibel" koennen (OpenHands
& Co.), bekommen den OpenAI-RVS-Wrapper neben sich gestellt
(xtts/openai_rvs_wrapper) — fuer sie ist der Cluster dann 127.0.0.1.

Schnittstelle fuer Drittprogramme: docs/LLM-SCHNITTSTELLE.md
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import time
import uuid
from typing import Optional

logger = logging.getLogger("aria-bridge.llm-cluster")

# Diese errorCodes heissen "die BOX war das Problem" → anderswo nochmal versuchen.
# Alles andere (Modell-Fehler, z.B. Kontext zu lang) waere ueberall gleich.
RETRY_CODES = {"timeout", "box_offline", "send_failed", "model_not_loaded"}
MAX_ATTEMPTS = 2


class LlmCluster:
    def __init__(self, bridge):
        self.b = bridge
        self.inflight: dict[str, int] = {}      # instanceId → laufende Anfragen (von uns)
        self._rr: dict[str, int] = {}
        self.timeout = float(os.getenv("CLUSTER_LLM_TIMEOUT_S", "600"))

    # ── Flotte ──────────────────────────────────────────────────────────

    def _online_llm(self) -> list:
        now = time.time()
        return [w for w in self.b._workers.values()
                if w.get("service") == "llm"
                and (now - w.get("last_seen", 0)) < self.b.WORKER_OFFLINE_S]

    def _is_online(self, instance_id: str) -> bool:
        return any(w.get("instanceId") == instance_id for w in self._online_llm())

    def models(self) -> list:
        """Jedes geladene Modell GENAU EINMAL, egal auf wie vielen Karten.
        ctx = das kleinste Kontextfenster unter seinen Karten (darauf muss sich
        ein Client verlassen koennen, egal wohin der Router verteilt)."""
        out: dict[str, dict] = {}
        for w in self._online_llm():
            slots = w.get("slots")
            if isinstance(slots, list):
                for sl in slots:
                    m = sl.get("model")
                    if not m or m not in (w.get("models") or []):
                        continue
                    e = out.setdefault(m, {"model": m, "ctx": 0, "cards": []})
                    e["cards"].append({"node": w.get("node", ""), "instanceId": w.get("instanceId", ""),
                                       "gpu": sl.get("gpu", ""), "gpuName": sl.get("gpuName", ""),
                                       "state": sl.get("state", "")})
                    c = int(sl.get("ctx") or 0)
                    if c and (not e["ctx"] or c < e["ctx"]):
                        e["ctx"] = c
            else:
                # Box mit altem Adapter (ohne Karten): nur die Modell-Liste
                for m in (w.get("models") or []):
                    e = out.setdefault(m, {"model": m, "ctx": 0, "cards": []})
                    e["cards"].append({"node": w.get("node", ""), "instanceId": w.get("instanceId", ""),
                                       "gpu": "", "gpuName": "", "state": ""})
        for e in out.values():
            e["available"] = len(e["cards"])
        return sorted(out.values(), key=lambda e: e["model"])

    def candidates(self, model: str) -> list:
        return [w for w in self._online_llm() if model in (w.get("models") or [])]

    def pick(self, model: str, exclude=()) -> Optional[str]:
        """Die Box mit den wenigsten laufenden Anfragen; Gleichstand reihum."""
        cands = [w["instanceId"] for w in self.candidates(model) if w["instanceId"] not in exclude]
        if not cands:
            return None
        low = min(self.inflight.get(c, 0) for c in cands)
        tied = sorted(c for c in cands if self.inflight.get(c, 0) == low)
        i = self._rr.get(model, 0) % len(tied)
        self._rr[model] = i + 1
        return tied[i]

    # ── Anfrage an eine bestimmte Box ───────────────────────────────────

    async def _send_to(self, target: str, req: dict, deadline: float) -> dict:
        """llm_request an GENAU diese Box, Antwort abwarten. Bricht frueh ab,
        wenn die Box unterwegs offline geht — dann kann route() woanders hin."""
        rid = str(uuid.uuid4())
        fut: asyncio.Future = asyncio.get_running_loop().create_future()
        self.b._pending_llm[rid] = fut
        self.inflight[target] = self.inflight.get(target, 0) + 1
        try:
            payload = {k: v for k, v in req.items() if v is not None}
            payload.update(requestId=rid, targetInstance=target)
            ok = await self.b._send_to_rvs({"type": "llm_request", "payload": payload,
                                            "timestamp": int(time.time() * 1000)})
            if not ok:
                return {"ok": False, "errorCode": "send_failed", "error": "RVS: Anfrage nicht gesendet"}
            while True:
                left = deadline - time.time()
                if left <= 0:
                    return {"ok": False, "errorCode": "timeout",
                            "error": f"Keine Antwort von {target} in der Frist"}
                try:
                    return await asyncio.wait_for(asyncio.shield(fut), timeout=min(5.0, left))
                except asyncio.TimeoutError:
                    if not self._is_online(target):
                        return {"ok": False, "errorCode": "box_offline",
                                "error": f"{target} ist waehrend der Anfrage offline gegangen"}
        finally:
            self.inflight[target] = max(0, self.inflight.get(target, 1) - 1)
            self.b._pending_llm.pop(rid, None)

    async def route(self, req: dict, timeout: Optional[float] = None) -> dict:
        """Eine Anfrage an den Cluster. req: model, messages, max_tokens,
        temperature, stop, tools. Antwort: wie llm_response der Box, plus
        attempts; bei Fehler errorCode + error (+ available).
        timeout: Gesamtfrist inkl. Ausfall-Ersatz. ARIA nutzt eine kurze (lokal
        ist ihr Schnellweg — dauert es zu lange, nimmt sie Claude)."""
        model = (req.get("model") or "").strip()
        messages = req.get("messages")
        if not isinstance(messages, list) or not messages:
            return {"ok": False, "errorCode": "bad_request", "error": "messages fehlt oder ist leer"}
        available = [m["model"] for m in self.models()]
        if not model:
            if len(available) == 1:
                model = available[0]       # eindeutig — dann darf es fehlen
            else:
                return {"ok": False, "errorCode": "bad_request", "available": available,
                        "error": "model fehlt. Geladen: " + (", ".join(available) or "nichts")}
        body = {"model": model, "messages": messages,
                "max_tokens": int(req.get("max_tokens") or 512),
                "temperature": float(req["temperature"]) if req.get("temperature") is not None else 0.7,
                "stop": req.get("stop") or None, "tools": req.get("tools") or None}
        limit = timeout or self.timeout
        deadline = time.time() + limit
        tried: list = []
        last: dict = {}
        for _ in range(MAX_ATTEMPTS):
            target = self.pick(model, exclude=tried)
            if not target:
                break
            tried.append(target)
            t0 = time.time()
            last = await self._send_to(target, body, deadline)
            if last.get("ok"):
                last["attempts"] = len(tried)
                return last
            logger.warning("[cluster] %s auf %s: %s (%s)%s", model, target, last.get("errorCode") or "Fehler",
                           (last.get("error") or "")[:120], " — versuche andere Box" if
                           last.get("errorCode") in RETRY_CODES else "")
            if last.get("errorCode") not in RETRY_CODES or time.time() >= deadline:
                break
        if not tried:
            return {"ok": False, "errorCode": "model_unavailable", "available": available,
                    "error": (f"Modell {model!r} ist gerade auf keiner KI-Box im Cluster geladen. "
                              f"Geladen: {', '.join(available) or 'nichts'}.")}
        last.setdefault("errorCode", "model_error")
        last["attempts"] = len(tried)
        if last["errorCode"] in RETRY_CODES and len(tried) >= 1:
            others = [c for c in self.candidates(model) if c["instanceId"] not in tried]
            if not others:
                last["error"] = (last.get("error") or "") + (
                    " — keine weitere Box mit diesem Modell erreichbar.")
        return last

    # ── RVS-Eingang ─────────────────────────────────────────────────────

    async def handle_rvs(self, mtype: str, payload: dict) -> bool:
        if mtype == "llm_cluster_models_request":
            await self.b._send_to_rvs({"type": "llm_cluster_models", "payload": {
                "requestId": payload.get("requestId", ""), "models": self.models()},
                "timestamp": int(time.time() * 1000)})
            return True
        if mtype == "llm_cluster_request":
            asyncio.create_task(self._rvs_request(payload))
            return True
        return False

    async def _rvs_request(self, payload: dict) -> None:
        rid = payload.get("requestId", "")
        res = await self.route(payload)
        out = {"requestId": rid, "ok": bool(res.get("ok"))}
        for k in ("content", "tool_calls", "model", "node", "gpu", "instanceId", "usage",
                  "elapsedMs", "attempts", "error", "errorCode", "available"):
            if res.get(k) is not None:
                out[k] = res[k]
        await self.b._send_to_rvs({"type": "llm_cluster_response", "payload": out,
                                   "timestamp": int(time.time() * 1000)})