Merge pull request #25 from ArshVermaGit/main

feat: complete implementation of CodeReview Environment, Models, and Scenarios

app.py

@@ -119,8 +119,8 @@ def submit_to_leaderboard(submission: SubmitScore):
     entries.append(new_entry)
     entries.sort(key=lambda x: x["score"], reverse=True)
     rank = entries.index(new_entry) + 1   # capture rank before slicing
-    leaderboard[submission.task_id] = entries[:10]
-    return {"status": "submitted", "rank": rank if rank <= 10 else None}
+    leaderboard[submission.task_id] = entries[:5]
+    return {"status": "submitted", "rank": rank if rank <= 5 else None}
 
 
 @app.websocket("/ws/events")

codereview_env/env.py

@@ -1,178 +1,177 @@
 from datetime import datetime, timezone
+from typing import List, Optional, Set
 from codereview_env.models import (
     TaskId, Action, Observation, StepResult, ResetResult,
-    ActionType, ActionRecord, EpisodeResult, FileChanged
+    ActionType, ActionRecord, EpisodeResult, Severity, GroundTruthIssue
 )
 from codereview_env.scenarios import get_scenario
-from codereview_env.graders.grader_utils import find_best_match
 from codereview_env.graders.bug_grader import grade_bug_detection
 from codereview_env.graders.security_grader import grade_security_audit
 from codereview_env.graders.arch_grader import grade_architectural_review
 
-
 class CodeReviewEnv:
+    MAX_NOISE_BUDGET = 5
     TASK_MAX_STEPS = {
-        TaskId.BUG_DETECTION:        10,
-        TaskId.SECURITY_AUDIT:       15,
+        TaskId.BUG_DETECTION: 10,
+        TaskId.SECURITY_AUDIT: 15,
         TaskId.ARCHITECTURAL_REVIEW: 20,
     }
+    SEVERITY_WEIGHTS = {
+        Severity.CRITICAL: 1.0,
+        Severity.HIGH: 0.8,
+        Severity.MEDIUM: 0.5,
+        Severity.LOW: 0.2,
+        Severity.INFO: 0.0,
+    }
 
     def __init__(self):
-        self._state = None
+        self.task_id: Optional[TaskId] = None
+        self.seed: int = 42
+        self.scenario = None
+        self.step_count: int = 0
+        self.noise_budget: int = self.MAX_NOISE_BUDGET
+        self.history: List[ActionRecord] = []
+        self.matched_issue_ids: Set[str] = set()
+        self.done: bool = False
+        self.terminated_reason: str = ""
+        self.episode_id: str = ""
 
     def reset(self, task_id: TaskId, seed: int = 42) -> ResetResult:
-        scenario = get_scenario(task_id, seed)
-        self._state = {
-            "task_id":      task_id,
-            "seed":         seed,
-            "scenario":     scenario,
-            "step_count":   0,
-            "noise_budget": 5,
-            "max_steps":    self.TASK_MAX_STEPS[task_id],
-            "history":      [],
-            "running_score": 0.0,
-            "done":         False,
-            "issues_found": set(),   # set of matched ground-truth issue IDs
-            "false_positives": []    # list of FP action bodies
-        }
+        self.scenario = get_scenario(task_id, seed)
+        self.task_id = task_id
+        self.seed = seed
+        self.step_count = 0
+        self.noise_budget = self.MAX_NOISE_BUDGET
+        self.history = []
+        self.matched_issue_ids = set()
+        self.done = False
+        self.terminated_reason = ""
+        
+        obs = self._build_observation()
         return ResetResult(
-            observation=self._build_obs(),
             task_id=task_id,
             seed=seed,
-            scenario_hash=scenario.hash
+            scenario_hash=self.scenario.hash,
+            observation=obs
         )
 
     def step(self, action: Action) -> StepResult:
-        if self._state is None or self._state["done"]:
-            raise RuntimeError("Episode is done or not initialized. Call reset().")
-
-        s = self._state
-        s["step_count"] += 1
-
-        # Record action in history (reward will be updated after calculation)
+        if self.done:
+            raise ValueError("Episode is already finished")
+
+        self.step_count += 1
+        reward = 0.0
+        
+        # Determine terminal state and reward
+        if action.action_type in (ActionType.APPROVE, ActionType.REQUEST_CHANGES):
+            self.done = True
+            self.terminated_reason = "terminal_action"
+            reward = 0.0 # Grader handles final score
+        
+        elif action.action_type == ActionType.FLAG_ISSUE:
+            match = None
+            for issue in self.scenario.ground_truth_issues:
+                if self._is_match(action, issue):
+                    match = issue
+                    break
+            
+            if match:
+                if match.id not in self.matched_issue_ids:
+                    reward = self.SEVERITY_WEIGHTS.get(match.severity, 0.0)
+                    self.matched_issue_ids.add(match.id)
+                else:
+                    # Already matched: penalty
+                    reward = -0.05
+                    self.noise_budget -= 1
+                    if self.noise_budget <= 0:
+                        self.done = True
+                        self.terminated_reason = "noise_exhausted"
+            else:
+                # False positive: penalty
+                reward = -0.05
+                self.noise_budget -= 1
+                if self.noise_budget <= 0:
+                    self.done = True
+                    self.terminated_reason = "noise_exhausted"
+        
+        # Max steps check
+        max_steps = self.TASK_MAX_STEPS.get(self.task_id, 10)
+        if not self.done and self.step_count >= max_steps:
+            self.done = True
+            self.terminated_reason = "max_steps"
+
+        # Record action
         record = ActionRecord(
             action_type=action.action_type,
             body=action.body,
             filename=action.filename,
             line_number=action.line_number,
-            severity=action.severity,
             category=action.category,
+            severity=action.severity,
             verdict=action.verdict,
-            reward=0.0,
+            reward=float(reward),
             timestamp=datetime.now(timezone.utc).isoformat()
         )
-        s["history"].append(record)
-
-        # Apply action logic and compute incremental reward delta
-        prev_score    = s["running_score"]
-        reward_delta  = self._apply_action(action)
-        s["running_score"] = prev_score + reward_delta
-
-        # Update the history record with the actual reward
-        record.reward = round(reward_delta, 4)
-
-        # Check termination
-        s["done"] = (
-            action.action_type in (ActionType.APPROVE, ActionType.REQUEST_CHANGES)
-            or s["step_count"] >= s["max_steps"]
-            or s["noise_budget"] <= 0
-        )
+        self.history.append(record)
 
         return StepResult(
-            observation=self._build_obs(),
-            reward=round(reward_delta, 4),   # ← incremental delta, NOT cumulative
-            done=s["done"],
-            info={
-                "step":               s["step_count"],
-                "cumulative_score":   round(s["running_score"], 4),
-                "noise_budget":       s["noise_budget"],
-                "issues_found_count": len(s["issues_found"]),
-            }
+            observation=self._build_observation(),
+            reward=float(reward),
+            done=self.done,
+            info={"terminated_reason": self.terminated_reason}
         )
 
-    def _build_obs(self) -> Observation:
-        s  = self._state
-        sc = s["scenario"]
+    def _is_match(self, action: Action, issue: GroundTruthIssue) -> bool:
+        if action.filename != issue.filename:
+            return False
+        if action.line_number is None:
+            return False
+        if abs(action.line_number - issue.line_number) > 3:
+            return False
+        if action.category != issue.category:
+            return False
+        
+        body_lower = (action.body or "").lower()
+        return any(kw.lower() in body_lower for kw in issue.keywords)
+
+    def _build_observation(self) -> Observation:
+        max_steps = self.TASK_MAX_STEPS.get(self.task_id, 10)
+        diff = "\n".join(f.patch for f in self.scenario.files_changed)
+        
         return Observation(
-            task_id=s["task_id"],
-            scenario_hash=sc.hash,
-            pr_title=sc.pr_title,
-            pr_description=sc.pr_description,
-            diff="\n".join([f.patch for f in sc.files_changed]),
-            files_changed=sc.files_changed,
-            step_count=s["step_count"],
-            max_steps=s["max_steps"],
-            noise_budget=s["noise_budget"],
-            max_noise_budget=5,
-            issues_flagged=len(s["issues_found"]),
-            done=s["done"]
+            task_id=self.task_id,
+            scenario_hash=self.scenario.hash,
+            pr_title=self.scenario.pr_title,
+            pr_description=self.scenario.pr_description,
+            diff=diff,
+            files_changed=self.scenario.files_changed,
+            step_count=self.step_count,
+            max_steps=max_steps,
+            noise_budget=self.noise_budget,
+            max_noise_budget=self.MAX_NOISE_BUDGET,
+            issues_flagged=len(self.matched_issue_ids),
+            done=self.done
         )
 
-    def _apply_action(self, action: Action) -> float:
-        """
-        Compute the incremental reward delta for this single action.
-        """
-        s  = self._state
-        sc = s["scenario"]
-
-        if action.action_type == ActionType.FLAG_ISSUE:
-            matched_gt = find_best_match(action, sc.ground_truth_issues, s["issues_found"])
-            if matched_gt:
-                s["issues_found"].add(matched_gt.id)
-                # Recalculate full grader score and return delta
-                new_score    = self._grade(sc, s)
-                reward_delta = new_score - s["running_score"]
-                return max(0.0, reward_delta)   # finding a real issue is always non-negative
-            else:
-                # False positive: consume noise budget and penalize
-                s["noise_budget"]     -= 1
-                s["false_positives"].append(action.body)
-                return -0.05
-
-        elif action.action_type in (ActionType.APPROVE, ActionType.REQUEST_CHANGES):
-            # Terminal: compute full final score delta
-            new_score    = self._grade(sc, s)
-            reward_delta = new_score - s["running_score"]
-            return reward_delta
-
-        # comment / ask_question — no reward signal
-        return 0.0
-
-    def _grade(self, sc, s) -> float:
-        """Route to the right grader based on task_id."""
-        if s["task_id"] == TaskId.BUG_DETECTION:
-            return grade_bug_detection(sc, s["history"])
-        elif s["task_id"] == TaskId.SECURITY_AUDIT:
-            return grade_security_audit(sc, s["history"])
-        else:
-            return grade_architectural_review(sc, s["history"])
-
     def get_final_result(self) -> EpisodeResult:
-        s  = self._state
-        sc = s["scenario"]
-
-        all_gt_ids = {gt.id for gt in sc.ground_truth_issues}
-        missed_ids = list(all_gt_ids - s["issues_found"])
-        final_score = self._grade(sc, s)
-
-        terminated_reason = ""
-        if s["done"]:
-            if s["noise_budget"] <= 0:
-                terminated_reason = "noise_exhausted"
-            elif s["history"] and s["history"][-1].action_type in (ActionType.APPROVE, ActionType.REQUEST_CHANGES):
-                terminated_reason = "terminal_action"
-            else:
-                terminated_reason = "max_steps"
-
+        if self.task_id == TaskId.BUG_DETECTION:
+            final_score = grade_bug_detection(self.scenario, self.history)
+        elif self.task_id == TaskId.SECURITY_AUDIT:
+            final_score = grade_security_audit(self.scenario, self.history)
+        else:
+            final_score = grade_architectural_review(self.scenario, self.history)
+            
         return EpisodeResult(
-            task_id=s["task_id"],
-            scenario_hash=sc.hash,
-            seed=s["seed"],
+            episode_id=self.episode_id,
+            task_id=self.task_id,
+            scenario_hash=self.scenario.hash,
+            seed=self.seed,
             final_score=round(final_score, 4),
-            steps_taken=s["step_count"],
-            issues_found=len(s["issues_found"]),
-            issues_total=len(sc.ground_truth_issues),
-            noise_penalties=5 - s["noise_budget"],
-            history=s["history"],
-            terminated_reason=terminated_reason
+            steps_taken=self.step_count,
+            issues_found=len(self.matched_issue_ids),
+            issues_total=len(self.scenario.ground_truth_issues),
+            noise_penalties=self.MAX_NOISE_BUDGET - self.noise_budget,
+            history=self.history,
+            terminated_reason=self.terminated_reason
         )
+

codereview_env/graders/arch_grader.py

@@ -1,39 +1,61 @@
 from typing import List
 from codereview_env.models import Scenario, ActionRecord, Category, ActionType, Verdict
-from codereview_env.graders.grader_utils import find_best_match
 
 def grade_architectural_review(scenario: Scenario, history: List[ActionRecord]) -> float:
-    """Grade architectural review: 0.6 * issue_score + 0.2 * verdict_score + 0.2 * quality_score."""
-    critical_issues = scenario.ground_truth_issues
-    flagged_actions = [r for r in history if r.category == Category.ARCHITECTURE]
+    if not history:
+        return 0.0
+        
+    flag_actions = [a for a in history if a.action_type == ActionType.FLAG_ISSUE]
     
-    # 1. Issue Detection (0.6)
-    found_count = 0
-    already_matched = set()
-    quality_bonus = 0.0
+    total_issues = len(scenario.ground_truth_issues)
+    if total_issues == 0:
+        return 0.0
+        
+    # 1. Match issues and calculate issue_score_avg
+    matched_gt_indices = set()
+    for i, truth in enumerate(scenario.ground_truth_issues):
+        if truth.category != Category.ARCHITECTURE:
+            continue
+            
+        for action in flag_actions:
+            # Match criteria: filename, line +- 5, category ARCHITECTURE, keyword match
+            if (action.filename == truth.filename and
+                action.line_number is not None and
+                abs(action.line_number - truth.line_number) <= 5 and
+                action.category == Category.ARCHITECTURE):
+                
+                body_lower = (action.body or "").lower()
+                if any(kw.lower() in body_lower for kw in truth.keywords):
+                    matched_gt_indices.add(i)
+                    break
+                    
+    issue_score_avg = len(matched_gt_indices) / total_issues
     
-    for action in flagged_actions:
-        match = find_best_match(action, critical_issues, already_matched)
-        if match:
-            found_count += 1
-            already_matched.add(match.id)
-            # Quality: + bonus if body is descriptive (> 80 chars)
-            if len(action.body) > 80:
-                quality_bonus += 0.05
+    # 2. Verdict Grading
+    terminal_action = None
+    for action in history:
+        if action.action_type in (ActionType.APPROVE, ActionType.REQUEST_CHANGES):
+            terminal_action = action
+            break # Use the first terminal action
+            
+    verdict_scores = []
+    for truth in scenario.ground_truth_issues:
+        if truth.required_verdict:
+            score = 1.0 if (terminal_action and terminal_action.verdict == truth.required_verdict) else 0.0
+            verdict_scores.append(score)
+            
+    verdict_avg = sum(verdict_scores) / len(verdict_scores) if verdict_scores else 0.0
     
-    issue_score = (found_count / len(critical_issues)) if critical_issues else 1.0
-    
-    # 2. Verdict Correctness (0.2)
-    verdict_score = 0.0
-    final_action = history[-1] if history else None
-    if final_action and final_action.action_type in (ActionType.APPROVE, ActionType.REQUEST_CHANGES):
-        # Check against the first required verdict in GT
-        required_verdicts = [gt.required_verdict for gt in critical_issues if gt.required_verdict]
-        if required_verdicts and final_action.verdict == required_verdicts[0]:
-            verdict_score = 1.0
-        elif not required_verdicts and final_action.verdict == Verdict.LGTM:
-            verdict_score = 1.0
+    # 3. Quality Score
+    max_body_len = 0
+    for action in flag_actions:
+        max_body_len = max(max_body_len, len(action.body or ""))
+        
+    quality_score = 0.0
+    if max_body_len > 20:
+        quality_score = min(1.0, max_body_len / 200)
+        
+    # 4. Final Weighted Calculation
+    final_score = 0.6 * issue_score_avg + 0.2 * verdict_avg + 0.2 * quality_score
+    return float(max(0.0, min(1.0, final_score)))
 
-    # 3. Final weighted calculation
-    score = 0.6 * issue_score + 0.2 * verdict_score + min(0.2, quality_bonus)
-    return round(min(1.0, score), 4)

codereview_env/graders/bug_grader.py

@@ -1,25 +1,66 @@
 from typing import List
-from codereview_env.models import Scenario, ActionRecord, Category
-from codereview_env.graders.grader_utils import find_best_match
+from codereview_env.models import Scenario, ActionRecord, Category, Severity, ActionType
 
 def grade_bug_detection(scenario: Scenario, history: List[ActionRecord]) -> float:
-    """Grade bug detection: 0.4 * coverage + 0.6 * precision."""
-    flagged_actions = [r for r in history if r.category == Category.BUG]
-    if not scenario.ground_truth_issues:
-        return 1.0 if not flagged_actions else 0.0
+    if not history:
+        return 0.0
+    
+    flag_actions = [a for a in history if a.action_type == ActionType.FLAG_ISSUE]
+    if not flag_actions:
+        return 0.0
+    
+    total_issues = len(scenario.ground_truth_issues)
+    if total_issues == 0:
+        return 0.0
         
-    found_ids = set()
-    matches = 0
-    already_matched = set()
+    matched_gt_indices = set()
+    used_action_indices = set()
+    issue_scores = []
     
-    for action in flagged_actions:
-        match = find_best_match(action, scenario.ground_truth_issues, already_matched)
-        if match:
-            matches += 1
-            found_ids.add(match.id)
-            already_matched.add(match.id)
+    for i, truth in enumerate(scenario.ground_truth_issues):
+        if truth.category != Category.BUG:
+            continue
+            
+        # Try to find a matching action for this ground truth issue
+        best_match_idx = -1
+        for j, action in enumerate(flag_actions):
+            if j in used_action_indices:
+                continue
+                
+            # Match criteria: filename, line +- 3, category BUG, >= 1 keyword
+            if (action.filename == truth.filename and
+                action.line_number is not None and
+                abs(action.line_number - truth.line_number) <= 3 and
+                action.category == Category.BUG):
+                
+                body_lower = (action.body or "").lower()
+                if any(kw.lower() in body_lower for kw in truth.keywords):
+                    best_match_idx = j
+                    break
+        
+        if best_match_idx != -1:
+            action = flag_actions[best_match_idx]
+            used_action_indices.add(best_match_idx)
+            matched_gt_indices.add(i)
+            
+            # Calculate issue score
+            sev_diff = abs(Severity.ordinal(truth.severity) - Severity.ordinal(action.severity))
+            sev_score = max(0.0, 1.0 - sev_diff * 0.3)
             
-    recall = len(found_ids) / len(scenario.ground_truth_issues)
-    precision = matches / len(flagged_actions) if flagged_actions else 0.0
+            body_lower = (action.body or "").lower()
+            match_count = sum(1 for kw in truth.keywords if kw.lower() in body_lower)
+            kw_score = match_count / len(truth.keywords)
+            
+            issue_score = 0.5 * kw_score + 0.5 * sev_score
+            issue_scores.append(issue_score)
+            
+    coverage = len(matched_gt_indices) / total_issues
+    avg_issue_score = sum(issue_scores) / len(issue_scores) if issue_scores else 0.0
     
-    return round(0.7 * recall + 0.3 * precision, 4)
+    total_flags = len(flag_actions)
+    unmatched_flags = total_flags - len(used_action_indices)
+    precision_penalty = unmatched_flags / max(1, total_flags)
+    
+    final_score = 0.4 * coverage + 0.6 * avg_issue_score - 0.1 * precision_penalty
+    return float(max(0.0, min(1.0, final_score)))
+

codereview_env/graders/security_grader.py

@@ -1,37 +1,56 @@
 from typing import List
-from codereview_env.models import Scenario, ActionRecord, Category, Severity
-from codereview_env.graders.grader_utils import find_best_match, keyword_overlap
-
-def severity_to_int(sev: Severity) -> int:
-    mapping = {
-        Severity.LOW: 1,
-        Severity.MEDIUM: 2,
-        Severity.HIGH: 3,
-        Severity.CRITICAL: 4
-    }
-    return mapping.get(sev, 0)
+from codereview_env.models import Scenario, ActionRecord, Category, Severity, ActionType
 
 def grade_security_audit(scenario: Scenario, history: List[ActionRecord]) -> float:
-    """Grade security audit: 0.7 * correct_severity + 0.3 * keyword_accuracy."""
-    flagged_actions = [r for r in history if r.category == Category.SECURITY]
-    if not scenario.ground_truth_issues:
-        return 1.0 if not flagged_actions else 0.0
+    if not history:
+        return 0.0
+        
+    flag_actions = [a for a in history if a.action_type == ActionType.FLAG_ISSUE]
+    if not flag_actions:
+        return 0.0
         
-    total_score = 0.0
-    already_matched = set()
+    matched_issue_scores = []
+    used_action_indices = set()
     
-    for action in flagged_actions:
-        match = find_best_match(action, scenario.ground_truth_issues, already_matched)
-        if match:
-            # Correct Severity: 1.0 if match, else penalty per level
-            sev_diff = abs(severity_to_int(action.severity) - severity_to_int(match.severity))
-            sev_score = max(0.0, 1.0 - (sev_diff * 0.3))
+    for truth in scenario.ground_truth_issues:
+        if truth.category != Category.SECURITY:
+            continue
+            
+        best_match_idx = -1
+        for j, action in enumerate(flag_actions):
+            if j in used_action_indices:
+                continue
+                
+            # Match criteria: filename, line +- 3, category SECURITY, >= 1 keyword
+            if (action.filename == truth.filename and
+                action.line_number is not None and
+                abs(action.line_number - truth.line_number) <= 3 and
+                action.category == Category.SECURITY):
+                
+                body_lower = (action.body or "").lower()
+                if any(kw.lower() in body_lower for kw in truth.keywords):
+                    best_match_idx = j
+                    break
+        
+        if best_match_idx != -1:
+            action = flag_actions[best_match_idx]
+            used_action_indices.add(best_match_idx)
             
-            # Keyword Accuracy
-            kw_score = keyword_overlap(action.body, match.keywords)
+            # Calculate issue score
+            sev_diff = abs(Severity.ordinal(truth.severity) - Severity.ordinal(action.severity))
+            sev_score = max(0.0, 1.0 - sev_diff * 0.3)
             
-            total_score += 0.7 * sev_score + 0.3 * kw_score
-            already_matched.add(match.id)
+            body_lower = (action.body or "").lower()
+            match_count = sum(1 for kw in truth.keywords if kw.lower() in body_lower)
+            kw_threshold = max(4, len(truth.keywords))
+            kw_score = match_count / kw_threshold
             
-    # Normalize by number of GT issues
-    return round(min(1.0, total_score / len(scenario.ground_truth_issues)), 4) if scenario.ground_truth_issues else 1.0
+            issue_score = 0.7 * sev_score + 0.3 * kw_score
+            matched_issue_scores.append(issue_score)
+            
+    if not matched_issue_scores:
+        return 0.0
+        
+    final_score = sum(matched_issue_scores) / len(matched_issue_scores)
+    return float(round(max(0.0, min(1.0, final_score)), 4))
+

openenv.yaml

@@ -45,22 +45,10 @@ tasks:
 
 grading:
   type: "deterministic"
-  score_range: [0.0, 1.0]
-  partial_credit: true
-  metrics:
-    - name: "recall"
-      description: "Fraction of ground-truth issues correctly identified"
-      weight: 0.7
-    - name: "precision"
-      description: "Fraction of flagged issues that are true positives"
-      weight: 0.3
-    - name: "severity_accuracy"
-      description: "How close the flagged severity is to ground truth (security task)"
-      weight: 0.7
-    - name: "keyword_accuracy"
-      description: "Semantic overlap between agent description and issue keywords"
-      weight: 0.3
-    - name: "verdict_correctness"
-      description: "Whether final approve/request_changes decision is correct (arch task)"
-      weight: 0.2
+  issue_matching:
+    coverage_weight: 0.4
+    precision_weight: 0.6
+  quality_scoring:
+    severity_weight: 0.7
+    keyword_weight: 0.3