Split monolithic tai.py into tai_pkg package, keep tai-orig.py as original
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.
This commit is contained in:
@@ -0,0 +1,246 @@
|
||||
"""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]
|
||||
Reference in New Issue
Block a user