YashR05/pullrequest-arena
0
1"""2inference.py — Baseline Inference Script for PullRequest Arena3 4Runs an AI agent (via OpenAI-compatible API) through all tasks in the5environment and logs results in the strict format required by OpenEnv.6 7Environment Variables:8 API_BASE_URL : Base URL for the OpenAI-compatible API endpoint9 MODEL_NAME : Model identifier to use for inference10 HF_TOKEN : HuggingFace / API token for authentication11 12Logging Format:13 [START] task=<task> env=<env> model=<model>14 [STEP] step=1 action=... reward=... done=false error=null15 [END] success=true steps=3 score=1.00 rewards=0.0,0.5,1.016"""17 18import json19import os20import sys21import time22 23from openai import OpenAI24 25from models import ReviewAction26 27class UnifiedEnv:28 def __init__(self):29 self.env_url = os.environ.get("OPENENV_BASE_URL", "")30 self.is_remote = bool(self.env_url)31 32 if self.is_remote:33 from client import PullRequestEnv34 import time35 import sys36 37 self.client = PullRequestEnv(base_url=self.env_url).sync()38 39 # Wait for server to boot (up to 60s)40 connected = False41 last_err = None42 for i in range(15):43 try:44 self.client.__enter__()45 connected = True46 break47 except Exception as e:48 last_err = e49 time.sleep(4)50 51 if not connected:52 print(f"[FATAL] Could not connect to remote env {self.env_url}: {last_err}")53 sys.exit(1)54 else:55 from server.pullrequest_environment import PullRequestEnvironment56 self.client = PullRequestEnvironment()57 58 def reset(self, task_id):59 if self.is_remote:60 res = self.client.reset(task_id=task_id)61 return res.observation62 else:63 return self.client.reset(task_id=task_id)64 65 def step(self, action_dict):66 action_obj = ReviewAction(**action_dict)67 if self.is_remote:68 res = self.client.step(action_obj)69 raw_reward = res.reward if res.reward is not None else 0.570 return res.observation, max(0.01, min(0.99, float(raw_reward))), res.done or False71 else:72 obs = self.client.step(action_obj)73 raw_reward = obs.reward if obs.reward is not None else 0.574 return obs, max(0.01, min(0.99, float(raw_reward))), obs.done or False75 76 def close(self):77 if self.is_remote:78 try:79 self.client.__exit__(None, None, None)80 except Exception:81 pass82 83 84# ---------------------------------------------------------------------------85# Constants86# ---------------------------------------------------------------------------87 88ENV_NAME = "pullrequest-arena"89VALID_ACTION_TYPES = {"approve", "request_changes", "comment", "suggest_fix"}90 91# System prompt instructing the LLM to act as a code reviewer.92SYSTEM_PROMPT = """You are an expert code reviewer. You will be given a pull request with a code diff to review.93 94Your job is to analyze the code and decide on one of the following actions:95- "approve": The code is correct and ready to merge.96- "request_changes": The code has bugs, errors, or security issues that must be fixed.97- "comment": You want to leave a comment or question about the code.98- "suggest_fix": You want to suggest a specific code improvement or refactoring.99- "submit_patch": You want to exactly submit a code patch solving the issue.100 101Respond ONLY with valid JSON in this exact format, no other text:102{103 "type": "<action_type>",104 "comment": "<your detailed review comment explaining the issue and how to fix it>",105 "patch": "<the diff format fix if type is submit_patch, otherwise empty string>"106}107 108Be thorough in your comment. Explain what the problem is and how to fix it."""109 110 111def build_review_prompt(observation: dict) -> str:112 """113 Build the user prompt from a PR observation.114 115 Args:116 observation: Dict with pr_title, pr_description, files_changed,117 code_diff, language, tests_passed, repository_context.118 119 Returns:120 Formatted prompt string for the LLM.121 """122 files = ", ".join(observation.get("files_changed", []))123 tests_status = "passing" if observation.get("tests_passed") else "failing"124 125 return f"""## Pull Request Review126 127**Title:** {observation['pr_title']}128**Description:** {observation['pr_description']}129**Files Changed:** {files}130**Language:** {observation['language']}131**Tests:** {tests_status}132**Repository Context:** {observation.get('repository_context', 'N/A')}133 134### Code Diff135```{observation['language']}136{observation['code_diff']}137```138 139Review this code diff carefully. Identify any bugs, security issues, style problems, or improvements. Respond with your action and comment as JSON."""140 141 142def parse_llm_response(response_text: str) -> dict:143 """144 Parse the LLM response into a structured action dict.145 146 Attempts to extract JSON from the response. Falls back to147 a heuristic parser if JSON parsing fails.148 149 Args:150 response_text: Raw text response from the LLM.151 152 Returns:153 Action dict with "type" and "comment" keys.154 """155 text = response_text.strip()156 157 # Try direct JSON parse158 try:159 action = json.loads(text)160 if isinstance(action, dict) and "type" in action and "comment" in action:161 action["type"] = action["type"].strip().lower()162 if action["type"] in VALID_ACTION_TYPES:163 return action164 except json.JSONDecodeError:165 pass166 167 # Try extracting JSON from markdown code block168 if "```" in text:169 for block in text.split("```"):170 block = block.strip()171 # Remove optional language identifier (e.g., "json")172 if block.startswith("json"):173 block = block[4:].strip()174 try:175 action = json.loads(block)176 if isinstance(action, dict) and "type" in action:177 action["type"] = action["type"].strip().lower()178 if action["type"] in VALID_ACTION_TYPES:179 action.setdefault("comment", "")180 return action181 except (json.JSONDecodeError, ValueError):182 continue183 184 # Try finding JSON object in the text with braces185 start = text.find("{")186 end = text.rfind("}") + 1187 if start != -1 and end > start:188 try:189 action = json.loads(text[start:end])190 if isinstance(action, dict) and "type" in action:191 action["type"] = action["type"].strip().lower()192 if action["type"] in VALID_ACTION_TYPES:193 action.setdefault("comment", "")194 return action195 except json.JSONDecodeError:196 pass197 198 # Heuristic fallback: detect action type from text199 text_lower = text.lower()200 detected_type = "comment" # default201 for action_type in ["request_changes", "suggest_fix", "approve", "comment"]:202 if action_type in text_lower or action_type.replace("_", " ") in text_lower:203 detected_type = action_type204 break205 206 return {"type": detected_type, "comment": text}207 208 209def call_llm(client: OpenAI, model: str, observation: dict) -> dict:210 """211 Call the LLM to review a pull request.212 213 Args:214 client: OpenAI client instance.215 model: Model identifier.216 observation: PR observation dict.217 218 Returns:219 Parsed action dict with "type" and "comment".220 """221 prompt = build_review_prompt(observation)222 223 response = client.chat.completions.create(224 model=model,225 messages=[226 {"role": "system", "content": SYSTEM_PROMPT},227 {"role": "user", "content": prompt},228 ],229 temperature=0.0, # Deterministic for reproducibility230 max_tokens=512,231 )232 233 response_text = response.choices[0].message.content or ""234 return parse_llm_response(response_text)235 236 237# ---------------------------------------------------------------------------238# Logging helpers (strict format)239# ---------------------------------------------------------------------------240 241def log_start(task_id: str, env_name: str, model: str) -> None:242 """Log the start of a task evaluation."""243 print(f"[START] task={task_id} env={env_name} model={model}")244 245 246def log_step(247 step: int,248 action: dict,249 reward: float,250 done: bool,251 error: str | None = None,252) -> None:253 """Log a single step in strict format."""254 action_str = json.dumps(action, separators=(",", ":")).replace("\n", "").replace("\r", "")255 done_str = str(done).lower()256 error_str = error if error else "null"257 print(258 f"[STEP] step={step} action={action_str} "259 f"reward={reward:.2f} done={done_str} error={error_str}"260 )261 262 263def log_end(264 success: bool,265 steps: int,266 rewards: list[float],267) -> None:268 """Log the end of a task evaluation."""269 success_str = str(success).lower()270 # Clamp each reward to (0, 1) for safety271 safe_rewards = [max(0.01, min(0.99, r)) for r in rewards]272 score = safe_rewards[-1] if safe_rewards else 0.01273 rewards_str = ",".join(f"{r:.2f}" for r in safe_rewards)274 print(275 f"[END] success={success_str} steps={steps} "276 f"score={score:.2f} rewards={rewards_str}"277 )278 279 280# ---------------------------------------------------------------------------281# Main282# ---------------------------------------------------------------------------283 284def run_inference() -> None:285 """286 Run the baseline inference agent through all tasks.287 288 Reads configuration from environment variables:289 API_BASE_URL : OpenAI-compatible API base URL290 MODEL_NAME : Model identifier291 HF_TOKEN : Authentication token292 """293 # --- Read configuration ---294 api_base_url = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1")295 model_name = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct")296 hf_token = os.environ.get("HF_TOKEN") or os.environ.get("OPENAI_API_KEY", "dummy_token_for_validation")297 298 # --- Initialize OpenAI client ---299 client = OpenAI(300 base_url=api_base_url,301 api_key=hf_token302 )303 304 # --- Initialize environment ---305 env = UnifiedEnv()306 task_ids = [str(i) for i in range(1, 20)]307 total_tasks = len(task_ids)308 309 print(f"PullRequest Arena — Baseline Inference")310 print(f"Model: {model_name}")311 print(f"API: {api_base_url}")312 print(f"Tasks: {total_tasks}")313 print("=" * 60)314 315 # --- Run agent through all tasks ---316 all_rewards: list[float] = []317 all_success: list[bool] = []318 start_time = time.time()319 320 for task_id in task_ids:321 observation_obj = env.reset(task_id=task_id)322 # convert OpenEnv observation model to dict323 observation = observation_obj.model_dump() if hasattr(observation_obj, "model_dump") else dict(observation_obj)324 325 log_start(task_id=task_id, env_name=ENV_NAME, model=model_name)326 327 step_num = 0328 episode_rewards: list[float] = []329 success = False330 error_msg = None331 332 try:333 # Build prompt and call LLM334 action = call_llm(client, model_name, observation)335 step_num += 1336 337 # Submit action to environment338 obs_obj, reward, done = env.step(action)339 episode_rewards.append(reward)340 341 log_step(342 step=step_num,343 action=action,344 reward=reward,345 done=done,346 error=None,347 )348 349 success = True350 351 except Exception as e:352 step_num += 1353 error_msg = str(e)354 episode_rewards.append(0.01)355 log_step(356 step=step_num,357 action={"type": "error", "comment": error_msg},358 reward=0.01,359 done=True,360 error=error_msg,361 )362 success = False363 364 log_end(365 success=success,366 steps=step_num,367 rewards=episode_rewards,368 )369 370 all_rewards.extend(episode_rewards)371 all_success.append(success)372 373 print() # Blank line between tasks374 375 env.close()376 377 # --- Aggregate results ---378 elapsed = time.time() - start_time379 total_score = sum(all_rewards) / len(all_rewards) if all_rewards else 0.01380 tasks_passed = sum(1 for s in all_success if s)381 382 print("=" * 60)383 print("SUMMARY")384 print(f" Tasks: {tasks_passed}/{total_tasks} completed")385 print(f" Avg Score: {total_score:.2f}")386 print(f" Total Time: {elapsed:.1f}s")387 print(f" All Rewards: {','.join(f'{r:.2f}' for r in all_rewards)}")388 print("=" * 60)389 390 391def main():392 import sys393 try:394 run_inference()395 except Exception as e:396 print(f"[ERROR] inference failed: {e}")397 import traceback398 traceback.print_exc()399 sys.exit(1)400 401if __name__ == "__main__":402 main()403