client.py

"""Smart Contract Audit Environment Client."""

from typing import Dict, Any

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

from models import AuditAction, AuditObservation, AuditState


class AuditEnv(EnvClient[AuditAction, AuditObservation, AuditState]):
    """
    Client for the Smart Contract Audit Environment.

    This client maintains a persistent WebSocket connection to the environment server,
    enabling efficient multi-step interactions with lower latency.
    """

    def _step_payload(self, action: AuditAction) -> Dict:
        """
        Convert AuditAction to JSON payload for step message.
        """
        return {
            "analysis": action.analysis,
            "vulnerabilities": action.vulnerabilities,
            "suggested_fixes": action.suggested_fixes,
        }

    def _parse_result(self, payload: Dict) -> StepResult[AuditObservation]:
        """
        Parse server response into StepResult[AuditObservation].
        """
        obs_data = payload.get("observation", {})
        observation = AuditObservation(
            contract_code=obs_data.get("contract_code", ""),
            contract_name=obs_data.get("contract_name", ""),
            task_id=obs_data.get("task_id", ""),
            task_description=obs_data.get("task_description", ""),
            hint=obs_data.get("hint", ""),
            feedback=obs_data.get("feedback", ""),
            reward=payload.get("reward", 0.0),
            done=payload.get("done", False),
        )

        return StepResult(
            observation=observation,
            reward=payload.get("reward", 0.0),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict) -> AuditState:
        """
        Parse server response into AuditState object.
        """
        return AuditState(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
            current_task=payload.get("current_task", ""),
            max_steps=payload.get("max_steps", 5),
            total_reward=payload.get("total_reward", 0.0)
        )