Update inference.py

inference.py

@@ -1,46 +1,78 @@
 """
-inference.py – PhishGuard-Env LLM Driver
-=========================================
-
-KEY CHANGES vs. previous version
-----------------------------------
-1. SUCCESS THRESHOLD RECALIBRATED for open-interval rewards:
-   Previous code used `success = final_score >= 0.7`. With rewards now in
-   (0.0, 1.0), a 0.7 threshold is still semantically correct (it sits between
-   R_PHISH_BEC_QUARANTINE=0.60 and R_MALWARE_QUARANTINE=0.75), but the
-   constant is now imported from grader.py as SUCCESS_THRESHOLD so it stays
-   in sync with any future reward-table changes.
-
-2. ERROR STEP REWARD UPDATED:
-   Transient-error steps previously appended 0.0 to the rewards list,
-   violating the open-interval contract. Now appends R_WRONG_PROCEDURE (0.10)
-   — a non-zero signal that reflects "the agent failed to produce an action".
-
-3. All previously documented bugs remain fixed:
-   - HF_TOKEN / OPENAI_API_KEY validation with clear error messages.
-   - base_url trailing-slash normalisation for openai >= 1.25.
-   - Transient errors retry in-place instead of resetting the episode.
-   - Final score uses calculate_overall_score(), not sum/10.
+inference.py – PhishGuard-Env  |  Hybrid LLM + HTTP Driver
+============================================================
+
+ARCHITECTURE — "Hybrid Agent" (resolves the Gemini vs GPT/Claude debate)
+--------------------------------------------------------------------------
+Both sides of the AI architecture debate were PARTIALLY correct:
+
+  ┌─────────────────────────────────────────────────────────────────────┐
+  │  GEMINI was right about:   The agent needs HF_TOKEN + API_BASE_URL  │
+  │                            to call the Qwen LLM "brain" for         │
+  │                            INFERENCE INTELLIGENCE (thinking).        │
+  │                                                                     │
+  │  GPT/CLAUDE was right about: The agent only needs requests.post()   │
+  │                            to hit the /step endpoint of env.py for  │
+  │                            ENVIRONMENT INTERACTION (acting).         │
+  │                                                                     │
+  │  UNIFIED SOLUTION:         Two separate clients, one agent:         │
+  │    llm_client  → OpenAI(base_url=HF_ROUTER)  for THINKING           │
+  │    ENV_BASE_URL → requests.post(ENV_BASE_URL) for ACTING            │
+  └─────────────────────────────────────────────────────────────────────┘
+
+Per-step flow
+-------------
+  1. POST  {ENV_BASE_URL}/reset        → receive first email observation
+  2. llm_client.chat.completions(obs)  → LLM "thinks" → returns action
+  3. POST  {ENV_BASE_URL}/step({action}) → env "acts"  → returns reward, done
+  4. Log [STEP] in OpenEnv format
+  5. Repeat 2-4 until done=True
+  6. Log [END] with final score
+
+Environment Variables
+---------------------
+  Required:
+    HF_TOKEN          : Hugging Face access token (hf_...)
+      OR
+    OPENAI_API_KEY    : Standard OpenAI key if using OpenAI router
+
+  Optional:
+    ENV_BASE_URL      : Base URL of the running env.py server
+                        Default: http://localhost:7860
+                        HF Spaces: https://<org>-<space-name>.hf.space
+    API_BASE_URL      : LLM inference router base URL
+                        Default: https://router.huggingface.co/v1
+    MODEL_NAME        : Inference model ID
+                        Default: Qwen/Qwen2.5-72B-Instruct
+    MAX_RETRIES       : LLM call retries before aborting a step (default 3)
+    RETRY_DELAY       : Seconds between LLM retries (default 2.0)
+    REQUEST_TIMEOUT   : HTTP timeout for env API calls in seconds (default 30)
+
+OpenEnv Logging Format (compliant)
+-----------------------------------
+  [START] task=phishing_triage env=phishguard_v1 model=<MODEL_NAME>
+  [STEP]  step=N action=ACTION reward=R health=H done=true|false error=null
+  [END]   success=true|false steps=N score=S rewards=r1,r2,...
 """
 
-import asyncio
-import os
-import requests
+from __future__ import annotations
+
 import json
-import textwrap
+import os
+import sys
 import time
+
+import requests
 from openai import OpenAI, AuthenticationError, APIConnectionError
 
-from env import env_instance
-from grader import calculate_overall_score, R_WRONG_PROCEDURE
+from grader import calculate_overall_score, R_WRONG_PROCEDURE, R_PERFECT
 
-# ── SUCCESS THRESHOLD ─────────────────────────────────────────────────────────
-# A run is considered successful if the average reward meets this bar.
-# 0.70 sits between "under-response on PHISH/BEC" (0.60) and
-# "correct malware containment" (0.75) — a meaningful quality gate.
-SUCCESS_THRESHOLD = 0.70
+# ══════════════════════════════════════════════════════════════════════════════
+# CONFIGURATION
+# ══════════════════════════════════════════════════════════════════════════════
 
-# ── CONFIGURATION ─────────────────────────────────────────────────────────────
+# ── API Key ───────────────────────────────────────────────────────────────────
+# Priority: HF_TOKEN → OPENAI_API_KEY → API_KEY
 _raw_key = (
     os.getenv("HF_TOKEN")
     or os.getenv("OPENAI_API_KEY")
@@ -52,81 +84,193 @@ if not _raw_key:
     print(
         "❌ ERROR: No API key found in environment variables.\n"
         "\n"
-        "  For Hugging Face:\n"
+        "  For Hugging Face Inference Router:\n"
         "    Linux/macOS : export HF_TOKEN='hf_...'\n"
         "    PowerShell  : $env:HF_TOKEN='hf_...'\n"
         "\n"
         "  For OpenAI:\n"
         "    Linux/macOS : export OPENAI_API_KEY='sk-...'\n"
-        "    PowerShell  : $env:OPENAI_API_KEY='sk-...'\n"
+        "    PowerShell  : $env:OPENAI_API_KEY='sk-...'\n",
+        file=sys.stderr,
     )
-    raise SystemExit(1)
+    sys.exit(1)
 
 API_KEY = _raw_key
 
+# ── LLM Router (the "brain" — for THINKING) ───────────────────────────────────
 # base_url must always end with "/" for openai >= 1.25.
-_raw_base = (os.getenv("https://huggingface.co/spaces/Scalar-hackathon/Phishing-env") or "https://router.huggingface.co/v1").rstrip("/") + "/"
-API_BASE_URL = _raw_base
+_raw_llm_base = (
+    os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
+).rstrip("/") + "/"
+LLM_BASE_URL = _raw_llm_base
+
+# ── Environment Server (the "world" — for ACTING) ─────────────────────────────
+# When running locally: http://localhost:7860
+# When env.py is deployed on HF Spaces: https://<org>-<space-name>.hf.space
+_raw_env_base = (
+    os.getenv("ENV_BASE_URL") or "http://localhost:7860"
+).rstrip("/")
+ENV_BASE_URL = _raw_env_base
+
+# ── Model and runtime settings ────────────────────────────────────────────────
+MODEL_NAME      = os.getenv("MODEL_NAME")    or "Qwen/Qwen2.5-72B-Instruct"
+MAX_RETRIES     = int(os.getenv("MAX_RETRIES",    "3"))
+RETRY_DELAY     = float(os.getenv("RETRY_DELAY",  "2.0"))   # seconds between LLM retries
+REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", "30")) # seconds for env HTTP calls
+
+# A run is successful if the average reward meets this bar.
+# Sits between R_PHISH_BEC_QUARANTINE (0.60) and R_MALWARE_QUARANTINE (0.75).
+SUCCESS_THRESHOLD = 0.70
+
+VALID_ACTIONS = frozenset({"MARK_SAFE", "MOVE_TO_SPAM", "QUARANTINE", "BLOCK_DOMAIN"})
 
-MODEL_NAME  = os.getenv("MODEL_NAME")  or "Qwen/Qwen2.5-72B-Instruct"
-MAX_RETRIES = int(os.getenv("MAX_RETRIES",  "3"))
-RETRY_DELAY = float(os.getenv("RETRY_DELAY", "2.0"))  # seconds between retries
 
-# ── CLIENT ────────────────────────────────────────────────────────────────────
+# ══════════════════════════════════════════════════════════════════════════════
+# LLM CLIENT  (THINKING layer)
+# ══════════════════════════════════════════════════════════════════════════════
+
 try:
-    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    llm_client = OpenAI(base_url=LLM_BASE_URL, api_key=API_KEY)
 except Exception as exc:
-    print(f"❌ Failed to initialise OpenAI client: {exc}")
-    raise SystemExit(1)
-
-# ── SYSTEM PROMPT ─────────────────────────────────────────────────────────────
-SYSTEM_PROMPT = textwrap.dedent("""
-    You are an expert SOC (Security Operations Center) Analyst responsible
-    for email triage in a corporate environment.
-
-    Analyse ALL fields of the provided email metadata:
-      - sender, subject, body, links, has_attachments
-      - spf_record    : SPF authentication result (pass / fail / softfail / none)
-      - dmarc_record  : DMARC policy result (pass / fail / none)
-      - urgency_level : how aggressively the email pressures the recipient
-      - confidence_hint : noisy SIEM / gateway / threat-intel signal — treat
-                          as one data-point, not ground truth
-
-    Choose EXACTLY ONE of these actions:
-      • MARK_SAFE     – Confirmed legitimate; deliver to inbox normally.
-      • MOVE_TO_SPAM  – Unsolicited bulk mail; no active security threat.
-      • QUARANTINE    – Suspicious; hold for deeper analyst investigation.
-      • BLOCK_DOMAIN  – Confirmed phishing / BEC source; block at perimeter.
-
-    Respond ONLY with a valid JSON object — no markdown fences, no preamble:
-    {
-      "action": "<one of the four actions above, in ALL_CAPS>",
-      "reasoning": "<one concise technical sentence explaining your decision>"
-    }
-""").strip()
-
-VALID_ACTIONS = {"MARK_SAFE", "MOVE_TO_SPAM", "QUARANTINE", "BLOCK_DOMAIN"}
-
-
-# ── LLM CALL WITH RETRY ───────────────────────────────────────────────────────
+    print(f"❌ Failed to initialise LLM client: {exc}", file=sys.stderr)
+    sys.exit(1)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# SYSTEM PROMPT
+# ══════════════════════════════════════════════════════════════════════════════
+
+SYSTEM_PROMPT = """
+You are an expert SOC (Security Operations Center) Analyst responsible
+for email triage in a corporate environment.
+
+Analyse ALL fields of the provided email metadata:
+  - sender          : full sender address — check for typosquatting
+  - subject         : urgency language, pressure tactics
+  - body            : instructions, links, attachment requests
+  - links           : check for URL shorteners, suspicious TLDs, mismatches
+  - has_attachments : True = elevated risk unless sender is fully trusted
+  - spf_record      : pass | softfail | fail | none
+  - dmarc_record    : pass | fail | none
+  - urgency_level   : low | medium | high | critical
+  - confidence_hint : noisy SIEM / gateway / threat-intel signal
+                      → treat as ONE data point, NOT ground truth
+
+SPF/DMARC adversarial note: some sophisticated attacks (BEC, supply-chain)
+will show spf_record=pass and dmarc_record=pass on compromised or lookalike
+domains. Never treat authentication pass as definitive clearance.
+
+Choose EXACTLY ONE of these actions:
+  MARK_SAFE     – Confirmed legitimate; deliver to inbox normally.
+  MOVE_TO_SPAM  – Unsolicited bulk mail; no active security threat.
+  QUARANTINE    – Suspicious; hold for deeper analyst investigation.
+  BLOCK_DOMAIN  – Confirmed phishing / BEC source; block at perimeter.
+
+Respond ONLY with a valid JSON object — no markdown fences, no preamble:
+{
+  "action": "<one of the four actions above, in ALL_CAPS>",
+  "reasoning": "<one concise technical sentence explaining your decision>"
+}
+""".strip()
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# ENVIRONMENT HTTP HELPERS  (ACTING layer)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def env_reset() -> dict:
+    """
+    POST {ENV_BASE_URL}/reset → returns the first email observation dict.
+
+    Raises
+    ------
+    RuntimeError if the server is unreachable or returns a non-200 status.
+    """
+    url = f"{ENV_BASE_URL}/reset"
+    try:
+        resp = requests.post(url, timeout=REQUEST_TIMEOUT)
+        resp.raise_for_status()
+        data = resp.json()
+        return data["observation"]
+    except requests.exceptions.ConnectionError as exc:
+        raise RuntimeError(
+            f"Cannot reach env.py server at {ENV_BASE_URL}.\n"
+            f"  → Is it running? Start with: uvicorn env:app --host 0.0.0.0 --port 7860\n"
+            f"  → For HF Spaces, set ENV_BASE_URL to your Space URL.\n"
+            f"  Original error: {exc}"
+        ) from exc
+    except Exception as exc:
+        raise RuntimeError(f"env_reset() failed: {exc}") from exc
+
+
+def env_step(action_str: str) -> tuple[dict | None, float, bool, dict]:
+    """
+    POST {ENV_BASE_URL}/step → returns (obs, reward, done, info).
+
+    Parameters
+    ----------
+    action_str : one of MARK_SAFE | MOVE_TO_SPAM | QUARANTINE | BLOCK_DOMAIN
+
+    Returns
+    -------
+    (observation, reward, done, info)
+    """
+    url = f"{ENV_BASE_URL}/step"
+    payload = {"action": action_str}
+    try:
+        resp = requests.post(url, json=payload, timeout=REQUEST_TIMEOUT)
+        resp.raise_for_status()
+        data    = resp.json()
+        obs     = data.get("observation")
+        reward  = float(data.get("reward", R_WRONG_PROCEDURE))
+        done    = bool(data.get("done", False))
+        info    = data.get("info", {})
+        return obs, reward, done, info
+    except requests.exceptions.ConnectionError as exc:
+        raise RuntimeError(
+            f"Cannot reach env.py server at {ENV_BASE_URL}: {exc}"
+        ) from exc
+    except Exception as exc:
+        raise RuntimeError(f"env_step() failed: {exc}") from exc
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# LLM CALL WITH RETRY  (THINKING layer)
+# ══════════════════════════════════════════════════════════════════════════════
+
 def call_llm_with_retry(obs: dict, step_num: int) -> str:
     """
     Call the LLM for one email observation and return a validated action string.
 
-    Retries up to MAX_RETRIES times on transient errors (JSON parse failures,
-    rate limits, timeouts). Re-raises immediately on unrecoverable errors
-    (AuthenticationError, APIConnectionError).
+    Strategy
+    --------
+    • Up to MAX_RETRIES attempts on transient errors (JSON parse, rate limit,
+      timeout, invalid action token).
+    • Re-raises immediately on unrecoverable errors (AuthenticationError,
+      APIConnectionError) — no point retrying these.
+
+    Parameters
+    ----------
+    obs      : email observation dict from the environment
+    step_num : 1-based step counter (for log messages only)
+
+    Returns
+    -------
+    str — one of MARK_SAFE | MOVE_TO_SPAM | QUARANTINE | BLOCK_DOMAIN
     """
-    last_error = None
+    last_error: Exception | None = None
 
     for attempt in range(1, MAX_RETRIES + 1):
         try:
-            completion = client.chat.completions.create(
+            completion = llm_client.chat.completions.create(
                 model=MODEL_NAME,
                 max_tokens=256,
                 messages=[
                     {"role": "system", "content": SYSTEM_PROMPT},
-                    {"role": "user",   "content": f"Email Data:\n{json.dumps(obs, indent=2)}"},
+                    {
+                        "role": "user",
+                        "content": f"Email Data:\n{json.dumps(obs, indent=2)}",
+                    },
                 ],
                 response_format={"type": "json_object"},
             )
@@ -138,51 +282,82 @@ def call_llm_with_retry(obs: dict, step_num: int) -> str:
             if action not in VALID_ACTIONS:
                 raise ValueError(
                     f"LLM returned an invalid action: '{action}'. "
-                    f"Must be one of {sorted(VALID_ACTIONS)}."
+                    f"Expected one of: {sorted(VALID_ACTIONS)}"
                 )
 
             return action
 
         except (AuthenticationError, APIConnectionError):
-            raise  # Unrecoverable — propagate immediately
+            raise  # Unrecoverable — propagate to caller immediately
 
         except Exception as exc:
             last_error = exc
             print(
-                f"  ⚠️  Step {step_num}, attempt {attempt}/{MAX_RETRIES} failed: {exc}",
+                f"  ⚠️  Step {step_num}, LLM attempt {attempt}/{MAX_RETRIES} "
+                f"failed: {exc}",
                 flush=True,
             )
             if attempt < MAX_RETRIES:
                 time.sleep(RETRY_DELAY)
 
     raise RuntimeError(
-        f"All {MAX_RETRIES} LLM call attempts failed for step {step_num}. "
+        f"All {MAX_RETRIES} LLM call attempts failed at step {step_num}. "
         f"Last error: {last_error}"
     )
 
 
-# ── MAIN EPISODE LOOP ─────────────────────────────────────────────────────────
-async def main() -> None:
+# ══════════════════════════════════════════════════════════════════════════════
+# MAIN EPISODE LOOP
+# ══════════════════════════════════════════════════════════════════════════════
+
+def main() -> None:
+    """
+    Run one full PhishGuard-Env episode using the Hybrid Agent pattern.
+
+    Flow
+    ----
+    1.  env_reset()                  → first email observation (HTTP)
+    2.  call_llm_with_retry(obs)     → triage action (LLM call)
+    3.  env_step(action)             → reward, next obs, done flag (HTTP)
+    4.  Log in OpenEnv format
+    5.  Repeat until done=True or safety-limit hit
+
+    Logging format (OpenEnv compliant)
+    -----------------------------------
+    [START] task=phishing_triage env=phishguard_v1 model=<MODEL>
+    [STEP]  step=N action=ACTION reward=R health=H done=true|false error=null
+    [END]   success=true|false steps=N score=S rewards=r1,r2,...
+    """
     rewards: list[float] = []
-    steps        = 0
-    success      = False
-    final_score  = R_WRONG_PROCEDURE  # Safe open-interval default
+    steps       = 0
+    success     = False
+    final_score = R_WRONG_PROCEDURE  # Open-interval safe default
 
     print(
         f"[START] task=phishing_triage env=phishguard_v1 model={MODEL_NAME}",
         flush=True,
     )
+    print(
+        f"        LLM router : {LLM_BASE_URL}\n"
+        f"        Env server : {ENV_BASE_URL}",
+        flush=True,
+    )
 
     try:
-        obs  = env_instance.reset()
+        # ── 1. Reset the environment (HTTP call) ──────────────────────────────
+        obs  = env_reset()
         done = False
 
+        # Safety cap: 10 scenarios + small buffer for transient step errors
         while not done and steps < 30:
             steps += 1
 
             try:
+                # ── 2. LLM decides action (THINKING) ─────────────────────────
                 action_str = call_llm_with_retry(obs, steps)
-                obs, reward, done, info = env_instance.step(action_str)
+
+                # ── 3. Submit action to environment (ACTING) ──────────────────
+                obs, reward, done, info = env_step(action_str)
                 rewards.append(reward)
 
                 print(
@@ -192,36 +367,58 @@ async def main() -> None:
                     flush=True,
                 )
 
+                # Emit the human-readable feedback from the grader
+                feedback = info.get("feedback", "")
+                if feedback:
+                    print(f"       {feedback}", flush=True)
+
             except (AuthenticationError, APIConnectionError) as exc:
-                # Unrecoverable API error — abort the entire run.
+                # Unrecoverable LLM API error — abort the entire run
                 print(
-                    f"[STEP] step={steps} action=ABORT reward={R_WRONG_PROCEDURE:.4f} "
-                    f"done=true error='{exc}'",
+                    f"[STEP] step={steps} action=ABORT "
+                    f"reward={R_WRONG_PROCEDURE:.4f} done=true "
+                    f"error='{exc}'",
                     flush=True,
                 )
                 print(
                     f"\n❌ Unrecoverable API error: {exc}\n"
-                    f"   Verify HF_TOKEN / OPENAI_API_KEY and API_BASE_URL.",
+                    f"   Check HF_TOKEN / OPENAI_API_KEY and API_BASE_URL.",
+                    file=sys.stderr,
+                )
+                rewards.append(R_WRONG_PROCEDURE)
+                done = True  # Force episode end
+
+            except RuntimeError as exc:
+                # Environment unreachable — abort
+                print(
+                    f"[STEP] step={steps} action=ABORT "
+                    f"reward={R_WRONG_PROCEDURE:.4f} done=true "
+                    f"error='{exc}'",
+                    flush=True,
                 )
+                print(f"\n❌ Environment error: {exc}", file=sys.stderr)
                 rewards.append(R_WRONG_PROCEDURE)
-                break
+                done = True
 
             except Exception as exc:
-                # Transient error: do NOT reset the episode.
-                # Append the minimum non-zero signal and keep the current obs.
+                # Transient / unexpected error:
+                # • Do NOT reset the episode.
+                # • Append minimum non-zero signal.
+                # • Keep the current obs and retry from the same email.
                 print(
-                    f"[STEP] step={steps} action=ERROR reward={R_WRONG_PROCEDURE:.4f} "
-                    f"done=false error='{exc}'",
+                    f"[STEP] step={steps} action=ERROR "
+                    f"reward={R_WRONG_PROCEDURE:.4f} done=false "
+                    f"error='{exc}'",
                     flush=True,
                 )
                 rewards.append(R_WRONG_PROCEDURE)
-                # obs is unchanged — loop retries from the same email
+                # obs is unchanged — next loop iteration retries same email
 
         final_score = calculate_overall_score(rewards)
         success     = final_score >= SUCCESS_THRESHOLD
 
     except Exception as exc:
-        print(f"❌ Execution Failure: {exc}")
+        print(f"❌ Execution failure: {exc}", file=sys.stderr)
         final_score = R_WRONG_PROCEDURE  # Open-interval safe default
 
     finally:
@@ -231,10 +428,17 @@ async def main() -> None:
             f"score={final_score:.4f} rewards={rewards_str}",
             flush=True,
         )
+        # Human-readable summary
+        print(
+            f"\n{'='*60}\n"
+            f"  Final Score   : {final_score:.4f}\n"
+            f"  Steps taken   : {steps}\n"
+            f"  Success       : {success}  (threshold ≥ {SUCCESS_THRESHOLD})\n"
+            f"  Rewards       : {[round(r, 2) for r in rewards]}\n"
+            f"{'='*60}",
+            flush=True,
+        )
 
 
 if __name__ == "__main__":
-    try:
-        asyncio.run(main())
-    except KeyboardInterrupt:
-        print("\nProcess interrupted by user.")
+    main()