common.py

"""Shared config for Shaktimaan: the agent pool, prompts, capabilities and graders."""
import json, os, re

# agent -> served model id, max concurrent requests we send, relative cost per token
POOL = {
    "gemma-3-1b":         dict(model="google/gemma-3-1b-it", concurrency=32, cost=0.02),
    "kimi-vl":            dict(model="Kimi-VL-A3B-Thinking", concurrency=64, cost=0.05),
    "nemotron-lightning": dict(model="nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", concurrency=64, cost=0.10),
    "glm-5.3-flash":      dict(model="zai-org/GLM-5.3-Flash", concurrency=32, cost=0.20),
    "qwen3.8-27b":        dict(model="Qwen/Qwen3.8-27B", concurrency=12, cost=0.50),
    "nemotron-ultra":     dict(model="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", concurrency=64, cost=1.00),
}
AGENTS = list(POOL)  # ordered cheapest -> most expensive
STRONGEST = "nemotron-ultra"
CONTEXT = {"gemma-3-1b": 32768, "kimi-vl": 8192, "nemotron-lightning": 131072, "glm-5.3-flash": 32768,
           "qwen3.8-27b": 32768, "nemotron-ultra": 262144}
VISION = {"kimi-vl"}  # only agent served with image input enabled
MAX_OUT = 8192


def endpoints() -> dict:
    """Agent -> OpenAI-compatible base URL (".../v1"), from $SHAKTIMAAN_ENDPOINTS: inline JSON or a path to a JSON file."""
    raw = os.environ["SHAKTIMAAN_ENDPOINTS"]
    return json.load(open(raw)) if os.path.exists(raw) else json.loads(raw)

SOLVER_SYSTEM = ("Solve the task. Be concise. End your response with a final line exactly of the form:\n"
                 "Final answer: <answer>\nFor multiple-choice questions, the answer is the option letter only.")
CODE_SYSTEM = ("Write a correct, self-contained Python solution. Put the complete code (imports, required function "
               "names and signatures) in a single ```python code block. Do not include tests or example usage.")

ROUTER_SYSTEM = (
    "You are Shaktimaan, an orchestrator. You do not answer the task yourself. You read the task and decide which "
    "agent(s) from the pool should handle it, minimising cost while maximising the chance of a correct answer.\n"
    "Agents (cheapest first): " + ", ".join(AGENTS) + ".\n"
    "Capabilities: only kimi-vl can see images; context limits: nemotron-ultra 262k tokens, nemotron-lightning 131k, "
    "kimi-vl 8k, all others 32k.\n"
    "Strategies: single = one agent answers; ensemble = several agents answer independently and the last one "
    "aggregates.\n"
    "Reply with JSON only: {\"domain\": ..., \"difficulty\": \"easy|medium|hard\", \"strategy\": \"single|ensemble\", "
    "\"agents\": [...]}"
)


def est_tokens(text: str) -> int:
    """Conservative token estimate (real tokenizers average ~3.5-4 chars/token on English)."""
    return len(text) // 3 + 1


def max_out(agent: str, prompt_tokens: int) -> int:
    return min(MAX_OUT, CONTEXT[agent] - prompt_tokens - 256)


def capable(agent: str, prompt: str, has_image: bool) -> bool:
    if has_image and agent not in VISION:
        return False
    return max_out(agent, est_tokens(prompt) + 1500 * has_image) >= 1024


def truncate(prompt: str, head: int = 5000, tail: int = 1500) -> str:
    return prompt if len(prompt) <= head + tail + 7 else prompt[:head] + "\n[...]\n" + prompt[-tail:]


def router_view(prompt: str, has_image: bool = False, n_tokens: int | None = None) -> str:
    """What the router reads: a size/attachment header plus the start and end of the task."""
    n = est_tokens(prompt) if n_tokens is None else n_tokens
    return f"[input ~{n // 1000}k tokens{'; 1 image attached' if has_image else ''}]\n" + truncate(prompt)


def strip_think(text: str) -> str:
    text = text or ""
    if "</think>" in text:
        text = text.rsplit("</think>", 1)[1]
    text = re.sub(r"◁think▷.*?◁/think▷", "", text, flags=re.S)  # Kimi-VL thinking markers
    return text.strip()


def extract_final(text: str) -> str:
    text = strip_think(text)
    m = re.findall(r"final answer\s*[:：]\s*(.+)", text, flags=re.I)
    if m:
        return m[-1].strip()
    m = re.findall(r"^\s*answer\s*[:：]\s*(.+)", text, flags=re.I | re.M)
    if m:
        return m[-1].strip()
    boxed = last_boxed(text)
    if boxed is not None:
        return boxed
    lines = [l for l in text.splitlines() if l.strip()]
    return lines[-1].strip() if lines else ""


def extract_code(text: str) -> str:
    text = strip_think(text)
    blocks = re.findall(r"```(?:python|py)?\s*\n(.*?)```", text, flags=re.S)
    return blocks[-1] if blocks else ""


def last_boxed(s: str):
    i = s.rfind("\\boxed{")
    if i < 0:
        return None
    i += len("\\boxed{")
    depth, j = 1, i
    while j < len(s) and depth:
        depth += {"{": 1, "}": -1}.get(s[j], 0)
        j += 1
    return s[i:j - 1] if depth == 0 else None


def _norm_math(s: str) -> str:
    s = s.strip().strip("$").strip()
    b = last_boxed(s)
    if b is not None:
        s = b
    for a, b2 in [("\\left", ""), ("\\right", ""), ("\\!", ""), ("\\,", ""), ("\\;", ""), ("dfrac", "frac"),
                  ("tfrac", "frac"), ("^\\circ", ""), ("^{\\circ}", ""), ("\\%", ""), ("%", ""), ("\\text{", "{"),
                  ("\\mbox{", "{"), (" ", ""), ("\\$", "")]:
        s = s.replace(a, b2)
    s = re.sub(r"\\frac(\d)(\d)", r"\\frac{\1}{\2}", s)
    s = s.rstrip(".")
    if s.startswith("{") and s.endswith("}") and s.count("{") == 1:
        s = s[1:-1]
    return s


def _num(s: str):
    s = s.replace(",", "").replace("$", "").strip()
    m = re.findall(r"-?\d+(?:\.\d+)?", s)
    return float(m[-1]) if m else None


def grade(kind: str, pred: str, gold: str) -> bool:
    """Grades extracted answers. Code is graded by execution (see gen.run_tests), not here."""
    if not pred:
        return False
    if kind == "mc":
        m = re.search(r"\b([A-J])\b", pred.upper().replace("(", " ").replace(")", " "))
        return bool(m) and m.group(1) == gold.upper()
    if kind == "num":
        p, g = _num(pred), _num(gold)
        return p is not None and g is not None and abs(p - g) < 1e-6 * max(1, abs(g))
    if kind == "relaxed":  # ChartQA convention: numbers within 5%, otherwise case-insensitive exact match
        pred = last_boxed(pred) or pred
        pred = re.sub(r"\\(?:mathbf|textbf|text|mathrm)\{([^}]*)\}", r"\1", pred).replace("\\(", "").replace("\\)", "")
        p, g = _num(pred), _num(gold)
        if g is not None and p is not None and re.fullmatch(r"[-\d.,%$ ]+", gold.strip()):
            return abs(p - g) <= 0.05 * abs(g) if g else p == 0
        return pred.strip().strip(".").lower() == gold.strip().lower()
    if kind == "math":
        p, g = _norm_math(pred), _norm_math(gold)
        if p == g:
            return True
        try:
            return abs(float(p) - float(g)) < 1e-6
        except ValueError:
            return False
    raise ValueError(kind)