update client.py

client.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""Code Output Assessment Environment Client."""
+"""AI Response Evaluation Environment Client."""
 
 from typing import Dict
 
@@ -19,61 +19,29 @@ class CodeAssessmentEnv(
     EnvClient[CodeAssessmentAction, CodeAssessmentObservation, State]
 ):
     """
-    Client for the Code Output Assessment Environment.
-
-    This client maintains a persistent WebSocket connection to the environment server,
-    enabling efficient multi-step interactions with lower latency.
-    Each client instance has its own dedicated environment session on the server.
+    Client for the AI Response Evaluation Environment.
 
     Example:
-        >>> # Connect to a running server
-        >>> with FirstRlProjEnv(base_url="http://localhost:8000") as client:
-        ...     result = client.reset()
-        ...     print(result.observation.problem_description)
-        ...     print(result.observation.test_case_input)
-        ...
-        ...     result = client.step(FirstRlProjAction(answer="8"))
-        ...     print(result.observation.is_correct)
-
-    Example with Docker:
-        >>> # Automatically start container and connect
-        >>> client = FirstRlProjEnv.from_docker_image("first_rl_proj:latest")
-        >>> try:
-        ...     result = client.reset()
-        ...     result = client.step(FirstRlProjAction(answer="8"))
-        ... finally:
-        ...     client.close()
+        >>> env = await CodeAssessmentEnv.from_docker_image("code_assessment_env:latest")
+        >>> result = await env.reset()
+        >>> print(result.observation.task_type)
+        >>> result = await env.step(CodeAssessmentAction(answer="incorrect, factual-error"))
     """
 
     def _step_payload(self, action: CodeAssessmentAction) -> Dict:
-        """
-        Convert CodeAssessmentAction to JSON payload for step message.
-
-        Args:
-            action: CodeAssessmentAction instance
-
-        Returns:
-            Dictionary representation suitable for JSON encoding
-        """
-        return {
-            "answer": action.answer,
-        }
+        return {"answer": action.answer}
 
     def _parse_result(self, payload: Dict) -> StepResult[CodeAssessmentObservation]:
-        """
-        Parse server response into StepResult[CodeAssessmentObservation].
-
-        Args:
-            payload: JSON response data from server
-
-        Returns:
-            StepResult with CodeAssessmentObservation
-        """
         obs_data = payload.get("observation", {})
         observation = CodeAssessmentObservation(
             problem_description=obs_data.get("problem_description", ""),
             difficulty=obs_data.get("difficulty", "easy"),
             test_case_input=obs_data.get("test_case_input", ""),
+            task_type=obs_data.get("task_type", "correctness_check"),
+            language=obs_data.get("language", "en"),
+            user_age=obs_data.get("user_age"),
+            user_mood=obs_data.get("user_mood"),
+            user_context=obs_data.get("user_context"),
             expected_output=obs_data.get("expected_output"),
             feedback=obs_data.get("feedback", ""),
             is_correct=obs_data.get("is_correct", False),
@@ -84,7 +52,6 @@ class CodeAssessmentEnv(
             reward=payload.get("reward"),
             metadata=obs_data.get("metadata", {}),
         )
-
         return StepResult(
             observation=observation,
             reward=payload.get("reward"),
@@ -92,15 +59,6 @@ class CodeAssessmentEnv(
         )
 
     def _parse_state(self, payload: Dict) -> State:
-        """
-        Parse server response into State object.
-
-        Args:
-            payload: JSON response from state request
-
-        Returns:
-            State object with episode_id and step_count
-        """
         return State(
             episode_id=payload.get("episode_id"),
             step_count=payload.get("step_count", 0),