refactor: Update inference.py to use AsyncOpenAI and make episode functions asynchronous

inference.py

@@ -29,7 +29,7 @@ import sys
 import time
 from typing import Any, Dict, List, Optional
 
-from openai import OpenAI
+from openai import AsyncOpenAI
 
 from server import Task1Environment, Task2Environment, Task3Environment
 from env.schemas import Action, ActionType
@@ -48,7 +48,7 @@ HF_TOKEN     = os.getenv("HF_TOKEN",     "")
 if not HF_TOKEN:
     raise RuntimeError("HF_TOKEN environment variable not set")
 
-client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)
+client = AsyncOpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)
 
 # Benchmark / environment identifier (constant for this env)
 ENV_BENCHMARK = "smart-contract-audit"
@@ -69,17 +69,17 @@ SUCCESS_SCORE_THRESHOLD = 0.5
 # Unified LLM call function
 # ─────────────────────────────────────────────────────────────────────────────
 
-def get_llm_response(
+async def get_llm_response(
     messages: List[Dict[str, str]],
     max_tokens: int = 200,
     temperature: float = 0.0,
 ) -> str:
     """
-    Call the Groq LLM with the given messages and parameters.
+    Call the LLM with the given messages and parameters.
     Returns the response content as a string.
     Raises an exception on failure (to be caught by the caller).
     """
-    completion = client.chat.completions.create(
+    completion = await client.chat.completions.create(
         model=MODEL_NAME,
         messages=messages,          # type: ignore
     )
@@ -137,7 +137,7 @@ def _t1_user_msg(obs: Dict[str, Any]) -> str:
     )
 
 
-def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str, Any]:
+async def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str, Any]:
     """Run one Task 1 episode; emit [START]/[STEP]/[END]."""
     r   = env.reset(seed=seed)
     obs = r.observation.model_dump()
@@ -156,7 +156,7 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
         for step in range(1, MAX_STEPS_T1 + 1):
             messages.append({"role": "user", "content": _t1_user_msg(obs)})
             try:
-                raw = get_llm_response(messages, max_tokens=200, temperature=0.0)
+                raw = await get_llm_response(messages, max_tokens=200, temperature=0.0)
                 error_msg = None
             except Exception as e:
                 raw = ""
@@ -194,8 +194,7 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
         "episode":          ep_num,
         "seed":             seed,
         "contract":         obs["contract_name"],
-        "grader_score":     grader_score,
-        "cumulative_reward": obs["cumulative_reward"],
+        "grader_score":     grader_score
     }
 
 
@@ -214,7 +213,7 @@ def _t2_user_msg(obs: Dict[str, Any]) -> str:
     )
 
 
-def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str, Any]:
+async def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str, Any]:
     """Run one Task 2 episode; emit [START]/[STEP]/[END]."""
     r   = env.reset(seed=seed)
     obs = r.observation.model_dump()
@@ -234,7 +233,7 @@ def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str,
         for step in range(1, MAX_STEPS_T2 + 1):
             messages.append({"role": "user", "content": _t2_user_msg(obs)})
             try:
-                raw = get_llm_response(messages, max_tokens=400, temperature=0.0)
+                raw = await get_llm_response(messages, max_tokens=400, temperature=0.0)
                 error_msg = None
             except Exception as e:
                 raw = ""
@@ -273,8 +272,7 @@ def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str,
         "seed":             seed,
         "contract":         obs["contract_name"],
         "function":         fn,
-        "grader_score":     grader_score,
-        "cumulative_reward": obs["cumulative_reward"],
+        "grader_score":     grader_score
     }
 
 
@@ -292,7 +290,7 @@ def _t3_user_msg(obs: Dict[str, Any]) -> str:
     )
 
 
-def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str, Any]:
+async def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str, Any]:
     """Run one Task 3 episode; emit [START]/[STEP]/[END]."""
     r   = env.reset(seed=seed)
     obs = r.observation.model_dump()
@@ -311,7 +309,7 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
         for step in range(1, MAX_STEPS_T3 + 1):
             messages.append({"role": "user", "content": _t3_user_msg(obs)})
             try:
-                raw = get_llm_response(messages, max_tokens=200, temperature=0.0)
+                raw = await get_llm_response(messages, max_tokens=200, temperature=0.0)
                 error_msg = None
             except Exception as e:
                 raw = ""
@@ -349,8 +347,7 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
         "episode":          ep_num,
         "seed":             seed,
         "contract":         obs["contract_name"],
-        "grader_score":     grader_score,
-        "cumulative_reward": obs["cumulative_reward"],
+        "grader_score":     grader_score
     }
 
 
@@ -358,12 +355,12 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
 # Task runners
 # ─────────────────────────────────────────────────────────────────────────────
 
-def run_task1(n: int = NUM_EPISODES) -> Dict[str, Any]:
+async def run_task1(n: int = NUM_EPISODES) -> Dict[str, Any]:
     print("\n" + "="*60, flush=True)
     print("TASK 1: Targeted Vulnerability Detection", flush=True)
     print("="*60, flush=True)
     env      = Task1Environment()
-    episodes = [_run_t1_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
+    episodes = [await _run_t1_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
     avg_s    = sum(e["grader_score"] for e in episodes) / n
     print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
     return {
@@ -373,14 +370,13 @@ def run_task1(n: int = NUM_EPISODES) -> Dict[str, Any]:
     }
 
 
-def run_task2(n: int = NUM_EPISODES) -> Dict[str, Any]:
+async def run_task2(n: int = NUM_EPISODES) -> Dict[str, Any]:
     print("\n" + "="*60, flush=True)
     print("TASK 2: Property Discovery", flush=True)
     print("="*60, flush=True)
     env      = Task2Environment()
-    episodes = [_run_t2_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
+    episodes = [await _run_t2_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
     avg_s    = sum(e["grader_score"] for e in episodes) / n
-    avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
     print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
     return {
         "task_id": "task2_property_discovery", "name": "Property Discovery",
@@ -389,14 +385,13 @@ def run_task2(n: int = NUM_EPISODES) -> Dict[str, Any]:
     }
 
 
-def run_task3(n: int = NUM_EPISODES) -> Dict[str, Any]:
+async def run_task3(n: int = NUM_EPISODES) -> Dict[str, Any]:
     print("\n" + "="*60, flush=True)
     print("TASK 3: Rule Checker", flush=True)
     print("="*60, flush=True)
     env      = Task3Environment()
-    episodes = [_run_t3_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
+    episodes = [await _run_t3_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
     avg_s    = sum(e["grader_score"] for e in episodes) / n
-    avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
     print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
     return {
         "task_id": "task3_rule_checker", "name": "Rule Checker",
@@ -412,17 +407,12 @@ def run_task3(n: int = NUM_EPISODES) -> Dict[str, Any]:
 async def main() -> None:
     """Async entry point (wraps sync env calls; asyncio.run() expected by caller)."""
     print("Smart Contract Audit RL Environment — Baseline Inference", flush=True)
-    print(f"Model: {MODEL_NAME} | Groq API", flush=True)
 
-    t1 = run_task1(NUM_EPISODES)
-    t2 = run_task2(NUM_EPISODES)
-    t3 = run_task3(NUM_EPISODES)
+    t1 = await run_task1(NUM_EPISODES)
+    t2 = await run_task2(NUM_EPISODES)
+    t3 = await run_task3(NUM_EPISODES)
 
-    results = {
-        "model":    MODEL_NAME,
-        "backend":  "groq",
-        "tasks":    [t1, t2, t3],
-    }
+    results: Dict[str, Any] = { "tasks": [t1, t2, t3] }
     overall = sum(t["avg_grader_score"] for t in results["tasks"]) / 3
     results["overall_avg_score"] = overall