inference.py

"""
Module: inference.py
Purpose: Root-level baseline inference script for Medical Triage Assistant.
Part of: Medical Triage Assistant — OpenEnv Round 1
Author: Team Squirrel

Overview:
    MANDATORY: Must be named inference.py and placed in project root.
    This script runs a baseline LLM agent against all 3 tasks in the
    TriageFlow environment. It reads API credentials from environment
    variables, uses the OpenAI client for LLM inference, and emits
    structured stdout logs in the exact [START]/[STEP]/[END] format
    required by the hackathon evaluator.

Dependencies:
    - openai: OpenAI client for LLM inference
    - os: Environment variable access
    - json: Action parsing

Usage:
    API_BASE_URL=<url> MODEL_NAME=<model> HF_TOKEN=<token> python inference.py
"""

import os
import sys
import json
import traceback

from openai import OpenAI

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from triage_flow.environment import TriageEnvironment
from triage_flow.graders import grade_task
from models import TriageAction, PriorityLevel


# ============================================================================
# Required Environment Variables
# ============================================================================

API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") or ""
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")

MAX_STEPS = 20  # Per task — well within 20 min limit
BENCHMARK = "triage-flow"
TASKS = ["basic-triage", "incomplete-records-triage", "mass-casualty-triage"]


# ============================================================================
# Stdout Logging — EXACT FORMAT, NO DEVIATION
# ============================================================================

def log_start(task: str, env: str, model: str):
    """Emit [START] line. Must be exactly one line."""
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error=None):
    """Emit [STEP] line. reward to 2 decimal places, done/success lowercase."""
    done_str = str(done).lower()
    error_str = str(error) if error else "null"
    print(
        f"[STEP]  step={step} action={action} reward={reward:.2f} "
        f"done={done_str} error={error_str}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: list):
    """Emit [END] line. Always emitted, even on exception."""
    success_str = str(success).lower()
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END]   success={success_str} steps={steps} score={score:.2f} "
        f"rewards={rewards_str}",
        flush=True,
    )


# ============================================================================
# LLM Agent
# ============================================================================

SYSTEM_PROMPT = """You are an AI triage assistant at a hospital emergency department intake desk.
Patients arrive one at a time. For each patient you must:

1. CLASSIFY the patient: immediate, urgent, less_urgent, non_urgent, or escalate
2. OUTPUT the full reordered priority queue of ALL patients seen so far

Classification guidelines:
- immediate: Life-threatening (chest pain + abnormal vitals, hemorrhagic shock, etc.)
- urgent: Serious but stable (low SpO2, respiratory distress, neurological signs)
- less_urgent: Needs care but can wait (lacerations, possible fractures, stable chronic conditions)
- non_urgent: Minor issues (scrapes, mild cold, anxiety-related symptoms)
- escalate: Patient has MISSING critical data (e.g., no vitals). Send to human. Patient is REMOVED from queue.

Queue ordering rules:
- All IMMEDIATE patients first, then URGENT, then LESS_URGENT, then NON_URGENT
- Within the same priority, order by clinical severity
- Escalated patients are REMOVED from the queue entirely

IMPORTANT: If a patient has missing vitals or incomplete critical data, you MUST classify as "escalate".

Respond with ONLY a JSON object:
{"classification": "<priority_or_escalate>", "reordered_queue": ["<patient_id_1>", "<patient_id_2>", ...]}

For escalation, the escalated patient should NOT appear in the queue."""


def build_user_prompt(observation_dict: dict) -> str:
    """Build the user prompt from the current observation."""
    patient = observation_dict.get("incoming_patient", {})
    current_queue = observation_dict.get("current_queue", [])
    step_num = observation_dict.get("step_number", 0)
    total = observation_dict.get("total_expected_patients", 0)
    feedback = observation_dict.get("previous_feedback", "")
    task = observation_dict.get("task_name", "")

    # Check for missing data
    missing_info = []
    if patient:
        if patient.get("vitals") is None:
            missing_info.append("VITALS ARE MISSING")
        if patient.get("history") is None:
            missing_info.append("HISTORY IS MISSING")
        if patient.get("allergies") is None:
            missing_info.append("ALLERGIES ARE MISSING")

    missing_str = "\n⚠️ WARNING: " + ", ".join(missing_info) if missing_info else ""

    prompt = f"""Task: {task}
Step: {step_num + 1} of {total}
Previous feedback: {feedback}

Current Queue: {json.dumps(current_queue)}

New Patient Arriving:
- ID: {patient.get('patient_id', 'N/A')}
- Age: {patient.get('age', 'N/A')}
- Chief Complaint: {patient.get('chief_complaint', 'N/A')}
- Symptoms: {', '.join(patient.get('symptoms', []))}
- Vitals: {json.dumps(patient.get('vitals')) if patient.get('vitals') else 'MISSING - CONSIDER ESCALATION'}
- History: {patient.get('history', 'MISSING')}
- Medications: {patient.get('medications', 'MISSING')}
- Allergies: {patient.get('allergies', 'MISSING')}
- Data Complete: {patient.get('info_complete', 'N/A')}{missing_str}

Classify this patient and output the full reordered queue. Respond with ONLY a JSON object."""

    return prompt


def parse_llm_action(response_text: str, current_queue: list, patient_id: str) -> TriageAction:
    """
    Parse the LLM's response into a TriageAction.

    Falls back to a safe default if parsing fails.
    """
    try:
        text = response_text.strip()
        # Handle markdown code blocks
        if "```json" in text:
            text = text.split("```json")[1].split("```")[0].strip()
        elif "```" in text:
            text = text.split("```")[1].split("```")[0].strip()

        # Find JSON object
        start = text.find("{")
        end = text.rfind("}") + 1
        if start >= 0 and end > start:
            text = text[start:end]

        data = json.loads(text)

        classification = data.get("classification", "less_urgent").lower().strip()
        queue = data.get("reordered_queue", current_queue)

        # Validate classification
        valid_classes = {"immediate", "urgent", "less_urgent", "non_urgent", "escalate"}
        if classification not in valid_classes:
            classification = "less_urgent"

        return TriageAction(
            classification=classification,
            reordered_queue=queue,
        )
    except Exception:
        # Fallback: classify as less_urgent, append to queue
        fallback_queue = list(current_queue) + [patient_id]
        return TriageAction(
            classification="less_urgent",
            reordered_queue=fallback_queue,
        )


def get_llm_action(client: OpenAI, observation_dict: dict, current_queue: list, patient_id: str) -> TriageAction:
    """Get an action from the LLM for the current observation."""
    user_prompt = build_user_prompt(observation_dict)

    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.1,
            max_tokens=300,
            stream=False,
        )
        text = (completion.choices[0].message.content or "").strip()
        if not text:
            text = '{"classification": "less_urgent", "reordered_queue": []}'
        return parse_llm_action(text, current_queue, patient_id)
    except Exception as exc:
        print(f"[DEBUG] Model request failed: {exc}", flush=True)
        # Fallback action
        fallback_queue = list(current_queue) + [patient_id]
        return TriageAction(
            classification="less_urgent",
            reordered_queue=fallback_queue,
        )


# ============================================================================
# Task Runner
# ============================================================================

def run_task(task_name: str, client: OpenAI):
    """
    Run a single task against the environment.

    Always emits [START] and [END] lines, even on exception.
    """
    env = TriageEnvironment()
    rewards = []
    steps_taken = 0
    score = 0.0
    success = False

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

    try:
        obs = env.reset(task_name=task_name)
        current_queue = list(obs.current_queue)

        for step in range(1, MAX_STEPS + 1):
            if obs.done:
                break

            patient = obs.incoming_patient or {}
            patient_id = patient.get("patient_id", "")

            obs_dict = {
                "incoming_patient": patient,
                "current_queue": current_queue,
                "step_number": obs.step_number,
                "total_expected_patients": obs.total_expected_patients,
                "previous_feedback": obs.previous_feedback or "",
                "task_name": obs.task_name,
            }

            action = get_llm_action(client, obs_dict, current_queue, patient_id)
            obs = env.step(action)

            reward = obs.reward if obs.reward is not None else 0.0
            done = obs.done
            error = None

            reward_val = round(reward, 2)
            rewards.append(reward_val)
            steps_taken = step

            # Update current queue from observation
            current_queue = list(obs.current_queue)

            # Build action string for logging
            action_str = f"{action.classification}({patient_id})[queue={len(action.reordered_queue)}]"

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

            if done:
                break

        # Compute final score using the grader
        state_dict = env.state.model_dump()
        score = grade_task(task_name, state_dict)
        score = round(score, 2)
        success = score >= 0.5

    except Exception as exc:
        print(f"[DEBUG] Task {task_name} error: {exc}", flush=True)
        traceback.print_exc(file=sys.stderr)
        score = 0.0
        success = False

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


# ============================================================================
# Main Entry Point
# ============================================================================

if __name__ == "__main__":
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    for task in TASKS:
        run_task(task, client)