refactor: Update API base URL and model name, enhance message handling with history tracking

inference.py

@@ -26,6 +26,7 @@ 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
@@ -40,8 +41,8 @@ from utils import T1_SYSTEM, T2_SYSTEM, T3_SYSTEM
 # ─────────────────────────────────────────────────────────────────────────────
 
 load_dotenv()
-API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
-MODEL_NAME   = os.getenv("MODEL_NAME",   "")
+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:
@@ -52,10 +53,6 @@ if not MODEL_NAME:
  
 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"
 
@@ -151,18 +148,13 @@ async def run_episode(
     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=task_id, env=ENV_BENCHMARK, model=MODEL_NAME)
 
-    messages: List[Dict[str, str]] = [
-        {"role": "system", "content": system_prompt}
-    ]
+    # 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
@@ -170,29 +162,39 @@ async def run_episode(
 
     try:
         for step in range(1, MAX_STEPS + 1):
-            messages.append({"role": "user", "content": user_msg_formatter(obs)})
+            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, max_tokens=max_tokens, temperature=0.0)
+                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 respoonse: " + str(e))
+                print("Error in parsing LLM response: " + 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
-
             step_rewards.append(r_val)
             steps_taken = step
             log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)
@@ -215,7 +217,6 @@ async def run_episode(
     }
     if extra_fields:
         result_dict.update(extra_fields(obs))
-
     return result_dict
 
 # ─────────────────────────────────────────────────────────────────────────────