Switching to old inference.py

inference.py

@@ -35,43 +35,34 @@ from openai import OpenAI
 from server import Task1Environment, Task2Environment, Task3Environment
 from env.schemas import Action, ActionType
 from utils import T1_SYSTEM, T2_SYSTEM, T3_SYSTEM
-from dotenv import load_dotenv
+from dotenv import dotenv_values
 
 # ─────────────────────────────────────────────────────────────────────────────
 # Configuration
 # ─────────────────────────────────────────────────────────────────────────────
 
-load_dotenv()  # Load from .env if available; otherwise rely on actual env vars
-API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
-MODEL_NAME   = os.getenv("MODEL_NAME",   "gpt-4o")
-HF_TOKEN     = os.getenv("HF_TOKEN")
+config = dotenv_values(".env")
+API_BASE_URL = config.get("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME   = config.get("MODEL_NAME",   "gpt-4o")
+HF_TOKEN     = config.get("HF_TOKEN",     "")
 
 if not HF_TOKEN:
     print("[WARN] HF_TOKEN not set — API calls may fail.", file=sys.stderr)
     exit(1)
 
-ENV_BENCHMARK          = "smart-contract-audit"
-NUM_EPISODES           = 2      # keep low on free tier; raise for full eval
-SEED_BASE              = 42
+# Benchmark / environment identifier (constant for this env)
+ENV_BENCHMARK = "smart-contract-audit"
 
-# Max LLM calls per episode (including the mandatory submit on last step).
-# Budget: free tier handles ~5-6 calls per episode before rate-limiting.
-MAX_STEPS_T1           = 5
-MAX_STEPS_T2           = 4
-MAX_STEPS_T3           = 4
+# Episodes per task
+NUM_EPISODES = 3
+SEED_BASE    = 42
 
-# How many steps before the end we start injecting "submit now" pressure.
-# E.g. PRESSURE_AT=2 means last 2 steps show a warning.
-PRESSURE_AT            = 2
+# Max steps per task
+MAX_STEPS_T1 = 15
+MAX_STEPS_T2 = 10
+MAX_STEPS_T3 = 12
 
-# Sliding-window size: how many recent (user, assistant) pairs to keep.
-# system prompt + 2 exchanges = ~800 tokens max — safe for free tier.
-HISTORY_WINDOW         = 2
-
-# Truncate action results to this many chars before inserting into the prompt.
-MAX_RESULT_CHARS       = 400
-
-# A grader_score >= this threshold → success=true in [END] line
+# A grader_score >= this is considered a "success" for the [END] line
 SUCCESS_SCORE_THRESHOLD = 0.5
 
 client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)
@@ -117,115 +108,30 @@ def log_end(
     )
 
 
-# ─────────────────────────────────────────────────────────────────────────────
-# Shared utilities
-# ─────────────────────────────────────────────────────────────────────────────
-
-def _truncate(text: str, limit: int = MAX_RESULT_CHARS) -> str:
-    """Truncate long action results to keep prompts small."""
-    if len(text) <= limit:
-        return text
-    return text[:limit] + f"... [truncated, {len(text) - limit} chars omitted]"
-
-
-def _sliding_messages(system: str, history: List[Dict[str, str]]) -> List[Dict[str, str]]:
-    """
-    Return system prompt + the last HISTORY_WINDOW (user, assistant) pairs.
-    This keeps total tokens bounded regardless of episode length.
-    """
-    # history = [..., user, assistant, user, assistant, ...]
-    # We want the last HISTORY_WINDOW complete pairs (2 messages each).
-    keep = HISTORY_WINDOW * 2
-    windowed = history[-keep:] if len(history) > keep else history
-    return [{"role": "system", "content": system}] + windowed
-
-
-def _call_llm(messages: List[Dict[str, str]], max_tokens: int = 150) -> tuple[str, Optional[str]]:
-    """Call the LLM; return (raw_response, error_string_or_None)."""
-    try:
-        resp = client.chat.completions.create(
-            model=MODEL_NAME, # type: ignore
-            messages=messages, # type: ignore
-            max_tokens=max_tokens,
-            temperature=0.0,
-        )
-        return resp.choices[0].message.content.strip(), None    # type: ignore
-    except Exception as e:
-        return "", str(e)[:80]
-
-
-def _parse_action(raw: str, fallback_at: ActionType,
-                  fallback_params: Dict[str, Any]) -> tuple[ActionType, Dict[str, Any]]:
-    """Parse LLM JSON response into (ActionType, params). Use fallback on failure."""
-    try:
-        parsed = json.loads(raw)
-        return ActionType(parsed["action"]), parsed.get("params", {})
-    except Exception:
-        return fallback_at, fallback_params
-
-
-def _pressure_suffix(steps_left: int) -> str:
-    """Return an urgent suffix when the step budget is nearly exhausted."""
-    if steps_left <= 0:
-        return (
-            "\n\n⚠️  FINAL STEP — you MUST submit your best answer RIGHT NOW.\n"
-            "Do not browse further. Emit a submit action immediately."
-        )
-    if steps_left <= PRESSURE_AT:
-        return (
-            f"\n\n⚠️  Only {steps_left} step(s) remaining. "
-            "You should submit your answer in the next step or two."
-        )
-    return ""
-
-
 # ─────────────────────────────────────────────────────────────────────────────
 # Task 1 — Targeted Vulnerability Detection
 # ─────────────────────────────────────────────────────────────────────────────
 
-def _t1_user(obs: Dict[str, Any], steps_left: int) -> str:
-    result = _truncate(obs.get("last_action_result") or "Episode just started.")
+def _t1_user_msg(obs: Dict[str, Any]) -> str:
     return (
-        f"Contract: {obs['contract_name']} | {obs['contract_description'][:80]}\n"
-        f"Step {obs['step_count']} | Reward: {obs['cumulative_reward']:.2f}\n"
-        f"Last action: {obs['last_action'] or 'None'}\n"
-        f"Result: {result}"
-        + _pressure_suffix(steps_left)
-    )
-
-
-def _t1_force_submit(obs: Dict[str, Any], history: List[Dict[str, str]]) -> tuple[ActionType, Dict[str, Any]]:
-    """
-    Build a forced submission from what we already know.
-    Strategy: ask the LLM one more time with an explicit 'submit NOW' mandate.
-    If that fails, fall back to a heuristic.
-    """
-    mandate = (
-        "Based on everything you have seen, submit your best answer NOW.\n"
-        "Respond ONLY with this JSON (fill in the values):\n"
-        '{"action":"submit","params":{"function_name":"<best_guess>","vulnerability_type":"<best_guess>"}}'
+        f"Contract: {obs['contract_name']}\n"
+        f"Description: {obs['contract_description']}\n"
+        f"Step: {obs['step_count']} | Reward so far: {obs['cumulative_reward']:.2f}\n\n"
+        f"Last action : {obs['last_action'] or 'None'}\n"
+        f"Last result : {obs['last_action_result'] or 'Episode just started.'}"
     )
-    messages = _sliding_messages(T1_SYSTEM, history) + [{"role": "user", "content": mandate}]
-    raw, _ = _call_llm(messages, max_tokens=80)
-    at, params = _parse_action(raw, ActionType.SUBMIT,
-                                {"function_name": "withdraw",
-                                 "vulnerability_type": "reentrancy"})
-    # Guarantee it's always a submit
-    if at != ActionType.SUBMIT:
-        at = ActionType.SUBMIT
-        if "function_name" not in params:
-            params["function_name"] = "withdraw"
-        if "vulnerability_type" not in params:
-            params["vulnerability_type"] = "reentrancy"
-    return at, params
 
 
 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()
+
     log_start(task="task1_vuln_detection", env=ENV_BENCHMARK, model=MODEL_NAME) # type: ignore
 
-    history:      List[Dict[str, str]] = []
+    messages: List[ChatCompletionMessageParam] = [ # type: ignore
+        {"role": "system", "content": T1_SYSTEM}
+    ]
     step_rewards: List[float] = []
     grader_score  = 0.0
     steps_taken   = 0
@@ -233,24 +139,31 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
 
     try:
         for step in range(1, MAX_STEPS_T1 + 1):
-            steps_left = MAX_STEPS_T1 - step
-            is_last    = (step == MAX_STEPS_T1)
-
-            if is_last:
-                # Never waste the last step on browsing — force a submission
-                at, params = _t1_force_submit(obs, history)
-            else:
-                user_msg = _t1_user(obs, steps_left)
-                history.append({"role": "user", "content": user_msg})
-                messages  = _sliding_messages(T1_SYSTEM, history)
-                raw, error_msg = _call_llm(messages)
-                history.append({"role": "assistant", "content": raw})
-                at, params = _parse_action(raw, ActionType.LIST_FUNCTIONS, {})
-
-            result   = env.step(Action(action_type=at, params=params))
-            obs      = result.observation.model_dump()
-            r_val    = result.reward.value
-            done     = result.done
+            messages.append({"role": "user", "content": _t1_user_msg(obs)})
+            try:
+                resp = client.chat.completions.create(
+                    model=MODEL_NAME, messages=messages, # type: ignore
+                    max_tokens=200, temperature=0.0,
+                )
+                raw = resp.choices[0].message.content.strip() # type: ignore
+                error_msg = None
+            except Exception as e:
+                raw = ""
+                error_msg = str(e)[:80]
+                print(f"[DEBUG] T1 LLM error ep={ep_num} step={step}: {e}", file=sys.stderr)
+
+            try:
+                parsed = json.loads(raw)
+                at     = ActionType(parsed["action"])
+                params = parsed.get("params", {})
+            except Exception:
+                at, params = ActionType.LIST_FUNCTIONS, {}
+
+            messages.append({"role": "assistant", "content": raw})
+            result = env.step(Action(action_type=at, params=params))
+            obs    = result.observation.model_dump()
+            r_val  = result.reward.value
+            done   = result.done
 
             step_rewards.append(r_val)
             steps_taken = step
@@ -258,18 +171,22 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
 
             if done:
                 v = r_val
-                grader_score = 0.999 if v >= 4.9 else (0.5 if v >= 0.9 else 0.0)
+                grader_score = 1.0 if v >= 4.9 else (0.5 if v >= 0.9 else 0.0)
                 break
 
-            if not is_last:
-                time.sleep(0.5)
+            time.sleep(0.3)
 
     finally:
         success = grader_score >= SUCCESS_SCORE_THRESHOLD
         log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)
 
-    return {"episode": ep_num, "seed": seed, "contract": obs["contract_name"],
-            "grader_score": grader_score, "cumulative_reward": obs["cumulative_reward"]}
+    return {
+        "episode":          ep_num,
+        "seed":             seed,
+        "contract":         obs["contract_name"],
+        "grader_score":     grader_score,
+        "cumulative_reward": obs["cumulative_reward"],
+    }
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -277,48 +194,29 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
 # ─────────────────────────────────────────────────────────────────────────────
 
 
-def _t2_user(obs: Dict[str, Any], steps_left: int) -> str:
-    extra  = obs.get("extra", {})
-    result = _truncate(obs.get("last_action_result") or "Episode just started.")
-    return (
-        f"Contract: {obs['contract_name']} | "
-        f"Function: {extra.get('target_function','?')} ({extra.get('target_signature','')})\n"
-        f"Step {obs['step_count']} | Reward: {obs['cumulative_reward']:.2f}\n"
-        f"Last action: {obs['last_action'] or 'None'}\n"
-        f"Result: {result}"
-        + _pressure_suffix(steps_left)
-    )
-
-
-def _t2_force_submit(obs: Dict[str, Any], history: List[Dict[str, str]]) -> tuple[ActionType, Dict[str, Any]]:
-    """Force a submit_property based on everything seen so far."""
+def _t2_user_msg(obs: Dict[str, Any]) -> str:
     extra = obs.get("extra", {})
-    fn    = extra.get("target_function", "this function")
-    mandate = (
-        f"You must now submit your best property for '{fn}'.\n"
-        "Write 2-3 sentences covering: what state changes, what is transferred, revert conditions.\n"
-        "Respond ONLY with:\n"
-        '{"action":"submit_property","params":{"property":"<your property here>"}}'
+    return (
+        f"Contract : {obs['contract_name']}\n"
+        f"Function : {extra.get('target_function', '?')}  "
+        f"({extra.get('target_signature', '')})\n"
+        f"Step: {obs['step_count']} | Reward so far: {obs['cumulative_reward']:.2f}\n\n"
+        f"Last action : {obs['last_action'] or 'None'}\n"
+        f"Last result :\n{obs['last_action_result'] or 'Episode just started.'}"
     )
-    messages = _sliding_messages(T2_SYSTEM, history) + [{"role": "user", "content": mandate}]
-    raw, _ = _call_llm(messages, max_tokens=200)
-    at, params = _parse_action(raw, ActionType.SUBMIT_PROPERTY, {})
-    if at != ActionType.SUBMIT_PROPERTY or not params.get("property", "").strip():
-        at     = ActionType.SUBMIT_PROPERTY
-        params = {"property": (
-            f"After a successful call to {fn}, the contract updates its internal state "
-            f"according to the function's logic. Reverts if input conditions are not met."
-        )}
-    return at, params
 
 
 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()
     fn  = obs["extra"].get("target_function", "?")
+
     log_start(task="task2_property_discovery", env=ENV_BENCHMARK, model=MODEL_NAME) # type: ignore
 
-    history:      List[Dict[str, str]] = []
+    messages: List[ChatCompletionMessageParam] = [ # type: ignore
+        {"role": "system", "content": T2_SYSTEM}
+    ]
     step_rewards: List[float] = []
     grader_score  = 0.0
     steps_taken   = 0
@@ -326,23 +224,31 @@ def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str,
 
     try:
         for step in range(1, MAX_STEPS_T2 + 1):
-            steps_left = MAX_STEPS_T2 - step
-            is_last    = (step == MAX_STEPS_T2)
-
-            if is_last:
-                at, params = _t2_force_submit(obs, history)
-            else:
-                user_msg = _t2_user(obs, steps_left)
-                history.append({"role": "user", "content": user_msg})
-                messages  = _sliding_messages(T2_SYSTEM, history)
-                raw, error_msg = _call_llm(messages, max_tokens=250)
-                history.append({"role": "assistant", "content": raw})
-                at, params = _parse_action(raw, ActionType.GET_FUNCTION_NATSPEC, {})
-
-            result   = env.step(Action(action_type=at, params=params))
-            obs      = result.observation.model_dump()
-            r_val    = result.reward.value
-            done     = result.done
+            messages.append({"role": "user", "content": _t2_user_msg(obs)})
+            try:
+                resp = client.chat.completions.create(
+                    model=MODEL_NAME, messages=messages, # type: ignore
+                    max_tokens=400, temperature=0.0,
+                )
+                raw = resp.choices[0].message.content.strip() # type: ignore
+                error_msg = None
+            except Exception as e:
+                raw = ""
+                error_msg = str(e)[:80]
+                print(f"[DEBUG] T2 LLM error ep={ep_num} step={step}: {e}", file=sys.stderr)
+
+            try:
+                parsed = json.loads(raw)
+                at     = ActionType(parsed["action"])
+                params = parsed.get("params", {})
+            except Exception:
+                at, params = ActionType.GET_FUNCTION_CODE, {}
+
+            messages.append({"role": "assistant", "content": raw})
+            result = env.step(Action(action_type=at, params=params))
+            obs    = result.observation.model_dump()
+            r_val  = result.reward.value
+            done   = result.done
 
             step_rewards.append(r_val)
             steps_taken = step
@@ -352,64 +258,48 @@ def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str,
                 grader_score = round(r_val / 5.0, 3) if r_val > 0 else 0.0
                 break
 
-            if not is_last:
-                time.sleep(0.5)
+            time.sleep(0.3)
 
     finally:
         success = grader_score >= SUCCESS_SCORE_THRESHOLD
         log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)
 
-    return {"episode": ep_num, "seed": seed,
-            "contract": obs["contract_name"], "function": fn,
-            "grader_score": grader_score, "cumulative_reward": obs["cumulative_reward"]}
+    return {
+        "episode":          ep_num,
+        "seed":             seed,
+        "contract":         obs["contract_name"],
+        "function":         fn,
+        "grader_score":     grader_score,
+        "cumulative_reward": obs["cumulative_reward"],
+    }
 
 
 # ─────────────────────────────────────────────────────────────────────────────
 # Task 3 — Rule Checker
 # ─────────────────────────────────────────────────────────────────────────────
 
-def _t3_user(obs: Dict[str, Any], steps_left: int) -> str:
-    extra  = obs.get("extra", {})
-    result = _truncate(obs.get("last_action_result") or "Episode just started.")
-    return (
-        f"Contract: {obs['contract_name']}\n"
-        f"Property: {extra.get('property_english', '(none)')[:200]}\n"
-        f"Step {obs['step_count']} | Reward: {obs['cumulative_reward']:.2f}\n"
-        f"Last action: {obs['last_action'] or 'None'}\n"
-        f"Result: {result}"
-        + _pressure_suffix(steps_left)
-    )
-
 
-def _t3_force_submit(obs: Dict[str, Any], history: List[Dict[str, str]]) -> tuple[ActionType, Dict[str, Any]]:
-    """Force a submit_function based on everything seen so far."""
-    prop = obs.get("extra", {}).get("property_english", "")
-    mandate = (
-        f"Property: {prop[:200]}\n"
-        "Based on everything you have seen, which function violates this property?\n"
-        "Respond ONLY with:\n"
-        '{"action":"submit_function","params":{"function_name":"<your_best_guess>"}}'
+def _t3_user_msg(obs: Dict[str, Any]) -> str:
+    extra = obs.get("extra", {})
+    return (
+        f"Contract  : {obs['contract_name']}\n"
+        f"Property  : {extra.get('property_english', '(none)')}\n"
+        f"Step: {obs['step_count']} | Reward so far: {obs['cumulative_reward']:.2f}\n\n"
+        f"Last action : {obs['last_action'] or 'None'}\n"
+        f"Last result :\n{obs['last_action_result'] or 'Episode just started.'}"
     )
-    messages = _sliding_messages(T3_SYSTEM, history) + [{"role": "user", "content": mandate}]
-    raw, _ = _call_llm(messages, max_tokens=80)
-    at, params = _parse_action(raw, ActionType.SUBMIT_FUNCTION, {})
-    if at != ActionType.SUBMIT_FUNCTION or not params.get("function_name", "").strip():
-        # Heuristic fallback: scan property text for a function name mention
-        fn_candidates = ["withdraw", "emergencyDrain", "buyTokens", "setPrice",
-                         "bid", "finalize", "stake", "claimRewards"]
-        prop_lower = prop.lower()
-        chosen = next((fn for fn in fn_candidates if fn.lower() in prop_lower), "withdraw")
-        at     = ActionType.SUBMIT_FUNCTION
-        params = {"function_name": chosen}
-    return at, params
 
 
 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()
+
     log_start(task="task3_rule_checker", env=ENV_BENCHMARK, model=MODEL_NAME) # type: ignore
 
-    history:      List[Dict[str, str]] = []
+    messages: List[ChatCompletionMessageParam] = [ # type: ignore
+        {"role": "system", "content": T3_SYSTEM}
+    ]
     step_rewards: List[float] = []
     grader_score  = 0.0
     steps_taken   = 0
@@ -417,23 +307,31 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
 
     try:
         for step in range(1, MAX_STEPS_T3 + 1):
-            steps_left = MAX_STEPS_T3 - step
-            is_last    = (step == MAX_STEPS_T3)
-
-            if is_last:
-                at, params = _t3_force_submit(obs, history)
-            else:
-                user_msg = _t3_user(obs, steps_left)
-                history.append({"role": "user", "content": user_msg})
-                messages  = _sliding_messages(T3_SYSTEM, history)
-                raw, error_msg = _call_llm(messages)
-                history.append({"role": "assistant", "content": raw})
-                at, params = _parse_action(raw, ActionType.GET_PROPERTY_SPECIFICATION, {})
-
-            result   = env.step(Action(action_type=at, params=params))
-            obs      = result.observation.model_dump()
-            r_val    = result.reward.value
-            done     = result.done
+            messages.append({"role": "user", "content": _t3_user_msg(obs)})
+            try:
+                resp = client.chat.completions.create(
+                    model=MODEL_NAME, messages=messages, # type: ignore
+                    max_tokens=200, temperature=0.0,
+                )
+                raw = resp.choices[0].message.content.strip()   # type: ignore
+                error_msg = None
+            except Exception as e:
+                raw = ""
+                error_msg = str(e)[:80]
+                print(f"[DEBUG] T3 LLM error ep={ep_num} step={step}: {e}", file=sys.stderr)
+
+            try:
+                parsed = json.loads(raw)
+                at     = ActionType(parsed["action"])
+                params = parsed.get("params", {})
+            except Exception:
+                at, params = ActionType.LIST_FUNCTIONS, {}
+
+            messages.append({"role": "assistant", "content": raw})
+            result = env.step(Action(action_type=at, params=params))
+            obs    = result.observation.model_dump()
+            r_val  = result.reward.value
+            done   = result.done
 
             step_rewards.append(r_val)
             steps_taken = step
@@ -441,18 +339,22 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
 
             if done:
                 v = r_val
-                grader_score = 0.999 if v >= 4.9 else (0.3 if v >= 0.999 else 0.0)
+                grader_score = 1.0 if v >= 4.9 else (0.3 if v >= 1.0 else 0.0)
                 break
 
-            if not is_last:
-                time.sleep(0.5)
+            time.sleep(0.3)
 
     finally:
         success = grader_score >= SUCCESS_SCORE_THRESHOLD
         log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)
 
-    return {"episode": ep_num, "seed": seed, "contract": obs["contract_name"],
-            "grader_score": grader_score, "cumulative_reward": obs["cumulative_reward"]}
+    return {
+        "episode":          ep_num,
+        "seed":             seed,
+        "contract":         obs["contract_name"],
+        "grader_score":     grader_score,
+        "cumulative_reward": obs["cumulative_reward"],
+    }
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -469,9 +371,11 @@ def run_task1(n: int = NUM_EPISODES) -> Dict[str, Any]:
     avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
     print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
     print(f"  Avg cum reward   : {avg_r:.2f}", flush=True)
-    return {"task_id": "task1_vuln_detection", "name": "Targeted Vulnerability Detection",
-            "status": "active", "num_episodes": n, "episodes": episodes,
-            "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r}
+    return {
+        "task_id": "task1_vuln_detection", "name": "Targeted Vulnerability Detection",
+        "status": "active", "num_episodes": n, "episodes": episodes,
+        "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r,
+    }
 
 
 def run_task2(n: int = NUM_EPISODES) -> Dict[str, Any]:
@@ -484,9 +388,11 @@ def run_task2(n: int = NUM_EPISODES) -> Dict[str, Any]:
     avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
     print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
     print(f"  Avg cum reward   : {avg_r:.2f}", flush=True)
-    return {"task_id": "task2_property_discovery", "name": "Property Discovery",
-            "status": "active", "num_episodes": n, "episodes": episodes,
-            "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r}
+    return {
+        "task_id": "task2_property_discovery", "name": "Property Discovery",
+        "status": "active", "num_episodes": n, "episodes": episodes,
+        "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r,
+    }
 
 
 def run_task3(n: int = NUM_EPISODES) -> Dict[str, Any]:
@@ -499,9 +405,11 @@ def run_task3(n: int = NUM_EPISODES) -> Dict[str, Any]:
     avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
     print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
     print(f"  Avg cum reward   : {avg_r:.2f}", flush=True)
-    return {"task_id": "task3_rule_checker", "name": "Rule Checker",
-            "status": "active", "num_episodes": n, "episodes": episodes,
-            "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r}
+    return {
+        "task_id": "task3_rule_checker", "name": "Rule Checker",
+        "status": "active", "num_episodes": n, "episodes": episodes,
+        "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r,
+    }
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -509,21 +417,18 @@ 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}", flush=True)
-    print(f"Base URL  : {API_BASE_URL}", flush=True)
-    print(f"Episodes  : {NUM_EPISODES} per task  |  "
-          f"Max steps: T1={MAX_STEPS_T1} T2={MAX_STEPS_T2} T3={MAX_STEPS_T3}", flush=True)
-    print(f"Hist window: last {HISTORY_WINDOW} exchanges  |  "
-          f"Result truncation: {MAX_RESULT_CHARS} chars", flush=True)
+    print(f"Model: {MODEL_NAME} | Base URL: {API_BASE_URL}", flush=True)
 
     t1 = run_task1(NUM_EPISODES)
     t2 = run_task2(NUM_EPISODES)
     t3 = run_task3(NUM_EPISODES)
 
     results = {
-        "model": MODEL_NAME, "base_url": API_BASE_URL,
-        "tasks": [t1, t2, t3],
+        "model":    MODEL_NAME,
+        "base_url": API_BASE_URL,
+        "tasks":    [t1, t2, t3],
     }
     overall = sum(t["avg_grader_score"] for t in results["tasks"]) / 3
     results["overall_avg_score"] = overall
@@ -541,4 +446,4 @@ async def main() -> None:
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    asyncio.run(main())