client.py

"""
WebSocket client for Code Debugging Challenge environment.
"""

from openenv.core.env_client import EnvClient
from openenv.core.client_types import StepResult, ResetResult
from .models import DebugAction, DebugObservation, DebugState


class DebugEnv(EnvClient[DebugAction, DebugObservation, DebugState]):
    """Client for interacting with Code Debugging Challenge environment."""
    
    def _step_payload(self, action: DebugAction) -> dict:
        """Convert action to JSON payload for server."""
        return {
            "action_type": action.action_type,
            "content": action.content
        }
    
    def _parse_result(self, data: dict) -> StepResult[DebugObservation]:
        """Parse step response from server into typed result."""
        obs_data = data["observation"]
        
        observation = DebugObservation(
            buggy_code=obs_data["buggy_code"],
            expected_output=obs_data["expected_output"],
            test_inputs=obs_data.get("test_inputs", []),
            current_output=obs_data.get("current_output"),
            error_message=obs_data.get("error_message"),
            attempts_remaining=obs_data["attempts_remaining"],
            hint=obs_data.get("hint"),
            success=obs_data.get("success", False)
        )
        
        return StepResult(
            observation=observation,
            reward=data["reward"],
            terminated=data["terminated"],
            truncated=data["truncated"],
            info=data.get("info", {})
        )
    
    def _parse_reset_result(self, data: dict) -> ResetResult[DebugObservation]:
        """Parse reset response from server into typed result."""
        obs_data = data["observation"]
        
        observation = DebugObservation(
            buggy_code=obs_data["buggy_code"],
            expected_output=obs_data["expected_output"],
            test_inputs=obs_data.get("test_inputs", []),
            attempts_remaining=obs_data.get("attempts_remaining", 5),
            success=False
        )
        
        return ResetResult(
            observation=observation,
            info=data.get("info", {})
        )
    
    def _parse_state(self, data: dict) -> DebugState:
        """Parse state response from server into typed state."""
        return DebugState(
            current_problem_index=data.get("current_problem_index", 0),
            attempts_made=data.get("attempts_made", 0),
            max_attempts=data.get("max_attempts", 5),
            score=data.get("score", 0.0),
            solved=data.get("solved", False),
            total_problems=data.get("total_problems", 7),
            episode_id=data.get("episode_id", "")
        )