agents/reward.py

"""
agents/reward.py - Reward function for CodeDebugger.
"""

import ast
import difflib

def calculate_reward(
    code: str,
    test_results: dict,
    original_buggy_code: str,
    iteration: int,
    execution_time: float,
    safety_check: dict,
    previous_tests_passed: int = 0,
    error_type: str = ""
) -> dict:
    """
    Calculate reward components and total score for a code fix attempt.
    """
    tests_passed = test_results.get("tests_passed", 0)
    tests_total = test_results.get("tests_total", 0)
    
    # ---------------------------------------------------------
    # 1. test_score (0-50)
    # ---------------------------------------------------------
    test_score = 0.0
    if tests_total > 0:
        test_score = (tests_passed / tests_total) * 40.0
        if tests_passed == tests_total:
            test_score += 10.0
    test_score = min(50.0, float(test_score))
    
    # ---------------------------------------------------------
    # 2. fix_quality (0-20, or negative if identical)
    # ---------------------------------------------------------
    fix_quality = 0.0
    added_lines = []
    
    if code.strip() == original_buggy_code.strip():
        fix_quality = -10.0
    else:
        orig_lines = original_buggy_code.splitlines()
        new_lines = code.splitlines()
        diff = list(difflib.unified_diff(orig_lines, new_lines, lineterm=""))
        
        changed_lines = 0
        for line in diff:
            if line.startswith(("+", "-")) and not line.startswith(("+++", "---")):
                changed_lines += 1
                if line.startswith("+"):
                    added_lines.append(line[1:])
                    
        if 1 <= changed_lines <= 3:
            fix_quality = 20.0
        elif 4 <= changed_lines <= 8:
            fix_quality = 10.0
        elif 9 <= changed_lines <= 15:
            fix_quality = 5.0
        else:
            fix_quality = 0.0
            
    # ---------------------------------------------------------
    # 3. format_score (0-15)
    # ---------------------------------------------------------
    format_score = 0.0
    
    try:
        orig_tree = ast.parse(original_buggy_code)
        orig_names = [node.name for node in orig_tree.body if isinstance(node, ast.FunctionDef)]
    except Exception:
        orig_names = []
        
    try:
        tree = ast.parse(code)
        format_score += 5.0  # AST parse succeeds
        
        new_names = [node.name for node in tree.body if isinstance(node, ast.FunctionDef)]
        if orig_names and new_names and orig_names[0] in new_names:
            format_score += 5.0  # Function name same as original
    except Exception:
        pass
        
    if '"""' in code or "'''" in code or "#" in code:
        format_score += 3.0
        
    if not ("\t" in code and "    " in code):
        format_score += 2.0
        
    format_score = min(15.0, format_score)
        
    # ---------------------------------------------------------
    # 4. speed_score (0-10)
    # ---------------------------------------------------------
    speed_score = 0.0
    if tests_total > 0 and tests_passed == tests_total:
        if iteration == 1:
            speed_score = 10.0
        elif iteration == 2:
            speed_score = 7.0
        elif iteration == 3:
            speed_score = 4.0
            
    # ---------------------------------------------------------
    # 5. safety_score (0-10)
    # ---------------------------------------------------------
    safety_score = 0.0
    if safety_check.get("safe") is True:
        safety_score += 10.0
    violations = len(safety_check.get("violations", []))
    safety_score -= (violations * 5.0)
    safety_score = max(0.0, min(10.0, safety_score))
    
    # ---------------------------------------------------------
    # 6. improvement_score (0-15)
    # ---------------------------------------------------------
    improvement_score = 0.0
    diff_passed = tests_passed - previous_tests_passed
    if diff_passed > 0:
        improvement_score = diff_passed * 5.0
    elif diff_passed < 0:
        improvement_score = -5.0
    else:
        improvement_score = 0.0
    improvement_score = min(15.0, improvement_score)
    
    # ---------------------------------------------------------
    # 7. anti_hack_penalty (negative or 0)
    # ---------------------------------------------------------
    anti_hack_penalty = 0.0
    stripped = code.strip()
    if stripped == "" or stripped == "pass":
        anti_hack_penalty -= 40.0
        
    if original_buggy_code:
        if len(code) < 0.2 * len(original_buggy_code):
            anti_hack_penalty -= 30.0
            
    # ---------------------------------------------------------
    # 8. process_bonus (0 or +3)
    # ---------------------------------------------------------
    process_bonus = 0.0
    if error_type and added_lines:
        added_text = "\n".join(added_lines).lower()
        
        # Heuristic keywords mapping for error_types
        keywords = {
            "IndexError": ["len", "[", "-1", "- 1", "range"],
            "TypeError": ["int", "str", "float", "strip", "split", "format"],
            "ValueError": ["if", "none", "try", "except"],
            "LogicError": ["==", "!=", ">", "<", ">=", "<=", "and", "or", "return", "+", "-", "is", "none", "extend", "append", "/", "//", "[", ":"],
            "KeyError": ["get", "in", "dict"],
            "InfiniteLoop": ["+", "-", "break", "return", "mid"]
        }
        
        kws = keywords.get(error_type, [error_type.lower()])
        if any(kw in added_text for kw in kws):
            process_bonus = 3.0

    # ---------------------------------------------------------
    # Calculate Total
    # ---------------------------------------------------------
    total = (
        test_score +
        fix_quality +
        format_score +
        speed_score +
        safety_score +
        improvement_score +
        anti_hack_penalty +
        process_bonus
    )
    
    total = max(0.0, float(total))
    
    return {
        "total": float(total),
        "test_score": float(test_score),
        "fix_quality": float(fix_quality),
        "format_score": float(format_score),
        "speed_score": float(speed_score),
        "safety_score": float(safety_score),
        "improvement_score": float(improvement_score),
        "anti_hack_penalty": float(anti_hack_penalty),
        "process_bonus": float(process_bonus)
    }