shashanks/medical_coding
0
1"""2Inference Script — Medical Coding Auditor Environment3=======================================================4MANDATORY ENVIRONMENT VARIABLES:5 API_BASE_URL The API endpoint for the LLM (default: "https://api.openai.com/v1").6 MODEL_NAME The model identifier to use for inference (default: "gpt-4.1-mini").7 HF_TOKEN Your Hugging Face / API key (used as API key for LLM calls). No default.8 LOCAL_IMAGE_NAME (optional) Docker image name if launching env via from_docker_image().9 ENV_BASE_URL (optional) Base URL of the running env server (default: http://localhost:7860).10 11STDOUT FORMAT (required by hackathon spec):12 [START] task=<task_id> env=medical_coding model=<model_name>13 [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>14 [END] success=<true|false> steps=<n> rewards=<r1,r2,...>15"""16 17import asyncio18import json19import os20import textwrap21from typing import Any, Dict, List, Optional22 23from openai import OpenAI24 25# ---------------------------------------------------------------------------26# Configuration27# ---------------------------------------------------------------------------28API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")29MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4.1-mini")30HF_TOKEN = os.getenv("HF_TOKEN")31 32if HF_TOKEN is None:33 raise ValueError("HF_TOKEN environment variable is required")34 35# Optional — if you use from_docker_image():36LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")37 38ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:7860")39BENCHMARK = "medical_coding"40 41TASK_IDS = [42 "easy_demographic",43 "medium_ncci_conflict",44 "medium_excludes1",45 "hard_specificity_untraceable",46 "expert_multi_error",47]48 49MAX_STEPS_PER_TASK: Dict[str, int] = {50 "easy_demographic": 10,51 "medium_ncci_conflict": 15,52 "medium_excludes1": 15,53 "hard_specificity_untraceable": 20,54 "expert_multi_error": 25,55}56 57SUCCESS_THRESHOLD = 0.558TEMPERATURE = 0.259MAX_TOKENS = 40060 61# ---------------------------------------------------------------------------62# Logging (required stdout format)63# ---------------------------------------------------------------------------64 65def log_start(task: str, env: str, model: str) -> None:66 print(f"[START] task={task} env={env} model={model}", flush=True)67 68 69def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:70 error_val = error if error else "null"71 done_val = str(done).lower()72 print(73 f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",74 flush=True,75 )76 77 78def log_end(success: bool, steps: int, rewards: List[float]) -> None:79 rewards_str = ",".join(f"{r:.2f}" for r in rewards)80 print(81 f"[END] success={str(success).lower()} steps={steps} rewards={rewards_str}",82 flush=True,83 )84 85 86# ---------------------------------------------------------------------------87# LLM prompting88# ---------------------------------------------------------------------------89 90SYSTEM_PROMPT = textwrap.dedent("""91You are an expert Medical Coding Auditor performing a pre-bill compliance review.92 93Your task: review proposed ICD-10-CM and CPT billing codes against the patient's94clinical note and demographics. Identify coding errors WITHOUT assigning new codes.95 96Available tools (respond with exactly ONE JSON object per turn):97 981. query_guideline — look up official coding guidelines for a specific code.99 {"action_type": "query_guideline", "code": "<CODE>"}100 1012. check_ncci_edits — check if two CPT codes have a CMS NCCI PTP bundling conflict.102 {"action_type": "check_ncci_edits", "code1": "<CPT1>", "code2": "<CPT2>"}103 1043. flag_error — record a confirmed coding error in the audit report.105 {"action_type": "flag_error", "code": "<CODE>", "error_type": "<TYPE>", "justification": "<REASON>"}106 error_type values:107 demographic_mismatch — code inapplicable to patient's sex or age108 excludes1_conflict — two mutually exclusive ICD-10-CM diagnosis codes109 ncci_edit — CMS NCCI PTP bundling violation110 specificity_error — wrong 7th character or insufficient code specificity111 untraceable_code — code does not exist in any official code set112 1134. ask_clarifying_question — ask the physician for missing clinical information when the note is ambiguous.114 {"action_type": "ask_clarifying_question", "question": "<YOUR QUESTION>"}115 Use when you need to confirm: patient sex/age, encounter type (initial vs. follow-up), or diagnosis history.116 1175. extract_evidence — highlight an exact text span from the clinical note as evidence for a coding decision.118 {"action_type": "extract_evidence", "evidence_text": "<EXACT SUBSTRING FROM NOTE>"}119 Use before flagging an error to ground your reasoning in specific documentation.120 1216. submit_audit — submit the completed audit report when you are done.122 {"action_type": "submit_audit"}123 124RULES:125- Only query or flag codes that appear in the proposed_codes list.126- Query each code's guidelines before flagging it.127- For CPT code pairs, use check_ncci_edits to find bundling conflicts.128- If query_guideline returns "CODE NOT FOUND" or "INVALID/UNTRACEABLE", flag it as untraceable_code.129- Check patient sex — maternity (O-codes) ONLY apply to female patients.130- Check 7th character on injury codes against the clinical note (initial vs. subsequent encounter).131- Use ask_clarifying_question if demographics or encounter type are unclear before flagging.132- Use extract_evidence to quote the exact clinical note text that supports each flag.133- Submit the audit only after reviewing all proposed codes.134- Respond with ONLY a valid JSON object — no markdown, no extra text.135""").strip()136 137 138def build_user_prompt(obs: Dict[str, Any], step: int, history: List[str]) -> str:139 proposed = obs.get("proposed_codes", {})140 codes_summary = "\n".join(141 f" {code}: {info.get('description', '')} [{info.get('code_type', '')}]"142 for code, info in proposed.items()143 )144 draft = obs.get("draft_report", [])145 draft_summary = json.dumps(draft, indent=2) if draft else " (none yet)"146 queried = obs.get("codes_queried", [])147 pairs_checked = obs.get("pairs_checked", [])148 last_result = obs.get("tool_result", "")[:700]149 150 history_block = "\n".join(history[-5:]) if history else "None"151 152 return textwrap.dedent(f"""153 === STEP {step} ===154 PATIENT: age={obs.get('patient_demographics', {}).get('age')}, sex={obs.get('patient_demographics', {}).get('sex')}155 CLINICAL NOTE:156 {obs.get('clinical_note', '')[:900]}157 158 PROPOSED CODES TO AUDIT:159 {codes_summary}160 161 Already queried: {queried}162 NCCI pairs checked: {pairs_checked}163 164 DRAFT AUDIT REPORT:165 {draft_summary}166 167 LAST TOOL RESULT:168 {last_result}169 170 RECENT HISTORY:171 {history_block}172 173 What is your next action? Reply with exactly one JSON object.174 """).strip()175 176 177def parse_action(text: str) -> Optional[Dict[str, Any]]:178 """Extract a JSON action from the LLM response."""179 text = text.strip()180 try:181 obj = json.loads(text)182 if isinstance(obj, dict) and "action_type" in obj:183 return obj184 except json.JSONDecodeError:185 pass186 start, end = text.find("{"), text.rfind("}") + 1187 if start >= 0 and end > start:188 try:189 obj = json.loads(text[start:end])190 if isinstance(obj, dict) and "action_type" in obj:191 return obj192 except json.JSONDecodeError:193 pass194 return None195 196 197def get_llm_action(198 client: OpenAI,199 obs: Dict[str, Any],200 step: int,201 history: List[str],202) -> Dict[str, Any]:203 """Call the LLM and parse its action."""204 user_prompt = build_user_prompt(obs, step, history)205 try:206 completion = client.chat.completions.create(207 model=MODEL_NAME,208 messages=[209 {"role": "system", "content": SYSTEM_PROMPT},210 {"role": "user", "content": user_prompt},211 ],212 temperature=TEMPERATURE,213 max_tokens=MAX_TOKENS,214 )215 text = (completion.choices[0].message.content or "").strip()216 action = parse_action(text)217 if action:218 return action219 print(f"[DEBUG] Could not parse action from LLM output: {text[:300]}", flush=True)220 except Exception as exc:221 print(f"[DEBUG] LLM error at step {step}: {exc}", flush=True)222 223 # Fallback: query unchecked codes, then submit224 queried = obs.get("codes_queried", [])225 unchecked = [c for c in obs.get("proposed_codes", {}) if c not in queried]226 if unchecked:227 return {"action_type": "query_guideline", "code": unchecked[0]}228 return {"action_type": "submit_audit"}229 230 231# ---------------------------------------------------------------------------232# WebSocket-based environment runner233# ---------------------------------------------------------------------------234 235async def run_task_websocket(236 ws_url: str,237 task_id: str,238 client: OpenAI,239) -> float:240 """241 Run one task episode over a WebSocket connection to maintain state.242 Returns the grader_score from the terminal observation (already in [0.0, 1.0]).243 """244 import websockets245 246 max_steps = MAX_STEPS_PER_TASK[task_id]247 rewards: List[float] = []248 steps_taken = 0249 score = 0.0250 success = False251 252 log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)253 254 try:255 async with websockets.connect(ws_url, max_size=10 * 1024 * 1024) as ws:256 # Reset with task_id257 await ws.send(json.dumps({"type": "reset", "data": {"task_id": task_id}}))258 raw = await ws.recv()259 response = json.loads(raw)260 obs_payload = response.get("data", {})261 obs = obs_payload.get("observation", obs_payload)262 done = obs_payload.get("done", False)263 264 history: List[str] = []265 266 for step in range(1, max_steps + 1):267 if done:268 break269 270 action_dict = get_llm_action(client, obs, step, history)271 action_str = json.dumps(action_dict)272 273 # Send step274 await ws.send(json.dumps({"type": "step", "data": action_dict}))275 raw = await ws.recv()276 step_response = json.loads(raw)277 step_data = step_response.get("data", {})278 279 obs = step_data.get("observation", obs)280 reward = step_data.get("reward") or 0.0281 done = step_data.get("done", False)282 error = obs.get("last_action_error") if isinstance(obs, dict) else None283 284 rewards.append(reward)285 steps_taken = step286 287 log_step(step=step, action=action_str, reward=reward, done=done, error=error)288 289 history.append(290 f"Step {step}: {action_str[:120]} → reward={reward:+.2f}"291 + (f" [ERROR: {error}]" if error else "")292 )293 294 if done:295 # Use grader_score from terminal observation as the score296 grader = obs.get("grader_score") if isinstance(obs, dict) else None297 if grader is not None:298 score = float(grader)299 break300 301 success = score >= SUCCESS_THRESHOLD302 303 except Exception as exc:304 print(f"[DEBUG] WebSocket task {task_id} error: {exc}", flush=True)305 306 finally:307 log_end(success=success, steps=steps_taken, rewards=rewards)308 309 return score310 311 312# ---------------------------------------------------------------------------313# HTTP-based fallback runner (stateless — demonstrates single-step interactions)314# ---------------------------------------------------------------------------315 316async def run_task_http(317 base_url: str,318 task_id: str,319 client: OpenAI,320) -> float:321 """322 HTTP-based stateless runner. Since the OpenEnv HTTP endpoints are stateless323 (each request creates a fresh environment), this runner passes the full324 state from the reset observation into a stateful local environment instance325 for demonstration.326 """327 # Import environment directly for stateful HTTP-like simulation328 import sys329 import os330 env_dir = os.path.dirname(os.path.abspath(__file__))331 if env_dir not in sys.path:332 sys.path.insert(0, env_dir)333 334 from server.environment import MedicalCodingEnvironment335 from models import MedicalCodingAction336 337 env = MedicalCodingEnvironment()338 max_steps = MAX_STEPS_PER_TASK[task_id]339 rewards: List[float] = []340 steps_taken = 0341 score = 0.0342 success = False343 344 log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)345 346 try:347 obs_obj = env.reset(task_id=task_id)348 obs = obs_obj.model_dump()349 done = obs_obj.done350 history: List[str] = []351 352 for step in range(1, max_steps + 1):353 if done:354 break355 356 action_dict = get_llm_action(client, obs, step, history)357 action_str = json.dumps(action_dict)358 359 try:360 action = MedicalCodingAction(**action_dict)361 step_obs = env.step(action)362 obs = step_obs.model_dump()363 reward = step_obs.reward or 0.0364 done = step_obs.done365 error = step_obs.last_action_error366 except Exception as e:367 reward = -0.2368 done = False369 error = str(e)370 print(f"[DEBUG] Step parse error: {e}", flush=True)371 372 rewards.append(reward)373 steps_taken = step374 375 log_step(step=step, action=action_str, reward=reward, done=done, error=error)376 377 history.append(378 f"Step {step}: {action_str[:120]} → reward={reward:+.2f}"379 + (f" [ERROR: {error}]" if error else "")380 )381 382 if done:383 # Use grader_score from terminal observation as the score384 if step_obs.grader_score is not None:385 score = float(step_obs.grader_score)386 break387 388 except Exception as exc:389 print(f"[DEBUG] HTTP task {task_id} error: {exc}", flush=True)390 391 finally:392 success = score >= SUCCESS_THRESHOLD393 log_end(success=success, steps=steps_taken, rewards=rewards)394 395 return score396 397 398# ---------------------------------------------------------------------------399# Main400# ---------------------------------------------------------------------------401 402async def main() -> None:403 """Run baseline inference across all three tasks."""404 client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)405 406 # If LOCAL_IMAGE_NAME is set, launch env from Docker image407 if LOCAL_IMAGE_NAME:408 from client import MedicalCodingEnv409 env = await MedicalCodingEnv.from_docker_image(LOCAL_IMAGE_NAME)410 base_url = env.base_url.rstrip("/")411 else:412 base_url = ENV_BASE_URL.rstrip("/")413 414 # Try WebSocket mode first, fall back to HTTP-local mode415 ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://") + "/ws"416 417 # Check if websockets library is available418 try:419 import websockets420 use_ws = True421 except ImportError:422 use_ws = False423 print("[DEBUG] websockets not installed, using HTTP/local mode", flush=True)424 425 task_scores: Dict[str, float] = {}426 427 for task_id in TASK_IDS:428 if use_ws:429 score = await run_task_websocket(ws_url, task_id, client)430 else:431 score = await run_task_http(base_url, task_id, client)432 task_scores[task_id] = score433 434 print("\n" + "=" * 60, flush=True)435 print("BASELINE SCORES SUMMARY", flush=True)436 print("=" * 60, flush=True)437 for task_id, score in task_scores.items():438 status = "PASS" if score >= SUCCESS_THRESHOLD else "FAIL"439 print(f" [{status}] {task_id}: {score:.3f}", flush=True)440 avg = sum(task_scores.values()) / len(task_scores)441 print(f" Average: {avg:.3f}", flush=True)442 print("=" * 60, flush=True)443 444 445if __name__ == "__main__":446 asyncio.run(main())447