Route Between a Small and a Large Model with Jev

"""Route Tallyboard support requests between a small and a large model, with Jev deciding.

The small model (GLM-5.3-flash) answers most billing questions correctly for a
fourteenth of the price. The large one (GLM-5.3) gets more of the hard ones
right. Jev reads each request and the policy, and your code picks the model.
"""
import json
import os
import re
from concurrent.futures import ThreadPoolExecutor

import jev
import llm

HERE = os.path.dirname(os.path.abspath(__file__))
SMALL = "z-ai/glm-5.3-flash"
LARGE = "z-ai/glm-5.3"

with open(os.path.join(HERE, "policy.md"), encoding="utf-8") as _f:
    POLICY = _f.read()

SYSTEM = (
    "You answer customer questions for Tallyboard using only the policy below. "
    "Work it out briefly, then end with one last line in the form ANSWER: <value>. "
    "For amounts give a plain number in dollars without the $ sign, rounded to 2 decimals. "
    "For yes/no questions answer yes or no. For dates answer like May 10.\n\n" + POLICY
)


def load_jsonl(path):
    with open(os.path.join(HERE, path), encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def load_outcomes(path="outcomes.jsonl"):
    """{id: {"small": {"ok", "cost", "ms"}, "large": {...}}}: each model's recorded
    share of correct answers over three runs, mean cost in dollars and median latency."""
    return {row["id"]: row for row in load_jsonl(path)}


def grade(got, want, kind):
    """True when the extracted answer matches the expected one."""
    if got is None:
        return False
    g = got.lower()
    if kind == "number":
        nums = re.findall(r"-?\d[\d,]*\.?\d*", g)
        if not nums:
            return want == 0 and ("no" in g or "none" in g)
        value = float(nums[0].replace(",", ""))
        return abs(value - want) <= max(0.02, abs(want) * 0.001)
    if kind == "yesno":
        return re.match(r"\W*" + want + r"\b", g) is not None
    if kind == "date":
        return want.lower() in g
    if want == "not allowed":
        return any(w in g for w in ("not allowed", "not possible", "cannot", "can't", "not"))
    if want == "Team and Business":
        return "team" in g and "business" in g and "starter" not in g
    return want.lower() in g


# ---------------------------------------------------------------- Step 1
def extract(text):
    """The value on the last "ANSWER:" line; without one, the last non-empty line; None for no text."""
    # TODO (Step 1): the text after the LAST "ANSWER:" line (strip spaces and * around it).
    # No ANSWER line: the last non-empty line. No text at all: None.
    raise NotImplementedError("Step 1: write extract()")


def answer(model, request, chat=llm.chat):
    """Ask `model` the request with the policy as the system prompt."""
    # TODO (Step 1): one chat(model, messages, reasoning_effort=...) call
    #   messages: a "system" message with SYSTEM, then a "user" message with the request
    #   reasoning_effort: "low" for SMALL, None for LARGE
    # return {"model", "answer": extract(reply["text"]), "text", "cost", "ms"}
    raise NotImplementedError("Step 1: write answer()")


# ---------------------------------------------------------------- Step 2
QUESTIONS = {
    # "combine" and "topic" arrive in Step 2
}


def assess(request, post=jev.post):
    """One Jev request about a support request. Returns {"combine", "topic", "cost"}."""
    raise NotImplementedError("assess() arrives in Step 2")


# ---------------------------------------------------------------- Step 3
def route(a, threshold=0.5):
    """LARGE when Jev thinks the answer combines rules, SMALL otherwise."""
    raise NotImplementedError("route() arrives in Step 3")


def assess_all(requests, post=jev.post):
    """{id: assessment} for every request, sixteen Jev requests at a time."""
    raise NotImplementedError("assess_all() arrives in Step 3")


def score(requests, outcomes, choose, assessments=None):
    """Replay the recorded outcomes for the model `choose(request, assessment)` picks.

    `assessments` is None for a router that does not ask Jev (then choose gets None
    and no Jev cost is added)."""
    raise NotImplementedError("score() arrives in Step 3")


# ---------------------------------------------------------------- Step 4
THRESHOLDS = [0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]


def sweep(requests, outcomes, assessments, thresholds=THRESHOLDS):
    """[(threshold, score report)] for the Jev router at each threshold."""
    raise NotImplementedError("sweep() arrives in Step 4")


def pick_threshold(rows, target_accuracy):
    """The cheapest threshold whose accuracy reaches the target, or None when none does.
    On a cost tie the higher threshold wins: same result here, fewer large calls elsewhere."""
    raise NotImplementedError("pick_threshold() arrives in Step 4")


# ---------------------------------------------------------------- Step 5
def choose_model(request, threshold, post=jev.post):
    """(model, reason, assessment). reason is "jev", or "jev_error" when Jev failed and the
    request goes to LARGE so a Jev outage costs money instead of answers."""
    raise NotImplementedError("choose_model() arrives in Step 5")


def handle(request, threshold, post=jev.post, chat=llm.chat):
    """Answer one request the way production would. When the small model gives no usable
    answer, ask the large one. Returns {"model", "answer", "reason", "escalated", "cost"}."""
    raise NotImplementedError("handle() arrives in Step 5")