Upload source.py with huggingface_hub

source.py

@@ -0,0 +1,157 @@
+"""ThirtySpokes KOTH routing agent — cost-aware cascade + optional verify.
+
+Contract required by the measured runtime:
+  build_agent(weights) -> agent
+  agent(prompt, call_model) -> answer
+
+`call_model(model, messages, params)` is the only channel to the owner-pinned pool.
+Answers must come from pool responses (grounding). Always prefer cheaper models and
+escalate only when heuristics say the ask is hard.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+
+# Match the owner's pool-reference generation settings (see docs/MINER.md).
+PARAMS = {"max_tokens": 16384, "reasoning": {"effort": "low"}}
+
+# Signals that LiveCodeBench / hard reasoning usually needs a stronger model.
+_HARD_MARKERS = (
+    "prove",
+    "proof",
+    "leetcode",
+    "write a function",
+    "write a program",
+    "implement",
+    "algorithm",
+    "complexity",
+    "def ",
+    "class ",
+    "```",
+    "stdin",
+    "stdout",
+    "test case",
+    "constraints",
+    "hard",
+)
+
+_CODE_MARKERS = (
+    "python",
+    "code",
+    "function",
+    "program",
+    "algorithm",
+    "compile",
+    "runtime error",
+)
+
+_WEAK_ANSWER = re.compile(
+    r"^\s*(i co(?:uld|n't)|sorry|as an ai|cannot|can't help|not sure|n/?a)\b",
+    re.I,
+)
+
+
+def _load_cfg(weights: bytes) -> dict:
+    try:
+        cfg = json.loads(weights.decode() if isinstance(weights, (bytes, bytearray)) else weights)
+    except Exception:
+        cfg = {}
+    if not isinstance(cfg, dict):
+        cfg = {}
+    # Defaults align with MockPool model ids for local smoke; override via weights.bin for mainnet.
+    cheap = cfg.get("cheap") or cfg.get("model") or "cheap"
+    mid = cfg.get("mid") or cheap
+    strong = cfg.get("strong") or mid
+    return {
+        "cheap": cheap,
+        "mid": mid,
+        "strong": strong,
+        "escalate_chars": int(cfg.get("escalate_chars", 400)),
+        "verify": bool(cfg.get("verify", True)),
+        "verify_on_code": bool(cfg.get("verify_on_code", True)),
+        "max_calls": int(cfg.get("max_calls", 3)),
+    }
+
+
+def _looks_hard(prompt: str, escalate_chars: int) -> bool:
+    p = prompt or ""
+    pl = p.lower()
+    if len(p) >= escalate_chars:
+        return True
+    if any(m in pl for m in _HARD_MARKERS):
+        return True
+    # Many coding tasks include multi-line specs / sample IO.
+    if pl.count("\n") >= 8 and any(m in pl for m in _CODE_MARKERS):
+        return True
+    return False
+
+
+def _looks_code(prompt: str) -> bool:
+    pl = (prompt or "").lower()
+    return any(m in pl for m in _CODE_MARKERS) or "```" in pl or "def " in pl
+
+
+def _is_weak(answer: str) -> bool:
+    if not answer or not str(answer).strip():
+        return True
+    text = str(answer).strip()
+    if len(text) < 8:
+        return True
+    if _WEAK_ANSWER.search(text):
+        return True
+    return False
+
+
+def _call(call_model, model: str, prompt: str) -> str:
+    return call_model(model, [{"role": "user", "content": prompt}], PARAMS)
+
+
+def build_agent(weights):
+    cfg = _load_cfg(weights)
+    cheap, mid, strong = cfg["cheap"], cfg["mid"], cfg["strong"]
+    escalate_chars = cfg["escalate_chars"]
+    do_verify = cfg["verify"]
+    verify_on_code = cfg["verify_on_code"]
+    max_calls = max(1, cfg["max_calls"])
+
+    def agent(prompt, call_model):
+        calls = 0
+        hard = _looks_hard(prompt, escalate_chars)
+        codey = _looks_code(prompt)
+
+        # 1) Cheap first on easy asks; mid/strong first on hard asks to avoid wasted cheap fails.
+        first = mid if hard else cheap
+        ans = _call(call_model, first, prompt)
+        calls += 1
+        chosen = first
+
+        # 2) Escalate if the cheap/mid answer looks weak or the ask is hard and we started cheap.
+        if calls < max_calls and (_is_weak(ans) or (hard and chosen == cheap)):
+            nxt = strong if chosen != strong else mid
+            if nxt != chosen:
+                ans = _call(call_model, nxt, prompt)
+                calls += 1
+                chosen = nxt
+
+        # 3) Optional verify: on code/hard tasks, ask a different model and keep the longer
+        # grounded pool response when the first looks thin. Still returns pool text only.
+        if (
+            do_verify
+            and calls < max_calls
+            and (hard or (verify_on_code and codey))
+            and chosen != strong
+        ):
+            alt = _call(call_model, strong, prompt)
+            calls += 1
+            if _is_weak(ans) or (not _is_weak(alt) and len(str(alt)) > len(str(ans)) * 1.15):
+                ans = alt
+
+        # Final grounding fallback: never return empty — one more strong call if needed.
+        if _is_weak(ans) and calls < max_calls and chosen != strong:
+            ans = _call(call_model, strong, prompt)
+
+        return ans
+
+    return agent