Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| inference.py - Code Debug Environment Baseline Agent | |
| Required env vars: API_BASE_URL, MODEL_NAME, HF_TOKEN | |
| Usage: | |
| python inference.py | |
| python inference.py --url https://Souravdanyal-code-debug-env.hf.space | |
| python inference.py --difficulty easy | |
| STDOUT FORMAT (required by evaluator): | |
| [START] task=<id> env=<benchmark> model=<model> | |
| [STEP] step=<n> action=<str> reward=<0.00> done=<true|false> error=<msg|null> | |
| [END] success=<true|false> steps=<n> rewards=<r1,r2,...> | |
| """ | |
| import os, sys, json, time, argparse, requests | |
| from openai import OpenAI | |
| from typing import List, Optional | |
| # ── Config ──────────────────────────────────────────────────────────────────── | |
| API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1") | |
| MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini") | |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") | |
| ENV_URL = os.environ.get("ENV_URL", "http://localhost:7860") | |
| BENCHMARK = "code-debug-env" | |
| MAX_STEPS = 5 | |
| client = OpenAI(api_key=HF_TOKEN or "dummy", base_url=API_BASE_URL) | |
| # ── Logging ─────────────────────────────────────────────────────────────────── | |
| def log_start(task_id, env, model): | |
| print(f"[START] task={task_id} env={env} model={model}", flush=True) | |
| def log_step(step, action, reward, done, error): | |
| print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error or 'null'}", flush=True) | |
| def log_end(success, steps, rewards): | |
| print(f"[END] success={str(success).lower()} steps={steps} rewards={','.join(f'{r:.2f}' for r in rewards)}", flush=True) | |
| # ── Env client ──────────────────────────────────────────────────────────────── | |
| def env_reset(url, difficulty): | |
| r = requests.post(f"{url}/reset", json={"difficulty": difficulty}, timeout=30) | |
| r.raise_for_status() | |
| return r.json() | |
| def env_step(url, fixed_code, explanation=None): | |
| payload = {"fixed_code": fixed_code} | |
| if explanation: | |
| payload["explanation"] = explanation | |
| r = requests.post(f"{url}/step", json=payload, timeout=30) | |
| r.raise_for_status() | |
| return r.json() | |
| # ── LLM ────────────────────────────────────────────────────────────────────── | |
| SYSTEM_PROMPT = """You are an expert Python debugging agent. Fix bugs in Python functions. | |
| RESPONSE FORMAT — strictly JSON only, no markdown: | |
| { | |
| "fixed_code": "<complete corrected Python function including imports>", | |
| "explanation": "<for hard tasks: explain the bug, root cause, and fix>" | |
| } | |
| RULES: | |
| - Return COMPLETE function with all imports (e.g. from collections import deque) | |
| - fixed_code must be valid Python | |
| - For hard tasks explanation MUST mention the algorithmic concept listed in instructions | |
| COMMON BUG PATTERNS: | |
| - List rotation RIGHT by k: correct is lst[-k:] + lst[:-k] NOT lst[k:] + lst[:k] | |
| - List rotation LEFT by k: correct is lst[k:] + lst[:k] | |
| - Graph/BFS missing visited set → infinite loop → add visited=set() | |
| - 0/1 Knapsack: must iterate BACKWARD: range(capacity, weight-1, -1) not forward | |
| - Binary search wrong boundary: return high not low, or high=n//2 | |
| - Off-by-one: lst[2] should be lst[1] for second element | |
| - Wrong operator: complement = target - n NOT target + n | |
| FOR HARD TASKS — explanation MUST include words from the instructions hint. | |
| Example: if instructions say "mention: iteration order" then write about iteration order. | |
| Example: if instructions say "mention: visited" then write about visited set. | |
| """ | |
| def call_llm(buggy_code, instructions, difficulty, feedback=None, attempt=1, prev_code=None): | |
| content = f"Difficulty: {difficulty}\nInstructions: {instructions}\n\nBuggy code:\n```python\n{buggy_code}\n```\n" | |
| if feedback and attempt > 1: | |
| content += f"\nPREVIOUS FIX FAILED. Feedback:\n{feedback}\n\nYour previous code:\n```python\n{prev_code or ''}\n```\n" | |
| content += "IMPORTANT: Your fix did not work. Look at the Expected vs Got values carefully.\n" | |
| content += "- If Got is a LEFT rotation but Expected is RIGHT: use lst[-k:] + lst[:-k]\n" | |
| content += "- If you see TimeoutError: add visited=set() for graph traversal\n" | |
| content += "- Try a COMPLETELY DIFFERENT approach.\n" | |
| if difficulty == "hard": | |
| # Extract keyword hints from instructions (e.g. "mention: visited, queue") | |
| import re | |
| hint_match = re.search(r'[Mm]ention[:\s]+([^.]+)', instructions) | |
| if hint_match: | |
| hints = hint_match.group(1).strip() | |
| content += f"\nFor your explanation, you MUST mention these concepts: {hints}\n" | |
| content += "Include a detailed explanation field — it counts for 30% of your reward.\n" | |
| try: | |
| resp = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": content}], | |
| max_tokens=1500, | |
| temperature=0.1 if attempt == 1 else 0.4, | |
| ) | |
| raw = resp.choices[0].message.content.strip() | |
| # Remove markdown fences | |
| if "```json" in raw: | |
| raw = raw.split("```json")[1].split("```")[0].strip() | |
| elif "```" in raw: | |
| raw = raw.split("```")[1].split("```")[0].strip() | |
| if raw.startswith("json"): | |
| raw = raw[4:].strip() | |
| # Find JSON object boundaries | |
| start = raw.find("{") | |
| end = raw.rfind("}") + 1 | |
| if start >= 0 and end > start: | |
| raw = raw[start:end] | |
| # Try direct parse first | |
| try: | |
| parsed = json.loads(raw) | |
| except json.JSONDecodeError: | |
| # Fix control characters by replacing literal newlines inside strings | |
| import re | |
| # Replace actual newlines within JSON string values with \n escape | |
| raw = re.sub(r'(?<!\\)\n', r'\\n', raw) | |
| raw = re.sub(r'(?<!\\)\t', r'\\t', raw) | |
| raw = re.sub(r'(?<!\\)\r', r'\\r', raw) | |
| try: | |
| parsed = json.loads(raw) | |
| except json.JSONDecodeError: | |
| # Last resort: extract fixed_code manually using regex | |
| code_match = re.search(r'"fixed_code"\s*:\s*"(.*?)"(?=\s*[,}])', raw, re.DOTALL) | |
| exp_match = re.search(r'"explanation"\s*:\s*"(.*?)"(?=\s*[,}])', raw, re.DOTALL) | |
| if code_match: | |
| code = code_match.group(1).encode().decode('unicode_escape') if '\\n' in code_match.group(1) else code_match.group(1) | |
| return {"fixed_code": code, "explanation": exp_match.group(1) if exp_match else None} | |
| raise | |
| return {"fixed_code": parsed.get("fixed_code", ""), "explanation": parsed.get("explanation")} | |
| except Exception as e: | |
| print(f"# LLM error: {e}", file=sys.stderr) | |
| return {"fixed_code": buggy_code, "explanation": None} | |
| # ── Episode ─────────────────────────────────────────────────────────────────── | |
| def run_episode(env_url, difficulty): | |
| data = env_reset(env_url, difficulty) | |
| obs = data["observation"] | |
| task_id = obs["task_id"] | |
| buggy_code = obs["buggy_code"] | |
| instructions = obs["instructions"] | |
| log_start(task_id, BENCHMARK, MODEL_NAME) | |
| rewards, steps_taken, success = [], 0, False | |
| last_feedback, last_code = None, None | |
| for attempt in range(1, MAX_STEPS + 1): | |
| steps_taken = attempt | |
| action = call_llm(buggy_code, instructions, difficulty, last_feedback, attempt, last_code) | |
| code = action["fixed_code"] | |
| last_code = code | |
| if not code or not code.strip(): | |
| log_step(attempt, "empty_submission", 0.0, False, "empty_code") | |
| rewards.append(0.0) | |
| continue | |
| try: | |
| result = env_step(env_url, code, action.get("explanation")) | |
| except Exception as e: | |
| log_step(attempt, "step_failed", 0.0, False, str(e)[:60]) | |
| rewards.append(0.0) | |
| continue | |
| reward = result.get("reward", 0.0) | |
| done = result.get("done", False) | |
| last_feedback = result.get("observation", {}).get("feedback", "") | |
| log_step(attempt, f"fix_{difficulty}_attempt{attempt}", reward, done, None) | |
| rewards.append(reward) | |
| if reward >= 1.0: | |
| success = True | |
| if done: | |
| break | |
| log_end(success, steps_taken, rewards) | |
| return success, steps_taken, rewards | |
| # ── Main ────────────────────────────────────────────────────────────────────── | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--url", default=ENV_URL) | |
| parser.add_argument("--difficulty", default=None, choices=["easy","medium","hard","all"]) | |
| args = parser.parse_args() | |
| url = args.url.rstrip("/") | |
| try: | |
| requests.get(f"{url}/health", timeout=10).raise_for_status() | |
| print(f"# Environment healthy at {url}", flush=True) | |
| except Exception as e: | |
| print(f"# Health check failed: {e}", file=sys.stderr) | |
| sys.exit(1) | |
| diffs = ["easy","medium","hard"] if args.difficulty in (None,"all") else [args.difficulty] | |
| all_rewards, successes = [], [] | |
| for d in diffs: | |
| ok, _, rewards = run_episode(url, d) | |
| all_rewards.extend(rewards) | |
| successes.append(ok) | |
| time.sleep(0.5) | |
| avg = round(sum(all_rewards)/len(all_rewards), 3) if all_rewards else 0.0 | |
| print(f"# SUMMARY: {sum(successes)}/{len(diffs)} tasks solved | avg_reward={avg}", flush=True) | |
| if __name__ == "__main__": | |
| main() |