File size: 9,116 Bytes
1f213fe
 
 
 
e48dce9
 
410e78d
 
 
 
 
 
 
 
 
1f213fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e48dce9
42466e3
 
1f213fe
e48dce9
1f213fe
 
 
 
 
 
 
 
e48dce9
 
 
1f213fe
 
 
 
74d0de6
 
 
1f213fe
 
 
 
 
 
74d0de6
 
 
1f213fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e48dce9
74d0de6
1f213fe
 
74d0de6
 
1f213fe
 
 
 
74d0de6
1f213fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74d0de6
1f213fe
 
 
 
 
 
 
 
e48dce9
 
 
 
 
 
 
 
 
 
 
 
1f213fe
 
 
 
74d0de6
 
 
 
 
 
1f213fe
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74d0de6
 
 
 
 
 
 
1f213fe
 
 
 
 
 
 
 
74d0de6
 
 
1f213fe
 
 
 
 
 
 
 
74d0de6
 
 
 
1f213fe
 
 
 
 
 
 
 
74d0de6
 
 
1f213fe
 
 
 
 
 
 
 
dab7b94
 
 
 
 
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
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
import copy
import random
from typing import Any, Dict, Optional

from openenv.core.env_server.interfaces import Environment
from models import PMOpsAction, PMOpsObservation, PMOpsState
from server.apps.ticketing import TicketingApp
from server.apps.codebase import CodebaseApp
from server.apps.chat import ChatApp
from server.world.org_generator import generate_org_config
from server.world.scenario_gen import generate_scenario
from server.tasks.triage_task import TriageTask
from server.tasks.incident_routing_task import IncidentRoutingTask
from server.tasks.release_notes_task import ReleaseNotesTask
from server.tasks.dep_update_task import DepUpdateTask

MAX_STEPS = 40
_TASK_TYPES = ["triage", "incident_routing", "release_notes", "dep_update"]
_DIFFICULTY_POOL = ["easy", "medium", "medium", "hard"]
_TASK_GRADERS = {
    "triage": TriageTask(),
    "incident_routing": IncidentRoutingTask(),
    "release_notes": ReleaseNotesTask(),
    "dep_update": DepUpdateTask(),
}


def _oracle_check(scenario: Dict[str, Any]) -> bool:
    expected = scenario.get("expected", {})
    if not expected:
        return False
    channel = expected.get("channel") or expected.get("channels")
    team = expected.get("team") or expected.get("teams_to_notify")
    return bool(channel) and (bool(team) or scenario["type"] == "release_notes")


class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        super().__init__()
        self._ticketing: Optional[TicketingApp] = None
        self._codebase: Optional[CodebaseApp] = None
        self._chat: Optional[ChatApp] = None
        self._org_config: Optional[Dict[str, Any]] = None
        self._scenario: Optional[Dict[str, Any]] = None
        self._step_count: int = 0
        self._done: bool = False

    def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs) -> PMOpsObservation:
        if seed is None:
            seed = random.randint(0, 2 ** 31)
        rng = random.Random(seed)
        difficulty = rng.choice(_DIFFICULTY_POOL)
        task_type = rng.choice(_TASK_TYPES)

        org: Optional[Dict[str, Any]] = None
        scenario: Optional[Dict[str, Any]] = None

        for attempt in range(10):
            org = generate_org_config(seed + attempt, difficulty)
            scenario = generate_scenario(task_type, org, seed + attempt)
            if _oracle_check(scenario):
                break

        if org is None or scenario is None:
            raise RuntimeError("Failed to initialize episode state")

        channels = list(org["oncall_channels"].values())
        noise = org.get("noise_channels", [])

        self._org_config = org
        self._scenario = scenario
        self._ticketing = TicketingApp(org)
        self._codebase = CodebaseApp(seed, org["services"])
        self._chat = ChatApp(channels, noise)
        self._step_count = 0
        self._done = False

        return PMOpsObservation(
            step=0,
            max_steps=MAX_STEPS,
            task_brief=scenario["brief"],
            last_action_result={
                "ok": True,
                "data": "Environment ready. Call meta.read_runbook to learn this org's conventions.",
            },
            app_state_deltas={"ticketing": [], "chat": [], "codebase": []},
            steps_remaining=MAX_STEPS,
            token_budget_remaining=12000,
            reward=0.0,
            done=False,
        )

    def step(self, action: PMOpsAction, timeout_s: Optional[float] = None, **kwargs) -> PMOpsObservation:
        if self._ticketing is None or self._codebase is None or self._chat is None or self._scenario is None:
            raise RuntimeError("Call reset() before step()")

        scenario = self._scenario

        if self._done:
            return PMOpsObservation(
                step=self._step_count,
                max_steps=MAX_STEPS,
                task_brief=scenario["brief"],
                last_action_result={"ok": False, "error": "Episode already finished"},
                app_state_deltas={"ticketing": [], "chat": [], "codebase": []},
                steps_remaining=0,
                token_budget_remaining=0,
                reward=0.0,
                done=True,
            )

        self._step_count += 1
        result = self._dispatch(action)

        done = action.action_type == "meta.finish" or self._step_count >= MAX_STEPS
        reward = 0.0
        if done:
            reward = self._grade()
            self._done = True

        side_effects = result.get("side_effects", [])
        deltas = {
            "ticketing": side_effects if any("ticket" in s for s in side_effects) else [],
            "chat": side_effects if any("message" in s for s in side_effects) else [],
            "codebase": [],
        }

        return PMOpsObservation(
            step=self._step_count,
            max_steps=MAX_STEPS,
            task_brief=scenario["brief"],
            last_action_result=result,
            app_state_deltas=deltas,
            steps_remaining=max(0, MAX_STEPS - self._step_count),
            token_budget_remaining=max(0, 12000 - self._step_count * 300),
            reward=reward,
            done=done,
        )

    @property
    def state(self) -> PMOpsState:
        return PMOpsState(
            org_config=self._org_config or {},
            task_config=self._scenario or {},
            ticketing=self._ticketing.snapshot() if self._ticketing else {},
            chat=self._chat.snapshot() if self._chat else {},
            codebase={},
            step_count=self._step_count,
            finished=self._done,
        )

    def _dispatch(self, action: PMOpsAction) -> Dict[str, Any]:
        at = action.action_type
        args = action.args or {}

        ticketing = self._ticketing
        codebase = self._codebase
        chat = self._chat
        if ticketing is None or codebase is None or chat is None:
            return {"ok": False, "error": "Environment not initialized. Call reset() before step()."}

        if at == "meta.noop":
            return {"ok": True, "data": "No operation."}

        if at == "meta.read_runbook":
            return {
                "ok": True,
                "data": {
                    "org_config": copy.deepcopy(self._org_config),
                    "hint": (
                        "Use label_taxonomy for valid ticket labels, "
                        "priority_levels for valid priorities, "
                        "team_map[service] to find the owning team, "
                        "oncall_channels[service] to find the channel to notify."
                    ),
                },
            }

        if at == "meta.finish":
            return {"ok": True, "data": "Episode finishing. Score will be computed."}

        if at.startswith("ticketing."):
            op = at.split(".", 1)[1]
            handlers = {
                "create_ticket":     ticketing.create_ticket,
                "update_ticket":     ticketing.update_ticket,
                "get_ticket":        ticketing.get_ticket,
                "list_tickets":      ticketing.list_tickets,
                "assign_ticket":     ticketing.assign_ticket,
                "comment_ticket":    ticketing.comment_ticket,
                "transition_ticket": ticketing.transition_ticket,
            }
            if op not in handlers:
                return {"ok": False, "error": f"Unknown ticketing action: {op}"}
            return handlers[op](args)

        if at.startswith("codebase."):
            op = at.split(".", 1)[1]
            handlers = {
                "list_commits": codebase.list_commits,
                "get_commit":   codebase.get_commit,
                "list_prs":     codebase.list_prs,
            }
            if op not in handlers:
                return {"ok": False, "error": f"Unknown codebase action: {op}"}
            return handlers[op](args)

        if at.startswith("chat."):
            op = at.split(".", 1)[1]
            handlers = {
                "post_message":  chat.post_message,
                "read_channel":  chat.read_channel,
                "list_channels": chat.list_channels,
                "search":        chat.search,
            }
            if op not in handlers:
                return {"ok": False, "error": f"Unknown chat action: {op}"}
            return handlers[op](args)

        return {"ok": False, "error": f"Unknown action_type: {at}"}

    def _grade(self) -> float:
        if self._ticketing is None or self._chat is None or self._org_config is None or self._scenario is None:
            return 0.0

        final_state = {
            "ticketing": self._ticketing.snapshot(),
            "chat": self._chat.snapshot(),
        }
        grader = _TASK_GRADERS.get(self._scenario["type"])
        if not grader:
            return 0.0
        return grader.grade(final_state, self._org_config, self._scenario)

    def close(self) -> None:
        self._ticketing = None
        self._codebase = None
        self._chat = None