inference.py

#!/usr/bin/env python3
"""
inference.py — MedTriage OpenEnv Baseline Inference Script
============================================================

Mandatory environment variables:
  API_BASE_URL   The API endpoint for the LLM
  MODEL_NAME     The model identifier
  HF_TOKEN       Your Hugging Face / API key
  MEDTRIAGE_URL  URL of the running MedTriage environment (default: http://localhost:7860)

STDOUT FORMAT (strictly followed):
  [START] task=<task_name> env=medtriage-env model=<model_name>
  [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
  [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>
"""

import asyncio
import json
import os
import textwrap
from typing import Any, Dict, List, Optional
import httpx
from openai import OpenAI

# ── Environment variables ─────────────────────────────────────────────────────
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
MEDTRIAGE_URL = os.getenv("MEDTRIAGE_URL", "http://localhost:7860").rstrip("/")

BENCHMARK = "medtriage-env"
MAX_STEPS = 3          # Single-turn tasks: reset → step → done
SUCCESS_THRESHOLD = 0.4
TEMPERATURE = 0.2
MAX_TOKENS = 512

TASKS = ["vital-triage", "differential-diagnosis", "treatment-safety"]

# ── Logging helpers ───────────────────────────────────────────────────────────

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    # Sanitize action string — no newlines
    action_clean = action.replace("\n", " ").replace("\r", "")[:200]
    print(
        f"[STEP] step={step} action={action_clean} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ── Prompt builders ───────────────────────────────────────────────────────────

SYSTEM_PROMPT = textwrap.dedent("""
    You are an expert emergency medicine physician with board certification in emergency medicine.
    You will be given a clinical task and patient information. You must respond with a valid JSON
    object containing your clinical decision. Be precise, evidence-based, and safety-conscious.
    Always check allergies and contraindications.
    
    Respond ONLY with a JSON object — no explanations, no markdown, no code blocks.
""").strip()

TASK1_PROMPT_TEMPLATE = textwrap.dedent("""
    TASK: Emergency Severity Index (ESI) Triage Classification
    
    Classify this patient using ESI levels:
    - ESI 1: Immediate life threat (airway compromise, pulselessness, severe shock)
    - ESI 2: High risk OR confused/lethargic OR severe pain/distress
    - ESI 3: Stable but needs multiple resources
    - ESI 4: Stable, needs one resource
    - ESI 5: No resources needed
    
    PATIENT:
    {patient_summary}
    
    Respond with JSON:
    {{"esi_level": <1-5>, "triage_reason": "<brief clinical justification>"}}
""").strip()

TASK2_PROMPT_TEMPLATE = textwrap.dedent("""
    TASK: Differential Diagnosis
    
    Generate a ranked differential diagnosis list (most likely first), identify red flags,
    and recommend additional diagnostic tests.
    
    PATIENT:
    {patient_summary}
    
    Respond with JSON:
    {{
      "diagnoses": ["<most likely>", "<2nd>", "<3rd>"],
      "red_flags": ["<finding1>", "<finding2>", ...],
      "recommended_tests": ["<test1>", "<test2>", ...]
    }}
""").strip()

TASK3_PROMPT_TEMPLATE = textwrap.dedent("""
    TASK: Treatment Planning (Safety-Critical)
    
    Propose a specific treatment for this patient. You MUST check:
    1. Patient allergies — NEVER prescribe allergens
    2. Drug-condition contraindications
    3. Renal/hepatic dose adjustments if applicable
    
    DIAGNOSIS CONTEXT: {diagnosis_context}
    
    PATIENT:
    {patient_summary}
    
    Respond with JSON:
    {{
      "diagnosis": "<working diagnosis>",
      "drug_name": "<specific drug or intervention>",
      "dose_mg": <numeric dose or null if not applicable>,
      "route": "<IV|PO|IM|Inhaled|Other>",
      "rationale": "<clinical reasoning including safety checks>"
    }}
""").strip()


def build_patient_summary(obs_data: Dict[str, Any]) -> str:
    """Convert observation data into a clear clinical summary string."""
    patient = obs_data.get("patient", {})
    vitals = patient.get("vitals", {})
    labs = patient.get("labs", [])
    
    lines = [
        f"Patient: {patient.get('age')}yo {patient.get('sex')}",
        f"Chief Complaint: {patient.get('chief_complaint', 'N/A')}",
        f"History: {patient.get('history', 'N/A')}",
        f"Allergies: {', '.join(patient.get('allergies', [])) or 'NKDA'}",
        f"Current Medications: {', '.join(patient.get('current_medications', [])) or 'None'}",
        f"Known Conditions: {', '.join(patient.get('known_conditions', [])) or 'None'}",
        "",
        "VITALS:",
        f"  HR: {vitals.get('heart_rate', 'N/A')} bpm",
        f"  BP: {vitals.get('systolic_bp', 'N/A')}/{vitals.get('diastolic_bp', 'N/A')} mmHg",
        f"  RR: {vitals.get('respiratory_rate', 'N/A')} breaths/min",
        f"  Temp: {vitals.get('temperature_c', 'N/A')} °C",
        f"  SpO2: {vitals.get('spo2_pct', 'N/A')}%",
        f"  GCS: {vitals.get('gcs', 'N/A')}",
        f"  Pain: {vitals.get('pain_score', 'N/A')}/10",
    ]
    
    if labs:
        lines.append("")
        lines.append("LABS:")
        for lab in labs:
            critical_marker = " ⚠ CRITICAL" if lab.get("is_critical") else ""
            lines.append(
                f"  {lab.get('test_name')}: {lab.get('value')} {lab.get('unit')} "
                f"(ref: {lab.get('reference_range')}){critical_marker}"
            )
    
    return "\n".join(lines)


def build_prompt(task: str, obs_data: Dict[str, Any]) -> str:
    """Build the appropriate model prompt based on task."""
    patient_summary = build_patient_summary(obs_data)
    
    if task == "vital-triage":
        return TASK1_PROMPT_TEMPLATE.format(patient_summary=patient_summary)
    elif task == "differential-diagnosis":
        return TASK2_PROMPT_TEMPLATE.format(patient_summary=patient_summary)
    elif task == "treatment-safety":
        context = obs_data.get("context", {}).get("diagnosis_context", "See clinical context above")
        return TASK3_PROMPT_TEMPLATE.format(
            patient_summary=patient_summary,
            diagnosis_context=context,
        )
    return f"Task: {task}\nPatient: {patient_summary}"


def get_model_action(client: OpenAI, task: str, obs_data: Dict[str, Any]) -> Dict[str, Any]:
    """Call the LLM and parse the action JSON."""
    user_prompt = build_prompt(task, obs_data)
    
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            stream=False,
        )
        raw = (completion.choices[0].message.content or "").strip()
        
        # Strip markdown code fences if present
        if raw.startswith("```"):
            raw = raw.split("```")[1]
            if raw.startswith("json"):
                raw = raw[4:]
        raw = raw.strip()
        
        action_dict = json.loads(raw)
        return action_dict
    except json.JSONDecodeError as e:
        print(f"[DEBUG] JSON parse error: {e}, raw={raw[:200]}", flush=True)
        # Return safe fallback for each task
        if task == "vital-triage":
            return {"esi_level": 3, "triage_reason": "Unable to parse model output — defaulting to ESI-3"}
        elif task == "differential-diagnosis":
            return {"diagnoses": ["Unknown"], "red_flags": [], "recommended_tests": []}
        else:
            return {"diagnosis": "Unknown", "drug_name": "IV Normal Saline",
                    "dose_mg": None, "route": "IV", "rationale": "Supportive care"}
    except Exception as e:
        print(f"[DEBUG] Model request failed: {e}", flush=True)
        if task == "vital-triage":
            return {"esi_level": 3, "triage_reason": "model error — default ESI-3"}
        elif task == "differential-diagnosis":
            return {"diagnoses": ["Unspecified"], "red_flags": [], "recommended_tests": []}
        else:
            return {"diagnosis": "Unknown", "drug_name": "Observation",
                    "dose_mg": None, "route": "Other", "rationale": "Model error"}


# ── HTTP helpers ──────────────────────────────────────────────────────────────

def http_reset(task: str) -> Dict[str, Any]:
    resp = httpx.post(f"{MEDTRIAGE_URL}/reset", json={"task": task}, timeout=30)
    resp.raise_for_status()
    return resp.json()


def http_step(action_dict: Dict[str, Any]) -> Dict[str, Any]:
    resp = httpx.post(f"{MEDTRIAGE_URL}/step", json=action_dict, timeout=30)
    resp.raise_for_status()
    return resp.json()


# ── Main episode runner ───────────────────────────────────────────────────────

def run_task(client: OpenAI, task: str) -> Dict[str, Any]:
    """Run a single task episode. Returns summary dict."""
    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False

    log_start(task=task, env=BENCHMARK, model=MODEL_NAME)

    try:
        # Reset
        reset_result = http_reset(task)
        obs_data = reset_result.get("observation", {})
        done = reset_result.get("done", False)

        # Single-turn task: one step per episode
        for step in range(1, MAX_STEPS + 1):
            if done:
                break

            action_dict = get_model_action(client, task, obs_data)
            action_str = json.dumps(action_dict, separators=(",", ":"))

            step_result = http_step(action_dict)
            reward = float(step_result.get("reward", 0.0))
            done = step_result.get("done", True)
            error = step_result.get("observation", {}).get("last_action_error")

            rewards.append(reward)
            steps_taken = step
            obs_data = step_result.get("observation", {})

            log_step(step=step, action=action_str, reward=reward, done=done, error=error)

            if done:
                break

        score = sum(rewards) / max(len(rewards), 1)
        score = min(max(score, 0.0), 1.0)
        success = score >= SUCCESS_THRESHOLD

    except Exception as e:
        print(f"[DEBUG] Task '{task}' error: {e}", flush=True)
        score = 0.0
        success = False

    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return {"task": task, "score": score, "success": success, "steps": steps_taken, "rewards": rewards}


def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
    
    print(f"[DEBUG] MedTriage inference starting", flush=True)
    print(f"[DEBUG] Server: {MEDTRIAGE_URL}", flush=True)
    print(f"[DEBUG] Model: {MODEL_NAME}", flush=True)
    print(f"[DEBUG] Tasks: {TASKS}", flush=True)
    
    all_results = []
    for task in TASKS:
        result = run_task(client, task)
        all_results.append(result)
        print("", flush=True)  # blank line between tasks

    # Summary
    avg_score = sum(r["score"] for r in all_results) / len(all_results)
    print(f"[DEBUG] === SUMMARY ===", flush=True)
    for r in all_results:
        status = "✓" if r["success"] else "✗"
        print(f"[DEBUG] {status} {r['task']}: score={r['score']:.3f}", flush=True)
    print(f"[DEBUG] Average score: {avg_score:.3f}", flush=True)


if __name__ == "__main__":
    main()