retoor
8df5cf7a45
tai.py (8643 lines) is now a thin facade over tai_pkg/ (33 feature modules, one class or feature per file). tai-orig.py preserves the pre-split single-file version byte for byte (sha256 verified against snapshot); import tai and ./tai.py behavior unchanged. Equivalence proof, all run on this tree immediately pre-commit: differential battery tai-orig.py vs facade 3718 checks 0 failures (surface, pure fuzz, tool dispatch, store/audit, release ops, fake-HOME install roundtrip, CLI runs); test_tai.py 261 OK; test_seal.py 16 OK; --version/--help byte-identical; dir(tai) zero removals. Fix included: opencode wire-protocol schema echo leaked into replies; opencode_parse now scans with JSONDecoder.raw_decode capped at 200 scans (13 OpencodeParseTests). Test count 261 = 237 HEAD + 8 pre-existing uncommitted ToolArgumentsTests (now committed with this change, disclosed) + 13 OpencodeParseTests + 3 EmbeddedStringsTests. Patch-target-only test migration (24 hunks), zero assertion changes. tai_pkg clean-regenerated after mid-verification tamper was found and reverted.
247 lines
10 KiB
Python
247 lines
10 KiB
Python
"""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]
|