""" Inference Script — Code Debug Environment ========================================== Runs the baseline LLM agent against all three tasks (fix_syntax, fix_logic, fix_algorithm) and emits the mandatory [START]/[STEP]/[END] log lines. Environment variables: API_BASE_URL API endpoint for the LLM (default: HF router) MODEL_NAME Model identifier (default: Qwen2.5-72B-Instruct) HF_TOKEN API key IMAGE_NAME Docker image name (if using from_docker_image) """ import asyncio import os import textwrap from typing import List, Optional from openai import OpenAI from code_debug_env import CodeDebugEnv, DebugAction IMAGE_NAME = os.getenv("IMAGE_NAME") API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct") BENCHMARK = "code_debug_env" TASKS = ["fix_syntax", "fix_logic", "fix_algorithm"] MAX_STEPS = 10 TEMPERATURE = 0.2 MAX_TOKENS = 1024 # ── logging helpers (match mandatory format exactly) ───────────────────── 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: error_val = error if error else "null" done_val = str(done).lower() safe_action = action.replace("\n", "\\n") print( f"[STEP] step={step} action={safe_action} " f"reward={reward:.2f} done={done_val} error={error_val}", flush=True, ) def log_end( success: bool, steps: int, score: float, rewards: List[float] ) -> None: rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={str(success).lower()} steps={steps} " f"score={score:.2f} rewards={rewards_str}", flush=True, ) # ── LLM interaction ───────────────────────────────────────────────────── SYSTEM_PROMPT = textwrap.dedent("""\ You are an expert Python debugger. You are given a buggy Python function, a description of what it should do, and test results showing which tests pass or fail. Your job is to output ONLY the corrected Python code — nothing else. Do not include explanations, markdown fences, or commentary. Output the complete function(s) so the code can be executed as-is. """) def build_user_prompt(obs_dict: dict) -> str: errors_section = "" if obs_dict.get("stderr"): errors_section = "## Errors\n" + obs_dict["stderr"] return textwrap.dedent(f"""\ ## Task {obs_dict['task_description']} ## Buggy Code (original) ``` {obs_dict['buggy_code']} ``` ## Current Code (your last submission) ``` {obs_dict['current_code']} ``` ## Test Results ({obs_dict['tests_passed']}/{obs_dict['tests_total']} passed) {obs_dict['test_results']} {errors_section} Reply with ONLY the corrected Python code. """) def call_llm(client: OpenAI, obs_dict: dict) -> str: user_prompt = build_user_prompt(obs_dict) try: resp = client.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, ) text = (resp.choices[0].message.content or "").strip() # Strip markdown fences if the model wraps the code if text.startswith("```"): lines = text.split("\n") lines = [l for l in lines if not l.startswith("```")] text = "\n".join(lines) return text except Exception as exc: print(f"[DEBUG] LLM call failed: {exc}", flush=True) return "" # ── main loop ──────────────────────────────────────────────────────────── async def run_task(client: OpenAI, env: CodeDebugEnv, task: str) -> None: rewards: List[float] = [] steps_taken = 0 score = 0.0 success = False log_start(task=task, env=BENCHMARK, model=MODEL_NAME) try: result = await env.reset() obs = result.observation obs_dict = { "task_description": obs.task_description, "buggy_code": obs.buggy_code, "current_code": obs.current_code, "test_results": obs.test_results, "tests_passed": obs.tests_passed, "tests_total": obs.tests_total, "stderr": obs.stderr, } for step in range(1, MAX_STEPS + 1): fixed_code = call_llm(client, obs_dict) if not fixed_code: fixed_code = obs_dict["current_code"] result = await env.step(DebugAction(code=fixed_code)) obs = result.observation reward = result.reward or 0.0 done = result.done rewards.append(reward) steps_taken = step log_step( step=step, action=fixed_code[:120], reward=reward, done=done, error=obs.stderr if obs.stderr else None, ) obs_dict = { "task_description": obs.task_description, "buggy_code": obs.buggy_code, "current_code": obs.current_code, "test_results": obs.test_results, "tests_passed": obs.tests_passed, "tests_total": obs.tests_total, "stderr": obs.stderr, } if done: break if obs_dict["tests_total"] > 0: score = obs_dict["tests_passed"] / obs_dict["tests_total"] score = min(max(score, 0.0), 1.0) success = score >= 0.5 finally: try: await env.close() except Exception as e: print(f"[DEBUG] env.close() error: {e}", flush=True) log_end( success=success, steps=steps_taken, score=score, rewards=rewards ) async def main() -> None: openai_client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) for task in TASKS: # Each reset() cycles to the next task automatically on the server. # When using Docker, pass CODE_DEBUG_TASK if you want a specific task. if IMAGE_NAME: env = await CodeDebugEnv.from_docker_image(IMAGE_NAME) else: base = os.getenv("CODE_DEBUG_BASE_URL", "http://localhost:8000") env = CodeDebugEnv(base_url=base) await run_task(openai_client, env, task) if __name__ == "__main__": asyncio.run(main())