Spaces:
Sleeping
Sleeping
File size: 6,334 Bytes
ee117a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | """
client.py — OpenEnv client for the Medical Diagnostic Environment.
This client enables training code to interact with the environment via WebSocket.
Provides both async and sync interfaces for flexibility.
"""
from typing import Optional
import asyncio
import json
try:
import websockets
except ImportError:
websockets = None
from openenv.core.env_client import EnvClient
from openenv.core.client_types import StepResult
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from models import DiagnosticAction, PatientObservation, ClinicalState
class DiagnosticEnv(EnvClient[DiagnosticAction, PatientObservation, ClinicalState]):
"""
Client for interacting with the Medical Diagnostic Environment.
Supports both async and sync usage:
Async (recommended for training):
async with DiagnosticEnv(base_url="...") as env:
obs = await env.reset()
obs = await env.step(DiagnosticAction(...))
Sync (for notebooks/simple scripts):
with DiagnosticEnv(base_url="...").sync() as env:
obs = env.reset()
obs = env.step(DiagnosticAction(...))
"""
@classmethod
async def from_docker_image(cls, image_name: Optional[str] = None, base_url: Optional[str] = None, **kwargs):
"""Create client connected to a running OpenEnv environment URL."""
if base_url is None:
base_url = os.getenv("ENV_URL", "ws://localhost:8000/ws")
return cls(base_url=base_url, **kwargs)
def _step_payload(self, action: DiagnosticAction) -> dict:
"""Convert action to JSON payload for server."""
return {
"action_type": action.action_type,
"question": action.question,
"test_name": action.test_name,
"diagnosis": action.diagnosis,
}
def _parse_result(self, payload: dict) -> StepResult:
"""Parse server response into StepResult."""
obs_data = payload.get("observation", {})
# Parse nested dictionaries if present
patient_data_revealed = obs_data.get("patient_data_revealed", {})
if isinstance(patient_data_revealed, str):
try:
patient_data_revealed = json.loads(patient_data_revealed)
except:
patient_data_revealed = {}
test_result = obs_data.get("test_result")
if isinstance(test_result, str):
try:
test_result = json.loads(test_result)
except:
test_result = None
observation = PatientObservation(
done=payload.get("done", False),
reward=payload.get("reward"),
message=obs_data.get("message", ""),
patient_response=obs_data.get("patient_response"),
test_result=test_result,
questions_asked=obs_data.get("questions_asked", []),
tests_completed=obs_data.get("tests_completed", []),
patient_data_revealed=patient_data_revealed,
steps_taken=obs_data.get("steps_taken", 0),
max_steps=obs_data.get("max_steps", 15),
)
return StepResult(
observation=observation,
reward=payload.get("reward"),
done=payload.get("done", False),
)
def _parse_state(self, payload: dict) -> ClinicalState:
"""Parse state response."""
patient_details = payload.get("patient_details", {})
if isinstance(patient_details, str):
try:
patient_details = json.loads(patient_details)
except:
patient_details = {}
return ClinicalState(
episode_id=payload.get("episode_id", ""),
step_count=payload.get("step_count", 0),
true_diagnosis=payload.get("true_diagnosis", ""),
patient_case=payload.get("patient_case", ""),
patient_details=patient_details,
difficulty=payload.get("difficulty", "easy"),
questions_asked=payload.get("questions_asked", []),
tests_completed=payload.get("tests_completed", []),
final_diagnosis_submitted=payload.get("final_diagnosis_submitted"),
final_accuracy=payload.get("final_accuracy", 0.0),
)
# ─────────────────────────────────────────────────────────────────────
# Sync wrapper for convenience
# ─────────────────────────────────────────────────────────────────────
def sync(self) -> "SyncDiagnosticEnv":
"""
Return synchronous wrapper for use in notebooks/simple scripts.
Usage:
with DiagnosticEnv(url).sync() as env:
obs = env.reset()
obs = env.step(action)
"""
return SyncDiagnosticEnv(self.base_url)
class SyncDiagnosticEnv:
"""Synchronous wrapper around async client."""
def __init__(self, base_url: str):
self.base_url = base_url
self._loop = asyncio.new_event_loop()
self._async_client = DiagnosticEnv(base_url)
def __enter__(self):
self._loop.run_until_complete(self._async_client.__aenter__())
return self
def __exit__(self, *args):
self._loop.run_until_complete(self._async_client.__aexit__(*args))
self._loop.close()
def reset(self, difficulty: str = "easy") -> PatientObservation:
"""Reset environment and start new episode."""
return self._loop.run_until_complete(
self._async_client.reset(difficulty=difficulty)
)
def step(self, action: DiagnosticAction) -> StepResult:
"""Take a step in the environment."""
return self._loop.run_until_complete(
self._async_client.step(action)
)
def state(self) -> ClinicalState:
"""Get current state (includes hidden information)."""
return self._loop.run_until_complete(
self._async_client.state()
)
|