Mist-ic commited on
Commit
b6ab5da
·
1 Parent(s): 641125f

Add client.py and __init__.py required by openenv push

Browse files
Files changed (2) hide show
  1. __init__.py +0 -0
  2. client.py +64 -0
__init__.py ADDED
File without changes
client.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SevZero Environment Client."""
2
+
3
+ from typing import Dict
4
+
5
+ from openenv.core import EnvClient
6
+ from openenv.core.client_types import StepResult
7
+ from openenv.core.env_server.types import State
8
+
9
+ from models import SevZeroAction, SevZeroObservation
10
+
11
+
12
+ class SevZeroEnv(EnvClient[SevZeroAction, SevZeroObservation, State]):
13
+ """
14
+ Client for the SevZero SRE Incident Response Environment.
15
+
16
+ Maintains a persistent WebSocket connection to the environment server,
17
+ enabling multi-step incident response episodes.
18
+
19
+ Example:
20
+ >>> with SevZeroEnv(base_url="http://localhost:7860") as client:
21
+ ... result = client.reset(task_id="easy", seed=42)
22
+ ... obs = result.observation
23
+ ... print(obs.global_slo_score)
24
+ ...
25
+ ... action = SevZeroAction(
26
+ ... action_type="inspect_logs",
27
+ ... params={"service_id": "order-service"}
28
+ ... )
29
+ ... result = client.step(action)
30
+ ... print(result.observation.logs)
31
+
32
+ Example with Docker:
33
+ >>> client = SevZeroEnv.from_docker_image("sevzero-env:latest")
34
+ >>> try:
35
+ ... result = client.reset(task_id="medium", seed=123)
36
+ ... action = SevZeroAction(action_type="noop", params={})
37
+ ... result = client.step(action)
38
+ ... finally:
39
+ ... client.close()
40
+ """
41
+
42
+ def _step_payload(self, action: SevZeroAction) -> Dict:
43
+ return {
44
+ "action_type": action.action_type,
45
+ "params": action.params,
46
+ }
47
+
48
+ def _parse_result(self, payload: Dict) -> StepResult[SevZeroObservation]:
49
+ obs_data = payload.get("observation", payload)
50
+ observation = SevZeroObservation(**{
51
+ k: v for k, v in obs_data.items()
52
+ if k in SevZeroObservation.model_fields
53
+ })
54
+ return StepResult(
55
+ observation=observation,
56
+ reward=payload.get("reward"),
57
+ done=payload.get("done", False),
58
+ )
59
+
60
+ def _parse_state(self, payload: Dict) -> State:
61
+ return State(
62
+ episode_id=payload.get("episode_id"),
63
+ step_count=payload.get("step_count", 0),
64
+ )