Update environment.py

environment.py

@@ -1,5 +1,6 @@
 from typing import Tuple, Dict, Any
 from models import Observation, Action, Reward, State
+from grader import grade_comment
 
 class CodeReviewEnv:
     def __init__(self, task: str = "easy"):
@@ -7,7 +8,7 @@ class CodeReviewEnv:
         self.reset()
 
     def set_task(self, task: str):
-        if task not in ["easy", "medium", "hard"]:
+        if task not in ["easy", "medium", "hard", "harder", "hardest"]:
             raise ValueError(f"Unknown task: {task}")
         self.task = task
 
@@ -18,17 +19,46 @@ class CodeReviewEnv:
         self.agent_comment = None
         self.done = False
 
+        # Task definitions with richer context
         if self.task == "easy":
-            self.pr_code = "def get_user(id):\n    return users[id]  # missing null check"
+            self.pr_title = "Fix missing null check in user lookup"
+            self.pr_description = "The current code does not handle missing user IDs. It raises a KeyError."
+            self.code_snippet = "def get_user(id):\n    return users[id]  # missing null check"
             self.comments = ["Looks good!", "Maybe add a comment?"]
+            self.expected_keywords = ["null", "key", "missing", "check", "exists", "handle"]
+            self.expert_comment = "Add a check to ensure the key exists before accessing the dictionary to avoid KeyError."
+
         elif self.task == "medium":
-            self.pr_code = "for i in range(len(items)):\n    process(items[i])\n# O(n^2) when it could be O(n)"
+            self.pr_title = "Improve loop efficiency"
+            self.pr_description = "The loop uses `range(len(items))` which is inefficient and less readable."
+            self.code_snippet = "for i in range(len(items)):\n    process(items[i])\n# O(n^2) when it could be O(n)"
             self.comments = ["Nice code"]
+            self.expected_keywords = ["enumerate", "for item in", "range", "inefficient", "optimize"]
+            self.expert_comment = "Use `for item in items:` for a more Pythonic and efficient loop."
+
         elif self.task == "hard":
-            self.pr_code = "def calculate_average(data):\n    total = sum(data)\n    return total / len(data)  # what if data is empty?"
+            self.pr_title = "Handle division by zero in average calculation"
+            self.pr_description = "The function crashes when the input list is empty."
+            self.code_snippet = "def calculate_average(data):\n    total = sum(data)\n    return total / len(data)  # what if data is empty?"
             self.comments = ["LGTM"]
-        else:
-            raise RuntimeError(f"Invalid task: {self.task}")
+            self.expected_keywords = ["empty", "zero", "length", "check", "handle", "exception"]
+            self.expert_comment = "Check if the list is empty and return a sensible default (e.g., 0) or raise a descriptive error."
+
+        elif self.task == "harder":
+            self.pr_title = "Fix race condition in counter increment"
+            self.pr_description = "Multiple threads increment a counter without synchronization, causing lost updates."
+            self.code_snippet = "counter = 0\ndef increment():\n    global counter\n    counter += 1\n# called from multiple threads"
+            self.comments = ["Simple code, should be fine?"]
+            self.expected_keywords = ["thread", "lock", "synchronization", "atomic", "race", "concurrent"]
+            self.expert_comment = "Use a threading.Lock to protect the counter increment, or use an atomic operation like `threading.atomic`."
+
+        else:  # hardest
+            self.pr_title = "Fix deadlock in database transaction"
+            self.pr_description = "Two threads acquire locks in opposite order, leading to potential deadlock."
+            self.code_snippet = "with lock1:\n    with lock2:\n        do_work()\n# another thread does lock2 then lock1"
+            self.comments = ["Seems okay?"]
+            self.expected_keywords = ["deadlock", "lock order", "acquire", "release", "trylock", "timeout"]
+            self.expert_comment = "Ensure all threads acquire locks in the same order to prevent deadlock. Consider using a timeout or a single lock."
 
         return self._get_observation()
 
@@ -41,8 +71,13 @@ class CodeReviewEnv:
 
         if action.action_type == "write_comment":
             self.agent_comment = action.comment_text or ""
-            reward = 0.2
-            quality_score = self._grade_comment(self.agent_comment)
+            reward = 0.2  # dense bonus for writing
+            # Semantic grader
+            quality_score = grade_comment(
+                self.agent_comment,
+                self.expected_keywords,
+                self.expert_comment
+            )
             reward += quality_score
             if reward > 1.0:
                 reward = 1.0
@@ -61,25 +96,11 @@ class CodeReviewEnv:
         obs = self._get_observation()
         return obs, Reward(value=reward), self.done, info
 
-    def _grade_comment(self, comment: str) -> float:
-        if self.task == "easy":
-            keywords = ["null", "key", "missing", "check", "exists", "handle"]
-            matched = sum(1 for kw in keywords if kw in comment.lower())
-            return min(1.0, matched / 3)
-        elif self.task == "medium":
-            keywords = ["enumerate", "for item in", "range", "inefficient", "optimize"]
-            matched = sum(1 for kw in keywords if kw in comment.lower())
-            return min(1.0, matched / 3)
-        elif self.task == "hard":
-            keywords = ["empty", "zero", "length", "check", "handle", "exception"]
-            matched = sum(1 for kw in keywords if kw in comment.lower())
-            return min(1.0, matched / 3)
-        else:
-            return 0.0
-
     def _get_observation(self) -> Observation:
         return Observation(
-            pr_code=self.pr_code,
+            pr_title=self.pr_title,
+            pr_description=self.pr_description,
+            code_snippet=self.code_snippet,
             comments=self.comments,
             agent_comment=self.agent_comment,
             step=self.step_count,
@@ -88,7 +109,9 @@ class CodeReviewEnv:
 
     def state(self) -> State:
         return State(
-            pr_code=self.pr_code,
+            pr_title=self.pr_title,
+            pr_description=self.pr_description,
+            code_snippet=self.code_snippet,
             comments=self.comments,
             agent_comment=self.agent_comment,
             step=self.step_count,