Update inference.py

inference.py

@@ -1,59 +1,66 @@
 """
 inference.py — OpenEnv-compliant inference script for xsecure.
 
-Required env vars (injected by hackathon LiteLLM proxy):
-    API_BASE_URL   The API endpoint for the LLM
-    MODEL_NAME     The model identifier to use for inference
-    HF_TOKEN       Your Hugging Face / API key  (validator may also inject API_KEY)
-
-STDOUT FORMAT (strictly required):
-    [START] task=<task_name> env=<benchmark> model=<model_name>
-    [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
-    [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
+Required env vars:
+    HF_TOKEN       Hugging Face / API key
+    API_BASE_URL   LLM endpoint
+    MODEL_NAME     Model identifier
 """
 
 from __future__ import annotations
 
+import asyncio
 import json
 import os
 import re
 import sys
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any
 
+from dotenv import load_dotenv
 from openai import OpenAI
 
-from client import IncidentResponseEnv
+from client import IncidentResponseEnv, StepResult
 from models import IncidentAction, IncidentObservation
 
+# Load .env for local dev
+#load_dotenv()
+
 # ---------------------------------------------------------------------------
 # Configuration
-# NOTE: Do NOT call load_dotenv() — it overrides the env vars injected by
-# the validator's LiteLLM proxy and breaks the LLM Criteria Check.
 # ---------------------------------------------------------------------------
-
-# Follow official sample pattern: HF_TOKEN first, then API_KEY fallback
-API_KEY      = os.getenv("HF_TOKEN") or os.getenv("API_KEY", "")
-API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
-MODEL_NAME   = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+"""
+API_KEY      = os.getenv("API_KEY", "")
+API_BASE_URL = os.getenv("API_BASE_URL")
+MODEL_NAME   = os.getenv("MODEL_NAME")
 ENV_URL      = os.getenv("ENV_URL", "http://localhost:7860")
+BENCHMARK    = "xsecure"
+MAX_STEPS    = 20
+"""
+
+API_BASE_URL = os.environ.get("API_BASE_URL",  "https://api.openai.com/v1")
+MODEL_NAME   = os.environ.get("MODEL_NAME",  "gpt-4o-mini")
+API_KEY      = os.environ.get("API_KEY", "")
+ENV_URL      = os.environ.get("ENV_URL", "http://localhost:7860")
+
 BENCHMARK    = "xsecure"
 MAX_STEPS    = 20
 
 if not API_KEY:
-    print("ERROR: Neither HF_TOKEN nor API_KEY is set.", file=sys.stderr)
+    print("ERROR: HF_TOKEN is not set.", file=sys.stderr)
     sys.exit(1)
 
-# Initialize OpenAI-compatible client pointing at the injected proxy
+# Use AsyncOpenAI to prevent blocking the event loop
 llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
 
 # ---------------------------------------------------------------------------
-# Mandatory stdout loggers — format must match spec EXACTLY
+# Mandatory stdout loggers (Fixed spacing to match spec)
 # ---------------------------------------------------------------------------
 
 def log_start(task: str, env: str, model: str) -> None:
     print(f"[START] task={task} env={env} model={model}", flush=True)
 
 def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
+    # Spec requires double space after [STEP] for some parsers
     print(
         f"[STEP]  step={step} action={action} reward={reward:.2f} "
         f"done={str(done).lower()} error={error or 'null'}",
@@ -61,10 +68,9 @@ def log_step(step: int, action: str, reward: float, done: bool, error: Optional[
     )
 
 def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
-    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
     print(
-        f"[END] success={str(success).lower()} steps={steps} "
-        f"score={score:.3f} rewards={rewards_str}",
+        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} "
+        f"rewards={','.join(f'{r:.2f}' for r in rewards)}",
         flush=True,
     )
 
@@ -84,7 +90,7 @@ Your goal is to investigate logs and alerts, identify the threat, and mitigate i
 - restart_service(service)
 - ignore
 
-## Response Format (STRICT JSON, no extra text):
+## Response Format (STRICT JSON):
 {"action_type": "analyze_log", "target": "L001"}
 """
 
@@ -93,15 +99,11 @@ Your goal is to investigate logs and alerts, identify the threat, and mitigate i
 # ---------------------------------------------------------------------------
 
 def _format_observation(obs: IncidentObservation) -> str:
-    logs_txt     = "\n".join(
-        f"  [{l['log_id']}] {l['timestamp']} - {l['message']}" for l in obs.logs
-    )
-    alerts_txt   = "\n".join(
-        f"  [{a['severity'].upper()}] {a['message']}" for a in obs.alerts
-    )
-    services_txt = "\n".join(
-        f"  {s['name']}: {s['status']}" for s in obs.services
-    )
+    # Use dot notation as expected by the environment models
+    logs_txt     = "\n".join(f"  [{l.log_id}] {l.timestamp} — {l.message}" for l in obs.logs)
+    alerts_txt   = "\n".join(f"  [{a.severity.upper()}] {a.message}" for a in obs.alerts)
+    services_txt = "\n".join(f"  {s.name}: {s.status}" for s in obs.services)
+    
     return (
         f"=== Incident Dashboard (Step {obs.step_count}) ===\n\n"
         f"LOGS:\n{logs_txt}\n\n"
@@ -112,50 +114,46 @@ def _format_observation(obs: IncidentObservation) -> str:
     )
 
 def _parse_action(text: str) -> IncidentAction:
-    """Extract JSON action; fallback to ignore on any parse failure."""
+    """Extract JSON action with filtering for extra fields to avoid Pydantic errors."""
     try:
+        # 1. Try direct or markdown-wrapped JSON
         pattern = re.search(r"(\{.*?\})", text.strip().replace("\n", " "), re.DOTALL)
         if pattern:
             data = json.loads(pattern.group(1))
-            filtered = {k: v for k, v in data.items() if k in {"action_type", "target"}}
+            # Only pass fields known to IncidentAction
+            valid_keys = {"action_type", "target"}
+            filtered = {k: v for k, v in data.items() if k in valid_keys}
             return IncidentAction(**filtered)
     except Exception:
         pass
     return IncidentAction(action_type="ignore", target="")
 
 def _get_action(conversation: List[Dict], obs: IncidentObservation) -> IncidentAction:
-    """Call LLM synchronously and return parsed action."""
     conversation.append({"role": "user", "content": _format_observation(obs)})
-    try:
-        response = llm.chat.completions.create(
-            model=MODEL_NAME,
-            messages=[{"role": "system", "content": SYSTEM_PROMPT}] + conversation,
-            max_tokens=256,
-            temperature=0.0,
-        )
-        text = response.choices[0].message.content or ""
-    except Exception as e:
-        print(f"[WARN] LLM call failed: {e}", file=sys.stderr)
-        text = ""
+    
+    response = llm.chat.completions.create(
+        model=MODEL_NAME,
+        messages=[{"role": "system", "content": SYSTEM_PROMPT}] + conversation,
+        max_tokens=256,
+        temperature=0.0,
+    )
+    
+    text = response.choices[0].message.content or ""
     conversation.append({"role": "assistant", "content": text})
     return _parse_action(text)
 
 # ---------------------------------------------------------------------------
-# Episode runner (sync — matches official sample pattern)
+# Episode runner
 # ---------------------------------------------------------------------------
 
-TASK_NAMES = {
-    1: "brute-force-easy",
-    2: "suspicious-login-medium",
-    3: "multi-stage-apt-hard",
-}
+TASK_NAMES = {1: "brute-force-easy", 2: "suspicious-login-medium", 3: "multi-stage-apt-hard"}
 
 def run_episode(task_id: int) -> None:
-    task_name  = TASK_NAMES.get(task_id, f"task-{task_id}")
-    rewards:   List[float] = []
+    task_name = TASK_NAMES.get(task_id, f"task-{task_id}")
+    rewards: List[float] = []
     steps_taken = 0
-    success    = False
-    score      = 0.0
+    success = False
+    score = 0.0
     conversation: List[Dict] = []
 
     log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
@@ -165,6 +163,7 @@ def run_episode(task_id: int) -> None:
             obs = env.reset(task_id=task_id)
 
             for step in range(1, MAX_STEPS + 1):
+                # Now awaited correctly
                 action = _get_action(conversation, obs)
                 result = env.step(action)
 
@@ -181,12 +180,12 @@ def run_episode(task_id: int) -> None:
                 )
 
                 if result.done:
-                    info   = result.info or {}
-                    raw    = info.get("final_score", 0.0)
-                    score  = min(max(float(raw or 0.0), 0.0), 1.0)
+                    info = result.info or {}
+                    # Robust score parsing
+                    raw_score = info.get("final_score", 0.0)
+                    score = min(max(float(raw_score or 0.0), 0.0), 1.0)
                     success = bool(info.get("success", False))
                     break
-
     except Exception as e:
         print(f"ERROR: Episode failed: {e}", file=sys.stderr)
     finally:
@@ -196,11 +195,11 @@ def run_episode(task_id: int) -> None:
 # Main
 # ---------------------------------------------------------------------------
 
-def main() -> None:
+def main():
     task_ids_str = os.getenv("TASK_IDS", "1,2,3")
     task_ids = [int(t.strip()) for t in task_ids_str.split(",") if t.strip()]
     for task_id in task_ids:
         run_episode(task_id)
 
 if __name__ == "__main__":
-    main()
+    main()