inference.py

"""Baseline inference script for the Hospital ED OpenEnv environment.

Hackathon requirements fulfilled by this file:

* File is named ``inference.py`` and lives in the repo root.
* Uses the OpenAI client for all LLM calls.
* Reads credentials from environment variables:
    - ``API_BASE_URL``     — LLM endpoint (default: HF router)
    - ``MODEL_NAME``       — model identifier (default: Qwen2.5-72B-Instruct)
    - ``HF_TOKEN`` / ``API_KEY`` — bearer token
    - ``LOCAL_IMAGE_NAME`` / ``IMAGE_NAME`` — Docker image name for
      :meth:`HospitalEdEnv.from_docker_image`
* Runs 3 tasks with a difficulty progression (easy → medium → hard):
    1. ``normal_day``     — steady-state ED operations (easy)
    2. ``surge``          — COVID-style respiratory surge (medium)
    3. ``mass_casualty``  — trauma burst, rapid triage (hard)
* Emits the exact ``[START] / [STEP] / [END]`` stdout format specified
  in the hackathon "Mandatory Additional Instructions" section.
* Produces per-task scores in [0, 1] (``task_score`` from the env
  metadata), plus a combined aggregate.

Runtime budget: well under 20 minutes on vCPU=2 / 8 GB memory. Each task
runs for at most 100 steps with a short system prompt and a fixed
``max_tokens=32`` response limit.
"""
from __future__ import annotations

import asyncio
import os
import re
import textwrap
import traceback
from typing import List, Optional

from openai import OpenAI

from client import HospitalEdEnv
from models import HospitalAction


# ---------------------------------------------------------------------
# Environment variables
# ---------------------------------------------------------------------
API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or "dummy"
IMAGE_NAME = (
    os.getenv("LOCAL_IMAGE_NAME")
    or os.getenv("IMAGE_NAME")
    or "hospital-ed-env:latest"
)

BENCHMARK = os.getenv("HOSPITAL_ED_BENCHMARK", "hospital_ed")

# Runtime caps — keeps total inference time well under 20 minutes even
# with a slow LLM. 100 env steps × 3 tasks × ~2 s/LLM call ≈ 10 min.
MAX_STEPS_PER_TASK = 100
TEMPERATURE = 0.2
MAX_TOKENS = 32  # we only need a single integer

# Tasks in ascending difficulty, as required by the hackathon rubric.
TASKS = ["normal_day", "surge", "mass_casualty"]


SYSTEM_PROMPT = textwrap.dedent(
    """
    You are the on-duty triage agent for a hospital emergency department.
    You control 20 general beds, 5 ICU beds, and 3 ventilators.

    Each timestep you pick exactly one discrete action from this table:

      0      No-op (wait)
      1-10   Assign waiting-queue patient [idx] to a GENERAL bed
      11-20  Assign waiting-queue patient [idx] to an ICU bed
      21-23  Use ventilator slot [idx] on the most severe ICU patient who
             does not already have one
      24-28  Transfer general bed [idx] patient to an ICU bed
      29-33  Discharge general bed [idx] patient
      34-38  Discharge ICU bed [idx] patient

    Severity scale: 1 = mild, 2 = moderate, 3 = severe, 4 = critical.
    Critical (4) patients die if they wait too long — prioritise them
    into the ICU first. Attach ventilators to severe respiratory cases.

    You will be given the current hospital state and a list of valid
    actions (a subset of 0-38). Reply with ONLY a single integer drawn
    from that valid list. No words, no punctuation, just the number.
    """
).strip()


# ---------------------------------------------------------------------
# Logging helpers — exact stdout format from the hackathon spec.
# ---------------------------------------------------------------------
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"
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"done={str(done).lower()} 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} "
        f"score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ---------------------------------------------------------------------
# LLM policy
# ---------------------------------------------------------------------
def _format_observation(obs, valid_actions: List[int]) -> str:
    """Format a HospitalObservation into a compact text prompt."""
    # Show only queue rows that have a waiting patient (severity > 0).
    queue_lines = []
    for i, row in enumerate(obs.waiting_queue):
        sev = int(row[0]) if len(row) > 0 else 0
        if sev == 0:
            continue
        cond = int(row[1]) if len(row) > 1 else 0
        wait = int(row[2]) if len(row) > 2 else 0
        queue_lines.append(f"  queue[{i}] sev={sev} cond={cond} wait={wait}")
    queue_block = "\n".join(queue_lines) if queue_lines else "  (empty)"

    stats = obs.stats or [0, 0, 0]
    treated = stats[0] if len(stats) > 0 else 0
    deaths = stats[1] if len(stats) > 1 else 0
    waiting = stats[2] if len(stats) > 2 else 0

    return textwrap.dedent(
        f"""
        Hospital state at timestep {obs.time_step}:
          General beds (severity, 0=empty): {list(obs.bed_occupancy)}
          ICU beds     (severity, 0=empty): {list(obs.icu_occupancy)}
          Ventilators (0=free, 1=in use):   {list(obs.ventilator_status)}
        Waiting queue ({waiting} patients):
        {queue_block}
        Cumulative: treated={treated}, deaths={deaths}

        Valid actions right now: {valid_actions}

        Reply with ONE integer from the valid list.
        """
    ).strip()


def _parse_action(
    text: str, valid_actions: List[int], fallback: int = 0
) -> int:
    """Parse the LLM response into an action id, clamped to valid_actions."""
    match = re.search(r"-?\d+", text or "")
    if match is None:
        return fallback
    try:
        a = int(match.group(0))
    except ValueError:
        return fallback
    if a in valid_actions:
        return a
    return fallback if fallback in valid_actions else valid_actions[0]


def _ask_llm(
    client: OpenAI, obs, valid_actions: List[int]
) -> int:
    """Query the LLM for the next action and return a valid integer id."""
    user_prompt = _format_observation(obs, valid_actions)
    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,
        )
        text = (completion.choices[0].message.content or "").strip()
    except Exception as exc:
        print(f"[DEBUG] LLM request failed: {exc}", flush=True)
        text = ""
    fallback = valid_actions[0] if valid_actions else 0
    return _parse_action(text, valid_actions, fallback=fallback)


# ---------------------------------------------------------------------
# Episode runner
# ---------------------------------------------------------------------
async def run_task(
    env: HospitalEdEnv,
    llm: OpenAI,
    task: str,
    seed: int,
) -> tuple[float, int, List[float], bool]:
    """Run one episode for ``task`` and return (score, steps, rewards, success)."""
    log_start(task=task, env=BENCHMARK, model=MODEL_NAME)

    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    last_error: Optional[str] = None
    obs = None

    try:
        result = await env.reset(seed=seed, task=task)
        obs = result.observation

        for step in range(1, MAX_STEPS_PER_TASK + 1):
            if result.done:
                break

            valid_actions = [i for i, m in enumerate(obs.action_mask) if m]
            if not valid_actions:
                valid_actions = [0]

            action_id = _ask_llm(llm, obs, valid_actions)

            try:
                result = await env.step(HospitalAction(action=action_id))
                obs = result.observation
                reward = float(result.reward or 0.0)
                done = bool(result.done)
                last_error = None
            except Exception as step_exc:  # pragma: no cover
                reward = 0.0
                done = True
                last_error = str(step_exc)

            rewards.append(reward)
            steps_taken = step
            log_step(
                step=step,
                action=f"discrete({action_id})",
                reward=reward,
                done=done,
                error=last_error,
            )

            if done:
                break

        # Final normalized task score is a direct Pydantic field on the
        # observation (metadata is stripped by OpenEnv's serializer).
        score = float(getattr(obs, "task_score", 0.0) or 0.0)
        score = max(0.0, min(1.0, score))
        success = score >= 0.5

    except Exception as exc:  # pragma: no cover
        last_error = f"{type(exc).__name__}: {exc}"
        traceback.print_exc()
    finally:
        log_end(
            success=success, steps=steps_taken, score=score, rewards=rewards
        )

    return score, steps_taken, rewards, success


# ---------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------
async def main() -> None:
    llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    # Spin up the OpenEnv Hospital ED server in a Docker container and
    # connect to it. The hackathon grader expects exactly this pattern.
    env = await HospitalEdEnv.from_docker_image(IMAGE_NAME)

    try:
        per_task_scores = []
        for i, task in enumerate(TASKS):
            score, _, _, _ = await run_task(env, llm, task=task, seed=42 + i)
            per_task_scores.append((task, score))
    finally:
        try:
            await env.close()
        except Exception as e:  # pragma: no cover
            print(f"[DEBUG] env.close() error: {e}", flush=True)

    # Aggregate summary printed after all [END] lines so judges can see
    # the composite score at a glance.
    if per_task_scores:
        mean_score = sum(s for _, s in per_task_scores) / len(per_task_scores)
        per_task_str = ", ".join(
            f"{name}={score:.3f}" for name, score in per_task_scores
        )
        print(
            f"[SUMMARY] tasks={len(per_task_scores)} "
            f"mean_score={mean_score:.3f} {per_task_str}",
            flush=True,
        )


if __name__ == "__main__":
    asyncio.run(main())