A reward is a score you can defend

import json
import re

import whileai as wai
from whileai.simulations import attach_labels


@wai.tool
def get_order(order_id: str) -> dict:
    """Look up an order by id."""
    ...


data = wai.simulate(
    wai.seeded_agent([get_order]),
    tools=[get_order],
    system_prompt="Help customers with orders.",
    simulator=False,
    mode="rl",
    repeats=4,
    repeat_policy="fixed",
    budget=64,
    seed=0,
)

ORDER_ID = re.compile(r"\bORD-\d+\b")
SUCCESS = {"ok", "created", "success", "done", "updated", "deleted"}
CLAIMED = re.compile(r"\b(done|went through|completed|succeeded|confirmed)\b", re.I)


# 1. A judge is any function from a row to a reward. This one is a program.
def judge(row):
    """1 when the agent looked the order up, led with it, and claimed nothing the tool did not say."""
    messages = row["messages"]
    asked = set(ORDER_ID.findall(next(m["content"] for m in messages if m["role"] == "user")))
    calls = [c for m in messages if m["role"] == "assistant" for c in m.get("tool_calls") or []]
    if not asked:  # no id to look up: ask for one, do not guess
        return {"reward": int(not calls)}
    call = (calls[0].get("function") or calls[0]) if calls else {}
    args = call.get("arguments") or {}
    args = json.loads(args) if isinstance(args, str) else args
    if call.get("name") != "get_order" or args.get("order_id") not in asked:
        return {"reward": 0}
    reply = messages[-1]["content"] if messages[-1]["role"] == "assistant" else ""
    if args["order_id"] not in re.split(r"(?<=[.!?])\s", reply.strip())[0]:
        return {"reward": 0}  # the first sentence is the result, not a preamble
    result = next((m["content"] for m in messages if m["role"] == "tool"), "{}")
    status = json.loads(result).get("status")
    return {"reward": int(status in SUCCESS or not CLAIMED.search(reply))}


scored = data.grade(judge=judge)
print(scored.pass_at)


# 2. A looser definition of the same job: the reply names the order, anywhere.
@wai.verifier
def names_the_order(candidate, reference, row):
    ids = [step["arguments"].get("order_id", "") for step in row["steps"]]
    return float(any(i and i in candidate for i in ids))


print(data.grade(judge=names_the_order).pass_at)

# 3. Check the judge against labels. Here the labels come from the stand-in's
#    answer key (its `seeded` field); in real life a person writes them, 50 to 200.
labels = [
    {"scenario_id": r["scenario_id"], "rollout_index": r["rollout_index"], "label": int(not r["seeded"])}
    for r in scored.rows[:40]
]
labeled, _ = attach_labels(scored.rows, labels, kind="human")
trust = wai.judge_trust(labeled, judge)
print("trusted:", trust["ok"])
print("agreement:", trust["agreement"]["agreement"], "kappa:", trust["agreement"]["kappa"])

pass@1 0.58 [0.38..0.77] | pass^4 (pass_pow_k) 0.31 [0.12..0.56] | pass@4 0.75 [0.50..0.94] | headroom 0.17 (16 groups, k=4)
pass@1 0.75 [0.50..0.94] | pass^4 (pass_pow_k) 0.75 [0.50..0.94] | pass@4 0.75 [0.50..0.94] | headroom 0.00 (16 groups, k=4)
trusted: False
agreement: 0.825 kappa: 0.6111