Blame for code-review-env/env/reward_engine.py

"""Reward engine for CodeReviewEnv.

Implements non-sparse, shaped rewards according to the master spec.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple

from env.graders.grader_easy import grade as grade_easy
from env.graders.grader_hard import grade as grade_hard
from env.graders.grader_medium import grade as grade_medium
from env.models import CodeReviewAction, GroundTruthBug, ReviewComment


@dataclass(frozen=True)
class RewardOutcome:
    """Outcome details from reward computation."""

    reward: float
    reason: str
    correctly_identified_bug_line: Optional[int]
    is_false_positive: bool
    is_red_herring_flag: bool
    is_duplicate: bool
    final_score: Optional[float]
    confidence_modifier: float = 0.0
    explanation_depth: Optional[str] = None


class RewardEngine:
    """Compute shaped rewards and final scores for a task episode."""

    def __init__(self, *, task_id: str, ground_truth: List[GroundTruthBug], max_steps: int) -> None:
        """Initialize the reward engine for a task."""

        self._task_id = task_id
        self._ground_truth = ground_truth
        self._max_steps = max_steps

    def _match_bug(self, line_number: int, filename: Optional[str] = None) -> Optional[GroundTruthBug]:
        """Find the closest ground-truth bug within +/-5 lines, preferring exact matches.

        Args:
            line_number: The line number to match against.
            filename: Optional filename for multi-file matching (Upgrade 4).
        """

        candidates: List[Tuple[int, GroundTruthBug]] = []
        for b in self._ground_truth:
            # If filename provided, only match bugs in that file
            if filename is not None and b.source_file is not None and b.source_file != filename:
                continue
            dist = abs(b.line_number - line_number)
            if dist <= 5:
                candidates.append((dist, b))
        if not candidates:
            # If filename was specified but no match, try all files (backward compatible)
            if filename is not None:
                return self._match_bug(line_number, filename=None)
            return None
        candidates.sort(key=lambda x: (x[0], x[1].line_number))
        return candidates[0][1]

    def _grade(self, comments: List[ReviewComment]) -> float:
        """Run the deterministic grader for this task."""

        if self._task_id == "easy":
            return grade_easy(comments, self._ground_truth)
        if self._task_id == "medium":
            return grade_medium(comments, self._ground_truth)
        if self._task_id == "hard":
            return grade_hard(comments, self._ground_truth)
        return 0.0

    def _evaluate_explanation_tiers(self, bug: GroundTruthBug, message: str) -> Tuple[bool, float, str]:
        """Upgrade 2: Evaluate explanation quality against tiered keywords.

        Args:
            bug: The matched ground-truth bug.
            message: The agent's comment message.

        Returns:
            Tuple of (should_register, reward_modifier, explanation_depth).
        """
        if bug.explanation_tiers is None:
            # Fall back to existing required_keywords logic
            return self._evaluate_required_keywords(bug, message)

        msg_lower = message.lower()
        tiers = bug.explanation_tiers

        tier3_keywords = tiers.get("tier3", [])
        tier2_keywords = tiers.get("tier2", [])
        tier1_keywords = tiers.get("tier1", [])

        has_tier3 = any(kw.lower() in msg_lower for kw in tier3_keywords) if tier3_keywords else False
        has_tier2 = any(kw.lower() in msg_lower for kw in tier2_keywords) if tier2_keywords else False
        has_tier1 = any(kw.lower() in msg_lower for kw in tier1_keywords) if tier1_keywords else False

        if has_tier3:
            # Deep explanation — full credit + bonus
            return True, 0.05, "deep"
        elif has_tier2:
            # Technical explanation — full credit, no bonus
            return True, 0.0, "technical"
        elif has_tier1:
            # Shallow mention — registered but with penalty
            return True, -0.05, "shallow"
        else:
            # No match at all — not registered, penalty
            return False, -0.10, "missing"

    def _evaluate_required_keywords(self, bug: GroundTruthBug, message: str) -> Tuple[bool, float, str]:
        """Original required_keywords logic for backward compatibility.

        Returns:
            Tuple of (should_register, reward_modifier, explanation_depth).
        """
        if not bug.required_keywords or not message:
            return True, 0.0, "technical"

        msg_lower = message.lower()
        has_keyword = any(kw.lower() in msg_lower for kw in bug.required_keywords)
        if has_keyword:
            return True, 0.0, "technical"
        else:
            return False, -0.10, "missing"

    def _compute_confidence_modifier(
        self,
        confidence: Optional[int],
        is_correct: bool,
        is_false_positive: bool,
        is_red_herring: bool,
    ) -> float:
        """Upgrade 1: Compute calibration modifier based on confidence level.

        Args:
            confidence: Agent's confidence value (0-100) or None.
            is_correct: Whether the bug was correctly matched.
            is_false_positive: Whether this was a false positive.
            is_red_herring: Whether this hit a red herring.

        Returns:
            Modifier to add to the base reward.
        """
        if confidence is None:
            return 0.0

        if 80 <= confidence <= 100:
            if is_correct and not is_false_positive and not is_red_herring:
                return 0.05  # High confidence + correct → bonus
            elif is_false_positive:
                return -0.10  # High confidence + false positive → extra penalty
            elif is_red_herring:
                return -0.10  # High confidence + red herring → extra penalty
        elif 50 <= confidence <= 79:
            return 0.0  # Medium confidence → no modifier
        elif 0 <= confidence <= 49:
            if is_correct and not is_false_positive and not is_red_herring:
                return -0.02  # Low confidence + correct → should know when it knows

        return 0.0

    def compute(
        self,
        action: CodeReviewAction,
        *,
        comments_so_far: List[ReviewComment],
        correctly_identified_bug_lines: set[int],
        step_number: int,
        steps_used_after_this: int,
    ) -> RewardOutcome:
        """Compute reward for an action.

        Args:
            action: Agent action.
            comments_so_far: Existing comments before applying this action.
            correctly_identified_bug_lines: Bug line numbers already credited.
            step_number: Current step number (1-indexed).
            steps_used_after_this: Step count used after applying this step (for efficiency bonus).

        Returns:
            RewardOutcome with reward and metadata.
        """

        # Handle inspect_file and inspect_lines actions
        if action.operation == "inspect_file":
            return RewardOutcome(
                reward=0.0,
                reason="Inspected file",
                correctly_identified_bug_line=None,
                is_false_positive=False,
                is_red_herring_flag=False,
                is_duplicate=False,
                final_score=None,
            )

        if action.operation == "inspect_lines":
            # Check if the inspected range contains a real bug line
            if action.start_line is not None and action.end_line is not None:
                for b in self._ground_truth:
                    if not b.is_red_herring and action.start_line <= b.line_number <= action.end_line:
                        if action.filename is None or b.source_file is None or action.filename == b.source_file:
                            return RewardOutcome(
                                reward=0.02,
                                reason="Inspected range contains a real bug",
                                correctly_identified_bug_line=None,
                                is_false_positive=False,
                                is_red_herring_flag=False,
                                is_duplicate=False,
                                final_score=None,
                            )
            return RewardOutcome(
                reward=0.0,
                reason="Inspected range contains no bugs",
                correctly_identified_bug_line=None,
                is_false_positive=False,
                is_red_herring_flag=False,
                is_duplicate=False,
                final_score=None,
            )

        if action.operation == "add_comment":
            if action.line_number is None:
                return RewardOutcome(
                    reward=-0.05,
                    reason="Invalid add_comment: missing line_number",
                    correctly_identified_bug_line=None,
                    is_false_positive=True,
                    is_red_herring_flag=False,
                    is_duplicate=False,
                    final_score=None,
                )

            matched = self._match_bug(action.line_number, filename=action.filename)
            if matched is None:
                # False positive
                conf_mod = self._compute_confidence_modifier(
                    action.confidence, is_correct=False,
                    is_false_positive=True, is_red_herring=False,
                )
                base_reward = -0.10 + conf_mod
                return RewardOutcome(
                    reward=base_reward,
                    reason="False positive: no ground-truth bug near commented line",
                    correctly_identified_bug_line=None,
                    is_false_positive=True,
                    is_red_herring_flag=False,
                    is_duplicate=False,
                    final_score=None,
                    confidence_modifier=conf_mod,
                )

            if matched.is_red_herring:
                conf_mod = self._compute_confidence_modifier(
                    action.confidence, is_correct=False,
                    is_false_positive=False, is_red_herring=True,
                )
                base_reward = -0.20 + conf_mod
                return RewardOutcome(
                    reward=base_reward,
                    reason="Flagged red herring",
                    correctly_identified_bug_line=None,
                    is_false_positive=False,
                    is_red_herring_flag=True,
                    is_duplicate=False,
                    final_score=None,
                    confidence_modifier=conf_mod,
                )

            if matched.line_number in correctly_identified_bug_lines:
                return RewardOutcome(
                    reward=-0.05,
                    reason="Duplicate comment on already-identified bug",
                    correctly_identified_bug_line=None,
                    is_false_positive=False,
                    is_red_herring_flag=False,
                    is_duplicate=True,
                    final_score=None,
                )

            base = 0.15
            sev_bonus = 0.05 if action.severity == matched.severity else 0.0
            cat_bonus = 0.05 if action.category == matched.category else 0.0

            # Use tiered evaluation if explanation_tiers is present
            should_register, semantic_modifier, explanation_depth = self._evaluate_explanation_tiers(
                matched, action.message or ""
            )

            reward = min(0.25, base + sev_bonus + cat_bonus) + semantic_modifier

            registered_line = matched.line_number if should_register else None

            # Apply confidence modifier AFTER all existing logic
            is_correct = registered_line is not None
            conf_mod = self._compute_confidence_modifier(
                action.confidence, is_correct=is_correct,
                is_false_positive=False, is_red_herring=False,
            )
            reward += conf_mod

            return RewardOutcome(
                reward=reward,
                reason="Correct proximity but missed semantic 'why'" if not should_register else "Correct bug proximity match",
                correctly_identified_bug_line=registered_line,
                is_false_positive=False,
                is_red_herring_flag=False,
                is_duplicate=False,
                final_score=None,
                confidence_modifier=conf_mod,
                explanation_depth=explanation_depth,
            )

        if action.operation == "approve":
            remaining_critical_or_major = [
                b
                for b in self._ground_truth
                if (not b.is_red_herring) and b.severity in {"critical", "major"} and b.line_number not in correctly_identified_bug_lines
            ]
            if remaining_critical_or_major:
                return RewardOutcome(
                    reward=-0.50,
                    reason="Approved while critical/major bugs remain unfound",
                    correctly_identified_bug_line=None,
                    is_false_positive=False,
                    is_red_herring_flag=False,
                    is_duplicate=False,
                    final_score=None,
                )
            return RewardOutcome(
                reward=0.10,
                reason="Approved with no critical/major bugs remaining",
                correctly_identified_bug_line=None,
                is_false_positive=False,
                is_red_herring_flag=False,
                is_duplicate=False,
                final_score=None,
            )

        if action.operation == "request_changes":
            if len(correctly_identified_bug_lines) > 0:
                return RewardOutcome(
                    reward=0.05,
                    reason="Requested changes with evidence",
                    correctly_identified_bug_line=None,
                    is_false_positive=False,
                    is_red_herring_flag=False,
                    is_duplicate=False,
                    final_score=None,
                )
            return RewardOutcome(
                reward=-0.05,
                reason="Requested changes without evidence",
                correctly_identified_bug_line=None,
                is_false_positive=False,
                is_red_herring_flag=False,
                is_duplicate=False,
                final_score=None,
            )

        if action.operation == "done":
            final_score = self._grade(comments_so_far)
            reward = float(final_score)
            if steps_used_after_this < int(0.6 * self._max_steps) and final_score > 0.8:
                reward += 0.10
            return RewardOutcome(
                reward=reward,
                reason="Final grading score",
                correctly_identified_bug_line=None,
                is_false_positive=False,
                is_red_herring_flag=False,
                is_duplicate=False,
                final_score=final_score,
            )

        return RewardOutcome(
            reward=-0.05,
            reason="Unknown operation",
            correctly_identified_bug_line=None,
            is_false_positive=True,
            is_red_herring_flag=False,
            is_duplicate=False,
            final_score=None,
        )