""" inference.py TraceRL OpenEnv Baseline Agent ============================================== MetaEnv Round 1 competition-compliant inference script. MANDATORY ENV VARS (injected by grading server): API_BASE_URL The LLM API endpoint (OpenAI-compatible) MODEL_NAME The model identifier for inference HF_TOKEN Hugging Face / API key STDOUT FORMAT (machine-parsed by automated validator): [START] task= env= model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= rewards= RULES: - reward / rewards always formatted to exactly 2 decimal places (f"{reward:.2f}") - done / success always lowercase boolean strings: true / false - error is raw last_action_error string, or null if none - One [START], one [STEP] per env.step(), one [END] always even on exception """ import os import sys import textwrap import difflib import re import time from typing import List, Optional from openai import OpenAI import unidiff from client import CodeFixerEnv from models import CodeFixerAction from dotenv import load_dotenv load_dotenv() import logging logger = logging.getLogger(__name__) def _enable_hermetic_runtime() -> None: """ Keep this process isolated from parent/global Python environment leakage. """ if os.getenv("HERMETIC_RUN", "1").strip().lower() not in {"1", "true", "yes"}: return os.environ["PYTHONPATH"] = "" os.environ["PYTHONNOUSERSITE"] = "1" os.environ["PYTHONDONTWRITEBYTECODE"] = "1" os.environ.pop("PYTHONHOME", None) _enable_hermetic_runtime() API_BASE_URL = os.getenv("API_BASE_URL", "https://integrate.api.nvidia.com/v1") API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") MODEL_NAME = os.getenv("MODEL_NAME", "qwen/qwen2.5-coder-32b-instruct") MAX_STEPS = int(os.getenv("MAX_STEPS", "10")) TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7")) MAX_TOKENS = int(os.getenv("MAX_TOKENS", "512")) SUCCESS_SCORE_THRESHOLD = float(os.getenv("SUCCESS_THRESHOLD", "0.5")) MAX_RETRIES = int(os.getenv("MAX_RETRIES", "3")) _DIFFICULTIES = ["easy", "medium", "hard"] TASK_NAME = (os.getenv("TRACERL_TASK") or "").strip().lower() BENCHMARK = os.getenv("TRACERL_BENCHMARK", "rl-code-fix-env") ENV_URL = os.getenv("ENV_URL", "http://localhost:8000") if "TASK_SOURCE" not in os.environ: os.environ["TASK_SOURCE"] = "swebench" if TASK_NAME and TASK_NAME not in _DIFFICULTIES: print( f"[DEBUG] Invalid TRACERL_TASK='{TASK_NAME}'. Falling back to all tasks.", flush=True, file=sys.stderr, ) TASK_NAME = "" llm = OpenAI(api_key=API_KEY, base_url=API_BASE_URL) SYSTEM_PROMPT = textwrap.dedent(""" You are an autonomous bug-fixing agent. You will be given a Python function or module that contains one or more bugs. Your task is to produce corrected Python code that fixes all bugs. Output format — you MUST follow this exactly: - Return ONLY the corrected code, nothing else. - Do NOT include markdown fences (```), prose, or explanations. - If the code has no bugs, output exactly: NO_CHANGE Code rules: - Preserve the original function signature exactly. - Do not add new imports unless strictly required by the fix. - Do not add print statements, debug code, or comments that were not in the original. - Apply the minimal fix needed. Do not refactor, rename, or reformat unrelated lines. Think step-by-step internally, but your final output must be ONLY corrected code or NO_CHANGE. """).strip() def _build_user_prompt(observation: dict, history: List[str], step: int) -> str: """Build the per-step prompt from the current environment observation.""" code = observation.get("code", "(no code provided)") test_score = observation.get("test_score", 0.0) total_tests = observation.get("total_tests", 1) error_logs = observation.get("logs", "") test_status = "PASSING" if test_score >= 1.0 else "FAILING" parts = [ f"Step: {step}", f"Tests Status: {test_status} (score {test_score:.2f}/{total_tests})", "", "=== Current Code ===", code, "", ] if error_logs: parts += ["=== Error Logs ===", error_logs, ""] # FIX: include actual previous patch attempts so LLM learns from history if history: parts += ["=== Previous Attempts (last 3) ==="] parts += history[-3:] parts += [""] parts += [ "Your task: provide the corrected full code for the snippet above.", "Return ONLY corrected code (or NO_CHANGE), no explanations or markdown.", ] return "\n".join(parts) 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: done_str = "true" if done else "false" error_str = error if error else "null" action_str = action.replace("\n", "\\n").replace("\r", "") print( f"[STEP] step={step} action={action_str} " f"reward={reward:.2f} done={done_str} error={error_str}", flush=True, ) def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: success_str = "true" if success else "false" rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={success_str} steps={steps} rewards={rewards_str}", flush=True, ) _env_client: Optional[CodeFixerEnv] = None def _get_client() -> CodeFixerEnv: global _env_client if _env_client is None: _env_client = CodeFixerEnv(base_url=ENV_URL) return _env_client def env_reset() -> dict: result = _get_client().reset() obs = result.observation return { "observation": { "code": obs.code, "logs": obs.logs, "test_score": obs.test_score, "total_tests": obs.total_tests, "steps": obs.steps, }, "done": getattr(obs, "done", False), } def env_step(action_type: str, payload: str = "") -> dict: try: result = _get_client().step(CodeFixerAction(type=action_type, payload=payload)) obs = result.observation return { "observation": { "code": obs.code, "logs": obs.logs, "test_score": obs.test_score, "total_tests": obs.total_tests, "steps": obs.steps, }, "reward": result.reward, "done": result.done, "error": None, } except Exception as exc: return { "observation": { "code": "", "logs": "", "test_score": 0.0, "total_tests": 1, "steps": 0, }, "reward": 0.0, "done": False, "error": str(exc), } def _strip_markdown_fences(text: str) -> str: cleaned = (text or "").strip() # Remove opening fence with optional language if cleaned.startswith("```"): cleaned = re.sub(r"^```[a-zA-Z0-9_-]*\s*\n", "", cleaned) cleaned = re.sub(r"\n```$", "", cleaned.strip()) return cleaned.strip() def _looks_like_unified_diff(text: str) -> bool: return ("--- " in text and "+++ " in text and "@@ " in text) def _is_valid_unified_diff(text: str) -> bool: try: unidiff.PatchSet(text) return True except Exception: return False def _build_unified_diff(original: str, revised: str) -> str: lines = difflib.unified_diff( original.splitlines(), revised.splitlines(), fromfile="a/code.py", tofile="b/code.py", lineterm="", ) return "\n".join(lines).strip() def _looks_like_python_code(text: str) -> bool: """Check if text looks like Python code (not a diff, not prose).""" text_lower = text.lower().strip() # Not a diff if _looks_like_unified_diff(text): return False # Not just "no change" or similar if text_lower in ("no_change", "no change", "no changes", "no changes made", "the code is correct", "already correct", "no bugs found", "no bug", "no bugs", "already fixed"): return False # Check for Python-like patterns (def, class, import, etc.) python_indicators = ['def ', 'class ', 'import ', 'from ', 'return ', 'if __name__', 'async def', 'self.', 'print('] return any(indicator in text for indicator in python_indicators) def _looks_like_output_value(text: str) -> bool: """Check if text is just an output value like [[3,1],[4,2]]""" text = text.strip() # Looks like a literal value/expression, not code if text.startswith('[') and text.endswith(']') and '\n' not in text: # Could be a list/dict literal if ',' in text or ':' in text: return True # Single word or simple expression if '\n' not in text and len(text.split()) <= 3: # Check if it's not a function definition if 'def ' not in text and 'class ' not in text: return True return False def _normalize_action(raw_action: str, original_code: str) -> str: """Normalize LLM output to unified diff format.""" cleaned = _strip_markdown_fences(raw_action) # Handle empty or NO_CHANGE cases if not cleaned: return "" cleaned_lower = cleaned.lower().strip() if cleaned_lower == "no_change": return "" # Check for common "no change" phrases no_change_phrases = [ "no change", "no changes", "no changes made", "the code is correct", "already correct", "no bugs found", "no bug", "no bugs", "already fixed", "no modifications needed", "the provided code is already correct", "code appears to be correct" ] if any(phrase in cleaned_lower for phrase in no_change_phrases): return "" # Check if output looks like a raw value (not code) - reject it if _looks_like_output_value(cleaned): print(f"[DEBUG] Rejected output that looks like a value: {cleaned[:50]}...", flush=True, file=sys.stderr) return "" # Check if it looks like a unified diff already if _looks_like_unified_diff(cleaned): return cleaned if _is_valid_unified_diff(cleaned) else "" # Check if it looks like Python code - convert to diff if _looks_like_python_code(cleaned): generated = _build_unified_diff(original_code, cleaned) return generated if _is_valid_unified_diff(generated) else "" # If nothing matches, return empty (invalid output) print(f"[DEBUG] Could not parse LLM output as code or diff: {cleaned[:100]}...", flush=True, file=sys.stderr) return "" def get_action(observation: dict, history: List[str], step: int) -> str: """ Call the LLM to produce the next action. Returns a unified diff patch string or empty string on failure. """ user_prompt = _build_user_prompt(observation, history, step) for attempt in range(MAX_RETRIES): try: completion = llm.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False, ) raw_action = (completion.choices[0].message.content or "").strip() return _normalize_action(raw_action, observation.get("code", "")) except Exception as exc: print( f"[DEBUG] LLM call failed at step {step} (attempt {attempt+1}/{MAX_RETRIES}): {exc}", flush=True, file=sys.stderr, ) if "429" in str(exc) and attempt < MAX_RETRIES - 1: # FIX: exponential backoff with longer base delay for 429s sleep_time = 5 * (2 ** attempt) print( f"[DEBUG] Rate limited. Sleeping {sleep_time}s before retry.", flush=True, file=sys.stderr, ) time.sleep(sleep_time) continue if attempt == MAX_RETRIES - 1: return "" return "" def main() -> None: forced_task = os.getenv("TRACERL_TASK") task_plan = [TASK_NAME] if forced_task else _DIFFICULTIES[:] try: for planned_task in task_plan: rewards: List[float] = [] history: List[str] = [] steps_taken: int = 0 success: bool = False score: float = 0.0 final_test_score: float = 0.0 episode_start = time.perf_counter() log_start(task=planned_task, env=BENCHMARK, model=MODEL_NAME) print(f"[DEBUG] planned_task={planned_task}", flush=True, file=sys.stderr) try: reset_resp = env_reset() observation = reset_resp.get("observation", {}) done = reset_resp.get("done", False) for step in range(1, MAX_STEPS + 1): step_start = time.perf_counter() if done: break # Get patch from LLM action = get_action(observation, history, step) # Step 1: apply the patch patch_resp = env_step("apply_patch", action) patch_obs = patch_resp.get("observation", {}) patch_error = patch_resp.get("error") or None if not patch_error and patch_obs.get("code"): observation = patch_obs # Step 2: always run tests after patching to get real reward signal test_resp = env_step("run_tests", "") next_obs = test_resp.get("observation", {}) reward = float(test_resp.get("reward", 0.0)) done = bool(test_resp.get("done", False)) test_error = test_resp.get("error") or None error = patch_error or test_error if not test_error and next_obs.get("code"): observation = next_obs test_score = float(next_obs.get("test_score", 0.0)) final_test_score = test_score total_tests = int(next_obs.get("total_tests", 1)) test_status = "PASS" if test_score >= 1.0 else "FAIL" raw_logs = (next_obs.get("logs") or "").strip() log_preview = raw_logs.replace("\n", " | ") if len(log_preview) > 300: log_preview = log_preview[:300] + "...(truncated)" print( f"[DEBUG] tests status={test_status} score={test_score:.2f}/{total_tests}", flush=True, file=sys.stderr, ) if log_preview: print(f"[DEBUG] test_logs={log_preview}", flush=True, file=sys.stderr) rewards.append(reward) steps_taken = step log_step(step=step, action=action, reward=reward, done=done, error=error) print( f"[DEBUG] step={step} elapsed_s={time.perf_counter() - step_start:.3f}", flush=True, file=sys.stderr, ) # FIX: store actual patch in history so LLM sees what it already tried history.append( f"Step {step}: patch={'(none)' if not action else action[:120]} | " f"test={test_status} reward={reward:.2f}" ) if done: success = final_test_score >= SUCCESS_SCORE_THRESHOLD break else: success = False score = min(max(final_test_score, 0.0), 1.0) if not done: success = score >= SUCCESS_SCORE_THRESHOLD except Exception as exc: print(f"[DEBUG] Episode crashed: {exc}", flush=True, file=sys.stderr) success = False score = min(max(final_test_score, 0.0), 1.0) finally: log_end(success=success, steps=steps_taken, score=score, rewards=rewards) print( f"[DEBUG] total_elapsed_s={time.perf_counter() - episode_start:.3f}", flush=True, file=sys.stderr, ) finally: if _env_client is not None: try: _env_client.close() except Exception: pass if __name__ == "__main__": main()