Update grader.py

grader.py

@@ -2,89 +2,91 @@
 grader.py – PhishGuard-Env  |  Triage Scoring & Task Registry
 ==============================================================
 
-WHAT CHANGED (and why)
+TASK IDENTITY CONTRACT
 -----------------------
-The OpenEnv validator's "Not enough tasks with graders" check works at the
-TASK-TYPE level, not the scenario level.  It counts distinct `task_id`
-values that appear in graded step responses, and it verifies upfront (via
-GET /metadata) that the environment declares at least 3 task types.
-
-Previous design used level IDs ("lv1"…"lv10") as task_id, which gave the
-validator 10 IDs each appearing once — not 5 semantic categories appearing
-multiple times.  The validator requires the latter.
-
-NEW DESIGN
-----------
-  task_id    : semantic category string, one of:
-                 "task_spam"      – unsolicited bulk mail
-                 "task_phishing"  – credential-harvesting / typosquat attacks
-                 "task_safe"      – legitimate emails (false-positive test)
-                 "task_malware"   – executable / macro payload delivery
-                 "task_bec"       – business email compromise / wire-fraud
-
-  scenario_id : individual level identifier ("lv1"…"lv10"), kept for logs.
-
-TASK_REGISTRY (exported)
-------------------------
-A dict consumed by env.py to:
-  1. Populate GET /metadata so the validator sees total_tasks ≥ 3 before
-     running any episode.
-  2. Drive grade_action() dispatch — all grading logic is keyed on task_id,
-     not on the raw email_type string, so the two are always in sync.
-
-REWARD SCALE CONTRACT (unchanged)
-----------------------------------
-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)
--------------------------------------------
+The OpenEnv validator counts "tasks with graders" by matching the
+`task_id` values declared in openenv.yaml against `task_id` values
+returned in /step responses at runtime.
+
+Both sides must use the SAME five semantic IDs:
+
+    task_spam       task_phishing       task_safe
+    task_malware    task_bec
+
+These are defined once in TASK_REGISTRY below and imported everywhere
+else — env.py, openenv.yaml is written to match them exactly.
+
+REWARD SCALE CONTRACT
+---------------------
+All rewards strictly in the open interval (0.0, 1.0).
+0 and 1 are NEVER returned — guaranteed by safe_score().
+
+    safe_score(raw) = 0.01 + 0.98 * raw,   raw in [0.0, 1.0]
+
+Full reward table
+-----------------
+  Outcome                                    raw     safe     Rationale
+  ─────────────────────────────────────────────────────────────────────
+  Perfect match                              1.00    0.99     Near-ideal
+  MALWARE → QUARANTINE (containment)         0.75    0.745    Correct call
+  PHISH/BEC → QUARANTINE (under-response)   0.60    0.598    Domain still active
+  SPAM → BLOCK_DOMAIN (over-escalation)     0.40    0.402    Disproportionate
+  SPAM → QUARANTINE (lighter escalation)    0.35    0.353    Wastes analyst queue
+  General wrong procedure                   0.10    0.108    Wrong but no harm
+  Business Disruption (SAFE → blocked)      0.05    0.059    Ops impact
+  Security Breach (threat → MARK_SAFE)      0.02    0.0296   Catastrophic
+
+Health-drain threshold
+----------------------
   HEALTH_DRAIN_THRESHOLD = 0.15
-  reward < 0.15  →  env.py deducts one life from the agent.
-  Covers: Security Breach (0.02), Business Disruption (0.05),
-          Wrong Procedure (0.10).
-  Cautious / partial-credit scores (≥ 0.35) never drain health.
+  env.py deducts one life when safe_score result < 0.15.
+  Covers: Breach (0.030), Disruption (0.059), Wrong (0.108).
+  Cautious/partial scores (≥ 0.35 raw → ≥ 0.353 safe) never drain health.
 """
 
 from typing import Tuple
 
-# ── Reward constants ───────────────────────────────────────��───────────────────
-# Defined once here; env.py imports them — no numeric literals elsewhere.
-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
+# ── safe_score ─────────────────────────────────────────────────────────────────
+# Linear bijection [0,1] → (0.01, 0.99).
+# Guarantees the open-interval contract at every call site.
+def safe_score(raw: float) -> float:
+    """Map a raw score in [0.0, 1.0] to a safe score strictly in (0.0, 1.0)."""
+    clamped = max(0.0, min(1.0, raw))
+    result  = 0.01 + 0.98 * clamped
+    assert 0.0 < result < 1.0, f"safe_score contract violated: {result}"
+    return round(result, 6)
+
+
+# ── Reward constants (raw values, passed through safe_score before use) ────────
+_R_PERFECT_RAW              = 1.00
+_R_MALWARE_QUARANTINE_RAW   = 0.75
+_R_PHISH_BEC_QUARANTINE_RAW = 0.60
+_R_SPAM_BLOCK_RAW           = 0.40
+_R_SPAM_QUARANTINE_RAW      = 0.35
+_R_WRONG_PROCEDURE_RAW      = 0.10
+_R_DISRUPTION_RAW           = 0.05
+_R_BREACH_RAW               = 0.02
+
+# Public safe-score constants (what env.py, inference.py actually use)
+R_PERFECT              = safe_score(_R_PERFECT_RAW)               # 0.99
+R_MALWARE_QUARANTINE   = safe_score(_R_MALWARE_QUARANTINE_RAW)    # 0.745
+R_PHISH_BEC_QUARANTINE = safe_score(_R_PHISH_BEC_QUARANTINE_RAW)  # 0.598
+R_SPAM_BLOCK           = safe_score(_R_SPAM_BLOCK_RAW)            # 0.402
+R_SPAM_QUARANTINE      = safe_score(_R_SPAM_QUARANTINE_RAW)       # 0.353
+R_WRONG_PROCEDURE      = safe_score(_R_WRONG_PROCEDURE_RAW)       # 0.108
+R_DISRUPTION           = safe_score(_R_DISRUPTION_RAW)            # 0.059
+R_BREACH               = safe_score(_R_BREACH_RAW)                # 0.0296
 
 HEALTH_DRAIN_THRESHOLD = 0.15   # env.py: reward < this → lose one life
 
 
 # ── Task Registry ──────────────────────────────────────────────────────────────
-# This is the single source of truth for task identity.
+# SINGLE SOURCE OF TRUTH for task identity.
 #
-# Keys   : task_id strings — what the OpenEnv validator counts.
-# Values : metadata dict consumed by GET /metadata and the grading dispatch.
-#
-#   email_types   : set of raw email type strings that map to this task_id.
-#                   Used by grade_action() to look up the correct grader.
-#   description   : human-readable summary for the /metadata endpoint.
-#   threat        : True  → dangerous email; graded as a threat task.
-#                   False → legitimate email; graded as a false-positive task.
+# Keys must exactly match the task `id` values in openenv.yaml.
+# The `email_types` set drives grade_action() dispatch.
+# `threat` determines Breach vs Disruption penalty direction.
 # ─────────────────────────────────────────────────────────────────────────────
-
 TASK_REGISTRY: dict[str, dict] = {
     "task_spam": {
         "email_types": {"SPAM"},
@@ -128,8 +130,7 @@ TASK_REGISTRY: dict[str, dict] = {
     },
 }
 
-# Reverse lookup: raw email_type string → task_id
-# Built once at import time; O(1) per grade_action call.
+# Reverse lookup: raw email_type string → task_id  (built once at import time)
 _TYPE_TO_TASK_ID: dict[str, str] = {
     etype: task_id
     for task_id, meta in TASK_REGISTRY.items()
@@ -139,11 +140,10 @@ _TYPE_TO_TASK_ID: dict[str, str] = {
 
 def resolve_task_id(email_type: str) -> str:
     """
-    Convert a raw email type string (e.g. "PHISH") to its task_id
-    (e.g. "task_phishing").
+    Convert a raw email type (e.g. "PHISH") → task_id (e.g. "task_phishing").
 
-    Raises ValueError for unrecognised types so misconfigured scenarios
-    fail loudly at development time rather than silently at benchmark time.
+    Raises ValueError for unknown types so misconfigured scenarios fail loudly
+    at startup rather than silently corrupting benchmark results.
     """
     key = email_type.strip().upper()
     if key not in _TYPE_TO_TASK_ID:
@@ -154,7 +154,51 @@ def resolve_task_id(email_type: str) -> str:
     return _TYPE_TO_TASK_ID[key]
 
 
-# ── Core grading function ──────────────────────────────────────────────────────
+# ── Per-task graders (one function per task_id) ────────────────────────────────
+# Each grader takes a list of per-step rewards for that task type and returns
+# a final safe_score. This is what openenv.yaml's grader field points to.
+
+def grade_task_spam(rewards: list[float]) -> float:
+    """Final score for all task_spam steps in an episode."""
+    return _aggregate(rewards)
+
+def grade_task_phishing(rewards: list[float]) -> float:
+    """Final score for all task_phishing steps in an episode."""
+    return _aggregate(rewards)
+
+def grade_task_safe(rewards: list[float]) -> float:
+    """Final score for all task_safe steps in an episode."""
+    return _aggregate(rewards)
+
+def grade_task_malware(rewards: list[float]) -> float:
+    """Final score for all task_malware steps in an episode."""
+    return _aggregate(rewards)
+
+def grade_task_bec(rewards: list[float]) -> float:
+    """Final score for all task_bec steps in an episode."""
+    return _aggregate(rewards)
+
+
+# Map task_id → its grader function (consumed by env.py /grade endpoint)
+TASK_GRADERS: dict[str, callable] = {
+    "task_spam":      grade_task_spam,
+    "task_phishing":  grade_task_phishing,
+    "task_safe":      grade_task_safe,
+    "task_malware":   grade_task_malware,
+    "task_bec":       grade_task_bec,
+}
+
+
+def _aggregate(rewards: list[float]) -> float:
+    """Average a list of per-step rewards, return as a safe_score."""
+    if not rewards:
+        return R_BREACH
+    raw_avg = sum(rewards) / len(rewards)
+    # rewards are already safe_scores; average is still in (0,1) but clamp defensively
+    return round(max(R_BREACH, min(R_PERFECT, raw_avg)), 6)
+
+
+# ── Core step grading function ─────────────────────────────────────────────────
 
 def grade_action(
     agent_output: str,
@@ -162,110 +206,74 @@ def grade_action(
     email_type: str,
 ) -> Tuple[float, str]:
     """
-    Grade one SOC triage decision and return a reward in the open interval
-    (0.0, 1.0) — the endpoints 0 and 1 are NEVER returned.
-
-    Parameters
-    ----------
-    agent_output    : Raw action string from the LLM (case-insensitive).
-    expected_output : Ground-truth correct action for this scenario.
-    email_type      : Raw threat category — PHISH | BEC | MALWARE | SPAM | SAFE.
-                      The function resolves this to a task_id internally.
+    Grade one SOC triage decision.
 
     Returns
     -------
     (reward, message)
-        reward  : float strictly in (0.0, 1.0).
+        reward  : float strictly in (0.0, 1.0) via safe_score().
         message : short verdict string for logs and agent feedback.
     """
     agent_action    = agent_output.strip().upper()
     expected_action = expected_output.strip().upper()
-    etype           = email_type.strip().upper()
 
-    # Resolve task_id for this email type (validates the type in the process).
-    task_id  = resolve_task_id(etype)
+    # resolve_task_id validates the type — raises ValueError on bad input
+    task_id  = resolve_task_id(email_type)
     is_threat = TASK_REGISTRY[task_id]["threat"]
 
-    # ── 1. PERFECT MATCH ──────────────────────────────────────────────────────
+    # ── 1. Perfect match ───────────────────────────────────────────────────
     if agent_action == expected_action:
         return R_PERFECT, "PERFECT_TRIAGE: Correct action taken"
 
-    # ── 2. SECURITY BREACH ────────────────────────────────────────────────────
-    # Threat email rubber-stamped as safe — worst possible outcome.
+    # ── 2. Security Breach — threat marked safe ────────────────────────────
     if is_threat and agent_action == "MARK_SAFE":
         return R_BREACH, "SECURITY_BREACH: Threat allowed into the network"
 
-    # ── 3. BUSINESS DISRUPTION ────────────────────────────────────────────────
-    # Legitimate email blocked or quarantined — costly false positive.
+    # ── 3. Business Disruption — safe email blocked ────────────────────────
     if not is_threat and agent_action in {"BLOCK_DOMAIN", "QUARANTINE", "MOVE_TO_SPAM"}:
         return R_DISRUPTION, "BUSINESS_DISRUPTION: Legitimate communication blocked"
 
-    # ── 4. PARTIAL CREDIT ─────────────────────────────────────────────────────
-
-    # MALWARE quarantined: textbook containment response.
-    # Branch fires only when scenario's `correct` field differs from QUARANTINE
-    # (kept for robustness against future scenario expansions).
+    # ── 4. Partial credit ───────────────────────────��──────────────────────
     if task_id == "task_malware" and agent_action == "QUARANTINE":
         return R_MALWARE_QUARANTINE, "CAUTIOUS: Malware isolated — strong containment"
 
-    # PHISH or BEC quarantined instead of domain-blocked.
-    # Threat is held, but the malicious domain stays active.
     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 domain fully blocked — disproportionate but not dangerous.
     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"
         )
 
-    # SPAM quarantined — lighter over-escalation, clogs analyst queue.
     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"
         )
 
-    # ── 5. GENERAL INCORRECT PROCEDURE ───────────────────────────────────────
+    # ── 5. General wrong procedure ─────────────────────────────────────────
     return R_WRONG_PROCEDURE, "INCORRECT_PROCEDURE: Decision does not match security policy"
 
 
-# ── Final score aggregation ────────────────────────────────────────────────────
+# ── Episode-level aggregation ──────────────────────────────────────────────────
 
-def calculate_overall_score(task_scores: list) -> float:
+def calculate_overall_score(task_scores: list[float]) -> float:
     """
-    Compute the final benchmark score from per-step rewards.
-
-    Result is clamped to (R_BREACH, R_PERFECT) to honour the open-interval
-    contract at the episode level as well as the step level.
-
-    Parameters
-    ----------
-    task_scores : list of floats, each in (0.0, 1.0).
-
-    Returns
-    -------
-    float in (0.0, 1.0) — never exactly 0 or 1.
+    Compute a final benchmark score from all per-step rewards.
+    Returns a float strictly in (0.0, 1.0).
     """
-    if not task_scores:
-        return R_BREACH  # Minimum signal value, not zero
-
-    raw_avg = sum(task_scores) / len(task_scores)
-    clamped = max(R_BREACH, min(R_PERFECT, raw_avg))
-    return round(clamped, 4)
+    return _aggregate(task_scores)
 
 
 def calculate_per_task_scores(
     task_score_map: dict[str, list[float]],
 ) -> dict[str, float]:
     """
-    Compute per-task-type average scores from a map of task_id → [rewards].
-
-    Used by GET /state to give the validator per-task breakdowns.
+    Compute per-task-type average scores.
 
     Parameters
     ----------
@@ -273,10 +281,10 @@ def calculate_per_task_scores(
 
     Returns
     -------
-    { task_id: averaged_score } — each value in (0.0, 1.0).
+    { task_id: score }  — each value strictly in (0.0, 1.0).
     """
     return {
-        task_id: calculate_overall_score(scores)
+        task_id: _aggregate(scores)
         for task_id, scores in task_score_map.items()
         if scores
     }