Update env.py

env.py

@@ -1,46 +1,6 @@
 """
 env.py – PhishGuard-Env | FastAPI Environment Server
 =======================================================
-
-ARCHITECTURE ROLE
------------------
-This file IS the environment. It runs as a persistent FastAPI server on
-Hugging Face Spaces (port 7860). The inference agent (inference.py) is a
-separate process that interacts with it exclusively through HTTP — it NEVER
-imports this module directly.
-
-Endpoints
----------
-  GET  /        → root metadata: lists task types, version, and task registry
-  POST /reset   → reset for a chosen difficulty level; returns first email observation
-  POST /step    → submit one triage action; returns (obs, reward, done, info)
-  GET  /state   → read-only snapshot of health, score, task index
-  GET  /health  → liveness probe for HF Spaces / load-balancers
-
-LEVEL DESIGN
-------------
-  easy   → lv1, lv2, lv3  (SPAM, PHISH, SAFE)
-  medium → lv4, lv5, lv6, lv7  (MALWARE, SAFE, BEC, PHISH)
-  hard   → lv8, lv9, lv10 (MALWARE, PHISH, BEC)
-
-TASK REGISTRY (3 graded tasks — satisfies OpenEnv validator)
--------------------------------------------------------------
-  task_phish    → PHISH scenarios (credential harvesting, typosquat)
-  task_bec      → BEC scenarios (wire fraud, supply-chain compromise)
-  task_internal → SAFE / MALWARE scenarios (internal mail & malware triage)
-
-FIX APPLIED
------------
-The OpenEnv validator calls POST /reset with an EMPTY body (no JSON at all).
-The previous endpoint signature was:
-    async def reset(request: ResetRequest)
-FastAPI requires a body when the parameter is not Optional, so it returned:
-    HTTP 422 Unprocessable Entity
-    {"detail":[{"type":"missing","loc":["body"],"msg":"Field required"}]}
-
-Fix: make the request body Optional. When None, default to level="easy".
-    async def reset(request: Optional[ResetRequest] = None)
-This is the ONLY change from the live repo.
 """
 
 from __future__ import annotations
@@ -80,8 +40,18 @@ from grader import (
     HEALTH_DRAIN_THRESHOLD,
     calculate_overall_score,
     grade_action,
+    TASK_REGISTRY,
 )
 
+# ── Single source of truth: email type → semantic task ID ────────────────────
+_TYPE_TO_TASK_ID: dict[str, str] = {
+    "SPAM":    "task_spam",
+    "PHISH":   "task_phishing",
+    "SAFE":    "task_safe",
+    "MALWARE": "task_malware",
+    "BEC":     "task_bec",
+}
+
 # ═════════════════════════════════════════════════════════════════════════════
 # SCENARIO DEFINITIONS (lv1 → lv10)
 # ═════════════════════════════════════════════════════════════════════════════
@@ -313,12 +283,13 @@ LEVEL_MAP: dict[str, list[str]] = {
 
 _SCENARIO_BY_ID: dict[str, dict] = {s["id"]: s for s in SCENARIOS}
 
+
 # ═════════════════════════════════════════════════════════════════════════════
 # REQUEST SCHEMA
 # ═════════════════════════════════════════════════════════════════════════════
 
 class ResetRequest(BaseModel):
-    level: str = "easy"   # "easy" | "medium" | "hard"
+    level: str = "easy"
 
 
 # ═════════════════════════════════════════════════════════════════════════════
@@ -344,8 +315,8 @@ class PhishGuardEnv(OpenEnv):
             raise ValueError(
                 f"Unknown level '{level}'. Valid choices: easy | medium | hard"
             )
-        ids = LEVEL_MAP[level]
-        subset = [dict(_SCENARIO_BY_ID[sid]) for sid in ids]
+        ids       = LEVEL_MAP[level]
+        subset    = [dict(_SCENARIO_BY_ID[sid]) for sid in ids]
         remaining = [s for s in SCENARIOS if s["id"] not in ids]
         subset.extend(remaining)
         random.shuffle(subset)
@@ -370,9 +341,14 @@ class PhishGuardEnv(OpenEnv):
         return first_task["data"]
 
     def step(self, action_str: str) -> tuple:
+
+        # ── Guard: episode already over ───────────────────────────────────────
+        # task_group MUST be present here — the /step FastAPI endpoint reads
+        # info["task_group"] and crashes with KeyError if it is missing.
         if self._is_over():
             return None, R_BREACH, True, {
                 "task_id":    None,
+                "task_group": None,       # ← CRITICAL: was missing, caused KeyError
                 "is_correct": False,
                 "health":     self.health,
                 "feedback":   "Episode already ended. Call /reset to start a new one.",
@@ -380,18 +356,12 @@ class PhishGuardEnv(OpenEnv):
                 "task_scores": list(self.task_scores),
             }
 
-        current_task = self.scenarios[self.current_task_idx]
-        _TYPE_TO_TASK_ID = {
-            "SPAM": "task_spam", "PHISH": "task_phishing", "SAFE":"task_safe", "MALWARE": "task_malware", "BEC":"task_bec",
-        }
-        task_id      = _TYPE_TO_TASK_ID.get(current_task["type"].upper(),current_task["id"])
-        semantic_task_id = {
-            "SPAM":    "task_spam",
-            "PHISH":   "task_phishing",
-            "SAFE":    "task_safe",
-            "MALWARE": "task_malware",
-            "BEC":     "task_bec"
-        }.get(current_task["type"].upper(), current_task["id"])
+        current_task     = self.scenarios[self.current_task_idx]
+        task_group       = current_task["level"]
+        # Use the module-level _TYPE_TO_TASK_ID — defined once, used everywhere
+        semantic_task_id = _TYPE_TO_TASK_ID.get(
+            current_task["type"].upper(), current_task["id"]
+        )
 
         reward, verdict_msg = grade_action(
             action_str,
@@ -404,7 +374,11 @@ class PhishGuardEnv(OpenEnv):
 
         log.info(
             "Step | level=%s | task=%s | action=%s | reward=%.4f | verdict=%s",
-            self.active_level, task_id, action_str.strip().upper(), reward, verdict_msg,
+            self.active_level,
+            semantic_task_id,
+            action_str.strip().upper(),
+            reward,
+            verdict_msg,
         )
 
         if reward < HEALTH_DRAIN_THRESHOLD:
@@ -438,8 +412,8 @@ class PhishGuardEnv(OpenEnv):
         )
 
         return obs, reward, done, {
-            "task_id":    semantic_task_id,
-            "task_group": current_task["level"],
+            "task_id":    semantic_task_id,   # "task_spam" not "lv1"
+            "task_group": task_group,
             "is_correct": reward >= R_PERFECT,
             "health":     self.health,
             "feedback":   feedback,
@@ -470,7 +444,7 @@ app = FastAPI(
     ),
     version="3.1.0",
     lifespan=lifespan,
-    docs_url="/docs",        # Swagger UI — required for OpenEnv validator inspection
+    docs_url="/docs",
     redoc_url="/redoc",
 )
 
@@ -483,15 +457,13 @@ app.add_middleware(
 )
 
 
+# ── Root metadata ─────────────────────────────────────────────────────────────
 @app.get("/", tags=["Meta"])
 async def root_metadata() -> dict:
     """
-    Root metadata endpoint — required by the OpenEnv Phase 2 validator.
-
-    Returns a JSON document listing all task types registered in this
-    environment so the validator can confirm there are ≥3 graded tasks.
+    Root metadata — lists all registered task types so the validator
+    can confirm at least 3 graded tasks exist before running the episode.
     """
-    from grader import TASK_REGISTRY
     return {
         "name":        "PhishGuard-Env",
         "version":     "3.1.0",
@@ -500,60 +472,39 @@ async def root_metadata() -> dict:
         "port":        7860,
         "tasks": [
             {
-                "id":          "task_phish",
-                "description": "Identify phishing emails (credential harvesting, typosquat domains) and block at perimeter.",
-                "grader":      "grader.grade_easy",
-            },
-            {
-                "id":          "task_bec",
-                "description": "Detect Business Email Compromise (wire fraud, CEO impersonation, supply-chain attacks).",
-                "grader":      "grader.grade_medium",
-            },
-            {
-                "id":          "task_internal",
-                "description": "Triage internal / malware emails: pass safe mail, quarantine malware payloads.",
-                "grader":      "grader.grade_hard",
-            },
-        ],
-        "task_registry": {
-            task_id: {
+                "id":          task_id,
                 "description": meta["description"],
-                "threat":      meta["threat"],
+                "grader":      "grader.grade_action",
             }
             for task_id, meta in TASK_REGISTRY.items()
-        },
-        "levels": list(LEVEL_MAP.keys()),
+        ],
+        "levels":    list(LEVEL_MAP.keys()),
         "endpoints": {
-            "reset": "POST /reset",
-            "step":  "POST /step",
-            "state": "GET /state",
-            "health":"GET /health",
-            "docs":  "GET /docs",
+            "reset":  "POST /reset",
+            "step":   "POST /step",
+            "state":  "GET /state",
+            "health": "GET /health",
+            "docs":   "GET /docs",
         },
     }
 
 
+# ── Liveness probe ─────────────────────────────────────────────────��──────────
 @app.get("/health", tags=["Meta"])
 async def health_probe() -> dict:
     return {"status": "ok", "env": "PhishGuard-Env", "version": "3.1.0"}
 
 
-# ── THE FIX IS HERE ───────────────────────────────────────────────────────────
-# The validator calls POST /reset with an empty body.
-# Making `request` Optional with a default of None means FastAPI no longer
-# requires a body — when nothing is sent, we just use level="easy".
+# ── Reset ─────────────────────────────────────────────────────────────────────
 @app.post("/reset", tags=["Environment"])
 async def reset(request: Optional[ResetRequest] = None) -> dict:
     """
     Reset the environment for a new episode.
-
-    Body (JSON) — fully optional, send {} or nothing at all:
-        { "level": "easy" | "medium" | "hard" }
-
-    Defaults to level="easy" when no body is provided.
+    Body is fully optional — validator calls this with no body at all.
+    Defaults to level='easy' when nothing is sent.
     """
     if request is None:
-        request = ResetRequest()   # defaults to level="easy"
+        request = ResetRequest()
 
     level = request.level.lower()
     if level not in LEVEL_MAP:
@@ -562,28 +513,19 @@ async def reset(request: Optional[ResetRequest] = None) -> dict:
             detail=f"Invalid level '{level}'. Must be one of: easy | medium | hard",
         )
 
-    obs          = _env.reset(level=level)
-    first_task   = _env.scenarios[_env.current_task_idx]
-    _TYPE_TO_SEMANTIC = {
-        "SPAM":    "task_spam",
-        "PHISH":   "task_phishing",
-        "SAFE":    "task_safe",
-        "MALWARE": "task_malware",
-        "BEC":     "task_bec",
-    }
-    _TYPE_TO_TASK_ID = {
-        "SPAM": "task_spam", "PHISH": "task_phishing",
-        "SAFE": "task_safe", "MALWARE": "task_malware", "BEC": "task_bec",
-    }
+    obs        = _env.reset(level=level)
+    first_task = _env.scenarios[_env.current_task_idx]
+
     return {
         "observation": obs,
-        "task_id": _TYPE_TO_TASK_ID.get(first_task["type"].upper(), first_task["id"]),
-        "task_group": first_task["level"],
-        "level": _env.active_level,
+        "task_id":     _TYPE_TO_TASK_ID.get(first_task["type"].upper(), first_task["id"]),
+        "task_group":  first_task["level"],
+        "level":       _env.active_level,
         "total_tasks": len(_env.scenarios),
     }
 
 
+# ── Step ──────────────────────────────────────────────────────────────────────
 @app.post("/step", tags=["Environment"])
 async def step(action: PhishAction) -> dict:
     obs, reward, done, info = _env.step(action.action)
@@ -598,6 +540,7 @@ async def step(action: PhishAction) -> dict:
     }
 
 
+# ── State ─────────────────────────────────────────────────────────────────────
 @app.get("/state", tags=["Environment"])
 async def state() -> dict:
     return {
@@ -611,6 +554,7 @@ async def state() -> dict:
     }
 
 
+# ── Entry point ───────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("env:app", host="0.0.0.0", port=7860, reload=False, log_level="info")
+    uvicorn.run("env:app", host="0.0.0.0", port=7860, reload=False, log_level="info")