run_model.sh

#!/usr/bin/env bash
# =============================================================================
# run_model.sh — pick a policy model by number, run the banking local-test.
#
# Thin wrapper around training_run.sh --localtest. All it does is set:
#     POLICY_MODEL         = the model you pick (with the fireworks_ai/ prefix)
#     USE_LEANMCP_GATEWAY  = 1 by default  (route inference through LeanMCP, logged)
# ...then call training_run.sh, so there is ONE source of truth for the run logic.
#
# Usage:
#     bash FIREWORKS_TRAINING/run_model.sh            # interactive: asks model + sim-user
#     bash FIREWORKS_TRAINING/run_model.sh 1          # pick #1; still asks about sim-user
#     bash FIREWORKS_TRAINING/run_model.sh 1 y        # pick #1, simulated user ON  (no prompts)
#     bash FIREWORKS_TRAINING/run_model.sh 1 n        # pick #1, simulated user OFF (no prompts)
#     bash FIREWORKS_TRAINING/run_model.sh 1 y 10 20  # pick #1, sim-user ON, 10 Easy + 20 Medium
#     USE_LEANMCP_GATEWAY=0 bash FIREWORKS_TRAINING/run_model.sh 1   # gateway OFF
#     ENABLE_USER_SIM=0    bash FIREWORKS_TRAINING/run_model.sh 1    # sim-user OFF (preset, no prompt)
#     USER_SIM_MODEL=fireworks_ai/...  bash FIREWORKS_TRAINING/run_model.sh 1   # override customer model
#
# Per repo convention, Claude does NOT run this — you do.
# =============================================================================
set -euo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"

# -----------------------------------------------------------------------------
# EDIT THIS LIST to add/remove models. Just the slug under accounts/fireworks/models/.
# (The fireworks_ai/accounts/fireworks/models/ prefix is added automatically.)
# -----------------------------------------------------------------------------
MODELS=(
  "minimax-m3"
  "minimax-m2p7"
  "glm-5p2"
  "glm-5p1"
  "kimi-k2p7-code"
  "kimi-k2p6"
  "qwen3p7-plus"
  "deepseek-v4-pro"
  "deepseek-v4-flash"
  "gpt-oss-120b"
  "gpt-oss-20b"
  "nemotron-3-ultra-nvfp4"
)

# Gateway ON by default; override with USE_LEANMCP_GATEWAY=0 in the environment.
GATEWAY="${USE_LEANMCP_GATEWAY:-1}"
# Simulated user: if preset via env (ENABLE_USER_SIM=0/1) we honor it; otherwise we
# ASK below (or take it from the 2nd positional arg). Empty = "ask".
USERSIM="${ENABLE_USER_SIM:-}"

# Task subset sizes: env preset > positional args 3 & 4 > interactive prompt.
# Default is a quick 5 Easy + 5 Medium subset (10 tasks). Set a big number for "all".
MAX_EASY="${MAX_EASY:-}"
MAX_MEDIUM="${MAX_MEDIUM:-}"

# -----------------------------------------------------------------------------
# Pick a model: from $1 if given, else show the numbered menu and read a number.
# -----------------------------------------------------------------------------
echo "== pick a policy model  (gateway: $([ "$GATEWAY" = 1 ] && echo ON || echo OFF)) =="
for i in "${!MODELS[@]}"; do
  printf "  %2d) %s\n" "$((i + 1))" "${MODELS[$i]}"
done
echo

choice="${1:-}"
if [ -z "$choice" ]; then
  read -rp "Select a model [1-${#MODELS[@]}]: " choice
fi

# Validate: must be an integer within range.
case "$choice" in
  ''|*[!0-9]*) echo "ERROR: '$choice' is not a number."; exit 1;;
esac
if [ "$choice" -lt 1 ] || [ "$choice" -gt "${#MODELS[@]}" ]; then
  echo "ERROR: pick a number between 1 and ${#MODELS[@]}."; exit 1
fi

SLUG="${MODELS[$((choice - 1))]}"
POLICY_MODEL="fireworks_ai/accounts/fireworks/models/${SLUG}"

# Resolve the simulated-user choice: env preset > 2nd arg (y/n) > interactive prompt.
sim_arg="${2:-}"
if [ -n "$USERSIM" ]; then
  :   # already set via ENABLE_USER_SIM env var; don't ask
elif [ -n "$sim_arg" ]; then
  case "$sim_arg" in
    [Yy]*|1) USERSIM=1;;
    [Nn]*|0) USERSIM=0;;
    *) echo "ERROR: 2nd arg must be y/n (include simulated user?)."; exit 1;;
  esac
else
  # ON is the faithful tau2 task (customer talks back); OFF = cheap one-shot check.
  read -rp "Include simulated user? [Y/n]: " ans
  case "$ans" in [Nn]*|0) USERSIM=0;; *) USERSIM=1;; esac
fi

# Resolve task subset sizes: env preset > 3rd/4th positional args > interactive prompt.
# Count how many tasks of each bucket the dataset actually has, so the prompts
# can show the max available. Falls back to '?' if the dataset isn't built yet.
DATASET="${EP_DATASET:-$HERE/data/banking_easy_medium.jsonl}"
if [ -f "$DATASET" ]; then
  N_EASY=$(grep -c '"bucket": "Easy"' "$DATASET" || true)
  N_MEDIUM=$(grep -c '"bucket": "Medium"' "$DATASET" || true)
else
  N_EASY="?"; N_MEDIUM="?"
fi
validate_count() {  # $1 = name, $2 = value
  case "$2" in
    ''|*[!0-9]*) echo "ERROR: $1 must be a non-negative integer (got '$2')."; exit 1;;
  esac
}
if [ -z "$MAX_EASY" ]; then
  if [ -n "${3:-}" ]; then
    MAX_EASY="$3"
  else
    read -rp "Max Easy tasks (max $N_EASY) [5]: " MAX_EASY
    MAX_EASY="${MAX_EASY:-5}"
  fi
fi
if [ -z "$MAX_MEDIUM" ]; then
  if [ -n "${4:-}" ]; then
    MAX_MEDIUM="$4"
  else
    read -rp "Max Medium tasks (max $N_MEDIUM) [5]: " MAX_MEDIUM
    MAX_MEDIUM="${MAX_MEDIUM:-5}"
  fi
fi
validate_count MAX_EASY "$MAX_EASY"
validate_count MAX_MEDIUM "$MAX_MEDIUM"

echo
echo "  selected      : $SLUG"
echo "  POLICY_MODEL  : $POLICY_MODEL"
echo "  gateway       : $([ "$GATEWAY" = 1 ] && echo "LeanMCP (needs LEANMCP_API_KEY)" || echo "direct Fireworks")"
echo "  simulated user: $([ "$USERSIM" = 1 ] && echo "ON (faithful multi-turn task)" || echo "OFF (one-shot smoke test)")"
echo "  task subset    : MAX_EASY=$MAX_EASY MAX_MEDIUM=$MAX_MEDIUM  (dataset has $N_EASY Easy / $N_MEDIUM Medium)"
echo

# Hand off to the real runner. It does the venv/PATH/preflight/non-interactive bits.
# These env vars are read by test_banking.py in the pytest subprocess:
#   POLICY_MODEL, USE_LEANMCP_GATEWAY, ENABLE_USER_SIM, MAX_EASY, MAX_MEDIUM
#   (and USER_SIM_MODEL / ENABLE_COMPLIANCE_GATE if you set them).
POLICY_MODEL="$POLICY_MODEL" USE_LEANMCP_GATEWAY="$GATEWAY" ENABLE_USER_SIM="$USERSIM" \
  MAX_EASY="$MAX_EASY" MAX_MEDIUM="$MAX_MEDIUM" \
  bash "$HERE/training_run.sh" --localtest