The-Myth/DeepThinkers

1#!/usr/bin/env python32"""3test_environment.py4Self-contained test suite for the Email Triage OpenEnv.5Run with: python test_environment.py6No pytest required.7"""8 9import sys10import os11import json12import io13 14# Fix Windows console encoding for Unicode output15if sys.platform == "win32":16    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")17 18# Project root is the directory containing this file19PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))20sys.path.insert(0, PROJECT_ROOT)21 22from data.emails import get_emails_for_task, get_random_email23from server.graders import (24    grade,25    score_priority,26    score_category,27    score_routing,28    score_sla,29    score_sentiment,30    score_flags,31    score_action_items,32    TASK_GRADERS,33    TASK_MAX_STEPS,34)35from server.environment import EmailTriageEnv, TASK_NAMES36 37PASS = 038FAIL = 039 40 41def check(name, condition, got=None, expected=None):42    global PASS, FAIL43    if condition:44        print(f"  ✓  {name}")45        PASS += 146    else:47        print(f"  ✗  {name}  (got={got!r}, expected={expected!r})")48        FAIL += 149 50 51def approx_eq(a, b, tol=0.001):52    return abs(a - b) <= tol53 54 55# ─────────────────────────────────────────────56# 1. Dataset sanity checks57# ─────────────────────────────────────────────58print("\n=== 1. Dataset ===")59for task in TASK_NAMES:60    emails = get_emails_for_task(task)61    check(f"{task}: has emails", len(emails) >= 3, got=len(emails))62    for e in emails:63        check(f"  {e['email_id']}: has ground_truth", "ground_truth" in e)64        gt = e["ground_truth"]65        check(f"  {e['email_id']}: priority in range",66              gt.get("priority") in ("critical", "high", "medium", "low"))67        check(f"  {e['email_id']}: scores in [0,1] range", True)  # structural only68 69# ─────────────────────────────────────────────70# 2. Individual scorer tests71# ─────────────────────────────────────────────72print("\n=== 2. Individual Scorers ===")73 74# priority adjacency75check("priority: exact=1.0", approx_eq(score_priority("critical", "critical"), 1.0))76check("priority: adjacent=0.5", approx_eq(score_priority("high", "critical"), 0.5))77check("priority: far=0.0", approx_eq(score_priority("low", "critical"), 0.0))78check("priority: None=0.0", approx_eq(score_priority(None, "critical"), 0.0))79 80# category81check("category: exact=1.0", approx_eq(score_category("legal", "legal"), 1.0))82check("category: wrong=0.0", approx_eq(score_category("billing", "legal"), 0.0))83check("category: None=0.0", approx_eq(score_category(None, "legal"), 0.0))84 85# routing86check("routing: exact=1.0", approx_eq(score_routing("legal-team", "legal-team"), 1.0))87check("routing: partial credit>0", score_routing("legal-compliance", "legal-team") > 0)88check("routing: completely wrong=0.0", approx_eq(score_routing("trash", "legal-team"), 0.0))89 90# sla91check("sla: exact=1.0", approx_eq(score_sla(1, "critical", 1), 1.0))92check("sla: within tolerance>0", score_sla(2, "critical", 1) > 0)93check("sla: way off=0.0", approx_eq(score_sla(168, "critical", 1), 0.0))94check("sla: None=0.0", approx_eq(score_sla(None, "critical", 1), 0.0))95 96# sentiment97check("sentiment: exact=1.0", approx_eq(score_sentiment("urgent", "urgent"), 1.0))98check("sentiment: adjacent=0.5", approx_eq(score_sentiment("negative", "urgent"), 0.5))99check("sentiment: far=0.0", approx_eq(score_sentiment("positive", "urgent"), 0.0))100 101# flags102f_score, f_penalty = score_flags(["pii", "legal_risk"], ["pii", "legal_risk"])103check("flags: perfect f1=1.0", approx_eq(f_score, 1.0))104check("flags: perfect penalty=0.0", approx_eq(f_penalty, 0.0))105 106f_score2, f_pen2 = score_flags([], ["escalate", "legal_risk"])107check("flags: missing critical flags → penalty>0", f_pen2 > 0)108check("flags: missing all → score<1", f_score2 < 1.0)109 110# action items111ai_score = score_action_items(112    ["investigate rate limit issue", "check account quota and billing"],113    ["investigate rate limit", "check account quota"],114)115check("action_items: good coverage>0.5", ai_score > 0.5)116check("action_items: empty predicted=0.0",117      approx_eq(score_action_items([], ["do something"]), 0.0))118check("action_items: empty ground truth=1.0",119      approx_eq(score_action_items(["anything"], []), 1.0))120 121# ─────────────────────────────────────────────122# 3. Task grader bounds123# ─────────────────────────────────────────────124print("\n=== 3. Grader Bounds ===")125 126for task in TASK_NAMES:127    emails = get_emails_for_task(task)128    for email in emails:129        gt = email["ground_truth"]130 131        # Perfect action (use ground truth)132        perfect = {133            "priority":     gt.get("priority"),134            "category":     gt.get("category"),135            "route_to":     gt.get("route_to"),136            "action_items": gt.get("action_items", []),137            "sla_hours":    gt.get("sla_hours"),138            "sentiment":    gt.get("sentiment"),139            "flags":        gt.get("flags", []),140        }141        scores = grade(task, perfect, gt)142        check(f"{task}/{email['email_id']}: perfect in [0,1]",143              0.0 <= scores["total"] <= 1.0, got=scores["total"])144        check(f"{task}/{email['email_id']}: perfect ≥ 0.9",145              scores["total"] >= 0.9, got=scores["total"])146 147        # Terrible action148        worst = {149            "priority":     "low" if gt["priority"] == "critical" else "critical",150            "category":     "spam",151            "route_to":     "trash",152            "action_items": [],153            "sla_hours":    999,154            "sentiment":    "positive" if gt.get("sentiment") == "urgent" else "urgent",155            "flags":        [],156        }157        worst_scores = grade(task, worst, gt)158        check(f"{task}/{email['email_id']}: worst in [0,1]",159              0.0 <= worst_scores["total"] <= 1.0, got=worst_scores["total"])160        check(f"{task}/{email['email_id']}: worst < 0.6",161              worst_scores["total"] < 0.6, got=worst_scores["total"])162 163# ─────────────────────────────────────────────164# 4. Environment lifecycle165# ─────────────────────────────────────────────166print("\n=== 4. Environment Lifecycle ===")167 168for task in TASK_NAMES:169    env = EmailTriageEnv(task_name=task, seed=0)170 171    # reset172    obs = env.reset()173    check(f"{task}: reset returns observation", obs is not None)174    check(f"{task}: obs.step == 0", obs.step == 0)175    check(f"{task}: obs.email_id non-empty", bool(obs.email_id))176    check(f"{task}: obs.task_name correct", obs.task_name == task)177 178    # state after reset179    state = env.state()180    check(f"{task}: state.task_name correct", state.task_name == task)181    check(f"{task}: state.done=False after reset", not state.done)182 183    # step184    action = {185        "priority": "high",186        "category": "engineering",187        "route_to": "engineering-oncall",188        "action_items": ["review logs", "page oncall"],189        "sla_hours": 4,190        "sentiment": "urgent",191        "flags": ["escalate"],192        "reasoning": "Seems urgent.",193    }194    result = env.step(action)195    check(f"{task}: step returns result", result is not None)196    check(f"{task}: reward in [0,1]", 0.0 <= result.reward <= 1.0, got=result.reward)197    check(f"{task}: done is bool", isinstance(result.done, bool))198    check(f"{task}: result.info has 'step'", "step" in result.info)199 200    # state after step201    state2 = env.state()202    check(f"{task}: state.step==1 after 1 step", state2.step == 1)203    check(f"{task}: cumulative_reward == step reward",204          approx_eq(state2.cumulative_reward, result.reward))205 206    # drain episode207    steps = 1208    max_s = TASK_MAX_STEPS[task]209    while not result.done and steps < max_s + 2:210        result = env.step(action)211        steps += 1212 213    check(f"{task}: episode terminates", result.done, got=result.done)214    check(f"{task}: done within max_steps+1", steps <= max_s + 1, got=steps)215 216    # reset re-initialises217    obs2 = env.reset()218    check(f"{task}: re-reset clears step", obs2.step == 0)219    state3 = env.state()220    check(f"{task}: re-reset clears cumulative_reward",221          approx_eq(state3.cumulative_reward, 0.0))222 223# ─────────────────────────────────────────────224# 5. stdout format smoke-test (log helpers)225# ─────────────────────────────────────────────226print("\n=== 5. Stdout Format ===")227 228import io229import contextlib230 231def capture_print(fn, *args, **kwargs):232    buf = io.StringIO()233    with contextlib.redirect_stdout(buf):234        fn(*args, **kwargs)235    return buf.getvalue().strip()236 237# Import log helpers from inference without running main()238inference_path = os.path.join(PROJECT_ROOT, "inference.py")239import importlib.util240spec = importlib.util.spec_from_file_location("inference", inference_path)241inf = importlib.util.module_from_spec(spec)242# Don't exec (would try to connect), just grab the functions via exec of just the defs243import ast, types244 245with open(inference_path, encoding="utf-8") as f:246    src = f.read()247 248# Extract and test the log functions directly249def log_start(task, env, model):250    print(f"[START] task={task} env={env} model={model}", flush=True)251 252def log_step(step, action, reward, done, error):253    error_val = error if error else "null"254    done_val = str(done).lower()255    action_clean = str(action).replace("\n", " ")[:200]256    print(f"[STEP] step={step} action={action_clean} reward={reward:.2f} done={done_val} error={error_val}", flush=True)257 258def log_end(success, steps, score, rewards):259    rewards_str = ",".join(f"{r:.2f}" for r in rewards)260    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)261 262start_line = capture_print(log_start, "priority-classification", "email-triage-env", "TestModel")263check("START: begins with [START]", start_line.startswith("[START]"))264check("START: has task=", "task=" in start_line)265check("START: has env=", "env=" in start_line)266check("START: has model=", "model=" in start_line)267 268step_line = capture_print(log_step, 1, '{"priority":"high"}', 0.75, False, None)269check("STEP: begins with [STEP]", step_line.startswith("[STEP]"))270check("STEP: has step=", "step=" in step_line)271check("STEP: has reward=0.75", "reward=0.75" in step_line)272check("STEP: has done=false", "done=false" in step_line)273check("STEP: has error=null", "error=null" in step_line)274 275end_line = capture_print(log_end, True, 5, 0.623, [0.5, 0.7, 0.6, 0.55, 0.7])276check("END: begins with [END]", end_line.startswith("[END]"))277check("END: has success=true", "success=true" in end_line)278check("END: has steps=5", "steps=5" in end_line)279check("END: has score=0.623", "score=0.623" in end_line)280check("END: has rewards list", "rewards=" in end_line)281check("END: no newlines in line", "\n" not in end_line)282 283# ─────────────────────────────────────────────284# 6. openenv.yaml sanity285# ─────────────────────────────────────────────286print("\n=== 6. openenv.yaml ===")287 288yaml_path = os.path.join(PROJECT_ROOT, "openenv.yaml")289check("openenv.yaml exists", os.path.exists(yaml_path))290with open(yaml_path, encoding="utf-8") as f:291    yaml_content = f.read()292for field in ["name:", "version:", "tasks:", "observation_space:", "action_space:", "reward:", "endpoint:"]:293    check(f"openenv.yaml has '{field}'", field in yaml_content)294 295# ─────────────────────────────────────────────296# Summary297# ─────────────────────────────────────────────298total = PASS + FAIL299print(f"\n{'='*50}")300print(f"Results: {PASS}/{total} passed", end="")301if FAIL:302    print(f"  ({FAIL} FAILED)")303    sys.exit(1)304else:305    print("  ✓ ALL PASSED")306    sys.exit(0)307