fix: add MAX_STEPS and rate limit delay to baseline

baseline.py

@@ -7,6 +7,7 @@ import json
 from groq import Groq
 from client import BugTriageClient
 from model import TriageAction
+import time
 
 # ── config ─────────────────────────────────────────────────
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
@@ -90,8 +91,9 @@ def main():
 
     with BugTriageClient() as env:
         obs = env.reset()
-
-        while not obs.done:
+        MAX_STEPS = 3
+        step_count = 0
+        while not obs.done and step_count < MAX_STEPS:
             task = obs.task_id
             print(f"\n── Task: {task.upper()} ──")
             print(f"  Bug: {obs.bug_report.title}")
@@ -112,6 +114,7 @@ def main():
 
             scores[task] = result.reward
             step_count += 1
+            time.sleep(2)
 
     print("\n" + "=" * 50)
     print("  BASELINE SCORES")

client.py

@@ -1,7 +1,6 @@
 # client.py
 import os
 import requests
-from dataclasses import asdict
 from typing import Optional
 from model import TriageAction, TriageObservation, BugReport
 
@@ -16,38 +15,18 @@ class StepResult:
 
 def _parse_observation(data: dict) -> TriageObservation:
     bug_data = data["bug_report"]
-    bug = BugReport(
-        id=bug_data["id"],
-        title=bug_data["title"],
-        body=bug_data["body"],
-        author=bug_data["author"],
-        labels_hint=bug_data.get("labels_hint", []),
-        comments=bug_data.get("comments", []),
-    )
+    bug = BugReport(**bug_data)
     return TriageObservation(
         bug_report=bug,
-        task_id=data["task_id"],
-        score=data["score"],
-        feedback=data["feedback"],
-        done=data["done"],
-        reward=data["reward"],
+        task_id=data.get("task_id", "easy"),
+        score=data.get("score", 0.0),
+        feedback=data.get("feedback", ""),
+        done=data.get("done", False),
+        reward=data.get("reward", 0.0),
     )
 
 
 class BugTriageClient:
-    """
-    HTTP REST client for BugTriageEnvironment.
-    Uses POST /reset and POST /step endpoints.
-
-    Usage:
-        with BugTriageClient() as env:
-            obs = env.reset()
-            while not obs.done:
-                action = TriageAction(...)
-                result = env.step(action)
-                obs = result.observation
-    """
-
     def __init__(self, base_url: Optional[str] = None):
         self.base_url = (
             base_url
@@ -57,27 +36,20 @@ class BugTriageClient:
         self.session.headers.update({"Content-Type": "application/json"})
 
     def reset(self) -> TriageObservation:
-        """Call POST /reset to start a new episode."""
-        response = self.session.post(
-            f"{self.base_url}/reset",
-            json={},
-            timeout=30,
-        )
+        response = self.session.post(f"{self.base_url}/reset", json={}, timeout=30)
         response.raise_for_status()
         data = response.json()
-        return _parse_observation(data["observation"])
+        obs_data = data.get("observation", data)
+        return _parse_observation(obs_data)
 
     def step(self, action: TriageAction) -> StepResult:
-        """Call POST /step with the triage action."""
-        payload = {"action": asdict(action)}
-        response = self.session.post(
-            f"{self.base_url}/step",
-            json=payload,
-            timeout=30,
-        )
+        from dataclasses import asdict
+        payload = {"action": action.dict()}
+        response = self.session.post(f"{self.base_url}/step", json=payload, timeout=30)
         response.raise_for_status()
         data = response.json()
-        obs = _parse_observation(data["observation"])
+        obs_data = data.get("observation", data)
+        obs = _parse_observation(obs_data)
         return StepResult(
             observation=obs,
             reward=data.get("reward", obs.reward) or 0.0,
@@ -86,11 +58,7 @@ class BugTriageClient:
         )
 
     def state(self) -> dict:
-        """Call GET /state to get current environment state."""
-        response = self.session.get(
-            f"{self.base_url}/state",
-            timeout=30,
-        )
+        response = self.session.get(f"{self.base_url}/state", timeout=30)
         response.raise_for_status()
         return response.json()