Update grader.py

grader.py

@@ -1,56 +1,83 @@
-
 """
 grader.py – PhishGuard-Env SOC Triage Scoring Logic
-=====================================================
+====================================================
 
-REWARD SCALE CONTRACT
----------------------
-All rewards are in the OPEN interval (0.0, 1.0).
-The endpoints 0 and 1 are NEVER returned. This is a hard invariant.
+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 the calculate_overall_score()
+clamp.
 
 Why open-interval?
-  • 1.0 would saturate the leaderboard and imply a theoretically perfect agent.
-  • 0.0 would be indistinguishable from a missing data-point in an RL pipeline.
-  • Every decision carries a non-zero signal so the training gradient never dies.
+  • 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.
 
-Full reward table
------------------
+REWARD TABLE
+────────────────────────────────────────────────────────────────────────
   Outcome                                    Reward   Rationale
-  ─────────────────────────────────────────────────────────────────────────────
-  Perfect match (any type)                    0.95    Near-ideal; headroom preserved
-  MALWARE → QUARANTINE (safe containment)     0.75    Correct call, no partial penalty
-  PHISH/BEC → QUARANTINE (under-response)    0.60    Stopped but domain still lives
-  SPAM → BLOCK_DOMAIN (over-escalation)      0.40    Too aggressive, not wrong
-  SPAM → QUARANTINE (lighter escalation)     0.35    Wastes analyst queue
-  General wrong procedure (no harm)          0.10    Wrong but no security/ops impact
-  Business Disruption (SAFE → blocked)       0.05    Ops cost, below drain threshold
-  Security Breach (threat → MARK_SAFE)       0.02    Catastrophic; minimum signal
-
-Health-drain threshold (consumed by env.py)
--------------------------------------------
+  ───────────────────────────────────────────────────────────────────────
+  Perfect match (any type)                    0.95    Near-ideal; headroom for 1.0 preserved
+  MALWARE → QUARANTINE  (correct containment) 0.75    Textbook isolation — full credit
+  PHISH/BEC → QUARANTINE (under-response)     0.60    Stopped but domain still active
+  SPAM → BLOCK_DOMAIN   (over-escalation)     0.40    Disproportionate but not harmful
+  SPAM → QUARANTINE     (lighter escalation)  0.35    Wastes analyst queue capacity
+  General wrong procedure (no harm)           0.10    Wrong, no security/ops impact
+  Business Disruption   (SAFE → blocked)      0.05    Operational cost; breaches threshold
+  Security Breach       (threat → MARK_SAFE)  0.02    Catastrophic; minimum non-zero signal
+
+HEALTH-DRAIN THRESHOLD
+────────────────────────────────────────────────────────────────────────
   HEALTH_DRAIN_THRESHOLD = 0.15
-  reward < 0.15  →  lose one life.
-  This covers Security Breach (0.02), Business Disruption (0.05), and
-  General Wrong Procedure (0.10). Cautious/partial-credit scores never
-  drain health, which is the intended design.
+
+  reward < 0.15  →  agent loses one life.  Covers:
+    Security Breach (0.02), Business Disruption (0.05), Wrong Procedure (0.10)
+
+  Cautious / partial-credit scores (≥ 0.35) NEVER drain health.
+
+VALID AGENT ACTIONS
+────────────────────────────────────────────────────────────────────────
+  MARK_SAFE      – deliver to inbox (use ONLY for confirmed-safe email)
+  MOVE_TO_SPAM   – bulk/unsolicited mail; no active threat
+  QUARANTINE     – hold for analyst review
+  BLOCK_DOMAIN   – perimeter block; for confirmed phishing / BEC sources
 """
 
+from __future__ import annotations
+
 from typing import Tuple
 
-# ── Reward constants ───────────────────────────────────────────────────────────
-# Change values here only — nowhere else in the codebase hard-codes these.
-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
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 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   # Perfect triage decision
+R_MALWARE_QUARANTINE   = 0.75   # MALWARE isolated — strong containment
+R_PHISH_BEC_QUARANTINE = 0.60   # PHISH/BEC held but 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
 
 # env.py compares `reward < HEALTH_DRAIN_THRESHOLD` to decide life-loss.
-HEALTH_DRAIN_THRESHOLD  = 0.15
+# Must remain above R_BREACH and R_DISRUPTION and R_WRONG_PROCEDURE,
+# and below R_SPAM_QUARANTINE so cautious calls never lose a life.
+HEALTH_DRAIN_THRESHOLD = 0.15
+
+# Convenience set used by grade_action() internal logic
+_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
+# ══════════════════════════════════════════════════════════════════════════════
 
 def grade_action(
     agent_output: str,
@@ -58,80 +85,124 @@ def grade_action(
     email_type: str,
 ) -> Tuple[float, str]:
     """
-    Grade one triage decision and return a reward in the open interval (0, 1).
+    Grade one SOC triage decision and return a reward in (0.0, 1.0).
 
     Parameters
     ----------
-    agent_output    : Raw action string from the LLM (case-insensitive).
+    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.
+        reward  : float strictly in (0.0, 1.0) — NEVER 0, NEVER 1.
         message : short verdict string for logs and agent feedback.
+
+    Decision Tree
+    -------------
+    1. PERFECT MATCH             → 0.95
+    2. SECURITY BREACH           → 0.02  (threat rubber-stamped into inbox)
+    3. BUSINESS DISRUPTION       → 0.05  (clean email blocked)
+    4. Cautious partial credit:
+       MALWARE → QUARANTINE      → 0.75
+       PHISH/BEC → QUARANTINE    → 0.60
+       SPAM → BLOCK_DOMAIN       → 0.40
+       SPAM → QUARANTINE         → 0.35
+    5. General wrong procedure   → 0.10  (catch-all; no harm done)
     """
     agent_action    = agent_output.strip().upper()
     expected_action = expected_output.strip().upper()
     etype           = email_type.strip().upper()
 
-    THREAT_TYPES = {"PHISH", "BEC", "MALWARE", "SPAM"}
+    # ── Normalise unexpected agent output ────────────────────────────────────
+    # If the LLM produces an unrecognised token, treat as wrong procedure
+    # rather than raising (avoids crashing 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",
+        )
 
-    # ── 1. PERFECT MATCH ──────────────────────────────────────────────────────
+    # ── 1. PERFECT MATCH ─────────────────────────────────────────────────────
     if agent_action == expected_action:
         return R_PERFECT, "PERFECT_TRIAGE: Correct action taken"
 
-    # ── 2. SECURITY BREACH — most severe ──────────────────────────────────────
-    # A dangerous email was rubber-stamped into the inbox.
-    if etype in THREAT_TYPES and agent_action == "MARK_SAFE":
-        return R_BREACH, "SECURITY_BREACH: Threat allowed into the network"
+    # ── 2. SECURITY BREACH — most severe outcome ──────────────────────────────
+    # A dangerous email was rubber-stamped as safe and delivered to the inbox.
+    # This covers ALL threat types (PHISH, BEC, MALWARE, SPAM).
+    if etype in _THREAT_TYPES and agent_action == "MARK_SAFE":
+        return (
+            R_BREACH,
+            f"SECURITY_BREACH: {etype} threat allowed into the network unimpeded",
+        )
 
-    # ── 3. BUSINESS DISRUPTION — severe false positive ────────────────────────
-    # A clean email was blocked or quarantined.
-    if etype == "SAFE" and agent_action in {"BLOCK_DOMAIN", "QUARANTINE", "MOVE_TO_SPAM"}:
-        return R_DISRUPTION, "BUSINESS_DISRUPTION: Legitimate communication blocked"
+    # ── 3. BUSINESS DISRUPTION — severe false-positive ───────────────────────
+    # A clean, legitimate email was blocked, quarantined, or spammed.
+    if etype == "SAFE" and agent_action in _BLOCKED_MOVES:
+        return (
+            R_DISRUPTION,
+            "BUSINESS_DISRUPTION: Legitimate communication was incorrectly blocked",
+        )
 
-    # ── 4. PARTIAL CREDIT — cautious but suboptimal ───────────────────────────
+    # ── 4. PARTIAL CREDIT — cautious but sub-optimal ─────────────────────────
 
-    # MALWARE → QUARANTINE: isolation is the textbook response.
-    # This branch fires only when the scenario's `correct` field differs
-    # (edge case for future scenario expansions — included for robustness).
+    # MALWARE → QUARANTINE: isolation is the textbook containment response.
+    # The scenario may require BLOCK_DOMAIN in future expansions; this branch
+    # handles that gracefully and still awards near-full credit.
     if etype == "MALWARE" and agent_action == "QUARANTINE":
-        return R_MALWARE_QUARANTINE, "CAUTIOUS: Malware isolated — strong containment"
+        return (
+            R_MALWARE_QUARANTINE,
+            "CAUTIOUS: Malware isolated via QUARANTINE — strong containment; "
+            "no further propagation risk",
+        )
 
-    # PHISH/BEC → QUARANTINE: threat is held but the domain keeps sending.
+    # PHISH/BEC → QUARANTINE: threat is held but the sending domain stays live.
+    # Preferred action is BLOCK_DOMAIN; QUARANTINE is under-response.
     if etype in {"PHISH", "BEC"} and agent_action == "QUARANTINE":
-        return R_PHISH_BEC_QUARANTINE, (
-            "UNDER_RESPONSE: Threat quarantined but domain still active — "
-            "consider BLOCK_DOMAIN"
+        return (
+            R_PHISH_BEC_QUARANTINE,
+            f"UNDER_RESPONSE: {etype} quarantined but source domain still active "
+            "— consider BLOCK_DOMAIN to prevent further delivery",
         )
 
-    # SPAM → BLOCK_DOMAIN: blocks the domain, but permanently for mere spam.
+    # SPAM → BLOCK_DOMAIN: disproportionate escalation for mere spam;
+    # burns perimeter block-list capacity on low-severity senders.
     if etype == "SPAM" and agent_action == "BLOCK_DOMAIN":
-        return R_SPAM_BLOCK, (
+        return (
+            R_SPAM_BLOCK,
             "OVER_ESCALATION: BLOCK_DOMAIN is disproportionate for SPAM — "
-            "prefer MOVE_TO_SPAM"
+            "prefer MOVE_TO_SPAM to avoid exhausting block-list resources",
         )
 
-    # SPAM → QUARANTINE: lighter over-escalation, still clogs the analyst queue.
+    # SPAM → QUARANTINE: lighter over-escalation; clogs the analyst review queue.
     if etype == "SPAM" and agent_action == "QUARANTINE":
-        return R_SPAM_QUARANTINE, (
+        return (
+            R_SPAM_QUARANTINE,
             "OVER_ESCALATION: QUARANTINE wastes analyst capacity for SPAM — "
-            "prefer MOVE_TO_SPAM"
+            "prefer MOVE_TO_SPAM",
         )
 
-    # ── 5. GENERAL INCORRECT PROCEDURE ───────────────────────────────────────
-    return R_WRONG_PROCEDURE, "INCORRECT_PROCEDURE: Decision does not match security policy"
+    # ── 5. GENERAL INCORRECT PROCEDURE (catch-all) ────────────────────────────
+    # Wrong action with no direct security or operational breach impact.
+    return (
+        R_WRONG_PROCEDURE,
+        f"INCORRECT_PROCEDURE: '{agent_action}' does not match security policy "
+        f"for {etype} email (expected: {expected_action})",
+    )
+
 
+# ═════════════════════════��════════════════════════════════════════════════════
+# CALCULATE_OVERALL_SCORE
+# ══════════════════════════════════════════════════════════════════════════════
 
 def calculate_overall_score(task_scores: list) -> float:
     """
-    Compute the final benchmark score from a list of per-task rewards.
+    Compute the final benchmark score from a list of per-step rewards.
 
-    The result is clamped to (R_BREACH, R_PERFECT) — matching the per-step
+    The result is clamped to [R_BREACH, R_PERFECT] — maintaining the
     open-interval contract — so downstream consumers always receive a float
-    that is strictly greater than 0 and strictly less than 1.
+    strictly greater than 0 and strictly less than 1.
 
     Parameters
     ----------
@@ -139,14 +210,19 @@ def calculate_overall_score(task_scores: list) -> float:
 
     Returns
     -------
-    float in (0.0, 1.0) — never exactly 0 or 1.
+    float in (R_BREACH, R_PERFECT) — always a valid open-interval value.
+
+    Edge cases
+    ----------
+    • Empty list → R_BREACH (minimum signal, not zero)
+    • Single-step list → that step's reward (clamped)
+    • All-perfect run → 0.95 (R_PERFECT), not 1.0
     """
     if not task_scores:
-        # No tasks completed: return the minimum signal value, not zero.
-        return R_BREACH
+        return R_BREACH  # No tasks completed: return minimum signal, not zero
 
     raw_avg = sum(task_scores) / len(task_scores)
 
-    # Clamp strictly within the open interval boundaries.
+    # Clamp strictly within the open-interval boundary constants.
     clamped = max(R_BREACH, min(R_PERFECT, raw_avg))
     return round(clamped, 4)