File size: 7,469 Bytes
2ac8bdd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce6b9af
2ac8bdd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce6b9af
2ac8bdd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce6b9af
2ac8bdd
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
"""
Senior Engineer Software PR Review Tasks
"""

from app.models import (
    Action, Observation, Reward, PullRequest, AgentAction, PRType
)
from data.pr_data import PULL_REQUESTS
import uuid
import random

class PRTypeTask:
    TASK_ID = "pr_type_classification"

    def __init__(self):
        self.current_idx = 0
        self.prs = []
        self.results = []

    def reset(self) -> Observation:
        self.current_idx = 0
        self.results = []
        self.prs = random.sample(PULL_REQUESTS, min(10, len(PULL_REQUESTS)))
        return self._make_obs()

    def step(self, action: Action):
        if self.current_idx >= len(self.prs):
            obs = Observation(task_id=self.TASK_ID, step=self.current_idx, episode_done=True)
            return obs, Reward(total=0.0), True, {}
        current = self.prs[self.current_idx]
        score = 1.0 if action.pr_type == current["true_pr_type"] else 0.0
        self.results.append({"score": score})
        self.current_idx += 1
        done = self.current_idx >= len(self.prs)
        return self._make_obs(done), Reward(total=score), done, {}

    def state(self): return {"task_id": self.TASK_ID, "step": self.current_idx}
    
    def grader_score(self):
        total = sum(r["score"] for r in self.results) / len(self.results) if self.results else 0.0
        return {"task_id": self.TASK_ID, "final_score": max(0.001, min(0.999, total)), "passed": total >= 0.7}

    def _make_obs(self, done=False):
        if done or self.current_idx >= len(self.prs):
            return Observation(task_id=self.TASK_ID, step=self.current_idx, episode_done=True)
        c = self.prs[self.current_idx]
        pr_obj = PullRequest(pr_id=c["pr_id"], title=c["title"], description=c["description"], diff=c["diff"], author=c["author"])
        return Observation(task_id=self.TASK_ID, step=self.current_idx, current_pr=pr_obj, valid_actions=[AgentAction.CLASSIFY_PR])


class PRBugIdentifyTask:
    TASK_ID = "pr_bug_identification"

    def __init__(self):
        self.current_idx = 0
        self.prs = []
        self.results = []

    def reset(self) -> Observation:
        self.current_idx = 0
        self.results = []
        # Filter to only PRs with actual bugs or security issues for this task
        issue_prs = [pr for pr in PULL_REQUESTS if pr["true_pr_type"] in [PRType.SECURITY, PRType.FEATURE, PRType.REFACTOR]]
        self.prs = random.sample(issue_prs, min(5, len(issue_prs)))
        return self._make_obs()

    def step(self, action: Action):
        if self.current_idx >= len(self.prs):
            obs = Observation(task_id=self.TASK_ID, step=self.current_idx, episode_done=True)
            return obs, Reward(total=0.0), True, {}
        current = self.prs[self.current_idx]
        
        # Heuristic grading: did they identify the core issue?
        score = 0.0
        if action.bug_description and len(action.bug_description) > 10:
            import re as _re
            desc = action.bug_description.lower()
            true_bug = current["true_bug_description"].lower()

            # Check if this is a "no bug" case
            no_bug_signals = ["none", "no bug", "correct", "correctly", "lgtm", "looks good", "addresses"]
            true_is_none = true_bug.startswith("none") or "no bug" in true_bug or "correctly" in true_bug

            if true_is_none:
                # Score 1.0 if agent also says no bug, else 0.0
                agent_says_no_bug = any(s in desc for s in no_bug_signals)
                score = 1.0 if agent_says_no_bug else 0.0
            else:
                # Strip punctuation from keywords for matching
                keywords = [_re.sub(r"[^a-z0-9_]", "", w) for w in true_bug.split() if len(w) > 4]
                keywords = [k for k in keywords if k]  # remove empty after strip
                if not keywords:
                    score = 1.0
                else:
                    matches = sum(1 for k in keywords if k in desc)
                    score = matches / len(keywords)
                    if score > 0.25: score = 1.0  # Lenient threshold
                
        self.results.append({"score": score})
        self.current_idx += 1
        done = self.current_idx >= len(self.prs)
        return self._make_obs(done), Reward(total=score), done, {}

    def state(self): return {"task_id": self.TASK_ID, "step": self.current_idx}
    
    def grader_score(self):
        total = sum(r["score"] for r in self.results) / len(self.results) if self.results else 0.0
        return {"task_id": self.TASK_ID, "final_score": max(0.001, min(0.999, total)), "passed": total >= 0.6}

    def _make_obs(self, done=False):
        if done or self.current_idx >= len(self.prs):
            return Observation(task_id=self.TASK_ID, step=self.current_idx, episode_done=True)
        c = self.prs[self.current_idx]
        pr_obj = PullRequest(pr_id=c["pr_id"], title=c["title"], description=c["description"], diff=c["diff"], author=c["author"])
        return Observation(task_id=self.TASK_ID, step=self.current_idx, current_pr=pr_obj, valid_actions=[AgentAction.IDENTIFY_BUG])


class PRReviewTask:
    TASK_ID = "pr_review_comment"

    def __init__(self):
        self.current_idx = 0
        self.prs = []
        self.results = []

    def reset(self) -> Observation:
        self.current_idx = 0
        self.results = []
        self.prs = random.sample(PULL_REQUESTS, min(5, len(PULL_REQUESTS)))
        return self._make_obs()

    def step(self, action: Action):
        if self.current_idx >= len(self.prs):
            obs = Observation(task_id=self.TASK_ID, step=self.current_idx, episode_done=True)
            return obs, Reward(total=0.0), True, {}
        current = self.prs[self.current_idx]
        score = 0.0
        if action.review_comment and len(action.review_comment) > 20:
            score += 0.5
            if "looks good" in action.review_comment.lower() or "lgtm" in action.review_comment.lower():
                if current["true_pr_type"] == PRType.BUG_FIX and current["true_bug_description"] == "None":
                    score += 0.5  # Correct to approve
                else:
                    score -= 0.5  # Incorrect to just approve a buggy PR
            else:
                # Some critique was offered
                score += 0.5
                
        # Ensure score within 0-1
        score = max(0.0, min(1.0, score))
        
        self.results.append({"score": score})
        self.current_idx += 1
        done = self.current_idx >= len(self.prs)
        return self._make_obs(done), Reward(total=score), done, {}

    def state(self): return {"task_id": self.TASK_ID, "step": self.current_idx}
    
    def grader_score(self):
        total = sum(r["score"] for r in self.results) / len(self.results) if self.results else 0.0
        return {"task_id": self.TASK_ID, "final_score": max(0.001, min(0.999, total)), "passed": total >= 0.6}

    def _make_obs(self, done=False):
        if done or self.current_idx >= len(self.prs):
            return Observation(task_id=self.TASK_ID, step=self.current_idx, episode_done=True)
        c = self.prs[self.current_idx]
        pr_obj = PullRequest(pr_id=c["pr_id"], title=c["title"], description=c["description"], diff=c["diff"], author=c["author"], true_bug_description=c["true_bug_description"])
        return Observation(task_id=self.TASK_ID, step=self.current_idx, current_pr=pr_obj, valid_actions=[AgentAction.REVIEW_PR])