training/smoke_test.py

"""Smoke test — runs locally without TRL/vLLM.

Tests in order:
  1. Dataset generation
  2. JSON extraction (edge cases)
  3. build_messages structure (no duplicates, correct alternation)
  4. Server startup + env connection
  5. One full heuristic episode (env API end-to-end)

Run from repo root:  python training/smoke_test.py
"""
import json
import os
import subprocess
import sys
import time

import requests

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

PASS = "\033[92m[PASS]\033[0m"
FAIL = "\033[91m[FAIL]\033[0m"
INFO = "\033[94m[INFO]\033[0m"


def check(label: str, cond: bool, detail: str = "") -> bool:
    if cond:
        print(f"{PASS} {label}")
    else:
        print(f"{FAIL} {label}" + (f" — {detail}" if detail else ""))
    return cond


# ---------------------------------------------------------------------------
# 1. Dataset
# ---------------------------------------------------------------------------
print("\n── 1. Dataset ──────────────────────────────────────────────────────")
from training.dataset import generate_triage_dataset, parse_seed_from_prompt

rows = generate_triage_dataset(n_episodes=5, base_seed=99)
check("generates 5 rows", len(rows) == 5)
check("each row has prompt+seed+difficulty", all(
    "prompt" in r and "seed" in r and "difficulty" in r for r in rows
))
seed_back = parse_seed_from_prompt(rows[0]["prompt"])
check("seed round-trips through prompt string", seed_back == rows[0]["seed"],
      f"got {seed_back}, expected {rows[0]['seed']}")
print(f"  sample prompt: {rows[0]['prompt'][:100]}...")


# ---------------------------------------------------------------------------
# 2. JSON extraction
# ---------------------------------------------------------------------------
print("\n── 2. JSON extraction ──────────────────────────────────────────────")
from training.rollout import extract_json_action, step_aware_fallback

cases = [
    # (description, input_text, expected_action_type)
    ("code block", 'Thinking...\n```json\n{"action_type": "meta.read_runbook", "args": {}}\n```', "meta.read_runbook"),
    ("raw JSON",   '{"action_type": "meta.finish", "args": {}}', "meta.finish"),
    ("regex hit",  'I will do this: {"action_type": "meta.noop", "args": {}} done.', "meta.noop"),
    ("no JSON",    "I cannot decide.", None),
    ("bad block",  '```json\n{broken\n```', None),
]
for desc, text, expected in cases:
    result = extract_json_action(text)
    got = result.get("action_type") if result else None
    check(f"extract: {desc}", got == expected, f"got {got!r}, expected {expected!r}")

fb = step_aware_fallback(0)
check("fallback step 0 → read_runbook", fb["action_type"] == "meta.read_runbook")
fb = step_aware_fallback(20)
check("fallback step 20 → noop", fb["action_type"] == "meta.noop")
fb = step_aware_fallback(38)
check("fallback step 38 → finish", fb["action_type"] == "meta.finish")

# reward_no_wrong_channels logic
from training.rewards import (
    WEIGHT_FINAL_SCORE, WEIGHT_NO_WRONG_CHANNELS,
    WEIGHT_VALID_JSON, WEIGHT_READ_RUNBOOK, WEIGHT_EFFICIENCY,
)
weights_sum = (WEIGHT_FINAL_SCORE + WEIGHT_NO_WRONG_CHANNELS +
               WEIGHT_VALID_JSON + WEIGHT_READ_RUNBOOK + WEIGHT_EFFICIENCY)
check("reward weights sum to 1.0", abs(weights_sum - 1.0) < 1e-9, f"sum={weights_sum}")


# ---------------------------------------------------------------------------
# 3. build_messages structure
# ---------------------------------------------------------------------------
print("\n── 3. build_messages ───────────────────────────────────────────────")
from training.rollout import build_messages

# Empty history (step 0)
msgs = build_messages([], "First observation")
check("step 0: starts with system", msgs[0]["role"] == "system")
check("step 0: ends with user", msgs[-1]["role"] == "user")
check("step 0: final user is current obs", msgs[-1]["content"] == "First observation")

# Simulate 3 turns
history = [
    {"obs_text": "obs_0", "completion": "comp_0", "is_runbook": True},
    {"obs_text": "obs_1", "completion": "comp_1", "is_runbook": False},
    {"obs_text": "obs_2", "completion": "comp_2", "is_runbook": False},
]
msgs = build_messages(history, "current_obs")

# Validate alternation: after system, must be user/asst/user/asst/.../user
roles = [m["role"] for m in msgs]
check("roles start with system", roles[0] == "system")
check("roles end with user", roles[-1] == "user")
pairs_ok = all(
    roles[i] == "user" and roles[i+1] == "assistant"
    for i in range(1, len(roles) - 2, 2)
)
check("strict user/asst alternation throughout", pairs_ok, str(roles))

# Verify current_obs appears exactly once as the last message
final_user_content = msgs[-1]["content"]
check("current_obs is last message content", final_user_content == "current_obs")
all_contents = [m["content"] for m in msgs]
check("current_obs not duplicated", all_contents.count("current_obs") == 1)

# Verify runbook is pinned (obs_0/comp_0 should appear even with many turns)
all_content_str = " ".join(all_contents)
check("runbook obs pinned", "obs_0" in all_content_str)
check("runbook completion pinned", "comp_0" in all_content_str)


# ---------------------------------------------------------------------------
# 4. Server startup
# ---------------------------------------------------------------------------
print("\n── 4. Server startup ───────────────────────────────────────────────")
REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ENV_URL = "http://localhost:8765"  # non-standard port to avoid conflicts

proc = subprocess.Popen(
    [sys.executable, "-m", "uvicorn", "server.app:app",
     "--host", "0.0.0.0", "--port", "8765"],
    cwd=REPO_DIR,
    stdout=subprocess.DEVNULL,
    stderr=subprocess.DEVNULL,
)

ready = False
for _ in range(20):
    try:
        r = requests.get(f"{ENV_URL}/", timeout=1)
        if r.status_code == 200:
            ready = True
            break
    except Exception:
        pass
    time.sleep(1)

check("server started on :8765", ready)
if not ready:
    proc.terminate()
    print(f"{FAIL} Cannot continue without server.")
    sys.exit(1)


# ---------------------------------------------------------------------------
# 5. Full heuristic episode (env API end-to-end)
# ---------------------------------------------------------------------------
print("\n── 5. End-to-end episode ───────────────────────────────────────────")
from openenv.core import GenericEnvClient

client = GenericEnvClient(base_url=ENV_URL).sync()
client.connect()

try:
    # reset with explicit seed (tests seed passthrough)
    res = client.reset(seed=rows[0]["seed"])
    obs = res.observation if hasattr(res, "observation") else res

    def _get(o, k):
        return getattr(o, k, None) if not isinstance(o, dict) else o.get(k)

    task_brief = _get(obs, "task_brief")
    check("reset returns task_brief", bool(task_brief), repr(task_brief)[:80])

    # Step 1: read runbook
    res = client.step({"action_type": "meta.read_runbook", "args": {}})
    obs = res.observation if hasattr(res, "observation") else res
    last = _get(obs, "last_action_result") or {}
    check("read_runbook ok=True", last.get("ok") is True)
    check("runbook contains org_config", "org_config" in str(last.get("data", "")))

    org = last.get("data", {}).get("org_config", {}) if isinstance(last.get("data"), dict) else {}
    labels = list(org.get("label_taxonomy", {}).values())
    priorities = org.get("priority_levels", [])
    teams = list(org.get("team_map", {}).values())
    channels = list(org.get("oncall_channels", {}).values())

    check("org has labels", len(labels) > 0)
    check("org has priorities", len(priorities) > 0)

    # Step 2: create ticket with correct label/priority
    res = client.step({
        "action_type": "ticketing.create_ticket",
        "args": {
            "summary": f"Triage: {task_brief[:60]}",
            "description": task_brief,
            "label": labels[0],
            "priority": priorities[1] if len(priorities) > 1 else priorities[0],
            "assignee": teams[0] if teams else "backend",
        },
    })
    obs = res.observation if hasattr(res, "observation") else res
    last = _get(obs, "last_action_result") or {}
    check("create_ticket ok=True", last.get("ok") is True, str(last.get("error")))
    ticket_id = (last.get("data") or {}).get("id") if isinstance(last.get("data"), dict) else None
    check("ticket_id returned", bool(ticket_id), repr(ticket_id))

    # Step 3: assign ticket
    if ticket_id and teams:
        res = client.step({
            "action_type": "ticketing.assign_ticket",
            "args": {"ticket_id": ticket_id, "team": teams[0]},
        })
        obs = res.observation if hasattr(res, "observation") else res
        last = _get(obs, "last_action_result") or {}
        check("assign_ticket ok=True", last.get("ok") is True, str(last.get("error")))

    # Step 4: post to oncall channel
    if channels:
        res = client.step({
            "action_type": "chat.post_message",
            "args": {"channel": channels[0], "text": f"Triage smoke test: {task_brief[:80]}"},
        })
        obs = res.observation if hasattr(res, "observation") else res
        last = _get(obs, "last_action_result") or {}
        check("post_message ok=True", last.get("ok") is True, str(last.get("error")))

    # Step 5: finish + check reward
    res = client.step({"action_type": "meta.finish", "args": {}})
    done = getattr(res, "done", None)
    reward = getattr(res, "reward", 0.0)
    check("episode done after finish", done is True, f"done={done}")
    check("reward is float in [0,1]", isinstance(reward, float) and 0.0 <= reward <= 1.0,
          f"reward={reward}")
    print(f"  episode reward: {reward:.3f}")

    # Verify make_rollout_func accepts max_steps param
    from training.rollout import make_rollout_func
    import inspect
    sig = inspect.signature(make_rollout_func)
    check("make_rollout_func has max_steps param", "max_steps" in sig.parameters)
    check("max_steps default is 15", sig.parameters["max_steps"].default == 15)

    # Verify no_wrong_channels_reward key in rollout output keys
    from training.rollout import rollout_once as _rollout_once
    sig2 = inspect.signature(_rollout_once)
    check("rollout_once has max_steps param", "max_steps" in sig2.parameters)

finally:
    client.close()

proc.terminate()

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
print("\n────────────────────────────────────────────────────────────────────")
print("Smoke test complete. If all checks passed, training pipeline is ready.")
print("Next: push repo to HF and run training/train.ipynb on A100 GPU Space.")