feat: Update inference script and Dockerfile for improved functionality

- Enhanced Dockerfile to include git-lfs and clone model directly into the container.
- Refactored inference logic to improve episode handling and logging.
- Introduced new constants for episode management and result truncation.

.gitignore

@@ -13,4 +13,5 @@ baseline_scores.json
 .pytest_cache/
 MySolution.md
 nltk_data
-uv.lock
+uv.lock
+eval_results.json

Dockerfile

@@ -4,20 +4,25 @@ FROM ${BASE_IMAGE}
 
 WORKDIR /app
 
-# Build tools needed for pandas/numpy/scikit-learn C extensions
+# Build tools
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
-        git build-essential \
+        git git-lfs build-essential curl \
     && rm -rf /var/lib/apt/lists/*
 
+# Enable git-lfs (required for HF model weights)
+RUN git lfs install
+
 # Copy source
 COPY . /app/env
 
-# Install directly into the base image's own Python — no venv, no --target
-# openenv is already provided by the base image, so exclude it
+# Install dependencies (excluding openenv)
 RUN grep -v 'openenv' /app/env/requirements.txt | pip install -r /dev/stdin
 
-# Download NLTK data — nltk is now properly installed, so this just works
+# Clone model directly into container (NO HF API at runtime)
+RUN git clone https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2 /app/model
+
+# Download NLTK data
 RUN python -m nltk.downloader \
         wordnet omw-1.4 stopwords punkt \
         averaged_perceptron_tagger_eng punkt_tab

inference.py

@@ -50,19 +50,28 @@ if not HF_TOKEN:
     print("[WARN] HF_TOKEN not set — API calls may fail.", file=sys.stderr)
     exit(1)
 
-# Benchmark / environment identifier (constant for this env)
-ENV_BENCHMARK = "smart-contract-audit"
+ENV_BENCHMARK          = "smart-contract-audit"
+NUM_EPISODES           = 1      # keep low on free tier; raise for full eval
+SEED_BASE              = 42
 
-# Episodes per task
-NUM_EPISODES = 3
-SEED_BASE    = 42
+# 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
 
-# Max steps per task
-MAX_STEPS_T1 = 15
-MAX_STEPS_T2 = 10
-MAX_STEPS_T3 = 12
+# 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
 
-# A grader_score >= this is considered a "success" for the [END] line
+# 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
 SUCCESS_SCORE_THRESHOLD = 0.5
 
 client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)
@@ -108,30 +117,115 @@ 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_msg(obs: Dict[str, Any]) -> str:
+def _t1_user(obs: Dict[str, Any], steps_left: int) -> str:
+    result = _truncate(obs.get("last_action_result") or "Episode just started.")
     return (
-        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.'}"
+        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>"}}'
+    )
+    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
 
-    messages: List[ChatCompletionMessageParam] = [ # type: ignore
-        {"role": "system", "content": T1_SYSTEM}
-    ]
+    history:      List[Dict[str, str]] = []
     step_rewards: List[float] = []
     grader_score  = 0.0
     steps_taken   = 0
@@ -139,31 +233,24 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
 
     try:
         for step in range(1, MAX_STEPS_T1 + 1):
-            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
+            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
 
             step_rewards.append(r_val)
             steps_taken = step
@@ -174,19 +261,15 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
                 grader_score = 1.0 if v >= 4.9 else (0.5 if v >= 0.9 else 0.0)
                 break
 
-            time.sleep(0.3)
+            if not is_last:
+                time.sleep(0.5)
 
     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"]}
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -194,29 +277,48 @@ def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str,
 # ─────────────────────────────────────────────────────────────────────────────
 
 
-def _t2_user_msg(obs: Dict[str, Any]) -> str:
-    extra = obs.get("extra", {})
+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']}\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.'}"
+        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."""
+    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>"}}'
     )
+    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
 
-    messages: List[ChatCompletionMessageParam] = [ # type: ignore
-        {"role": "system", "content": T2_SYSTEM}
-    ]
+    history:      List[Dict[str, str]] = []
     step_rewards: List[float] = []
     grader_score  = 0.0
     steps_taken   = 0
@@ -224,31 +326,23 @@ def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str,
 
     try:
         for step in range(1, MAX_STEPS_T2 + 1):
-            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
+            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
 
             step_rewards.append(r_val)
             steps_taken = step
@@ -258,48 +352,64 @@ 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
 
-            time.sleep(0.3)
+            if not is_last:
+                time.sleep(0.5)
 
     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_msg(obs: Dict[str, Any]) -> str:
-    extra = obs.get("extra", {})
+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)')}\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.'}"
+        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>"}}'
+    )
+    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
 
-    messages: List[ChatCompletionMessageParam] = [ # type: ignore
-        {"role": "system", "content": T3_SYSTEM}
-    ]
+    history:      List[Dict[str, str]] = []
     step_rewards: List[float] = []
     grader_score  = 0.0
     steps_taken   = 0
@@ -307,31 +417,23 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
 
     try:
         for step in range(1, MAX_STEPS_T3 + 1):
-            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
+            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
 
             step_rewards.append(r_val)
             steps_taken = step
@@ -342,19 +444,15 @@ def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str,
                 grader_score = 1.0 if v >= 4.9 else (0.3 if v >= 1.0 else 0.0)
                 break
 
-            time.sleep(0.3)
+            if not is_last:
+                time.sleep(0.5)
 
     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"]}
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -371,11 +469,9 @@ 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]:
@@ -388,11 +484,9 @@ 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]:
@@ -405,11 +499,9 @@ 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}
 
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -417,18 +509,21 @@ 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} | Base URL: {API_BASE_URL}", 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)
 
     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

org_inference.py

@@ -0,0 +1,449 @@
+"""
+inference.py
+------------
+Inference script — Smart Contract Audit RL Environment.
+
+Implements agents for all three tasks using the OpenAI-compatible client.
+Emits mandatory structured stdout in the OpenEnv format.
+
+MANDATORY ENV VARS:
+  API_BASE_URL   LLM API endpoint   (default: https://api.openai.com/v1)
+  MODEL_NAME     Model identifier   (default: gpt-4o-mini)
+  HF_TOKEN       API key / HF token
+
+MANDATORY STDOUT FORMAT (per episode):
+  [START] task=<id> env=smart-contract-audit model=<model>
+  [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<str|null>
+  [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
+
+Usage:
+  python inference.py
+
+Output:
+  Structured stdout per episode, plus baseline_scores.json summary.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import time
+from typing import Any, Dict, List, Optional
+
+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 dotenv_values
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Configuration
+# ─────────────────────────────────────────────────────────────────────────────
+
+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)
+
+# Benchmark / environment identifier (constant for this env)
+ENV_BENCHMARK = "smart-contract-audit"
+
+# Episodes per task
+NUM_EPISODES = 3
+SEED_BASE    = 42
+
+# Max steps per task
+MAX_STEPS_T1 = 15
+MAX_STEPS_T2 = 10
+MAX_STEPS_T3 = 12
+
+# 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)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# 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,
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Task 1 — Targeted Vulnerability Detection
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _t1_user_msg(obs: Dict[str, Any]) -> str:
+    return (
+        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.'}"
+    )
+
+
+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
+
+    messages: List[ChatCompletionMessageParam] = [ # type: ignore
+        {"role": "system", "content": T1_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_T1 + 1):
+            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
+            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
+
+            if done:
+                v = r_val
+                grader_score = 1.0 if v >= 4.9 else (0.5 if v >= 0.9 else 0.0)
+                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"],
+        "grader_score":     grader_score,
+        "cumulative_reward": obs["cumulative_reward"],
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Task 2 — Property Discovery
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def _t2_user_msg(obs: Dict[str, Any]) -> str:
+    extra = obs.get("extra", {})
+    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.'}"
+    )
+
+
+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
+
+    messages: List[ChatCompletionMessageParam] = [ # type: ignore
+        {"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:
+                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
+            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
+
+            if done:
+                grader_score = round(r_val / 5.0, 3) if r_val > 0 else 0.0
+                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,
+        "cumulative_reward": obs["cumulative_reward"],
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Task 3 — Rule Checker
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+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.'}"
+    )
+
+
+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
+
+    messages: List[ChatCompletionMessageParam] = [ # type: ignore
+        {"role": "system", "content": T3_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_T3 + 1):
+            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
+            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
+
+            if done:
+                v = r_val
+                grader_score = 1.0 if v >= 4.9 else (0.3 if v >= 1.0 else 0.0)
+                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"],
+        "grader_score":     grader_score,
+        "cumulative_reward": obs["cumulative_reward"],
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Task runners
+# ─────────────────────────────────────────────────────────────────────────────
+
+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)]
+    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)
+    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,
+    }
+
+
+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)]
+    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)
+    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,
+    }
+
+
+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)]
+    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)
+    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,
+    }
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Main
+# ─────────────────────────────────────────────────────────────────────────────
+
+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} | 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],
+    }
+    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}: {t['avg_grader_score']:.3f}", flush=True)
+    print(f"\n  Overall avg grader score: {overall:.3f}", 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())