test_airline.py

#!/usr/bin/env python3
"""Eval Protocol entrypoint for τ²-bench AIRLINE RFT.

Airline twin of test_banking.py. EP discovers the `@evaluation_test`-decorated
function below and from it knows how to roll out one airline task and score it:

    dataset row ──► policy model takes a turn ──► MCP server runs an airline tool
                  ▲                                          │
                  └────────────  (repeat until stop)  ◄──────┘
                                        │
                                        ▼
                              this file's reward fn → one scalar 0..1

Three things this file does:
  1. dataset_adapter : JSONL row -> EvaluationRow            (airline_to_evaluation_row)
  2. @evaluation_test : declare HOW to roll out (model, MCP server, parallelism, …)
  3. reward          : score the FINISHED rollout (delegated to score_airline.py)

Heavy lifting lives elsewhere:
  - ENVIRONMENT + TOOLS -> FIREWORKS_TRAINING/airline_mcp/server.py
  - REWARD MATH         -> FIREWORKS_TRAINING/score_airline.py   (ONE self-contained file)

Run order (see FIREWORKS_TRAINING/RUNBOOK.md):
    bash FIREWORKS_TRAINING/run_airline_model.sh          # local eval, no training
    bash FIREWORKS_TRAINING/training_airline.sh --go      # launch RFT on Fireworks
"""
import json   # noqa: F401  (parity with test_banking; harmless if unused)
import os
import socket  # probe for a free MCP server port (see MCP_PORT below)
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, HERE)
# The reward: ONE self-contained file, no reward_functions/ dependency.
from score_airline import score_evaluation_row  # noqa: E402

from eval_protocol.models import EvaluationRow, Message, InputMetadata    # noqa: E402
from eval_protocol.pytest import evaluation_test, MCPGymRolloutProcessor  # noqa: E402

# ===========================================================================
# WHICH MODEL PLAYS THE AGENT (the "policy")
# ===========================================================================
# Default policy for LOCAL TEST = gpt-oss-20b on serverless (no deployment/cold
# start). For real training, training_airline.sh `--base-model` selects what's tuned.
POLICY_MODEL = os.environ.get(
    "POLICY_MODEL",
    "fireworks_ai/accounts/fireworks/models/gpt-oss-20b",
)

# Policy sampling knobs (env-overridable). TEMPERATURE=0 + SEED=<int> for debug.
TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.8"))
SEED = os.environ.get("SEED")  # str | None

# Which JSONL of airline tasks to roll out. Default = airline_easy (the user asked
# to start with Easy only). Override with EP_DATASET to point elsewhere.
DATASET = os.environ.get("EP_DATASET", "FIREWORKS_TRAINING/data/airline_easy.jsonl")

# Turn the multi-turn SIMULATED USER on/off (default ON = faithful tau2 task).
ENABLE_USER_SIM = os.environ.get("ENABLE_USER_SIM", "1") == "1"
# Which model plays the simulated CUSTOMER (falls back to the row's recommendation).
USER_SIM_MODEL = os.environ.get(
    "USER_SIM_MODEL", "fireworks_ai/accounts/fireworks/models/gpt-oss-120b"
)


# Optional per-difficulty SUBSET (default: use the WHOLE dataset). Cap how many
# tasks of each bucket get rolled out via MAX_<BUCKET> env vars, e.g.
#     MAX_EASY=5   -> a 5-task smoke run.
def _bucket_caps() -> dict:
    caps = {}
    for name, value in os.environ.items():
        if name.startswith("MAX_") and value.strip().isdigit():
            caps[name[4:].upper()] = int(value)   # MAX_EASY -> "EASY" -> 5
    return caps


def _subset_by_bucket(rows: list[dict]) -> list[dict]:
    caps = _bucket_caps()
    if not caps:
        return rows
    seen: dict[str, int] = {}
    kept: list[dict] = []
    for r in rows:
        key = str(r.get("bucket") or "").strip().upper().replace(" ", "_")
        cap = caps.get(key)
        if cap is None:
            kept.append(r)                      # no cap for this bucket -> keep all
        elif seen.get(key, 0) < cap:
            kept.append(r)
            seen[key] = seen.get(key, 0) + 1
    print(f"[test_airline] bucket subset {caps} -> {len(kept)}/{len(rows)} tasks")
    return kept


# ===========================================================================
# OPTIONAL: route INFERENCE through the LeanMCP gateway (for logging)
# ===========================================================================
LEANMCP_GATEWAY_BASE = "https://aigateway.leanmcp.com/v1/fireworks"
USE_LEANMCP_GATEWAY = os.environ.get("USE_LEANMCP_GATEWAY", "0") == "1"


def _completion_params() -> list[dict]:
    """LiteLLM call settings for the AGENT (policy) model. EP passes this straight
    to LiteLLM for every agent turn. Returns a LIST (we use exactly one)."""
    p = {
        "model": POLICY_MODEL,
        "temperature": TEMPERATURE,   # 0.8 = exploration -> in-group variance -> gradient
        "max_tokens": 4096,           # tight cap; small models ramble
    }
    if SEED not in (None, ""):
        p["seed"] = int(SEED)
    if USE_LEANMCP_GATEWAY:
        key = os.environ.get("LEANMCP_API_KEY")
        if not key:
            raise RuntimeError(
                "USE_LEANMCP_GATEWAY=1 but LEANMCP_API_KEY not in the environment. "
                "Export it, or set USE_LEANMCP_GATEWAY=0 to call Fireworks directly."
            )
        p["api_base"] = LEANMCP_GATEWAY_BASE
        p["api_key"] = key
    return [p]


def airline_to_evaluation_row(rows: list[dict]) -> list[EvaluationRow]:
    """ADAPTER (step 1): our JSONL rows -> EP EvaluationRow objects.

    Per task we set:
      • messages: the SEED conversation (the user-prompt/persona template).
      • input_metadata: rides along through the rollout so the reward fn can read
        back the task id (row_id) + dataset_info (bucket, user_simulation, criteria).
    """
    rows = _subset_by_bucket(rows)   # optional MAX_EASY caps (no-op by default)
    out: list[EvaluationRow] = []
    for r in rows:
        us = r.get("user_simulation") or {}
        # The simulated user's LLM call (vendor.tau2 UserSimulator -> generate() ->
        # litellm) does NOT go through _completion_params(); it uses these llm_args
        # directly (manager.py passes them to UserSimulator, which spreads them into
        # generate(**llm_args) -> acompletion(**kwargs)). So when gateway routing is
        # on we must inject api_base+api_key HERE too, or the customer turns would
        # bypass LeanMCP and hit Fireworks directly. This mirrors the agent (policy)
        # injection in _completion_params(): with the gateway ON, BOTH sides go via LeanMCP.
        user_llm_args = {"temperature": us.get("temperature", 0.0)}
        if USE_LEANMCP_GATEWAY:
            key = os.environ.get("LEANMCP_API_KEY")
            if not key:
                raise RuntimeError(
                    "USE_LEANMCP_GATEWAY=1 but LEANMCP_API_KEY not in the environment "
                    "(needed to route the simulated user through the gateway)."
                )
            user_llm_args["api_base"] = LEANMCP_GATEWAY_BASE  # route customer turns via LeanMCP
            user_llm_args["api_key"] = key                    # LeanMCP key, NOT Fireworks
        user_simulation = {
            "enabled": ENABLE_USER_SIM,
            "system_prompt": us.get("instructions", ""),          # the customer persona
            "llm": us.get("recommended_user_model") or USER_SIM_MODEL,
            "llm_args": user_llm_args,
        }
        out.append(
            EvaluationRow(
                messages=[Message(role="system", content=r.get("user_prompt_template", ""))],
                input_metadata=InputMetadata(
                    row_id=r["id"],   # task id -> reward loads the matching tau2 Task
                    dataset_info={
                        "bucket": r.get("bucket"),
                        "environment_context": r.get("environment_context"),
                        "user_simulation": user_simulation,
                        "evaluation_criteria": r.get("evaluation_criteria"),
                        "max_turns_hint": (r.get("meta") or {}).get("max_turns_hint", 20),
                    },
                ),
            )
        )
    return out


# WHICH PORT the airline MCP server binds to (same scheme as test_banking): EP's
# processor hardcodes 9700 and hard-fails if busy, so pick a free one ourselves —
# start at MCP_PORT (default 9700) and walk forward until one is free. This lets
# banking + airline (or two models) run concurrently.
def _find_free_port(start: int, max_tries: int = 50) -> int:
    for candidate in range(start, start + max_tries):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(1)
            if s.connect_ex(("localhost", candidate)) != 0:   # nobody listening -> free
                if candidate != start:
                    print(f"⚠️ Port {start} is in use — MCP server will use {candidate} instead")
                return candidate
    raise RuntimeError(f"No free port in range {start}-{start + max_tries - 1}")


MCP_PORT = _find_free_port(int(os.environ.get("MCP_PORT", "9700")))


# ===========================================================================
# THE EP TEST DECLARATION (step 2)
# ===========================================================================
@evaluation_test(
    input_dataset=[DATASET],
    dataset_adapter=airline_to_evaluation_row,
    completion_params=_completion_params(),
    rollout_processor=MCPGymRolloutProcessor(),
    rollout_processor_kwargs={"port": MCP_PORT},
    server_script_path="FIREWORKS_TRAINING/airline_mcp/server.py",
    mode="pointwise",
    passed_threshold=0.4,
    num_runs=2,                  # simulated user is stochastic -> average two rollouts
    max_concurrent_rollouts=16,
)
def test_airline(row: EvaluationRow) -> EvaluationRow:
    """REWARD HOOK (step 3): score ONE finished rollout. Delegates everything to
    score_airline.score_evaluation_row, which rebuilds tau2 reward_info from the
    trajectory, prints it (SHOW_REWARD_INFO), and runs the simple reward math."""
    return score_evaluation_row(row)


# ===========================================================================
# Running this file directly is just a cheap "does the dataset load?" check.
# Real runs go through `ep local-test` (run_airline_model.sh) or `create rft`.
# ===========================================================================
if __name__ == "__main__":
    path = os.path.join(ROOT, DATASET)
    if os.path.exists(path):
        n = sum(1 for _ in open(path))
        print(f"dataset OK: {n} rows in {DATASET}")
    else:
        print(f"dataset MISSING: {DATASET}")