File size: 1,259 Bytes
3807ea3
287d681
3807ea3
287d681
3807ea3
287d681
 
 
 
3807ea3
 
 
 
 
287d681
 
3807ea3
 
287d681
3807ea3
287d681
3807ea3
 
287d681
3807ea3
287d681
3807ea3
287d681
 
3807ea3
 
287d681
 
3807ea3
 
287d681
 
3807ea3
 
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
"""CyberSecurity_OWASP OpenEnv client."""

from __future__ import annotations

from typing import Any

from openenv.core import EnvClient
from openenv.core.client_types import StepResult

from .models import (
    CyberSecurityOWASPAction,
    CyberSecurityOWASPObservation,
    CyberSecurityOWASPState,
)


class CyberSecurityOWASPEnv(
    EnvClient[CyberSecurityOWASPAction, CyberSecurityOWASPObservation, CyberSecurityOWASPState]
):
    """WebSocket client for the CyberSecurity_OWASP environment."""

    def _step_payload(self, action: CyberSecurityOWASPAction) -> dict[str, Any]:
        return action.model_dump()

    def _parse_result(self, payload: dict[str, Any]) -> StepResult[CyberSecurityOWASPObservation]:
        obs_data = payload.get("observation", {})
        observation = CyberSecurityOWASPObservation(**obs_data)
        return StepResult(
            observation=observation,
            reward=payload.get("reward", observation.reward),
            done=payload.get("done", observation.done),
        )

    def _parse_state(self, payload: dict[str, Any]) -> CyberSecurityOWASPState:
        return CyberSecurityOWASPState(**payload)


# Backward-compatible alias from generated scaffold.
CybersecurityOwaspEnv = CyberSecurityOWASPEnv