fix round 1

env/env.py

@@ -1,5 +1,6 @@
 import requests
 from env.deception import deploy_honeypot, fake_database, block_attacker
+from models import Observation, Reward
 
 SERVER = "http://127.0.0.1:7860"
 
@@ -19,7 +20,7 @@ class DeceptionEnv:
         logs = requests.get(f"{SERVER}/logs").json()
         self._state = logs
 
-        return self._state
+        return Observation(**self._state)
 
     def step(self, action):
 
@@ -34,7 +35,7 @@ class DeceptionEnv:
         # Detect brute force
         if action == "detect_attack":
             if failed_logins > 3:
-                reward += 0.2
+                reward += 0.15
             else:
                 reward -= 0.05
 
@@ -42,51 +43,57 @@ class DeceptionEnv:
         if action == "detect_attack":
             for r in requests_log:
                 if isinstance(r, dict) and r.get("type") == "port_scan":
-                    reward += 0.2
+                    reward += 0.15
                     break
 
         # Detect SQL injection
         if action == "detect_attack":
             for r in requests_log:
                 if isinstance(r, dict) and r.get("type") == "sql_injection":
-                    reward += 0.2
+                    reward += 0.15
                     break
 
         # Detect directory traversal
         if action == "detect_attack":
             for r in requests_log:
                 if isinstance(r, dict) and r.get("type") == "directory_traversal":
-                    reward += 0.2
+                    reward += 0.15
                     break
 
         # Deploy honeypot
         elif action == "deploy_honeypot":
             deploy_honeypot()
-            reward += 0.3
+            reward += 0.30
 
         # Fake database
         elif action == "fake_database":
             fake_database()
-            reward += 0.2
+            reward += 0.20
 
         # Block attacker
         elif action == "block_ip":
             if logs.get("suspicious_ips"):
                 ip = logs["suspicious_ips"][0]
                 block_attacker(ip)
-                reward += 0.5
+                reward += 0.50
                 self.done = True
 
         # Episode boundary
         if self.current_step >= self.max_steps:
             self.done = True
 
+        # Clamp reward
+        reward = min(max(reward, 0.0), 1.0)
+
         self._state = logs
 
-        return self._state, reward, self.done, {}
+        observation = Observation(**self._state)
+        reward_obj = Reward(reward=reward, done=self.done)
+
+        return observation, reward_obj.reward, reward_obj.done, {}
 
     def state(self):
-        return self._state
+        return Observation(**self._state)
 
     def action_space(self):
         return [
@@ -94,4 +101,4 @@ class DeceptionEnv:
             "deploy_honeypot",
             "fake_database",
             "block_ip"
-        ]
+        ]

inference.py

@@ -9,11 +9,44 @@ random.seed(42)
 from env.env import DeceptionEnv
 from env.attacker import simulate_attack
 
+# Import graders
+from tasks.easy.grader import grade as easy_grade
+from tasks.medium.grader import grade as medium_grade
+from tasks.hard.grader import grade as hard_grade
+
+
 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 
+def choose_action(client, state):
+
+    try:
+        response = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a cybersecurity AI agent. Choose one action from: detect_attack, deploy_honeypot, fake_database, block_ip"
+                },
+                {
+                    "role": "user",
+                    "content": f"Current state: {state}"
+                }
+            ],
+            max_tokens=10,
+            temperature=0.0
+        )
+
+        action = response.choices[0].message.content.strip()
+
+    except Exception:
+        action = "detect_attack"
+
+    return action
+
+
 def run_task(task_name):
 
     try:
@@ -24,7 +57,7 @@ def run_task(task_name):
         )
 
         env = DeceptionEnv()
-        env.reset()
+        state = env.reset()
 
         print(
             f"[START] task={task_name} env=ai-deception-openenv model={MODEL_NAME}",
@@ -46,30 +79,12 @@ def run_task(task_name):
             except Exception:
                 pass
 
-            # Required OpenAI call
-            try:
-                client.chat.completions.create(
-                    model=MODEL_NAME,
-                    messages=[{"role": "user", "content": "choose action"}],
-                    timeout=10
-                )
-            except Exception:
-                pass
-
-            # Task logic
-            if task_name == "easy":
-                action = "detect_attack" if step < 3 else "deploy_honeypot"
+            # Model chooses action
+            action = choose_action(client, state)
 
-            elif task_name == "medium":
-                action = "detect_attack" if step == 1 else "deploy_honeypot"
-
-            else:  # hard
-                if step == 1:
-                    action = "detect_attack"
-                elif step == 2:
-                    action = "deploy_honeypot"
-                else:
-                    action = "block_ip"
+            # fallback
+            if action not in env.action_space():
+                action = "detect_attack"
 
             state, reward, done, _ = env.step(action)
 
@@ -86,10 +101,16 @@ def run_task(task_name):
                 break
 
         steps = len(rewards)
-        score = sum(rewards) / steps if steps > 0 else 0.0
-        score = min(max(score, 0.05), 0.95)
 
-        success = score >= 0.3
+        # Use graders
+        if task_name == "easy":
+            score = easy_grade(rewards)
+        elif task_name == "medium":
+            score = medium_grade(rewards)
+        else:
+            score = hard_grade(rewards)
+
+        success = score >= 0.5
 
         print(
             f"[END] success={str(success).lower()} steps={steps} "
@@ -110,4 +131,4 @@ run_task("medium")
 run_task("hard")
 
 # Keep alive briefly
-time.sleep(180)
+time.sleep(180)

models.py

@@ -1,13 +1,13 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from typing import List, Optional, Dict, Any
 
 
 class Observation(BaseModel):
     failed_logins: int
     port_scans: int
-    suspicious_ips: List[str]
+    suspicious_ips: List[str] = Field(default_factory=list)
     total_requests: Optional[int] = 0
-    attack_types: Optional[List[str]] = []
+    attack_types: List[str] = Field(default_factory=list)
 
 
 class Action(BaseModel):
@@ -17,4 +17,4 @@ class Action(BaseModel):
 class Reward(BaseModel):
     reward: float
     done: bool
-    info: Optional[Dict[str, Any]] = {}
+    info: Dict[str, Any] = Field(default_factory=dict)

tasks/easy/grader.py

@@ -1,9 +1,9 @@
-def grade(reward):
+def grade(rewards):
+    score = sum(rewards) / len(rewards)
 
-    if reward >= 0.4:
+    if score >= 0.4:
         return 1.0
-
-    elif reward >= 0.2:
+    elif score >= 0.2:
         return 0.5
-
-    return 0.0
+    else:
+        return 0.0

tasks/hard/grader.py

@@ -1,12 +1,9 @@
-def grade(reward):
+def grade(rewards):
+    score = sum(rewards) / len(rewards)
 
-    if reward >= 0.8:
+    if score >= 0.6:
         return 1.0
-
-    elif reward >= 0.5:
+    elif score >= 0.4:
         return 0.7
-
-    elif reward >= 0.3:
-        return 0.4
-
-    return 0.0
+    else:
+        return 0.3

tasks/medium/grader.py

@@ -1,9 +1,9 @@
-def grade(reward):
+def grade(rewards):
+    score = sum(rewards) / len(rewards)
 
-    if reward >= 0.6:
+    if score >= 0.5:
         return 1.0
-
-    elif reward >= 0.3:
-        return 0.5
-
-    return 0.0
+    elif score >= 0.3:
+        return 0.6
+    else:
+        return 0.2