import time import re from env.executor import CodeExecutor from agents.reward import calculate_reward from agents.critic import run_all_anti_hack_checks from fastapi import FastAPI import threading class CodeDebuggerEnv: def __init__(self, max_iterations=5): self.executor = CodeExecutor(timeout_seconds=10) self.max_iterations = max_iterations self.problem = None self.iteration = 0 self.history = [] self.prev_tests_passed = 0 def reset(self, problem: dict) -> dict: """Reset for new problem. Returns initial observation.""" self.problem = problem self.iteration = 0 self.history = [] self.prev_tests_passed = 0 return { "problem_id": problem["id"], "difficulty": problem["difficulty"], "description": problem["description"], "buggy_code": problem["buggy_code"], "error_type": problem.get("error_type", ""), "test_cases": problem["test_cases"], "iteration": 0 } def step(self, action: str) -> tuple: """ action = submitted fixed code string Process: 1. Increment iteration 2. Safety check on submitted code 3. If unsafe: skip tests, return penalty reward 4. Run all tests via executor 5. Run anti-hack checks 6. Calculate reward 7. Store in history 8. Return (observation, reward_dict, done, info) """ self.iteration += 1 # 2. Safety check safety = self.executor.check_code_safety(action) func_name = extract_function_name(self.problem["buggy_code"]) if not safety["safe"]: # 3. If unsafe: skip tests test_results = { "tests_passed": 0, "tests_total": len(self.problem["test_cases"]), "pass_rate": 0.0, "results": [], "execution_time_total": 0.0 } execution_time = 0.0 anti_hack = {"all_passed": True, "total_penalty": 0.0, "checks": {}} else: # 4. Run all tests via executor start_t = time.time() test_results = self.executor.run_all_tests( code=action, function_name=func_name, test_cases=self.problem["test_cases"] ) execution_time = time.time() - start_t # 5. Run anti-hack checks anti_hack = run_all_anti_hack_checks( code=action, original_code=self.problem["buggy_code"], test_cases=self.problem["test_cases"], test_results=test_results, execution_time=test_results.get("execution_time_total", execution_time) ) # 6. Calculate reward reward_dict = calculate_reward( code=action, test_results=test_results, original_buggy_code=self.problem["buggy_code"], iteration=self.iteration, execution_time=test_results.get("execution_time_total", execution_time), safety_check=safety, previous_tests_passed=self.prev_tests_passed, error_type=self.problem.get("error_type", "") ) # incorporate critic penalties into reward critic_penalty = anti_hack.get("total_penalty", 0.0) reward_dict["anti_hack_penalty"] += critic_penalty reward_dict["total"] += critic_penalty reward_dict["total"] = max(0.0, reward_dict["total"]) # clamp to >= 0 # 7. Store in history self.history.append({ "iteration": self.iteration, "code": action, "reward": reward_dict, "test_results": test_results }) self.prev_tests_passed = test_results.get("tests_passed", 0) # done = True if solved or out of iterations done = False if test_results.get("pass_rate", 0.0) == 1.0: done = True elif self.iteration >= self.max_iterations: done = True # 8. Return observation = { "problem_id": self.problem["id"], "iteration": self.iteration, "code": action, "test_results": test_results, "safety": safety, "anti_hack": anti_hack } info = { "solved": test_results.get("pass_rate", 0.0) == 1.0, "safety_violations": safety.get("violations", []), "anti_hack_failed": not anti_hack.get("all_passed", True) } return observation, reward_dict, done, info def render(self) -> str: """ Returns formatted display string showing: - Problem title [difficulty] - Iteration N/max - Test results table - Reward breakdown - Issues list - Best reward so far """ if not self.problem: return "Environment not initialized." title = self.problem.get("title", "Unknown") diff = self.problem.get("difficulty", "unknown").upper() lines = [ f"=== {title} [{diff}] ===", f"Iteration: {self.iteration} / {self.max_iterations}" ] if not self.history: lines.append("No actions taken yet.") return "\n".join(lines) last = self.history[-1] tr = last["test_results"] rw = last["reward"] passed = tr.get("tests_passed", 0) total = tr.get("tests_total", 0) rate = tr.get("pass_rate", 0.0) lines.append(f"\nTest Results: {passed}/{total} passed ({rate*100:.1f}%)") lines.append("\nReward Breakdown:") for k, v in rw.items(): if k != "total": lines.append(f" {k}: {v:.2f}") lines.append(f" --> TOTAL: {rw['total']:.2f}") issues = [] if not last.get("safety", {}).get("safe", True): issues.append(f"Unsafe code detected: {last['safety'].get('violations')}") if not last.get("anti_hack", {}).get("all_passed", True): for name, chk in last["anti_hack"].get("checks", {}).items(): if not chk.get("passed", True): issues.append(f"Anti-hack failure ({name}): {chk.get('message')}") if issues: lines.append("\nIssues:") for issue in issues: lines.append(f" - {issue}") best_reward = max(h["reward"]["total"] for h in self.history) lines.append(f"\nBest reward so far: {best_reward:.2f}") return "\n".join(lines) def extract_function_name(code: str) -> str: """Use regex to find 'def function_name(' in code""" match = re.search(r"def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(", code) if match: return match.group(1) return "unknown" # FastAPI app app = FastAPI(title="CodeDebugger RL Environment") env_instance = CodeDebuggerEnv() @app.post("/reset") def api_reset(problem: dict): return env_instance.reset(problem) @app.post("/step") def api_step(payload: dict): action = payload.get("code", "") obs, reward, done, info = env_instance.step(action) return {"observation": obs, "reward": reward, "done": done, "info": info} @app.get("/render") def api_render(): return {"display": env_instance.render()} @app.get("/health") def health(): return {"status": "ok", "version": "1.0.0"}