"""Model health tracking, scoring, and roster probing.""" import json import os import time import urllib.error import urllib.request from tai_pkg.claude import Claude from tai_pkg.codex import Codex from tai_pkg.consts import Consts from tai_pkg.gemini import Gemini from tai_pkg.grok import Grok from tai_pkg.muse import Muse from tai_pkg.opencode import Opencode from tai_pkg.text import Text __all__ = ["ModelHealth", "ModelScoring"] class ModelHealth: def __init__(self, home): self.path = os.path.join(home, Consts.MODEL_HEALTH_FILE) if home else "" self.data = {"models": {}, "lists": {}} if self.path: try: with open(self.path, "r", encoding="utf-8") as handle: raw = json.load(handle) if isinstance(raw, dict): if isinstance(raw.get("models"), dict): self.data["models"] = raw["models"] if isinstance(raw.get("lists"), dict): self.data["lists"] = raw["lists"] except (OSError, ValueError): pass def _save(self): if not self.path: return try: tmp = self.path + ".tmp" with open(tmp, "w", encoding="utf-8") as handle: json.dump(self.data, handle) os.replace(tmp, self.path) except OSError: pass def entry(self, provider, model): return self.data["models"].get("%s|%s" % (provider or "default", model or "")) or {} def weight(self, provider, model): return ModelScoring.model_weight(self.entry(provider, model)) def circuit_open(self, provider, model): return (self.entry(provider, model).get("until") or 0.0) > time.time() def record(self, provider, model, success, latency_ms=None, tokens_per_second=None): if not model: return key = "%s|%s" % (provider or "default", model or "") entry = self.data["models"].setdefault(key, {"s": 0, "f": 0, "r": 0.0, "lat": 0.0, "tps": 0.0, "cf": 0, "until": 0.0, "used": 0.0, "n": 0}) entry["n"] += 1 entry["used"] = time.time() if success: reward = Consts.MODEL_NEUTRAL_REWARD if latency_ms is None and tokens_per_second is None else ModelScoring.model_outcome_reward(latency_ms, tokens_per_second) entry["s"] += 1 entry["cf"] = 0 entry["until"] = 0.0 entry["r"] += reward if latency_ms is not None: entry["lat"] += latency_ms if tokens_per_second is not None: entry["tps"] += tokens_per_second else: entry["f"] += 1 entry["cf"] += 1 if entry["cf"] >= Consts.MODEL_CIRCUIT_THRESHOLD: entry["until"] = time.time() + Consts.MODEL_CIRCUIT_COOLDOWN self._save() def cached_list(self, name, ttl=Consts.OPENROUTER_FREE_CACHE_TTL): row = self.data["lists"].get(name) if not isinstance(row, dict): return None models = row.get("models") if not models or (time.time() - (row.get("at") or 0)) > ttl: return None return [str(item) for item in models if str(item).strip()] def store_list(self, name, models): cleaned = [str(item).strip() for item in models if str(item).strip()] if not cleaned: return self.data["lists"][name] = {"at": time.time(), "models": cleaned} self._save() def snapshot(self): rows = [] for key, entry in self.data["models"].items(): provider, _, model = key.partition("|") ok = entry.get("s") or 0 rows.append({ "provider": provider, "model": model, "weight": round(ModelScoring.model_weight(entry), 4), "ok": ok, "fail": entry.get("f") or 0, "avg_ms": round((entry.get("lat") or 0.0) / ok) if ok and entry.get("lat") else None, "avg_tps": round((entry.get("tps") or 0.0) / ok, 1) if ok and entry.get("tps") else None, "circuit_open": (entry.get("until") or 0.0) > time.time(), "used": entry.get("used") or 0.0, }) rows.sort(key=lambda row: (-row["weight"], row["provider"], row["model"])) return rows class ModelScoring: """ModelScoring.""" @staticmethod def model_speed_reward(tokens_per_second): if not tokens_per_second or tokens_per_second <= 0: return Consts.MODEL_NEUTRAL_REWARD return tokens_per_second / (tokens_per_second + Consts.MODEL_SPEED_REF_TPS) @staticmethod def model_latency_reward(latency_ms): if not latency_ms or latency_ms <= 0: return 1.0 return 1.0 / (1.0 + latency_ms / Consts.MODEL_LATENCY_REF_MS) @staticmethod def model_outcome_reward(latency_ms=None, tokens_per_second=None): return ModelScoring.model_speed_reward(tokens_per_second) * ModelScoring.model_latency_reward(latency_ms) @staticmethod def model_weight(entry): entry = entry or {} attempts = (entry.get("s") or 0) + (entry.get("f") or 0) return ((entry.get("r") or 0.0) + Consts.MODEL_NEUTRAL_REWARD * Consts.MODEL_WEIGHT_PRIOR) / (attempts + Consts.MODEL_WEIGHT_PRIOR) @staticmethod def fetch_model_ids(models_url, key="", extra_headers=None, timeout=Consts.MODEL_LIST_TIMEOUT): headers = dict(extra_headers or {}) if key: headers["Authorization"] = "Bearer " + key try: request = urllib.request.Request(models_url, headers=headers, method="GET") with urllib.request.urlopen(request, timeout=timeout) as response: payload = json.loads(response.read().decode("utf-8", "replace")) except (OSError, ValueError): return None rows = payload.get("data") if isinstance(payload, dict) else None if not isinstance(rows, list): return None ids = [str(row.get("id")).strip() for row in rows if isinstance(row, dict) and str(row.get("id") or "").strip()] return ids or None @staticmethod def embedded_failure_message(data): if not isinstance(data, dict): return "" error = data.get("error") if isinstance(error, dict): message = error.get("message") or error.get("type") or "" return str(message).strip() if isinstance(error, str) and error.strip(): return error.strip() if data.get("type") == "error": return str(data.get("message") or "upstream error").strip() return "" @staticmethod def prefer_muse(models): return sorted(models, key=lambda item: 0 if Consts.MUSE_LABEL in item.lower() else 1) @staticmethod def openrouter_free_tier(cand): model = (cand.get("model") or "").lower() muse = 0 if Consts.MUSE_LABEL in model else 1 if cand.get("provider") == Consts.OPENCODE_LABEL: return muse if cand.get("provider") == Consts.DEVPLACE_LABEL: return 2 + muse if cand.get("provider") == Consts.OPENROUTER_LABEL and Text.is_free_model(model): return 4 + muse return 6 @staticmethod def models_report(chat): lines = ["label/model | weight ok/fail avg-ms avg-tps status"] seen = set() for cand in chat.ordered_candidates(): key = (cand["provider"], cand["model"]) if key in seen: continue seen.add(key) entry = chat.health.entry(cand["provider"], cand["model"]) ok = entry.get("s") or 0 fail = entry.get("f") or 0 avg_ms = str(round((entry.get("lat") or 0.0) / ok)) if ok and entry.get("lat") else "-" avg_tps = str(round((entry.get("tps") or 0.0) / ok, 1)) if ok and entry.get("tps") else "-" if cand["chat_url"] == Consts.MUSE_CHAT_URL: status = Muse.muse_health() elif chat.health.circuit_open(cand["provider"], cand["model"]): status = "circuit open" elif not ok and not fail: status = "untested" else: status = "ready" lines.append("%s/%s | %.4f %d/%d %s %s %s" % (cand["provider"], cand["model"], ModelScoring.model_weight(entry), ok, fail, avg_ms, avg_tps, status)) return "\n".join(lines) @staticmethod def probe_backend(models_url, key): if models_url == Consts.MUSE_MODELS_URL: return Muse.muse_health() if models_url == Consts.GROK_MODELS_URL: return Grok.grok_health() if models_url == Consts.CLAUDE_MODELS_URL: return Claude.claude_health() if models_url == Consts.CODEX_MODELS_URL: return Codex.codex_health() if models_url == Consts.GEMINI_MODELS_URL: return Gemini.gemini_health() if models_url == Consts.OPENCODE_MODELS_URL: return "ok, %s" % Opencode.opencode_binary() if Opencode.opencode_binary() else "opencode CLI not found" request = urllib.request.Request(models_url, headers={"Authorization": "Bearer " + key}) started = time.time() try: with urllib.request.urlopen(request, timeout=6) as response: try: payload = json.loads(response.read().decode("utf-8", "replace")) except ValueError: return "bad response" data = payload.get("data", []) if isinstance(payload, dict) else [] count = len(data) if isinstance(data, list) else 0 return "ok, %d models, %dms" % (count, (time.time() - started) * 1000) except urllib.error.HTTPError as exc: if exc.code == 401: return "needs key" return "http %d" % exc.code except OSError as exc: return "fail: %s" % Text.short_error(exc)[:60] except Exception as exc: return "fail: %s" % Text.short_error(exc)[:60]