Update inference.py

inference.py

@@ -1,73 +1,71 @@
-import asyncio
-import os
-import textwrap
-from typing import List, Optional
-from openai import OpenAI
-from my_env_v4 import MyEnvV4Action, MyEnvV4Env
-
-# Environment Configuration
-API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
-MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
-TASK_NAME = "email-triage"
-BENCHMARK = "mit-manipal-v4"
-MAX_STEPS = 3
-SUCCESS_THRESHOLD = 0.5
-
-SYSTEM_PROMPT = """
-You are an Email Security Agent. Triage the following email based on sender, headers, and body content.
-Digital Seduction Rules:
-- 'INBOX': Official domains (.edu, .gov) and passed security headers.
-- 'SPAM': Marketing, gambling, or generic lottery win claims.
-- 'QUARANTINE': Phishing, high-urgency threats, suspicious links (.net, .co), or failed headers (SPF/DMARC Fail).
-
-REPLY WITH EXACTLY ONE WORD: 'INBOX', 'SPAM', or 'QUARANTINE'.
-"""
-
-def log_start():
-    print(f"[START] task={TASK_NAME} env={BENCHMARK} model={MODEL_NAME}", flush=True)
-
-def log_step(step, action, reward, done):
-    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error=null", flush=True)
-
-def log_end(success, steps, score, rewards):
-    r_str = ",".join(f"{r:.2f}" for r in rewards)
-    print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={r_str}", flush=True)
-
-async def main():
-    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-    env = MyEnvV4Env() # Local instance for testing, can use from_docker_image if needed
-    
-    rewards = []
-    log_start()
-    
-    try:
-        result = await env.reset()
-        for step in range(1, MAX_STEPS + 1):
-            if result.done: break
-            
-            obs = result.observation
-            prompt = f"Sender: {obs.sender}\nSubject: {obs.subject}\nBody: {obs.body}\nHeaders: {obs.headers}"
-            
-            # OpenAI Call
-            response = client.chat.completions.create(
-                model=MODEL_NAME,
-                messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}],
-                max_tokens=10,
-                temperature=0.0 # Deterministic for testing
-            )
-            action_text = response.choices[0].message.content.strip().upper()
-            
-            result = await env.step(MyEnvV4Action(message=action_text))
-            rewards.append(result.reward)
-            
-            log_step(step, action_text, result.reward, result.done)
-            if result.done: break
-
-        total_score = sum(rewards) / MAX_STEPS
-        log_end(total_score >= SUCCESS_THRESHOLD, len(rewards), total_score, rewards)
-    finally:
-        await env.close()
-
-if __name__ == "__main__":
+import asyncio
+import os
+from typing import List
+from openai import OpenAI
+from env import MyEnvV4Env
+from models import MyEnvV4Action
+
+# Environment Configuration
+API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
+API_KEY = os.getenv("HF_TOKEN") or ""
+MODEL_NAME = os.getenv("MODEL_NAME") or "gemini-2.5-flash-preview-09-2025"
+TASK_NAME = "security-mail-triage"
+
+SYSTEM_PROMPT = """
+You are an Advanced Email Security Agent. Analyze the metadata (headers, SPF/DKIM), URLs, and content.
+Categories:
+- INBOX: Trusted academic/official domains, passed auth, clean history.
+- SPAM: Mass marketing, generic lottery/sales, usually safe but unwanted.
+- QUARANTINE: Phishing, spear-phishing, credential theft, high-urgency threats, typo-squatted domains.
+
+Rules:
+1. Examine 'raw_headers' and 'auth_results'.
+2. Inspect 'urls' for low reputation or high age.
+3. Provide reasoning first, then your decision.
+
+Respond in JSON format:
+{
+  "reasoning": "Explain your logic here...",
+  "message": "INBOX|SPAM|QUARANTINE"
+}
+"""
+
+async def main():
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    env = MyEnvV4Env()
+    
+    rewards = []
+    print(f"[START] Testing Security Triage Environment...")
+    
+    result = await env.reset()
+    step_idx = 1
+    
+    while not result.done:
+        obs = result.observation
+        prompt = f"Sender: {obs.sender}\nSubject: {obs.subject}\nBody: {obs.body}\nHeaders: {obs.raw_headers}\nURLs: {obs.urls}"
+        
+        try:
+            response = client.chat.completions.create(
+                model=MODEL_NAME,
+                messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}],
+                response_format={"type": "json_object"},
+                temperature=0.0
+            )
+            import json
+            data = json.loads(response.choices[0].message.content)
+            
+            action = MyEnvV4Action(message=data["message"], reasoning=data["reasoning"])
+            result = await env.step(action)
+            rewards.append(result.reward)
+            
+            print(f"[STEP {step_idx}] Action: {action.message} | Reward: {result.reward:.2f}")
+            step_idx += 1
+        except Exception as e:
+            print(f"[ERROR] Step {step_idx}: {e}")
+            break
+
+    score = sum(rewards) / len(rewards) if rewards else 0
+    print(f"[END] Final Score: {score:.3f}")
+
+if __name__ == "__main__":
     asyncio.run(main())