Update env.py

env.py

@@ -1,77 +1,83 @@
-import asyncio
-from typing import Optional
-from types import SimpleNamespace
-from openenv.core.env_server import Environment
-from models import MyEnvV4Observation, MyEnvV4Action
-
-class MyEnvV4Env(Environment):
-    def __init__(self):
-        super().__init__()
-        # Realistic Dataset with Digital Seduction/Phishing markers
-        self.dataset = [
-            {
-                "sender": "dean.office@manipal.edu",
-                "subject": "B.Tech Lab Exam Schedule",
-                "body": "Please find the attached PDF for the upcoming CSE lab exams.",
-                "headers": ["SPF: Pass", "DKIM: Pass"],
-                "label": "INBOX"
-            },
-            {
-                "sender": "verify-account@security-amazon.net",
-                "subject": "Urgent: Your account is locked!",
-                "body": "Digital Seduction Alert: High urgency used. Click http://bit.ly/fake-link to unlock.",
-                "headers": ["SPF: Fail", "DMARC: Fail"],
-                "label": "QUARANTINE"
-            },
-            {
-                "sender": "prize@lottery-winner.co",
-                "subject": "Congratulations! You won $10,000",
-                "body": "Reply with your bank details to claim your cash prize immediately.",
-                "headers": ["SPF: Neutral"],
-                "label": "SPAM"
-            }
-        ]
-        self.current_step = 0
-
-    async def reset(self):
-        self.current_step = 0
-        return self._get_result()
-
-    def _get_result(self, reward=0.0, done=False):
-        if self.current_step >= len(self.dataset):
-            obs = MyEnvV4Observation(
-                sender="N/A", subject="N/A", body="N/A", 
-                headers=[], echoed_message="End of Data"
-            )
-            return SimpleNamespace(observation=obs, reward=reward, done=True)
-        
-        data = self.dataset[self.current_step]
-        obs = MyEnvV4Observation(
-            sender=data["sender"],
-            subject=data["subject"],
-            body=data["body"],
-            headers=data["headers"],
-            echoed_message=f"Step {self.current_step + 1}"
-        )
-        return SimpleNamespace(observation=obs, reward=reward, done=done)
-
-    async def step(self, action: MyEnvV4Action):
-        if self.current_step >= len(self.dataset):
-            return self._get_result(done=True)
-
-        correct_label = self.dataset[self.current_step]["label"]
-        # Exact match reward logic for 0.0 - 1.0 range
-        reward = 1.0 if action.message.strip().upper() == correct_label else 0.0
-        
-        self.current_step += 1
-        done = self.current_step >= len(self.dataset)
-        
-        return self._get_result(reward=reward, done=done)
-
-    async def close(self):
-        pass
-
-    @classmethod
-    async def from_docker_image(cls, image_name: str):
-        """Simulated helper for local/containerized runs."""
-        return cls()
+import asyncio
+from typing import Optional
+from types import SimpleNamespace
+from openenv.core.env_server import Environment
+from models import MyEnvV4Observation, MyEnvV4Action
+
+class MyEnvV4Env(Environment):
+    def __init__(self):
+        super().__init__()
+        # Realistic Dataset with Digital Seduction/Phishing markers
+        self.dataset = [
+            {
+                "sender": "dean.office@manipal.edu",
+                "subject": "B.Tech Lab Exam Schedule",
+                "body": "Please find the attached PDF for the upcoming CSE lab exams.",
+                "headers": ["SPF: Pass", "DKIM: Pass"],
+                "label": "INBOX"
+            },
+            {
+                "sender": "verify-account@security-amazon.net",
+                "subject": "Urgent: Your account is locked!",
+                "body": "Digital Seduction Alert: High urgency used. Click http://bit.ly/fake-link to unlock.",
+                "headers": ["SPF: Fail", "DMARC: Fail"],
+                "label": "QUARANTINE"
+            },
+            {
+                "sender": "prize@lottery-winner.co",
+                "subject": "Congratulations! You won $10,000",
+                "body": "Reply with your bank details to claim your cash prize immediately.",
+                "headers": ["SPF: Neutral"],
+                "label": "SPAM"
+            }
+        ]
+        self.current_step = 0
+
+    async def reset(self):
+        self.current_step = 0
+        return self._get_result()
+
+    def _get_result(self, reward=0.0, done=False):
+        if self.current_step >= len(self.dataset):
+            obs = MyEnvV4Observation(
+                sender="N/A", subject="N/A", body="N/A", 
+                headers=[], echoed_message="End of Data"
+            )
+            return SimpleNamespace(observation=obs, reward=reward, done=True)
+        
+        data = self.dataset[self.current_step]
+        obs = MyEnvV4Observation(
+            sender=data["sender"],
+            subject=data["subject"],
+            body=data["body"],
+            headers=data["headers"],
+            echoed_message=f"Step {self.current_step + 1}"
+        )
+        return SimpleNamespace(observation=obs, reward=reward, done=done)
+
+    async def step(self, action: MyEnvV4Action):
+        if self.current_step >= len(self.dataset):
+            return self._get_result(done=True)
+
+        correct_label = self.dataset[self.current_step]["label"]
+        # Exact match reward logic for 0.0 - 1.0 range
+        reward = 1.0 if action.message.strip().upper() == correct_label else 0.0
+        
+        self.current_step += 1
+        done = self.current_step >= len(self.dataset)
+        
+        return self._get_result(reward=reward, done=done)
+
+    async def close(self):
+        pass
+
+    @classmethod
+    async def from_docker_image(cls, image_name: str):
+        """Simulated helper for local/containerized runs."""
+        return cls()
+
+# Create the environment instance
+env = MyEnvV4Env()
+
+# This is the 'app' that Docker and Uvicorn are looking for
+app = env.app