File size: 1,623 Bytes
7f42e9d
 
 
64ba330
 
 
7f42e9d
 
 
64ba330
7f42e9d
 
64ba330
7f42e9d
 
 
64ba330
7f42e9d
 
 
 
 
 
 
 
64ba330
 
 
7f42e9d
 
 
 
 
64ba330
7f42e9d
 
 
 
 
64ba330
 
 
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
from __future__ import annotations
from typing import Any, Dict

# Standard OpenEnv imports
from openenv_core import EnvClient 
from openenv_core.client_types import StepResult

from models import IncidentAction, IncidentObservation, IncidentState

class IncidentResponseEnv(EnvClient[IncidentAction, IncidentObservation, IncidentState]):
    """
    Client for the xsecure Incident Response environment.
    Inherits all core methods (reset, step, state) from HTTPEnvClient.
    """

    def _step_payload(self, action: IncidentAction) -> Dict[str, Any]:
        """Serialize action for the /step endpoint."""
        return {
            "action_type": action.action_type,
            "target":      action.target,
        }

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult:
        """Deserialize HTTP response into a typed StepResult."""
        obs_data = payload.get("observation", {})
        
        # Leverage Pydantic's ability to handle the dictionary directly
        obs = IncidentObservation(**obs_data)

        return StepResult(
            observation=obs,
            reward=payload.get("reward", 0.0),
            done=payload.get("done", False),
            #info=payload.get("info", {}),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> IncidentState:
        """Deserialize /state response into IncidentState."""
        data = payload.get("state", payload)
        # Filters keys to match the Pydantic model fields
        valid_fields = set(IncidentState.model_fields.keys())
        return IncidentState(**{k: v for k, v in data.items() if k in valid_fields})