README.md

---
title: TriageFlow
emoji: 🏥
colorFrom: red
colorTo: blue
sdk: docker
app_port: 8000
tags:
  - openenv
pinned: false
---

# TriageFlow: AI-Assisted Patient Intake Prioritization

[![OpenEnv](https://img.shields.io/badge/OpenEnv-compatible-brightgreen)](https://github.com/meta-pytorch/OpenEnv)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Docker](https://img.shields.io/badge/Docker-ready-blue)](https://www.docker.com/)

---

## Motivation

When emergency departments are overwhelmed, the initial intake assessment — deciding who needs immediate attention and who can safely wait — becomes a critical cognitive load for clinical staff.

TriageFlow simulates the intake desk of a busy hospital emergency department. Patients arrive **one at a time**. For each patient, the AI agent must:

1. **Classify** the patient's urgency (IMMEDIATE / URGENT / LESS_URGENT / NON_URGENT / ESCALATE)
2. **Output the full reordered priority queue** — the correct ordering of all patients seen so far

This dual-output design tests both clinical judgment (is this patient critical?) and queue management (where does this patient rank against everyone else?).

**Why this matters:** Training AI agents to assist with intake prioritization — flagging obvious high-acuity cases and managing queue state — could reduce cognitive load on clinical staff. This environment provides a safe, reproducible benchmark for evaluating such agents.

**What the agent simulates:** The role of a triage nurse or intake coordinator, NOT a doctor. This is an operational, administrative, safety-critical workflow — not diagnosis.

---

## Environment Overview

### Episode Lifecycle

```
reset(task_name) → first patient presented with empty queue
  ↓
agent observes incoming patient + current queue state
  ↓
agent outputs: classification + full reordered priority queue
  ↓
step(action) → reward computed (classification accuracy + queue correctness + escalation judgement)
  ↓
next patient arrives → repeat
  ↓
all patients seen → episode ends
  ↓
grader scores full trajectory → score in [0.0, 1.0]
```

### What the Agent Can Do
- **Classify patients:** IMMEDIATE, URGENT, LESS_URGENT, NON_URGENT
- **Escalate patients** with missing critical data (sends to human, removes from queue)
- **Reorder the priority queue** after each classification

### What the Agent Must Never Do
- Issue a medical diagnosis
- Prescribe or recommend treatments
- Make clinical judgments about disease causation

---

## Observation Space

| Field | Type | Description |
|---|---|---|
| `incoming_patient` | `dict` | The patient who just arrived (symptoms, vitals, age, complaint, history) |
| `current_queue` | `List[str]` | Patient IDs in current priority order (agent's queue so far) |
| `step_number` | `int` | Current step in the episode |
| `total_expected_patients` | `int` | Total patients the agent will see |
| `previous_feedback` | `str` | Feedback on last classification |
| `task_name` | `str` | Name of the active task |
| `done` | `bool` | Whether the episode has ended |
| `reward` | `float` | Reward for the last action |

**Hidden from agent (grader-only):** Per-step answer keys, full agent trajectory, internal scoring components.

---

## Action Space

Each step, the agent outputs a single JSON object:

```json
{
  "classification": "immediate|urgent|less_urgent|non_urgent|escalate",
  "reordered_queue": ["P001", "P003", "P002"]
}
```

| Field | Type | Description |
|---|---|---|
| `classification` | `str` | Priority level for the incoming patient, or `"escalate"` if data is missing |
| `reordered_queue` | `List[str]` | Full queue of all non-escalated patients in priority order |

**Rules:**
- IMMEDIATE patients go first, then URGENT, LESS_URGENT, NON_URGENT
- Escalated patients are **removed** from the queue entirely
- The queue should contain all patients classified so far (excluding escalated ones)

---

## Reward Function

Three scoring components per step, weighted:

| Component | Weight | What It Measures |
|---|---|---|
| **Classification Accuracy** | 50% | Did the agent assign the correct priority label? |
| **Queue Ordering** | 30% | Does the agent's queue match the ideal queue? (position-weighted) |
| **Escalation Judgment** | 20% | Did the agent correctly identify missing data? |

### Classification Penalties (Asymmetric)

| Mistake | Severity |
|---|---|
| Correct classification | Full credit |
| Off by one level (e.g., URGENT instead of IMMEDIATE) | Small penalty |
| Off by two levels (e.g., LESS_URGENT instead of IMMEDIATE) | Heavy penalty |
| IMMEDIATE classified as NON_URGENT | **Maximum penalty** (patient could die) |
| NON_URGENT classified as IMMEDIATE | Mild penalty (wastes resources but safe) |
| Correctly escalating incomplete patient | Full credit |
| Failing to escalate missing data (guessing) | Heavy penalty |
| Escalating complete patient unnecessarily | Moderate penalty |

**Philosophy:** Under-triaging dangerous patients is punished far harder than over-triaging safe patients.

---

## Tasks

### Task 1: Basic Triage (Easy) — 3 patients
- **Scenario:** 3 patients with complete records and clearly distinct urgency levels
- **Patients:** Chest pain with cardiac history (IMMEDIATE), sprained ankle (NON_URGENT), persistent fever with cough (URGENT)
- **Challenge:** Classify correctly and build the priority queue — no escalations needed
- **Max Steps:** 10

### Task 2: Incomplete Records Triage (Medium) — 5 patients (3 + 1 escalation + 1 re-entry)
- **Scenario:** One patient arrives with missing vitals → must be **escalated**. They return later with complete data as a new patient.
- **Challenge:** Escalation judgment + correct re-classification after data completion
- **Max Steps:** 15

### Task 3: Mass Casualty Triage (Hard) — 10 patients (8 + 1 escalation + 1 re-entry)
- **Scenario:** Building collapse. Multiple patients at similar urgency levels. Drug interactions (warfarin + head injury). Missing data requiring escalation.
- **Challenge:** Ranking 8+ patients correctly, escalation judgment under pressure, intra-priority ordering
- **Max Steps:** 20

---

## Grading Formula

```
final_score = (classification_accuracy × 0.5) + (queue_ordering × 0.3) + (escalation_judgment × 0.2)
```

All scores are clipped to `[0.0, 1.0]`. Deterministic: same state always produces the same score.

---

## Setup

### Prerequisites
- Python 3.10+
- Docker (for containerized deployment)

### Installation

```bash
git clone <repo-url>
cd triage-flow
pip install -r requirements.txt
```

---

## Running Locally

### Start the server:
```bash
ENABLE_WEB_INTERFACE=true uvicorn server.app:app --host 0.0.0.0 --port 8000
```

### Run the baseline inference:
```bash
# Set environment variables (Linux/Mac)
export API_BASE_URL="https://router.huggingface.co/v1"
export MODEL_NAME="meta-llama/Llama-3.1-8B-Instruct"
export HF_TOKEN="your-hf-token"

# Windows CMD
set API_BASE_URL=https://router.huggingface.co/v1
set MODEL_NAME=meta-llama/Llama-3.1-8B-Instruct
set HF_TOKEN=your-hf-token

# Run inference
python inference.py
```

### Validate:
```bash
openenv validate
```

---

## Docker

```bash
# Build
docker build -t triage-flow .

# Run
docker run -p 8000:8000 triage-flow

# Test
curl -X POST http://localhost:8000/reset -H "Content-Type: application/json" -d '{}'
```

---

## Baseline Scores (Llama 3.1 8B Instruct)

| Task | Score | Notes |
|---|---|---|
| basic-triage | **0.85** | Correctly identifies IMMEDIATE chest pain case |
| incomplete-records-triage | **0.76** | Successfully escalates patient with missing vitals |
| mass-casualty-triage | **0.86** | Handles escalation + 8-patient queue ordering |

---

## Limitations

- Patient cases are synthetic and simplified — real triage involves much more nuanced assessment
- The environment does not model patient deterioration over time
- Vitals are static snapshots, not continuous monitoring data
- The queue ordering is evaluated against a single "ideal" ordering — real triage may have multiple valid orderings
- Graders focus on priority accuracy and queue correctness, not clinical reasoning quality
- The step budget abstraction does not capture real-time constraints

---

## Repository Structure

```
triage-flow/
├── inference.py              # Baseline inference script (project root)
├── models.py                 # Pydantic models: TriageAction, TriageObservation, TriageState
├── openenv.yaml              # OpenEnv spec metadata
├── Dockerfile                # Container build
├── README.md                 # This file
├── requirements.txt          # Python dependencies
├── pyproject.toml            # Package config
├── client.py                 # OpenEnv client for remote access
├── triage_flow/              # Main environment package
│   ├── __init__.py
│   ├── environment.py        # Core env: reset(), step(), state
│   ├── tasks.py              # Task definitions with per-step answer keys
│   ├── graders.py            # Deterministic 3-component grading
│   └── reward.py             # Per-step reward (classification + queue + escalation)
├── server/
│   ├── __init__.py
│   ├── app.py                # FastAPI server (OpenEnv create_app)
│   ├── ui.py                 # Custom Gradio clinical dashboard
│   └── triage_flow_environment.py  # Server-side environment re-export
└── .github/
    └── workflows/
        └── sync_to_hub.yml   # GitHub Actions → HF Spaces sync
```

---

## License

MIT

---

*Built by Team Squirrel for OpenEnv Round 1 Hackathon.*