fix: update OpenAI client usage and improve error handling in task execution

inference.py

@@ -3,7 +3,7 @@ import json
 import asyncio
 import re
 from typing import List, Dict, Any
-from openai import AsyncOpenAI
+from openai import OpenAI
 
 from src.env import QuasarEnv
 from src.models import QuasarAction
@@ -18,10 +18,8 @@ def log_step(step: int, action: str, reward: float, done: bool, error: str = Non
 def log_end(success: bool, steps: int, score: float, rewards: List[float]):
     print(f"[END] success={success} steps={steps} score={score} rewards={rewards}", flush=True)
 
-async def run_task(client: AsyncOpenAI, task_name: str, model_name: str):
+async def run_task(client: OpenAI, task_name: str, model_name: str):
     env = QuasarEnv(task_name=task_name)
-    
-    history: List[str] = []
     rewards: List[float] = []
     steps_taken = 0
     score = 0.0
@@ -44,36 +42,33 @@ You MUST return strictly valid JSON matching this schema:
 {
   "command": "allow_packet" | "flag_packet" | "isolate_ip" | "pass",
   "target_id": "string or null"
-}
-Do not include markdown blocks or any other text."""
+}"""
 
         for step in range(1, max_steps + 1):
-            if state.done:
-                break
+            if state.done: break
 
             obs_dict = state.observation.model_dump()
             user_message = f"Current State: {json.dumps(obs_dict)}\nWhat is your action?"
 
             try:
-                response = await client.chat.completions.create(
+                response = client.chat.completions.create(
                     model=model_name,
                     messages=[
                         {"role": "system", "content": system_prompt},
                         {"role": "user", "content": user_message}
                     ],
-                    temperature=0.0
+                    temperature=0.0 
                 )
-
-                raw_action = response.choices[0].message.content or ""
+                
+                raw_action = response.choices[0].message.content
                 json_match = re.search(r'\{.*?\}', raw_action, re.DOTALL)
 
                 if json_match:
                     action_dict = json.loads(json_match.group(0))
+                    action_obj = QuasarAction(**action_dict)
+                    error = None
                 else:
-                    raise ValueError("No valid JSON object found in LLM response.")
-
-                action_obj = QuasarAction(**action_dict)
-                error = None
+                    raise ValueError("No valid JSON object found.")
                 
             except Exception as e:
                 action_obj = QuasarAction(command="pass", target_id=None)
@@ -81,44 +76,35 @@ Do not include markdown blocks or any other text."""
                 error = f"LLM parsing error: {str(e)}"
 
             state = await env.step(action_obj)
-            
             reward = state.reward.score if state.reward else 0.0
             done = state.done
 
             rewards.append(reward)
             steps_taken = step
-
             log_step(step=step, action=action_dict, reward=reward, done=done, error=error)
 
-            if done:
-                break
+            if done: break
 
-        score = rewards[-1] if rewards else 0.0
-        success = score >= 0.7
+        score = rewards[-1] if rewards else 0.0 
+        success = score >= 0.7 
 
     finally:
         log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
 
 
 async def main():
-    api_base = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
-    api_key = os.environ.get("HF_TOKEN") or os.environ.get("OPENAI_API_KEY") 
-    model_name = os.environ.get("MODEL_NAME", "gpt-3.5-turbo")
-
-    if not api_key:
-        print("WARNING: HF_TOKEN or OPENAI_API_KEY environment variable not set. LLM calls will fail.")
-
-    client = AsyncOpenAI(base_url=api_base, api_key=api_key)
+    API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
+    MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-3.5-turbo")
+    HF_TOKEN = os.environ.get("HF_TOKEN")
+    
+    api_key = HF_TOKEN or os.environ.get("OPENAI_API_KEY")
 
-    tasks = [
-        "task_1_volumetric_flood",
-        "task_2_contextual_injection",
-        "task_3_stealth_poisoning"
-    ]
+    client = OpenAI(base_url=API_BASE_URL, api_key=api_key)
 
+    tasks = ["task_1_volumetric_flood", "task_2_contextual_injection", "task_3_stealth_poisoning"]
     for task in tasks:
         print(f"\n--- Running Baseline for {task.upper()} ---")
-        await run_task(client, task, model_name)
+        await run_task(client, task, MODEL_NAME)
 
 if __name__ == "__main__":
     asyncio.run(main())

server/app.py

@@ -5,10 +5,8 @@ from openenv.core.env_server import create_app
 from src.env import QuasarEnv
 from src.models import QuasarAction, QuasarObservation
 
-# Initialize the OpenEnv FastAPI wrapper
 app = create_app(QuasarEnv, QuasarAction, QuasarObservation)
 
-# Add a root endpoint for human judges and HF health checks
 @app.get("/", response_class=HTMLResponse)
 async def root():
     return """