Upload folder using huggingface_hub

client_env.py

@@ -1,22 +1,12 @@
-"""
-AegisGym WebSocket client — concrete subclass of openenv EnvClient.
-
-Usage:
-    client = AegisGymWsClient()
-    sync_client = client.sync()
-    obs = sync_client.reset()
-    obs = sync_client.step({...})
-"""
 from typing import Any, Dict
 from openenv.core.env_client import EnvClient
 from openenv.core.sync_client import SyncEnvClient
 from server.models import AuditAction, AuditObservation
 
-HF_SPACE_WSS = "wss://armaan020-aegisgym.hf.space"
-
+HF_SPACE_WSS = "wss://armaan020-aegisopenenv.hf.space"
 
 class AegisGymWsClient(EnvClient):
-    """Concrete EnvClient implementation for the AegisGym HF Space."""
+    """Concrete EnvClient implementation for the AegisOpenEnv HF Space."""
 
     def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]:
         """Convert an action dict into the WS step payload."""
@@ -24,44 +14,12 @@ class AegisGymWsClient(EnvClient):
 
     def _parse_result(self, payload: Dict[str, Any]) -> Any:
         """Parse reset/step response from the server into usable result."""
-        return payload  # keep as dict; training code accesses .observation, .reward, .done
+        return payload
 
     def _parse_state(self, payload: Dict[str, Any]) -> Any:
         """Parse the state endpoint response."""
         return payload
 
-
 def get_sync_client(ws_url: str = HF_SPACE_WSS) -> SyncEnvClient:
     """Return a synchronous wrapper over the WebSocket client."""
     return AegisGymWsClient(base_url=ws_url).sync()
-
-
-def test_live_env():
-    print("=== Testing Live AegisGym Space (wss) ===")
-    client = get_sync_client()
-
-    print("reset() ...")
-    result = client.reset()
-    print(f"  result type: {type(result)}")
-    print(f"  keys: {list(result.keys()) if isinstance(result, dict) else dir(result)}")
-
-    action = AuditAction(
-        action_type="FLAG",
-        target_id="ACC-BL-001",
-        regulation_citation="EU-AI-Act-Art-57"
-    ).model_dump()
-
-    print("step(FLAG ACC-BL-001) ...")
-    result = client.step(action)
-    print(f"  reward={result.get('reward') if isinstance(result, dict) else getattr(result,'reward',None)}")
-    print(f"  done={result.get('done') if isinstance(result, dict) else getattr(result,'done',None)}")
-
-    print("state() ...")
-    s = client.state()
-    print(f"  State: {s}")
-
-    print("\n=== Live environment OK! ===")
-
-
-if __name__ == "__main__":
-    test_live_env()

inference.py

@@ -1,7 +1,3 @@
-"""
-AegisGym Baseline Inference Script
-Requirement: Must use OpenAI client and env vars for credentials.
-"""
 import os
 import json
 from openai import OpenAI
@@ -9,21 +5,14 @@ from client_env import get_sync_client
 from server.models import AuditAction
 
 # ─── Config (Required by Meta OpenEnv) ──────────────────────────────────────
-API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
-MODEL_NAME   = os.getenv("MODEL_NAME", "gpt-4o")
+API_BASE_URL = os.getenv("API_BASE_URL", "https://openrouter.ai/api/v1")
+MODEL_NAME   = os.getenv("MODEL_NAME", "stepfun/step-3.5-flash:free")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
-ENV_URL      = os.getenv("ENV_URL", "https://armaan020-aegisgym.hf.space")
-
-# OpenRouter recommends specific headers for their API
-extra_headers = {
-    "HTTP-Referer": "https://huggingface.co/spaces/armaan020/AegisGym",
-    "X-Title": "AegisGym OpenEnv Submission"
-}
+ENV_URL      = os.getenv("ENV_URL", "https://armaan020-aegisopenenv.hf.space")
 
 client = OpenAI(
     api_key=OPENAI_API_KEY, 
-    base_url=API_BASE_URL,
-    default_headers=extra_headers if "openrouter" in API_BASE_URL.lower() else None
+    base_url=API_BASE_URL
 )
 
 SYSTEM_PROMPT = """You are a high-performance financial auditor AI.
@@ -48,6 +37,7 @@ def run_baseline(num_episodes=10):
         try:
             obs_payload = env.reset()
             obs = obs_payload.get("observation", {})
+            target_id = obs.get("account_metadata", {}).get("target_id", "N/A")
             
             user_msg = (
                 f"Audit the transaction.\n"
@@ -56,7 +46,6 @@ def run_baseline(num_episodes=10):
                 f"Account: {obs.get('account_metadata', {})}"
             )
             
-            # OpenAI API Call
             response = client.chat.completions.create(
                 model=MODEL_NAME,
                 messages=[
@@ -68,6 +57,11 @@ def run_baseline(num_episodes=10):
             
             content = response.choices[0].message.content
             action_data = json.loads(content)
+            
+            # Robust Fallbacks
+            if "target_id" not in action_data: action_data["target_id"] = target_id
+            if "action_type" not in action_data: action_data["action_type"] = "APPROVE"
+            
             print(f"  Target: {action_data.get('target_id')} | Action: {action_data.get('action_type')}")
             
             result = env.step(action_data)
@@ -78,7 +72,7 @@ def run_baseline(num_episodes=10):
             episodes_run += 1
             
         except Exception as e:
-            print(f"  Error in episode: {e}")
+            print(f"  Error in episode {i+1}: {e}")
             continue
 
     print(f"\n--- AegisGym Reproducibility Report ---")

server/server.py

@@ -1,5 +1,5 @@
 import random
-from typing import Dict, Any, Tuple
+from typing import Dict, Any, Tuple, List
 from openenv.core.env_server import Environment
 from .models import AuditAction, AuditObservation, AuditState
 from .grader import Grader
@@ -63,7 +63,7 @@ class AegisOpenEnv(Environment):
 
     def _set_next_scenario(self):
         tiers = ["easy", "medium", "hard"]
-        self.current_tier = tiers[self.step_count % 3]
+        self.current_tier = tiers[self.step_count % len(tiers)]
         
         if self.current_tier == "easy":
             self.current_target_id = random.choice(["ACC-BL-001", "ACC-CLEAN-01"])