fix: Update output instructions in prompts.py to enforce case sensitivity and structure, inference.py made DRY, PROJECT works

inference.py

@@ -7,7 +7,7 @@ Implements agents for all three tasks using the Groq client.
 Emits mandatory structured stdout in the OpenEnv format.
 
 MANDATORY ENV VARS:
-  GROQ_API_KEY   Groq API key (required)
+  HF_TOKEN       Hugging Face Token (required)
   MODEL_NAME     Model identifier   (default: openai/gpt-oss-20b)
 
 MANDATORY STDOUT FORMAT (per episode):
@@ -26,15 +26,14 @@ import asyncio
 import json
 import os
 import sys
-import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Callable, Awaitable, Union
 
 from openai import AsyncOpenAI
+from dotenv import load_dotenv
 
 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
 
 # ─────────────────────────────────────────────────────────────────────────────
 # Configuration
@@ -47,24 +46,32 @@ HF_TOKEN     = os.getenv("HF_TOKEN",     "")
 
 if not HF_TOKEN:
     raise RuntimeError("HF_TOKEN environment variable not set")
-
+ 
 client = AsyncOpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)
 
+# from groq import AsyncGroq
+# GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+# client = AsyncGroq(api_key=GROQ_API_KEY)
+
 # Benchmark / environment identifier (constant for this env)
 ENV_BENCHMARK = "smart-contract-audit"
 
 # Episodes per task
-NUM_EPISODES = 3
+NUM_EPISODES = 5
 SEED_BASE    = 42
 
-# Max steps per task
-MAX_STEPS_T1 = 15
-MAX_STEPS_T2 = 15
-MAX_STEPS_T3 = 15
+# Max steps per task (same for all tasks)
+MAX_STEPS = 35
 
 # A grader_score >= this is considered a "success" for the [END] line
 SUCCESS_SCORE_THRESHOLD = 0.5
 
+# Throttle concurrent LLM calls
+SEMAPHORE = asyncio.Semaphore(3)
+
+# Timeout for each LLM request
+LLM_TIMEOUT = 20
+
 # ─────────────────────────────────────────────────────────────────────────────
 # Unified LLM call function
 # ─────────────────────────────────────────────────────────────────────────────
@@ -79,12 +86,19 @@ async def get_llm_response(
     Returns the response content as a string.
     Raises an exception on failure (to be caught by the caller).
     """
-    completion = await client.chat.completions.create(
-        model=MODEL_NAME,
-        messages=messages,          # type: ignore
-    )
-    return completion.choices[0].message.content.strip()  # type: ignore
-
+    try:
+        async with SEMAPHORE:
+            completion = await asyncio.wait_for(
+                client.chat.completions.create(
+                    model=MODEL_NAME,
+                    messages=messages,  # type: ignore
+                ),
+                timeout=LLM_TIMEOUT,
+            )
+        return completion.choices[0].message.content.strip()  # type: ignore
+
+    except asyncio.TimeoutError:
+        raise RuntimeError("LLM request timed out")
 
 # ─────────────────────────────────────────────────────────────────────────────
 # Mandatory stdout helpers
@@ -95,12 +109,7 @@ def log_start(task: str, env: str, model: str) -> None:
     print(f"[START] task={task} env={env} model={model}", flush=True)
 
 
-def log_step(
-    step: int,
-    action: str,
-    reward: float,
-    done: bool,
-    error: Optional[str] = None,
+def log_step( step: int, action: str, reward: float, done: bool, error: Optional[str] = None,
 ) -> None:
     """Emit a [STEP] line — one per env.step() call."""
     error_val = error if error else "null"
@@ -111,12 +120,7 @@ def log_step(
     )
 
 
-def log_end(
-    success: bool,
-    steps: int,
-    score: float,
-    rewards: List[float],
-) -> None:
+def log_end( success: bool, steps: int, score: float, rewards: List[float]) -> None:
     """Emit the [END] line — one per episode, always emitted."""
     rewards_str = ",".join(f"{r:.2f}" for r in rewards)
     print(
@@ -125,163 +129,112 @@ def log_end(
         flush=True,
     )
 
-
 # ─────────────────────────────────────────────────────────────────────────────
-# Task 1 — Targeted Vulnerability Detection
+# Generic episode runner
 # ─────────────────────────────────────────────────────────────────────────────
 
-def _t1_user_msg(obs: Dict[str, Any]) -> str:
-    return (
-        f"Last action : {obs['last_action'] or 'None'}\n"
-        f"Last result : {obs['last_action_result'] or 'Episode just started.'}"
-    )
-
-
-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)
+async def run_episode(
+    env: Union[Task1Environment, Task2Environment, Task3Environment],
+    seed: int,
+    ep_num: int,
+    *,
+    task_id: str,
+    system_prompt: str,
+    user_msg_formatter: Callable[[Dict[str, Any]], str],
+    max_tokens: int = 200,
+    default_action: ActionType = ActionType.LIST_FUNCTIONS,
+    extra_fields: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
+) -> Dict[str, Any]:
+    """
+    Run one episode with the given environment and task-specific parameters.
+    Emits [START]/[STEP]/[END] lines and returns a dict with episode results.
+    """
+    r = env.reset(seed=seed)
     obs = r.observation.model_dump()
 
-    log_start(task="task1_vuln_detection", env=ENV_BENCHMARK, model=MODEL_NAME)
+    log_start(task=task_id, env=ENV_BENCHMARK, model=MODEL_NAME)
 
     messages: List[Dict[str, str]] = [
-        {"role": "system", "content": T1_SYSTEM}
+        {"role": "system", "content": system_prompt}
     ]
     step_rewards: List[float] = []
-    grader_score  = 0.0
-    steps_taken   = 0
+    grader_score = 0.0
+    steps_taken = 0
     error_msg: Optional[str] = None
 
     try:
-        for step in range(1, MAX_STEPS_T1 + 1):
-            messages.append({"role": "user", "content": _t1_user_msg(obs)})
+        for step in range(1, MAX_STEPS + 1):
+            messages.append({"role": "user", "content": user_msg_formatter(obs)})
             try:
-                raw = await get_llm_response(messages, max_tokens=200, temperature=0.0)
+                raw = await get_llm_response(messages, max_tokens=max_tokens, temperature=0.0)
                 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)
+                print(f"[DEBUG] {task_id} LLM error ep={ep_num} step={step}: {e}", file=sys.stderr)
 
             try:
                 parsed = json.loads(raw)
-                at     = ActionType(parsed["action"])
+                at = ActionType(parsed["action"])
                 params = parsed.get("params", {})
-            except Exception:
-                at, params = ActionType.LIST_FUNCTIONS, {}
+            except Exception as e:
+                at, params = default_action, {}
+                print("Error in parsing LLM respoonse: " + str(e))
 
             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
+            obs = result.observation.model_dump()
+            r_val = result.reward.value
+            done = result.done
 
             step_rewards.append(r_val)
             steps_taken = step
+            print(raw, at.value, r_val)
             log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
 
             if done:
                 grader_score = r_val
                 break
 
-            time.sleep(0.3)
+            await asyncio.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
+    result_dict = {
+        "episode": ep_num,
+        "seed": seed,
+        "grader_score": grader_score,
+        "contract": obs.get("contract_name", ""),
     }
+    if extra_fields:
+        result_dict.update(extra_fields(obs))
 
+    return result_dict
 
 # ─────────────────────────────────────────────────────────────────────────────
-# Task 2 — Property Discovery
+# Task-specific user message formatters and extra field extractors
 # ─────────────────────────────────────────────────────────────────────────────
 
+def t1_user_msg(obs: Dict[str, Any]) -> str:
+    return (
+        f"Last action : {obs['last_action'] or 'None'}\n"
+        f"Last result : {obs['last_action_result'] or 'Episode just started.'}"
+    )
 
-def _t2_user_msg(obs: Dict[str, Any]) -> str:
+def t2_user_msg(obs: Dict[str, Any]) -> str:
     extra = obs.get("extra", {})
     return (
         f"Target Function : {extra.get('target_function', '?')}  "
-        # f"({extra.get('target_signature', '')})\n"
         f"Last action : {obs['last_action'] or 'None'}\n"
         f"Last result :\n{obs['last_action_result'] or 'Episode just started.'}"
     )
 
+def t2_extra_fields(obs: Dict[str, Any]) -> Dict[str, Any]:
+    return {"function": obs.get("extra", {}).get("target_function", "?")}
 
-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()
-    fn  = obs["extra"].get("target_function", "?")
-
-    log_start(task="task2_property_discovery", env=ENV_BENCHMARK, model=MODEL_NAME)
-
-    messages: List[Dict[str, str]] = [
-        {"role": "system", "content": T2_SYSTEM}
-    ]
-    step_rewards: List[float] = []
-    grader_score  = 0.0
-    steps_taken   = 0
-    error_msg: Optional[str] = None
-
-    try:
-        for step in range(1, MAX_STEPS_T2 + 1):
-            messages.append({"role": "user", "content": _t2_user_msg(obs)})
-            try:
-                raw = await get_llm_response(messages, max_tokens=400, temperature=0.0)
-                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
-            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
-
-            if done:
-                grader_score = r_val
-                break
-
-            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
-    }
-
-
-# ─────────────────────────────────────────────────────────────────────────────
-# Task 3 — Rule Checker
-# ─────────────────────────────────────────────────────────────────────────────
-
-
-def _t3_user_msg(obs: Dict[str, Any]) -> str:
+def t3_user_msg(obs: Dict[str, Any]) -> str:
     extra = obs.get("extra", {})
     return (
         f"Verify Property  : {extra.get('property_english', '(none)')}\n"
@@ -289,139 +242,117 @@ def _t3_user_msg(obs: Dict[str, Any]) -> str:
         f"Last result :\n{obs['last_action_result'] or 'Episode just started.'}"
     )
 
+# ─────────────────────────────────────────────────────────────────────────────
+# Generic task runner
+# ─────────────────────────────────────────────────────────────────────────────
 
-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()
-
-    log_start(task="task3_rule_checker", env=ENV_BENCHMARK, model=MODEL_NAME)
-
-    messages: List[Dict[str, str]] = [
-        {"role": "system", "content": T3_SYSTEM}
+async def run_task(
+    task_id: str,
+    task_name: str,
+    env_class: type,
+    system_prompt: str,
+    user_msg_formatter: Callable[[Dict[str, Any]], str],
+    max_tokens: int = 200,
+    default_action: ActionType = ActionType.LIST_FUNCTIONS,
+    extra_fields: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
+    num_episodes: int = NUM_EPISODES,
+) -> Dict[str, Any]:
+    """Run multiple episodes for a given task and return aggregated results."""
+    print("\n" + "=" * 60, flush=True)
+    print(f"TASK: {task_name}", flush=True)
+    print("=" * 60, flush=True)
+
+    env = env_class()
+    tasks = [
+        run_episode(
+            env,
+            seed=SEED_BASE + i,
+            ep_num=i + 1,
+            task_id=task_id,
+            system_prompt=system_prompt,
+            user_msg_formatter=user_msg_formatter,
+            max_tokens=max_tokens,
+            default_action=default_action,
+            extra_fields=extra_fields,
+        )
+        for i in range(num_episodes)
     ]
-    step_rewards: List[float] = []
-    grader_score  = 0.0
-    steps_taken   = 0
-    error_msg: Optional[str] = None
-
-    try:
-        for step in range(1, MAX_STEPS_T3 + 1):
-            messages.append({"role": "user", "content": _t3_user_msg(obs)})
-            try:
-                raw = await get_llm_response(messages, max_tokens=200, temperature=0.0)
-                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
-            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
-
-            if done:
-                grader_score = r_val
-                break
-
-            time.sleep(0.3)
-
-    finally:
-        success = grader_score >= SUCCESS_SCORE_THRESHOLD
-        log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)
+    episodes = await asyncio.gather(*tasks)
+    avg_score = sum(e["grader_score"] for e in episodes) / num_episodes
 
+    print(f"\n  Avg grader score : {avg_score:.4f}", flush=True)
     return {
-        "episode":          ep_num,
-        "seed":             seed,
-        "contract":         obs["contract_name"],
-        "grader_score":     grader_score
+        "task_id": task_id,
+        "name": task_name,
+        "status": "active",
+        "num_episodes": num_episodes,
+        "episodes": episodes,
+        "avg_grader_score": avg_score,
     }
 
-
 # ─────────────────────────────────────────────────────────────────────────────
-# Task runners
+# Task-specific runners (thin wrappers for clarity)
 # ─────────────────────────────────────────────────────────────────────────────
 
 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 = [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 {
-        "task_id": "task1_vuln_detection", "name": "Targeted Vulnerability Detection",
-        "status": "active", "num_episodes": n, "episodes": episodes,
-        "avg_grader_score": avg_s
-    }
-
+    return await run_task(
+        task_id="task1_vuln_detection",
+        task_name="Targeted Vulnerability Detection",
+        env_class=Task1Environment,
+        system_prompt=T1_SYSTEM,
+        user_msg_formatter=t1_user_msg,
+        max_tokens=200,
+        default_action=ActionType.LIST_FUNCTIONS,
+        num_episodes=n,
+    )
 
 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 = [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
-    print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
-    return {
-        "task_id": "task2_property_discovery", "name": "Property Discovery",
-        "status": "active", "num_episodes": n, "episodes": episodes,
-        "avg_grader_score": avg_s
-    }
-
+    return await run_task(
+        task_id="task2_property_discovery",
+        task_name="Property Discovery",
+        env_class=Task2Environment,
+        system_prompt=T2_SYSTEM,
+        user_msg_formatter=t2_user_msg,
+        max_tokens=400,
+        default_action=ActionType.GET_FUNCTION_CODE,
+        extra_fields=t2_extra_fields,
+        num_episodes=n,
+    )
 
 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 = [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
-    print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
-    return {
-        "task_id": "task3_rule_checker", "name": "Rule Checker",
-        "status": "active", "num_episodes": n, "episodes": episodes,
-        "avg_grader_score": avg_s
-    }
-
+    return await run_task(
+        task_id="task3_rule_checker",
+        task_name="Rule Checker",
+        env_class=Task3Environment,
+        system_prompt=T3_SYSTEM,
+        user_msg_formatter=t3_user_msg,
+        max_tokens=200,
+        default_action=ActionType.LIST_FUNCTIONS,
+        num_episodes=n,
+    )
 
 # ─────────────────────────────────────────────────────────────────────────────
 # Main
 # ─────────────────────────────────────────────────────────────────────────────
 
 async def main() -> None:
-    """Async entry point (wraps sync env calls; asyncio.run() expected by caller)."""
+    """Async entry point."""
     print("Smart Contract Audit RL Environment — Baseline Inference", flush=True)
 
     t1 = await run_task1(NUM_EPISODES)
     t2 = await run_task2(NUM_EPISODES)
     t3 = await run_task3(NUM_EPISODES)
 
-    results: Dict[str, Any] = { "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
 
-    print("\n" + "="*60, flush=True)
+    print("\n" + "=" * 60, flush=True)
     print("BASELINE SUMMARY", flush=True)
-    print("="*60, flush=True)
+    print("=" * 60, flush=True)
     for t in results["tasks"]:
         print(f"  ✅ {t['name']:40s}: {t['avg_grader_score']:.3f}", flush=True)
-    print(f"\n  Overall avg grader score: {overall:.3f}", flush=True)
+    print(f"\n  Overall avg grader score: {overall:.4f}", flush=True)
 
     with open("baseline_scores.json", "w") as f:
         json.dump(results, f, indent=2)

utils/prompts.py

@@ -32,7 +32,8 @@ Common vulnerabilities in contracts:
 - denial of service
 
 Submit immediately once confident.
-Output: JSON only. No text.
+Output: JSON only. No text. FOLLOW EXACT STRCUTURE OF ACTIONS GIVEN ANY CHANGE WILL LEAD TO
+INVALID ACTION. It's case-sensitive as well.
 """
 
 T2_SYSTEM = """You are a Solidity formal methods engineer.
@@ -71,7 +72,8 @@ Format:
 
 Submit immediately once confident.
 
-Output: JSON only.
+Output: JSON only. No text. FOLLOW EXACT STRCUTURE OF ACTIONS GIVEN ANY CHANGE WILL LEAD TO
+INVALID ACTION. It's case-sensitive as well.
 """
 
 T3_SYSTEM = """You are a Solidity security auditor.
@@ -108,5 +110,6 @@ Example Violation heuristics:
 Select the function that clearly breaks the property.  
 Submit immediately once confident.
 
-Output: JSON only.
+Output: JSON only. No text. FOLLOW EXACT STRCUTURE OF ACTIONS GIVEN ANY CHANGE WILL LEAD TO
+INVALID ACTION. It's case-sensitive as well.
 """