Updated client.py

client.py

@@ -1,48 +1,20 @@
-"""
-client.py — xsecure environment client.
-Extends openenv-core's HTTPEnvClient — WebSocket, sync wrapper,
-from_hub() and from_docker_image() all come for free.
-
-Usage (async):
-    async with IncidentResponseEnv(base_url="http://localhost:7860") as env:
-        obs = await env.reset(task_id=1)
-        result = await env.step(IncidentAction(action_type="analyze_log", target="L001"))
-
-Usage (sync):
-    with IncidentResponseEnv(base_url="http://localhost:7860").sync() as env:
-        obs = env.reset(task_id=1)
-        result = env.step(IncidentAction(action_type="analyze_log", target="L001"))
-
-Usage (Docker — auto-pulls and runs):
-    env = await IncidentResponseEnv.from_docker_image("xsecure:latest")
-    async with env:
-        obs = await env.reset(task_id=1)
-"""
-
 from __future__ import annotations
-
-from dataclasses import asdict
 from typing import Any, Dict
 
-try:
-    from openenv_core.http_env_client import HTTPEnvClient
-    from openenv_core.types import StepResult
-except ImportError:
-    from core.http_env_client import HTTPEnvClient
-    from core.types import StepResult
+# Standard OpenEnv imports
+from openenv_core import EnvClient 
+from openenv_core.client_types import StepResult
 
 from models import IncidentAction, IncidentObservation, IncidentState
 
-
-class IncidentResponseEnv(HTTPEnvClient[IncidentAction, IncidentObservation]):
+class IncidentResponseEnv(EnvClient[IncidentAction, IncidentObservation, IncidentState]):
     """
     Client for the xsecure Incident Response environment.
-    Inherits reset(), step(), state(), sync(), from_hub(), from_docker_image()
-    from openenv-core's HTTPEnvClient.
+    Inherits all core methods (reset, step, state) from HTTPEnvClient.
     """
 
     def _step_payload(self, action: IncidentAction) -> Dict[str, Any]:
-        """Serialize action to JSON dict for the HTTP /step endpoint."""
+        """Serialize action for the /step endpoint."""
         return {
             "action_type": action.action_type,
             "target":      action.target,
@@ -51,43 +23,20 @@ class IncidentResponseEnv(HTTPEnvClient[IncidentAction, IncidentObservation]):
     def _parse_result(self, payload: Dict[str, Any]) -> StepResult:
         """Deserialize HTTP response into a typed StepResult."""
         obs_data = payload.get("observation", {})
-
-        # Re-hydrate nested dataclass lists
-        from models import LogEntry, AlertEntry, ServiceStatus
-        logs     = [LogEntry(**l)     for l in obs_data.get("logs", [])]
-        alerts   = [AlertEntry(**a)   for a in obs_data.get("alerts", [])]
-        services = [ServiceStatus(**s) for s in obs_data.get("services", [])]
-
-        obs = IncidentObservation(
-            logs=logs,
-            alerts=alerts,
-            services=services,
-            active_users=obs_data.get("active_users", []),
-            step_count=obs_data.get("step_count", 0),
-            reward=obs_data.get("reward", 0.0),
-            done=obs_data.get("done", False),
-            info=obs_data.get("info", {}),
-            last_action_result=obs_data.get("last_action_result", ""),
-        )
+        
+        # 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", {}),
+            #info=payload.get("info", {}),
         )
 
     def _parse_state(self, payload: Dict[str, Any]) -> IncidentState:
         """Deserialize /state response into IncidentState."""
         data = payload.get("state", payload)
-        return IncidentState(**{
-            k: v for k, v in data.items()
-            if k in IncidentState.__dataclass_fields__
-        })
-
-
-# ---------------------------------------------------------------------------
-# Backward-compatible StepResult re-export
-# ---------------------------------------------------------------------------
-
-__all__ = ["IncidentResponseEnv", "StepResult"]
+        # 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})