rishil2005/github-issue-triage-env

1from __future__ import annotations
2
3import json
4from pathlib import Path
5from typing import Any
6
7from models import AgentAction, IssueObservation
8
9
10def _compute_reward(action: AgentAction, expected_action: str, expected_label: str | None) -> int:
11    if expected_action == "RequestMoreInfo":
12        return 2 if action.action_type == "RequestMoreInfo" else -1
13    if action.action_type == "AddLabel" and action.label == expected_label:
14        return 1
15    return -1
16
17
18class IssueEnvironment:
19    def __init__(self, dataset_path: str | Path = "dataset.json") -> None:
20        self.dataset_path = Path(dataset_path)
21        self._issues: list[dict[str, Any]] = []
22        self._index = 0
23
24    def reset(self) -> IssueObservation:
25        with open(self.dataset_path, encoding="utf-8") as fh:
26            self._issues = json.load(fh)
27        if not self._issues:
28            raise ValueError("dataset.json is empty.")
29        self._index = 0
30        return self._current_observation()
31
32    def step(self, action: AgentAction) -> tuple[IssueObservation | None, int, bool, dict[str, Any]]:
33        current = self._issues[self._index]
34        reward = _compute_reward(
35            action=action,
36            expected_action=current["expected_action"],
37            expected_label=current.get("expected_label"),
38        )
39
40        info: dict[str, Any] = {
41            "issue_id": current["issue_id"],
42            "expected_action": current["expected_action"],
43            "expected_label": current.get("expected_label"),
44            "agent_action": action.action_type,
45            "agent_label": action.label,
46            "reward": reward,
47        }
48
49        self._index += 1
50        done = self._index >= len(self._issues)
51        next_obs = None if done else self._current_observation()
52        return next_obs, reward, done, info
53
54    def close(self) -> None:
55        return None
56
57    def _current_observation(self) -> IssueObservation:
58        issue = self._issues[self._index]
59        return IssueObservation(
60            issue_id=issue["issue_id"],
61            title=issue["title"],
62            body=issue["body"],
63        )
64