Shumaque/code-review-agent
0
1"""2Inference Script for Code Review Agent3===================================4MANDATORY5- Before submitting, ensure the following variables are defined in your environment configuration:6 API_BASE_URL The API endpoint for the LLM.7 MODEL_NAME The model identifier to use for inference.8 HF_TOKEN Your Hugging Face / API key.9 10- Defaults are set only for API_BASE_URL and MODEL_NAME11 (and should reflect your active inference setup):12 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")13 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")14 15- The inference script must be named `inference.py` and placed in the root directory of the project16- Participants must use OpenAI Client for all LLM calls using above variables17 18STDOUT FORMAT19- The script must emit exactly three line types to stdout, in this order:20 21 [START] task=<task_name> env=<benchmark> model=<model_name>22 [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>23 [END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>24 25 Rules:26 - One [START] line at episode begin.27 - One [STEP] line per step, immediately after env.step() returns.28 - One [END] line after env.close(), always emitted (even on exception).29 - reward and rewards are formatted to 2 decimal places.30 - done and success are lowercase booleans: true or false.31 - error is the raw last_action_error string, or null if none.32 - All fields on a single line with no newlines within a line.33 - Each tasks should return score in [0, 1]34 35 Example:36 [START] task=detect_logic_bug env=code-review-agent model=Qwen/Qwen2.5-72B-Instruct37 [STEP] step=1 action=submit_findings(count=2) reward=0.70 done=true error=null38 [END] success=true steps=1 score=0.700 rewards=0.7039"""40 41import os42import json43from typing import List, Optional44from dotenv import load_dotenv45 46# Load environment variables from .env file47load_dotenv()48 49from openai import OpenAI50 51API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")52API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")53MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")54BENCHMARK = "code-review-agent"55MAX_STEPS = 556SUCCESS_SCORE_THRESHOLD = 0.5 # F1 score threshold for success57 58if API_KEY is None:59 raise ValueError("HF_TOKEN or API_KEY environment variable is required")60 61client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)62 63SYSTEM_PROMPT = """You are a senior software engineer performing a code review.64You will receive a unified diff where EVERY line is prefixed with its line number like "LINE 1: ...".65Your job is to identify ALL bugs, security vulnerabilities, and significant issues.66 67CRITICAL INSTRUCTIONS:681. Line Numbers: Use EXACTLY the "LINE X" number shown on the left of each buggy line.69 Example: If you see "LINE 017: + buggy_code_here" then report line_number: 1770 Do NOT count lines yourself - just read the LINE number directly from the left side.712. Keywords: Use EXACT terms in descriptions - "sql injection", "hardcoded", "off-by-one", "path traversal", "command injection", "pickle", "deserialization", "ceiling division"723. Precision: Only report REAL issues - don't hallucinate or report non-issues734. Categories: Use "bug", "security", "style", or "performance"745. Severity: Use "critical", "high", "medium", "low", or "info"75 76COMMON PATTERNS TO IDENTIFY:77- SQL Injection: f-string or string formatting in SQL queries78- Hardcoded Secrets: API keys, passwords, JWT secrets in code79- Command Injection: subprocess.run with shell=True and user input80- Path Traversal: Using unsanitized filenames in file paths81- Unsafe Deserialization: pickle.loads() on untrusted data82- Off-by-one Errors: Wrong index math like (page * size) instead of ((page-1) * size)83- Missing Edge Cases: Integer division instead of ceiling division84 85Respond ONLY with a valid JSON object matching this exact schema:86{87 "findings": [88 {89 "line_number": <integer, 1-indexed line in diff where issue appears>,90 "severity": "<critical|high|medium|low|info>",91 "category": "<bug|security|style|performance>",92 "description": "<clear explanation using exact keywords>"93 }94 ]95}96 97EXAMPLES:98Example 1 - SQL Injection:99Diff shows: sql = f"SELECT * FROM users WHERE name LIKE '%{query}%'"100Finding: {"line_number": 5, "severity": "critical", "category": "security", "description": "SQL injection vulnerability - unsanitized user input in f-string query allows SQL injection"}101 102Example 2 - Off-by-one Error:103Diff shows: start = page * page_size (should be (page - 1) * page_size)104Finding: {"line_number": 3, "severity": "high", "category": "bug", "description": "Off-by-one error in pagination - should use (page - 1) * page_size instead of page * page_size"}105 106Example 3 - Hardcoded Secret:107Diff shows: SECRET_KEY = "hardcoded_jwt_secret_abc123"108Finding: {"line_number": 10, "severity": "critical", "category": "security", "description": "Hardcoded JWT secret key exposed in source code - should use environment variable"}109 110Example 4 - Missing Ceiling Division:111Diff shows: return total_items // page_size (should be (total_items + page_size - 1) // page_size)112Finding: {"line_number": 8, "severity": "medium", "category": "bug", "description": "Missing ceiling division in total_pages calculation - use (total_items + page_size - 1) // page_size to include last partial page"}113 114Rules:115- Count line numbers from the FIRST line of the diff output shown116- Use the exact keywords listed above in descriptions117- Do NOT include findings for things that are fine118- Return only the JSON, no preamble, no markdown fences"""119 120 121def log_start(task: str, env: str, model: str) -> None:122 print(f"[START] task={task} env={env} model={model}", flush=True)123 124 125def log_step(126 step: int, action: str, reward: float, done: bool, error: Optional[str]127) -> None:128 error_val = error if error else "null"129 done_val = str(done).lower()130 print(131 f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",132 flush=True,133 )134 135 136def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:137 rewards_str = ",".join(f"{r:.2f}" for r in rewards)138 print(139 f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",140 flush=True,141 )142 143 144def review_diff(diff: str, file_path: str, pr_context: str, task_name: str) -> dict:145 # Number every line sequentially so AI can read exact line numbers146 numbered_diff = "\n".join(147 f"LINE {i + 1:3}: {line}" for i, line in enumerate(diff.split("\n"))148 )149 150 user_content = f"""PR Context: {pr_context}151File: {file_path}152Task: {task_name}153 154Diff (use the LINE numbers on the left as your line_number values):155{numbered_diff}"""156 157 try:158 response = client.chat.completions.create(159 model=MODEL_NAME,160 messages=[161 {"role": "system", "content": SYSTEM_PROMPT},162 {"role": "user", "content": user_content},163 ],164 temperature=0.0,165 )166 raw = response.choices[0].message.content.strip()167 # Strip markdown fences if model added them despite instructions168 if raw.startswith("```"):169 raw = raw.split("```")[1]170 if raw.startswith("json"):171 raw = raw[4:]172 return json.loads(raw.strip())173 except Exception as exc:174 print(f"[DEBUG] Model request failed: {exc}", flush=True)175 return {"findings": []}176 177 178def run_task(task_name: str, diff: str, file_path: str, pr_context: str) -> dict:179 from server.code_review_environment import CodeReviewEnv180 from server.models import CodeReviewAction, Finding181 from server.models import Severity, Category182 183 env = CodeReviewEnv(task_name=task_name)184 obs = env.reset()185 186 rewards: List[float] = []187 steps_taken = 0188 score = 0.0189 success = False190 last_error = None191 192 log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)193 194 try:195 result = review_diff(obs.diff, obs.file_path, obs.pr_context, task_name)196 findings = []197 for f in result.get("findings", []):198 try:199 findings.append(200 Finding(201 line_number=int(f["line_number"]),202 severity=Severity(f["severity"]),203 category=Category(f["category"]),204 description=str(f["description"]),205 )206 )207 except Exception:208 continue209 210 action = CodeReviewAction(findings=findings)211 obs = env.step(action)212 reward = obs.reward or 0.0213 done = obs.done214 error = None215 216 rewards.append(reward)217 steps_taken = 1218 action_str = f"submit_findings(count={len(findings)})"219 220 log_step(221 step=steps_taken, action=action_str, reward=reward, done=done, error=error222 )223 224 # Clamp score to (0, 1) range - hackathon requires strictly between 0 and 1225 score = min(0.999, max(0.001, reward))226 success = score >= SUCCESS_SCORE_THRESHOLD227 228 except Exception as exc:229 last_error = str(exc).replace("\n", " ")230 log_step(step=1, action="error", reward=0.0, done=True, error=last_error)231 232 finally:233 try:234 env.close()235 except Exception as e:236 print(f"[DEBUG] env.close() error: {e}", flush=True)237 log_end(success=success, steps=steps_taken, score=score, rewards=rewards)238 239 return {240 "task": task_name,241 "score": score,242 "success": success,243 "steps": steps_taken,244 }245 246 247if __name__ == "__main__":248 from tasks.seeds import TASK_REGISTRY249 250 all_results = []251 for task_name, seed in TASK_REGISTRY.items():252 result = run_task(task_name, seed.diff, seed.file_path, seed.pr_context)253 all_results.append(result)254 255 print("\n--- BASELINE SCORES ---", flush=True)256 total = 0.0257 for r in all_results:258 print(f" {r['task']}: {r['score']:.4f} | success={r['success']}", flush=True)259 total += r["score"]260 print(f" AVERAGE: {total / len(all_results):.4f}", flush=True)261