sankar-raul/ICD-10-code-predictor-env

1"""Client for the Medical Coding Assistant environment."""2 3from typing import Any4 5from openenv.core import EnvClient6from openenv.core.client_types import StepResult7 8from .models import MedicalCodingAction, MedicalCodingObservation, MedicalCodingState9 10 11class MedicalCodingAssistantEnv(12    EnvClient[MedicalCodingAction, MedicalCodingObservation, MedicalCodingState]13):14    """Persistent client for the Medical Coding Assistant environment."""15 16    def _step_payload(self, action: MedicalCodingAction) -> dict[str, Any]:17        return action.model_dump()18 19    def _parse_result(self, payload: dict[str, Any]) -> StepResult[MedicalCodingObservation]:20        obs_data = payload.get("observation", {})21        observation = MedicalCodingObservation(22            task_id=obs_data.get("task_id", ""),23            difficulty=obs_data.get("difficulty", "easy"),24            objective=obs_data.get("objective", ""),25            encounter_text=obs_data.get("encounter_text", ""),26            allowed_codes=obs_data.get("allowed_codes", []),27            revealed_hints=obs_data.get("revealed_hints", []),28            current_primary_code=obs_data.get("current_primary_code", ""),29            current_secondary_codes=obs_data.get("current_secondary_codes", []),30            current_needs_review=obs_data.get("current_needs_review", False),31            attempts_remaining=obs_data.get("attempts_remaining", 0),32            progress_score=obs_data.get("progress_score", 0.0),33            grader_feedback=obs_data.get("grader_feedback", []),34            reward_breakdown=obs_data.get("reward_breakdown", {}),35            done=payload.get("done", False),36            reward=payload.get("reward"),37            metadata=obs_data.get("metadata", {}),38        )39        return StepResult(40            observation=observation,41            reward=payload.get("reward"),42            done=payload.get("done", False),43        )44 45    def _parse_state(self, payload: dict[str, Any]) -> MedicalCodingState:46        return MedicalCodingState(47            episode_id=payload.get("episode_id"),48            step_count=payload.get("step_count", 0),49            current_task_id=payload.get("current_task_id", ""),50            difficulty=payload.get("difficulty", "easy"),51            current_primary_code=payload.get("current_primary_code", ""),52            current_secondary_codes=payload.get("current_secondary_codes", []),53            current_needs_review=payload.get("current_needs_review", False),54            best_score=payload.get("best_score", 0.0),55            hints_used=payload.get("hints_used", 0),56            repeated_actions=payload.get("repeated_actions", 0),57            completed=payload.get("completed", False),58        )59