""" inference.py ------------ Inference script — Smart Contract Audit RL Environment. Implements agents for all three tasks using the Groq client. Emits mandatory structured stdout in the OpenEnv format. MANDATORY ENV VARS: HF_TOKEN Hugging Face Token (required) MODEL_NAME Model identifier (default: openai/gpt-oss-20b) MANDATORY STDOUT FORMAT (per episode): [START] task= env=smart-contract-audit model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= score=<0.000> rewards= Usage: python inference.py Output: Structured stdout per episode, plus baseline_scores.json summary. """ import asyncio import json import os import sys from collections import deque 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 # ───────────────────────────────────────────────────────────────────────────── # Configuration # ───────────────────────────────────────────────────────────────────────────── load_dotenv() API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1/") MODEL_NAME = os.getenv("MODEL_NAME", "CohereLabs/tiny-aya-fire:cohere") HF_TOKEN = os.getenv("HF_TOKEN", "") if not HF_TOKEN: raise RuntimeError("HF_TOKEN environment variable not set") if not MODEL_NAME: raise RuntimeError("MODEL_NAME not set") client = AsyncOpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL) # Benchmark / environment identifier (constant for this env) ENV_BENCHMARK = "smart-contract-audit" # Episodes per task NUM_EPISODES = 5 SEED_BASE = 42 # 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 # ───────────────────────────────────────────────────────────────────────────── async def get_llm_response( messages: List[Dict[str, str]], max_tokens: int = 200, temperature: float = 0.0, ) -> str: """ 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). """ 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 # ───────────────────────────────────────────────────────────────────────────── def log_start(task: str, env: str, model: str) -> None: """Emit the [START] line — one per episode.""" 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, ) -> None: """Emit a [STEP] line — one per env.step() call.""" error_val = error if error else "null" print( f"[STEP] step={step} action={action} " f"reward={reward:.2f} done={str(done).lower()} error={error_val}", flush=True, ) 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( f"[END] success={str(success).lower()} steps={steps} " f"score={score:.3f} rewards={rewards_str}", flush=True, ) def _clamp(reward: float) -> float: return max(0.001, min(0.999, reward)) # ───────────────────────────────────────────────────────────────────────────── # Generic episode runner # ───────────────────────────────────────────────────────────────────────────── 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]: r = env.reset(seed=seed) obs = r.observation.model_dump() log_start(task=task_id, env=ENV_BENCHMARK, model=MODEL_NAME) # Keep only the last 2 user-assistant pairs (4 messages). history: deque = deque(maxlen=4) step_rewards: List[float] = [] grader_score = 0.0 steps_taken = 0 error_msg: Optional[str] = None try: for step in range(1, MAX_STEPS + 1): user_msg_content = user_msg_formatter(obs) user_message = {"role": "user", "content": user_msg_content} history.append(user_message) # Always prepend the system prompt so it survives deque eviction messages_for_llm = [ {"role": "system", "content": system_prompt}, *list(history), ] try: raw = await get_llm_response(messages_for_llm, max_tokens=max_tokens, temperature=0.0) error_msg = None except Exception as e: raw = "" error_msg = str(e)[:80] print(f"[DEBUG] {task_id} LLM error ep={ep_num} step={step}: {e}", file=sys.stderr) # Append the assistant reply so the next step sees the full turn history.append({"role": "assistant", "content": raw}) try: parsed = json.loads(raw) at = ActionType(parsed["action"]) params = parsed.get("params", {}) except Exception as e: at, params = default_action, {} print("Error in parsing LLM response: " + str(e)) 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 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) result_dict = { "episode": ep_num, "seed": seed, "grader_score": _clamp(grader_score), "contract": obs.get("contract_name", ""), } if extra_fields: result_dict.update(extra_fields(obs)) return result_dict # ───────────────────────────────────────────────────────────────────────────── # 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: extra = obs.get("extra", {}) return ( f"Target Function : {extra.get('target_function', '?')} " 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", "?")} def t3_user_msg(obs: Dict[str, Any]) -> str: extra = obs.get("extra", {}) return ( f"Verify Property : {extra.get('property_english', '(none)')}\n" f"Last action : {obs['last_action'] or 'None'}\n" f"Last result :\n{obs['last_action_result'] or 'Episode just started.'}" ) # ───────────────────────────────────────────────────────────────────────────── # Generic task runner # ───────────────────────────────────────────────────────────────────────────── 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) ] episodes = await asyncio.gather(*tasks) avg_score = sum(e["grader_score"] for e in episodes) / num_episodes print(f"\n Avg grader score : {_clamp(avg_score):.4f}", flush=True) return { "task_id": task_id, "name": task_name, "status": "active", "num_episodes": num_episodes, "episodes": episodes, "avg_grader_score": _clamp(avg_score), } # ───────────────────────────────────────────────────────────────────────────── # Task-specific runners (thin wrappers for clarity) # ───────────────────────────────────────────────────────────────────────────── async def run_task1(n: int = NUM_EPISODES) -> Dict[str, Any]: 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]: 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]: 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.""" 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]} overall = sum(t["avg_grader_score"] for t in results["tasks"]) / 3 results["overall_avg_score"] = overall print("\n" + "=" * 60, flush=True) print("BASELINE SUMMARY", flush=True) print("=" * 60, flush=True) for t in results["tasks"]: print(f" ✅ {t['name']:40s}: {_clamp(t['avg_grader_score']):.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) print("\n Scores written to baseline_scores.json", flush=True) if __name__ == "__main__": asyncio.run(main())