Update grader.py

grader.py

@@ -1,134 +1,106 @@
 """
-grader.py – PhishGuard-Env SOC Triage Scoring Logic
-====================================================
-
-REWARD CONTRACT  →  OPEN INTERVAL (0.0, 1.0)
----------------------------------------------
-All rewards are STRICTLY greater than 0 and STRICTLY less than 1.
-The endpoints 0 and 1 are NEVER returned.  This is a hard invariant
-enforced by the constant table below and by calculate_overall_score().
-
-Why open-interval?
-  • 1.0 saturates the leaderboard and implies a theoretically perfect agent.
-  • 0.0 is indistinguishable from a missing data-point in an RL pipeline.
-  • Every decision carries a non-zero gradient signal so training never dies.
-
-REWARD TABLE
-────────────────────────────────────────────────────────────────────────────
-  Constant                 Value   Outcome / Rationale
-  ─────────────────────────────────────────────────────────────────────────
-  R_PERFECT                0.95    Exact match — near-ideal; headroom for 1.0
-  R_MALWARE_QUARANTINE     0.75    MALWARE → QUARANTINE (textbook isolation)
-  R_PHISH_BEC_QUARANTINE   0.60    PHISH/BEC → QUARANTINE (domain still live)
-  R_SPAM_BLOCK             0.40    SPAM → BLOCK_DOMAIN (over-escalation)
-  R_SPAM_QUARANTINE        0.35    SPAM → QUARANTINE (lighter over-escalation)
-  R_WRONG_PROCEDURE        0.10    Wrong action, no direct security/ops harm
-  R_DISRUPTION             0.05    SAFE email blocked — operational cost
-  R_BREACH                 0.02    Threat allowed into inbox — catastrophic
-
-  R_PARTIAL is exported as an alias for R_MALWARE_QUARANTINE (0.75) so
-  inference.py and env.py can import a single mid-range signal constant
-  without hard-coding a numeric literal.
-
-HEALTH-DRAIN THRESHOLD
-────────────────────────────────────────────────────────────────────────────
+grader.py – PhishGuard-Env | Triage Scoring & Task Registry
+==============================================================
+
+REWARD SCALE CONTRACT
+----------------------------------
+All rewards strictly in the open interval (0.0, 1.0) — 0 and 1 never returned.
+
+  Outcome                                    Reward   Rationale
+  ─────────────────────────────────────────────────────────────────────────────
+  Perfect match (any task type)               0.95    Near-ideal; headroom kept
+  MALWARE → QUARANTINE (safe containment)     0.75    Correct containment call
+  PHISH/BEC → QUARANTINE (under-response)    0.60    Held but domain still active
+  SPAM → BLOCK_DOMAIN (over-escalation)      0.40    Too aggressive, not wrong
+  SPAM → QUARANTINE (lighter escalation)     0.35    Wastes analyst queue
+  General wrong procedure (no active harm)   0.10    Wrong but not dangerous
+  Business Disruption (SAFE email blocked)   0.05    Ops impact, below drain floor
+  Security Breach (threat → MARK_SAFE)       0.02    Catastrophic; minimum signal
+
+Health-drain threshold (consumed by env.py)
+-------------------------------------------
   HEALTH_DRAIN_THRESHOLD = 0.15
-
-  reward < 0.15  →  agent loses one life.  Covers:
-    R_BREACH (0.02), R_DISRUPTION (0.05), R_WRONG_PROCEDURE (0.10)
-
-  Cautious / partial-credit scores (≥ 0.35) NEVER drain health.
-
-PASS_THRESHOLD
-────────────────────────────────────────────────────────────────────────────
-  PASS_THRESHOLD = 0.50
-
-  The minimum overall_score an agent must achieve across a full run to be
-  considered a passing benchmark result.  Imported by inference.py so the
-  bar is defined in exactly one place.
-
-LEVEL CONTEXT (from env.py)
-────────────────────────────────────────────────────────────────────────────
-  easy   → lv1 (SPAM),    lv2 (PHISH),   lv3 (SAFE)
-  medium → lv4 (MALWARE), lv5 (SAFE),    lv6 (BEC),   lv7 (PHISH)
-  hard   → lv8 (MALWARE), lv9 (PHISH),   lv10 (BEC)
-
-  email_type values in play: SPAM | PHISH | BEC | MALWARE | SAFE
-
-VALID AGENT ACTIONS
-────────────────────────────────────────────────────────────────────────────
-  MARK_SAFE      – deliver to inbox          (use ONLY for confirmed-safe)
-  MOVE_TO_SPAM   – bulk / unsolicited mail   (no active threat)
-  QUARANTINE     – hold for analyst review
-  BLOCK_DOMAIN   – perimeter block           (confirmed phishing / BEC source)
-
-DECISION TREE  (grade_action)
-────────────────────────────────────────────────────────────────────────────
-  1. Unrecognised action token               → R_WRONG_PROCEDURE
-  2. action == correct                       → R_PERFECT
-  3. Threat type (incl. SPAM) + MARK_SAFE    → R_BREACH      (Security Breach)
-  4. SAFE type + blocking move               → R_DISRUPTION  (Business Disruption)
-  5. MALWARE → QUARANTINE                    → R_MALWARE_QUARANTINE
-  6. PHISH/BEC → QUARANTINE                  → R_PHISH_BEC_QUARANTINE
-  7. SPAM → BLOCK_DOMAIN                     → R_SPAM_BLOCK
-  8. SPAM → QUARANTINE                       → R_SPAM_QUARANTINE
-  9. Catch-all wrong procedure               → R_WRONG_PROCEDURE
-
-BUG FIX (v1.0.2 → v1.0.3)
-────────────────────────────────────────────────────────────────────────────
-  _THREAT_TYPES previously excluded SPAM.  This meant lv1 SPAM + MARK_SAFE
-  returned R_WRONG_PROCEDURE (0.10) instead of R_BREACH (0.02) — a
-  security-critical email type was not penalised as a breach.
-  Fix: SPAM added to _THREAT_TYPES so MARK_SAFE on any threat drains health.
+  reward < 0.15  →  env.py deducts one life from the agent.
 """
 
-from __future__ import annotations
-
 from typing import Tuple
 
+# ── Reward constants ───────────────────────────────────────────────────────────
+R_PERFECT               = 0.95
+R_MALWARE_QUARANTINE    = 0.75
+R_PHISH_BEC_QUARANTINE  = 0.60
+R_SPAM_BLOCK            = 0.40
+R_SPAM_QUARANTINE       = 0.35
+R_WRONG_PROCEDURE       = 0.10
+R_DISRUPTION            = 0.05
+R_BREACH                = 0.02
+
+HEALTH_DRAIN_THRESHOLD  = 0.15   # env.py: reward < this → lose one life
+
+# ── Task Registry ──────────────────────────────────────────────────────────────
+TASK_REGISTRY: dict[str, dict] = {
+    "task_spam": {
+        "email_types": {"SPAM"},
+        "description": (
+            "Identify unsolicited bulk/commercial mail and route it to spam "
+            "without over-escalating to a full domain block."
+        ),
+        "threat": True,
+    },
+    "task_phishing": {
+        "email_types": {"PHISH"},
+        "description": (
+            "Detect credential-harvesting and typosquat phishing attacks; "
+            "block the offending domain at the perimeter."
+        ),
+        "threat": True,
+    },
+    "task_safe": {
+        "email_types": {"SAFE"},
+        "description": (
+            "Recognise legitimate internal and partner emails and pass them "
+            "through without causing business disruption."
+        ),
+        "threat": False,
+    },
+    "task_malware": {
+        "email_types": {"MALWARE"},
+        "description": (
+            "Identify emails delivering executable payloads or macro-enabled "
+            "attachments and quarantine them for analysis."
+        ),
+        "threat": True,
+    },
+    "task_bec": {
+        "email_types": {"BEC"},
+        "description": (
+            "Detect business email compromise and supply-chain attacks, "
+            "including cases where SPF/DMARC authentication passes."
+        ),
+        "threat": True,
+    },
+}
+
+# Reverse lookup: raw email_type string → task_id
+_TYPE_TO_TASK_ID: dict[str, str] = {
+    etype: task_id
+    for task_id, meta in TASK_REGISTRY.items()
+    for etype in meta["email_types"]
+}
+
+
+def resolve_task_id(email_type: str) -> str:
+    """Convert raw email type (e.g. 'PHISH') to task_id (e.g. 'task_phishing')."""
+    key = email_type.strip().upper()
+    if key not in _TYPE_TO_TASK_ID:
+        raise ValueError(
+            f"Unknown email_type '{email_type}'. "
+            f"Valid types: {sorted(_TYPE_TO_TASK_ID)}"
+        )
+    return _TYPE_TO_TASK_ID[key]
+
 
-# ══════════════════════════════════════════════════════════════════════════════
-# REWARD CONSTANTS
-# ══════════════════════════════════════════════════════════════════════════════
-# All numeric reward values are defined ONCE here.
-# env.py and inference.py import these — neither file hard-codes numbers.
-
-R_PERFECT              = 0.95   # Exact triage match — near-ideal signal
-R_MALWARE_QUARANTINE   = 0.75   # MALWARE isolated via QUARANTINE
-R_PHISH_BEC_QUARANTINE = 0.60   # PHISH/BEC held but source domain still live
-R_SPAM_BLOCK           = 0.40   # SPAM → BLOCK_DOMAIN (over-escalation)
-R_SPAM_QUARANTINE      = 0.35   # SPAM → QUARANTINE   (lighter over-escalation)
-R_WRONG_PROCEDURE      = 0.10   # Wrong action; no security / operational harm
-R_DISRUPTION           = 0.05   # Business Disruption — SAFE email blocked
-R_BREACH               = 0.02   # Security Breach — threat allowed through
-
-# ── Convenience alias ─────────────────────────────────────────────────────────
-# R_PARTIAL = mid-range cautious signal used by inference.py / reporting layers.
-# Maps to R_MALWARE_QUARANTINE (0.75) — the highest partial-credit reward.
-R_PARTIAL = R_MALWARE_QUARANTINE
-
-# ── Benchmark pass threshold ──────────────────────────────────────────────────
-# Imported by inference.py so the pass/fail bar is defined in one place.
-PASS_THRESHOLD = 0.50
-
-# ── Health-drain threshold ────────────────────────────────────────────────────
-# env.py compares `reward < HEALTH_DRAIN_THRESHOLD` to decide life-loss.
-# Must sit above R_WRONG_PROCEDURE (0.10) and below R_SPAM_QUARANTINE (0.35)
-# so cautious over-escalations never drain health.
-HEALTH_DRAIN_THRESHOLD = 0.15
-
-# ── Internal lookup sets ──────────────────────────────────────────────────────
-# BUG FIX: SPAM is now included in _THREAT_TYPES.
-# Previously SPAM was omitted, so SPAM + MARK_SAFE returned R_WRONG_PROCEDURE
-# (0.10) instead of the correct R_BREACH (0.02).
-_THREAT_TYPES  = frozenset({"PHISH", "BEC", "MALWARE", "SPAM"})
-_BLOCKED_MOVES = frozenset({"BLOCK_DOMAIN", "QUARANTINE", "MOVE_TO_SPAM"})
-_VALID_ACTIONS = frozenset({"MARK_SAFE", "MOVE_TO_SPAM", "QUARANTINE", "BLOCK_DOMAIN"})
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# GRADE_ACTION
-# ══════════════════════════════════════════════════════════════════════════════
+# ── Core grading function ──────────────────────────────────────────────────────
 
 def grade_action(
     agent_output: str,
@@ -136,162 +108,119 @@ def grade_action(
     email_type: str,
 ) -> Tuple[float, str]:
     """
-    Grade one SOC triage decision and return a reward in (0.0, 1.0).
-
-    Parameters
-    ----------
-    agent_output    : Raw action string from the LLM agent (case-insensitive).
-    expected_output : Ground-truth correct action for this scenario.
-    email_type      : Threat category — PHISH | BEC | MALWARE | SPAM | SAFE.
-
-    Returns
-    -------
-    (reward, message)
-        reward  : float strictly in (0.0, 1.0) — NEVER 0, NEVER 1.
-        message : short verdict string for logs and agent feedback.
-
-    Examples (aligned with env.py scenarios)
-    -----------------------------------------
-    lv1  SPAM  / MOVE_TO_SPAM  + MOVE_TO_SPAM  → 0.95  PERFECT_TRIAGE
-    lv1  SPAM  / MOVE_TO_SPAM  + MARK_SAFE     → 0.02  SECURITY_BREACH  ← BUG FIX
-    lv1  SPAM  / MOVE_TO_SPAM  + QUARANTINE    → 0.35  OVER_ESCALATION
-    lv2  PHISH / BLOCK_DOMAIN  + BLOCK_DOMAIN  → 0.95  PERFECT_TRIAGE
-    lv2  PHISH / BLOCK_DOMAIN  + QUARANTINE    → 0.60  UNDER_RESPONSE
-    lv2  PHISH / BLOCK_DOMAIN  + MARK_SAFE     → 0.02  SECURITY_BREACH
-    lv3  SAFE  / MARK_SAFE     + MARK_SAFE     → 0.95  PERFECT_TRIAGE
-    lv3  SAFE  / MARK_SAFE     + QUARANTINE    → 0.05  BUSINESS_DISRUPTION
-    lv4  MALWARE/ QUARANTINE   + QUARANTINE    → 0.95  PERFECT_TRIAGE
-    lv4  MALWARE/ QUARANTINE   + BLOCK_DOMAIN  → 0.10  INCORRECT_PROCEDURE
-    lv6  BEC   / QUARANTINE    + MARK_SAFE     → 0.02  SECURITY_BREACH
-    lv10 BEC   / BLOCK_DOMAIN  + QUARANTINE    → 0.60  UNDER_RESPONSE
+    Grade one SOC triage decision.
+    Returns (reward, message) where reward is strictly in (0.0, 1.0).
     """
     agent_action    = agent_output.strip().upper()
     expected_action = expected_output.strip().upper()
     etype           = email_type.strip().upper()
 
-    # ── Step 1: Reject unrecognised tokens ───────────────────��───────────────
-    # Treat unknown output as wrong procedure rather than raising so a single
-    # malformed response does not crash the entire episode.
-    if agent_action not in _VALID_ACTIONS:
-        return (
-            R_WRONG_PROCEDURE,
-            f"INVALID_ACTION: '{agent_action}' is not a recognised triage action — "
-            f"must be one of: {', '.join(sorted(_VALID_ACTIONS))}",
-        )
+    task_id   = resolve_task_id(etype)
+    is_threat = TASK_REGISTRY[task_id]["threat"]
 
-    # ── Step 2: Perfect match ─────────────────────────────────────────────────
     if agent_action == expected_action:
         return R_PERFECT, "PERFECT_TRIAGE: Correct action taken"
 
-    # ── Step 3: Security Breach — most severe outcome ────────────────────────
-    # Any threat type (PHISH, BEC, MALWARE, SPAM) rubber-stamped as safe.
-    # Covers: lv1 SPAM→MARK_SAFE, lv2 PHISH→MARK_SAFE, lv6 BEC→MARK_SAFE, etc.
-    if etype in _THREAT_TYPES and agent_action == "MARK_SAFE":
-        return (
-            R_BREACH,
-            f"SECURITY_BREACH: {etype} threat delivered to inbox unimpeded — "
-            "catastrophic failure; source remains active",
-        )
-
-    # ── Step 4: Business Disruption — severe false-positive ──────────────────
-    # A clean, legitimate email (SAFE) was blocked, quarantined, or spammed.
-    # Covers: lv3 SAFE→QUARANTINE, lv5 SAFE→BLOCK_DOMAIN, etc.
-    if etype == "SAFE" and agent_action in _BLOCKED_MOVES:
-        return (
-            R_DISRUPTION,
-            "BUSINESS_DISRUPTION: Legitimate communication was incorrectly blocked — "
-            "operational impact; sender trust degraded",
-        )
+    if is_threat and agent_action == "MARK_SAFE":
+        return R_BREACH, "SECURITY_BREACH: Threat allowed into the network"
 
-    # ── Step 5: Partial credit — cautious but sub-optimal ────────────────────
+    if not is_threat and agent_action in {"BLOCK_DOMAIN", "QUARANTINE", "MOVE_TO_SPAM"}:
+        return R_DISRUPTION, "BUSINESS_DISRUPTION: Legitimate communication blocked"
 
-    # MALWARE → QUARANTINE
-    # Isolation is textbook containment; BLOCK_DOMAIN is preferred for confirmed
-    # malware sources but QUARANTINE still prevents propagation.
-    if etype == "MALWARE" and agent_action == "QUARANTINE":
-        return (
-            R_MALWARE_QUARANTINE,
-            "CAUTIOUS: Malware isolated via QUARANTINE — strong containment; "
-            "no further propagation risk detected",
-        )
+    if task_id == "task_malware" and agent_action == "QUARANTINE":
+        return R_MALWARE_QUARANTINE, "CAUTIOUS: Malware isolated — strong containment"
 
-    # PHISH/BEC → QUARANTINE
-    # Threat is held but the sending domain remains live and can re-deliver.
-    # Preferred action is BLOCK_DOMAIN to sever the attack vector entirely.
-    if etype in {"PHISH", "BEC"} and agent_action == "QUARANTINE":
-        return (
-            R_PHISH_BEC_QUARANTINE,
-            f"UNDER_RESPONSE: {etype} quarantined but source domain still active — "
-            "prefer BLOCK_DOMAIN to prevent further phishing delivery attempts",
+    if task_id in {"task_phishing", "task_bec"} and agent_action == "QUARANTINE":
+        return R_PHISH_BEC_QUARANTINE, (
+            "UNDER_RESPONSE: Threat quarantined but domain still active — "
+            "consider BLOCK_DOMAIN"
         )
 
-    # SPAM → BLOCK_DOMAIN
-    # Disproportionate escalation; burns perimeter block-list capacity on a
-    # low-severity bulk sender.  Prefer MOVE_TO_SPAM.
-    if etype == "SPAM" and agent_action == "BLOCK_DOMAIN":
-        return (
-            R_SPAM_BLOCK,
+    if task_id == "task_spam" and agent_action == "BLOCK_DOMAIN":
+        return R_SPAM_BLOCK, (
             "OVER_ESCALATION: BLOCK_DOMAIN is disproportionate for SPAM — "
-            "prefer MOVE_TO_SPAM to preserve block-list resources",
+            "prefer MOVE_TO_SPAM"
         )
 
-    # SPAM → QUARANTINE
-    # Lighter over-escalation; clogs the analyst review queue with bulk mail.
-    if etype == "SPAM" and agent_action == "QUARANTINE":
-        return (
-            R_SPAM_QUARANTINE,
-            "OVER_ESCALATION: QUARANTINE wastes analyst capacity on SPAM — "
-            "prefer MOVE_TO_SPAM for bulk unsolicited mail",
+    if task_id == "task_spam" and agent_action == "QUARANTINE":
+        return R_SPAM_QUARANTINE, (
+            "OVER_ESCALATION: QUARANTINE wastes analyst capacity for SPAM — "
+            "prefer MOVE_TO_SPAM"
         )
 
-    # ── Step 6: General incorrect procedure (catch-all) ──────────────────────
-    # Wrong action with no direct security breach or operational disruption.
-    # e.g. MALWARE → BLOCK_DOMAIN when QUARANTINE is expected (lv4, lv8).
-    return (
-        R_WRONG_PROCEDURE,
-        f"INCORRECT_PROCEDURE: '{agent_action}' does not match security policy "
-        f"for {etype} email (expected: {expected_action}) — "
-        "review triage guidelines",
-    )
+    return R_WRONG_PROCEDURE, "INCORRECT_PROCEDURE: Decision does not match security policy"
 
 
-# ══════════════════════════════════════════════════════════════════════════════
-# CALCULATE_OVERALL_SCORE
-# ══════════════════════════════════════════════════════════════════════════════
+# ── Episode-level score aggregation ───────────────────────────────────────────
 
 def calculate_overall_score(task_scores: list) -> float:
-    """
-    Compute the final benchmark score from a list of per-step rewards.
-
-    The result is clamped to [R_BREACH, R_PERFECT] — maintaining the
-    open-interval contract — so downstream consumers always receive a float
-    strictly greater than 0 and strictly less than 1.
-
-    Parameters
-    ----------
-    task_scores : list of floats, each in (0.0, 1.0).
-
-    Returns
-    -------
-    float in [R_BREACH, R_PERFECT] — always a valid open-interval value.
-
-    Edge cases
-    ----------
-    • Empty list    → R_BREACH (minimum non-zero signal; not zero)
-    • Single step   → that step's reward, clamped to [R_BREACH, R_PERFECT]
-    • All perfect   → R_PERFECT (0.95), never 1.0
-
-    Level score examples (env.py alignment)
-    ----------------------------------------
-    easy   (3 tasks): [0.95, 0.95, 0.95]       → 0.95
-    medium (4 tasks): [0.95, 0.02, 0.95, 0.60] → 0.63
-    hard   (3 tasks): [0.02, 0.05, 0.95]        → 0.34
-    """
+    """Compute final benchmark score from per-step rewards. Returns float in (0,1)."""
     if not task_scores:
-        return R_BREACH  # No tasks completed: minimum signal, not zero
-
+        return R_BREACH
     raw_avg = sum(task_scores) / len(task_scores)
-
-    # Clamp strictly within the open-interval boundary constants.
     clamped = max(R_BREACH, min(R_PERFECT, raw_avg))
     return round(clamped, 4)
+
+
+def calculate_per_task_scores(
+    task_score_map: dict[str, list[float]],
+) -> dict[str, float]:
+    """Compute per-task-type average scores. Returns {task_id: score}."""
+    return {
+        task_id: calculate_overall_score(scores)
+        for task_id, scores in task_score_map.items()
+        if scores
+    }
+
+
+# ── Named graders — ONE PER DIFFICULTY LEVEL ──────────────────────────────────
+#
+# WHY THESE EXIST:
+#   openenv.yaml declares 3 tasks: easy, medium, hard.
+#   Each task's `grader` field must point to a real callable in this file.
+#   The validator does:  import grader; callable(grader.grade_easy)  → True
+#   Without these functions, it finds nothing → "Not enough tasks with graders"
+#
+# WHAT THEY DO:
+#   Each receives the list of per-step rewards collected during that level's
+#   scenarios and returns a final aggregated score in (0.0, 1.0).
+#
+# LINE NUMBERS WHERE OPENENV.YAML REFERENCES THESE:
+#   openenv.yaml line 7  → grader: "grader.grade_easy"
+#   openenv.yaml line 11 → grader: "grader.grade_medium"
+#   openenv.yaml line 15 → grader: "grader.grade_hard"
+# ─────────────────────────────────────────────────────────────────────────────
+
+def grade_easy(task_scores: list) -> float:
+    """
+    Grader for easy-level scenarios (lv1–lv3: SPAM, PHISH, SAFE).
+    Referenced by openenv.yaml task id 'easy'.
+    Returns float strictly in (0.0, 1.0).
+    """
+    return calculate_overall_score(task_scores)
+
+
+def grade_medium(task_scores: list) -> float:
+    """
+    Grader for medium-level scenarios (lv4–lv7: MALWARE, SAFE, BEC, PHISH).
+    Referenced by openenv.yaml task id 'medium'.
+    Returns float strictly in (0.0, 1.0).
+    """
+    return calculate_overall_score(task_scores)
+
+
+def grade_hard(task_scores: list) -> float:
+    """
+    Grader for hard-level scenarios (lv8–lv10: MALWARE, PHISH, BEC).
+    Referenced by openenv.yaml task id 'hard'.
+    Returns float strictly in (0.0, 1.0).
+    """
+    return calculate_overall_score(task_scores)
+
+
+# Registry map — consumed by env.py for the /grade endpoint
+GRADERS = {
+    "easy":   grade_easy,
+    "medium": grade_medium,
+    "hard":   grade_hard,
+}