score_airline.py

#!/usr/bin/env python3
"""Airline RFT reward — ONE self-contained file: reward_info JSON -> a scalar 0..1.

This is the airline twin of score_from_reward_info.py, but DELIBERATELY simpler and
STANDALONE: it imports NOTHING from FIREWORKS_TRAINING/reward_functions/. The whole
reward is the tiny `simple_reward()` function below — no gated partial credit, no
compliance gates, no curriculum CAP. Read reward_info, do simple arithmetic, return
a score.

What is `reward_info`?  It is the dict tau2's evaluator produces for one finished
rollout. The shape you will actually see (printed for you when SHOW_REWARD_INFO=1):

    {
      "reward": 1.0,                      # tau2's own 0/1 verdict for the whole task
      "db_check":   {"db_match": true},   # did the final DB match the gold DB?
      "action_checks": [                  # one entry per REQUIRED gold action
        {"action": {"name": "cancel_reservation",
                    "arguments": {"reservation_id": "EHGLP3"}},
         "action_match": true},
        ...
      ],
      "communicate_checks": [...],        # info the agent had to tell the user
      "reward_basis": ["DB", "COMMUNICATE"],
      "info": {...}
    }

The reward (see `simple_reward`):
  * tau2 already solved it (reward >= 0.999)      -> 1.0
  * otherwise -> the AVERAGE of whichever of these signals are present:
        - db_check.db_match                         (0 or 1)
        - fraction of action_checks that matched    (matched / N)
  * nothing to score on                            -> fall back to tau2's reward

Three entry points (same names as the banking scorer, so the rest of the
pipeline is drop-in):
  simple_reward(reward_info) -> (score, reason)      # the pure reward math
  should_early_stop(tool_calls, max_repeat=8) -> bool  # loop guard
  score_evaluation_row(row) -> EvaluationRow         # EP @evaluation_test adapter

Per repo convention, Claude does NOT run this. To eyeball the math offline:
    python3 FIREWORKS_TRAINING/score_airline.py     # runs the self-test at the bottom
"""
from __future__ import annotations
import os
import sys
import json
import collections

# Pretty-print the raw reward_info JSON for each scored rollout (the user asked to
# "just show me the reward info json"). ON by default; SHOW_REWARD_INFO=0 to silence.
SHOW_REWARD_INFO = os.environ.get("SHOW_REWARD_INFO", "1") == "1"
# How many identical (tool, args) calls before the loop guard fires.
REP_FULL = int(os.environ.get("REP_FULL", "8"))


# ---------------------------------------------------------------------------
# THE REWARD  (this is the whole thing — simple arithmetic on reward_info)
# ---------------------------------------------------------------------------
def simple_reward(reward_info: dict) -> tuple[float, str]:
    """reward_info dict -> (score in [0,1], human-readable reason).

    No gates, no caps, no curriculum — just:
      full tau2 solve          -> 1.0
      else                     -> mean of (db_match, matched_actions / N)
      no signal                -> tau2's own reward (or 0.0)
    """
    old = reward_info.get("reward")
    if isinstance(old, (int, float)) and old >= 0.999:
        return 1.0, "full solve (tau2 reward == 1.0)"

    parts: list[float] = []

    acs = reward_info.get("action_checks") or []
    if acs:
        matched = sum(1 for a in acs if a.get("action_match"))
        parts.append(matched / len(acs))

    db = reward_info.get("db_check") or {}
    if "db_match" in db:
        parts.append(1.0 if db.get("db_match") else 0.0)

    if not parts:
        score = float(old or 0.0)
        return score, f"no action/db checks; using tau2 reward={score:.2f}"

    score = sum(parts) / len(parts)
    matched = sum(1 for a in acs if a.get("action_match"))
    reason = (
        f"partial: actions={matched}/{len(acs)} "
        f"db_match={int(bool(db.get('db_match')))} -> {score:.3f}"
    )
    return score, reason


# ---------------------------------------------------------------------------
# LOOP GUARD — kill degenerate "call the same tool forever" rollouts early
# ---------------------------------------------------------------------------
def should_early_stop(tool_calls: list[tuple[str, dict]], max_repeat: int = REP_FULL) -> bool:
    """True once ANY identical (tool, args) call has been made `max_repeat` times.

    Wire into the rollout processor's per-turn hook: when it returns True, end the
    rollout and score the (short) trajectory instead of looping to 300 messages."""
    if not tool_calls:
        return False
    counts = collections.Counter(
        (n, json.dumps(a, sort_keys=True, default=str)) for n, a in tool_calls
    )
    return max(counts.values()) >= max_repeat


# ===========================================================================
# EP ADAPTER — rebuild tau2 reward_info from the rollout, then score it
# ===========================================================================
# The MCP server returns reward 0 per step (it has no task), so we reconstruct a
# tau2 SimulationRun from the rollout messages, load the real airline Task by id,
# and run tau2's evaluate_simulation. Its RewardInfo is the `reward_info` that
# simple_reward() above consumes.
_DOMAIN = "airline"
# ALL = faithful tau2 reward (env*action*communicate*nl). NL assertions call an LLM
# judge; set AIRLINE_EVAL_TYPE=action (or =env) for a cheap, fully-offline reward.
_EVAL_TYPE = os.environ.get("AIRLINE_EVAL_TYPE", "all")

_TASKS_BY_ID = None  # lazy cache: task_id -> tau2 Task


def _load_task(task_id):
    """Load the airline Task by id (cached). task_split_name=None -> all tasks."""
    global _TASKS_BY_ID
    if _TASKS_BY_ID is None:
        from tau2.domains.airline.environment import get_tasks
        _TASKS_BY_ID = {t.id: t for t in get_tasks(None)}
    return _TASKS_BY_ID.get(task_id)


def _row_meta(row):
    """(task_id, bucket) from the EvaluationRow's input_metadata (set by the
    dataset_adapter in test_airline.py: row_id + dataset_info)."""
    im = getattr(row, "input_metadata", None)
    task_id = getattr(im, "row_id", None)
    di = getattr(im, "dataset_info", None) or {}
    if not isinstance(di, dict):
        di = getattr(di, "__dict__", {}) or {}
    return task_id, di.get("bucket")


def _norm_msg(m) -> dict:
    """EP message (dict or pydantic) -> plain dict with the fields we read."""
    if isinstance(m, dict):
        return m
    return {
        "role": getattr(m, "role", None),
        "content": getattr(m, "content", None),
        "tool_calls": getattr(m, "tool_calls", None),
        "tool_call_id": getattr(m, "tool_call_id", None),
    }


def _norm_args(args) -> dict:
    """Tool-call arguments arrive as a JSON string (OpenAI/EP) or already a dict."""
    if isinstance(args, str):
        try:
            return json.loads(args)
        except Exception:
            return {}
    return args or {}


def _to_tau2_messages(messages):
    """Convert EP/OpenAI-format messages -> tau2 Message objects for the evaluator."""
    from tau2.data_model.message import (
        AssistantMessage, SystemMessage, ToolCall, ToolMessage, UserMessage,
    )
    out = []
    for raw in messages or []:
        m = _norm_msg(raw)
        role, content = m.get("role"), m.get("content")
        if role == "system":
            out.append(SystemMessage(role="system", content=content or ""))
        elif role == "user":
            out.append(UserMessage(role="user", content=content or ""))
        elif role == "assistant":
            tcs = []
            for tc in (m.get("tool_calls") or []):
                if isinstance(tc, dict):
                    fn = tc.get("function") or {}
                    name = fn.get("name") if isinstance(fn, dict) else None
                    raw_args = fn.get("arguments") if isinstance(fn, dict) else None
                    tc_id = tc.get("id", "")
                else:
                    fn = getattr(tc, "function", tc)
                    name = getattr(fn, "name", None)
                    raw_args = getattr(fn, "arguments", None)
                    tc_id = getattr(tc, "id", "")
                tcs.append(ToolCall(id=tc_id or "", name=name or "",
                                    arguments=_norm_args(raw_args), requestor="assistant"))
            out.append(AssistantMessage(role="assistant", content=content,
                                        tool_calls=tcs or None))
        elif role == "tool":
            out.append(ToolMessage(id=m.get("tool_call_id", "") or "", role="tool",
                                   content=content, requestor="assistant"))
    return out


def _extract_reward_info(row) -> dict:
    """Compute tau2 `reward_info` for the finished rollout via evaluate_simulation."""
    import time
    from tau2.data_model.simulation import SimulationRun, TerminationReason
    from tau2.evaluator.evaluator import EvaluationType, evaluate_simulation

    task_id, _ = _row_meta(row)
    task = _load_task(task_id)
    if task is None:
        raise NotImplementedError(
            f"score_evaluation_row: no airline task found for id={task_id!r} — "
            "cannot compute reward_info (check the dataset row id matches a task)."
        )

    messages = _to_tau2_messages(getattr(row, "messages", None) or [])
    now = time.strftime("%Y-%m-%dT%H:%M:%S")
    sim = SimulationRun(
        id=str(getattr(row, "rollout_id", None) or task_id),
        task_id=task_id,
        start_time=now, end_time=now, duration=0.0,
        termination_reason=TerminationReason.AGENT_STOP,
        messages=messages,
    )

    def _run(eval_type):
        # Airline get_environment takes NO retrieval_variant (unlike banking) -> no env_kwargs.
        return evaluate_simulation(
            simulation=sim, task=task, evaluation_type=eval_type,
            solo_mode=False, domain=_DOMAIN,
        )

    try:
        return _run(EvaluationType(_EVAL_TYPE)).model_dump()
    except Exception as e:
        # ENV/DB evaluator replays mutating tool calls and compares responses; our
        # MCP server wraps results, so strict equality can raise. Degrade to ACTION
        # (grades tool choices/args, no env replay) so one row can't redden the suite.
        if EvaluationType(_EVAL_TYPE) == EvaluationType.ACTION:
            raise
        ri = _run(EvaluationType.ACTION).model_dump()
        if isinstance(ri.get("info"), dict):
            ri["info"]["eval_fallback"] = f"{type(e).__name__}: env replay failed -> ACTION only"
        else:
            ri["info"] = {"eval_fallback": f"{type(e).__name__}: env replay failed -> ACTION only"}
        return ri


def _extract_tool_calls(row) -> list[tuple[str, dict]]:
    """The agent's (tool_name, arguments) calls, in order, from the trajectory."""
    calls: list[tuple[str, dict]] = []
    for raw in getattr(row, "messages", None) or []:
        m = _norm_msg(raw)
        for tc in (m.get("tool_calls") or []):
            if isinstance(tc, dict):
                fn = tc.get("function") if isinstance(tc.get("function"), dict) else tc
                calls.append((fn.get("name"), _norm_args(fn.get("arguments"))))
            else:
                fn = getattr(tc, "function", tc)
                calls.append((getattr(fn, "name", None), _norm_args(getattr(fn, "arguments", None))))
    return calls


def score_evaluation_row(row):
    """EP @evaluation_test body: score the finished rollout, attach the result.

    1. rebuild tau2 reward_info from the rollout,
    2. (optionally) PRINT the reward_info JSON so you can see exactly what's scored,
    3. run simple_reward() -> (score, reason),
    4. set row.evaluation_result = EvaluateResult(score, reason)."""
    from eval_protocol.models import EvaluateResult

    reward_info = _extract_reward_info(row)
    task_id, _ = _row_meta(row)

    if SHOW_REWARD_INFO:
        print(f"\n===== reward_info for task {task_id} =====")
        print(json.dumps(reward_info, indent=2, default=str))
        print("==========================================\n")

    score, reason = simple_reward(reward_info)
    row.evaluation_result = EvaluateResult(score=score, reason=reason)
    return row


# ---------------------------------------------------------------------------
# Offline self-test (no EP, no model) — validates the simple reward math
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    full = {"reward": 1.0, "db_check": {"db_match": True},
            "action_checks": [{"action": {"name": "cancel_reservation"}, "action_match": True}]}
    half = {"reward": 0.0, "db_check": {"db_match": False},
            "action_checks": [{"action": {"name": "cancel_reservation"}, "action_match": True}]}
    zero = {"reward": 0.0, "db_check": {"db_match": False},
            "action_checks": [{"action": {"name": "cancel_reservation"}, "action_match": False}]}
    nodb = {"reward": 0.0, "action_checks": [{"action": {"name": "x"}, "action_match": True}]}
    empty = {"reward": 0.0}

    print("full solve            :", simple_reward(full))    # -> 1.0
    print("db miss, action match :", simple_reward(half))    # -> mean(1, 0) = 0.5
    print("all missed            :", simple_reward(zero))    # -> 0.0
    print("action only, no db    :", simple_reward(nodb))    # -> 1.0
    print("no checks at all      :", simple_reward(empty))   # -> tau2 reward 0.0