HiberNET/drug-interaction-checker
0
1# Drug Interaction Environment — Complete Codebase Explanation2 3## 📋 Executive Summary4 5This is a **clinical pharmacology RL environment** built on OpenEnv that trains and evaluates LLM agents on their ability to identify dangerous **drug-drug interactions** from patient medication lists. The agent acts as a clinical pharmacist, flagging pairs of drugs that interact dangerously, then receives rewards based on accuracy of the severity classification and recommended clinical action.6 7**Framework**: OpenEnv (OpenAI-style environment interface) + FastAPI server + LLM inference client8 9---10 11## 🏗️ Overall Architecture12 13```14┌─────────────────────────────────────────────────────────────────┐15│ FASTAPI SERVER (app.py) │16│ ┌─────────────┐ ┌──────────────┐ ┌────────────────────────┐ │17│ │ /reset │ │ /step │ │ /state │ │18│ │ Initialize │ │ Process │ │ Get episode internal │ │19│ │ new episode │ │ LLM action │ │ state (for grading) │ │20│ └─────────────┘ └──────────────┘ └────────────────────────┘ │21│ ▲ ▲ ▲ │22│ │ │ │ │23│ ┌────────▼────────────────▼────────────────────▼──────────┐ │24│ │ DrugInteractionEnvironment (drug_interaction_env.py) │ │25│ │ │ │26│ │ • reset(task_level) → Initialize episode │ │27│ │ • step(action) → Process one action │ │28│ │ • validate(action) → Check action validity │ │29│ │ • calculate_reward() → Score the action │ │30│ │ • state / observation → Multi-level state tracking │ │31│ └──────────────────────────────────────────────────────────┘ │32│ ▲ ▲ ▲ │33│ │ │ │ │34│ ┌────────┴─────────┬──────┴──────────┬────────┴─────────────┐ │35│ │ PATIENTS │ DRUG_DATABASE │ MODELS (Pydantic) │ │36│ │ ├─ easy (1 int) │ 35 interactions │ ├─ Action │ │37│ │ ├─ medium (3) │ 12 severe │ ├─ Observation │ │38│ │ └─ hard (5) │ 13 moderate │ ├─ State │ │39│ │ │ 10 mild │ └─ PredictionEntry │ │40│ └─────────────────┴─────────────────┴──────────────────────┘ │41└─────────────────────────────────────────────────────────────────┘42 ▲ ▲43 │ │44 │ HTTP (JSON) │45 │ │46 │ LLM Client (inference.py) │47 │ • OpenAI API calls │48 │ • Structured system prompt │49 │ • Multi-turn conversation loop │50 │ • Episode grading (grader.py) │51 │ │52 └────────────────────────────────────────────53```54 55---56 57## 🔧 Core Components58 59### 1. **DrugInteractionEnvironment** (`drug_interaction_environment.py`)60 61The main RL environment implementing the OpenAI Gym-like interface:62 63#### **State Variables** (initialized in `__init__`)64```python65self.task_level: str # "easy", "medium", or "hard"66self.patient: dict # Patient profile (age, conditions, meds)67self.ground_truth_keys: set # Canonical interaction pairs to find68self.attempted_keys: set # All pairs agent has tried flagging69self.identified_pairs: set # Correct pairs identified by agent70self.perfectly_completed: set # Pairs with perfect severity + action71self.predictions: dict # Detailed prediction records72self.flags_raised: list # Log of all actions taken73self.step_count: int # Current step number74self.max_steps: int # Episode step budget75self.episode_reward: float # Cumulative reward76self.done: bool # Episode terminal state77```78 79#### **Key Methods**80 81**`reset(task_level: str) → dict`**82- Loads a patient scenario from `PATIENTS[task_level]`83- Derives `ground_truth_keys` by scanning all medication pairs against `DRUG_INTERACTIONS`84- Sets `max_steps = len(ground_truth_keys) * 3` (generous budget)85- Returns initial observation dict86 87**`validate(action: dict) → bool | None`**88- Returns `True` if pair is in DRUG_INTERACTIONS (valid)89- Returns `False` if drug names invalid OR pair not in database (phantom)90- Returns `None` if pair already attempted (duplicate)91- **Normalization**: `key = tuple(sorted([drug_a.lower(), drug_b.lower()]))`92 93**`calculate_reward(key, severity, action) → float`**94Reward breakdown:95| Component | Score | Condition |96|----|----|----|97| **Base (pair identified)** | +0.4 | Always for valid pair |98| **Severity match** | +0.2 | Predicted == ground truth |99| **Severity mismatch (GT severe)** | -0.2 | Predicted ≠ severe (when GT is severe) |100| **Severity mismatch (GT moderate)** | -0.1 | Predicted ≠ moderate |101| **Severity mismatch (GT mild)** | -0.05 | Predicted ≠ mild |102| **Action match** | +0.2 | Predicted action == ground truth |103| **Action mismatch** | -0.1 | Predicted action ≠ ground truth |104 105**Perfect reward**: 0.4 + 0.2 + 0.2 = **0.8** (when severity and action both correct)106 107**`step(action: dict) → (observation, reward, done, state)`**108Main game loop:1091. If action is `DONE`, apply termination penalty and end episode1102. Validate the flagged pair1113. If invalid/phantom: reward = -0.31124. If duplicate: reward = -0.051135. If valid: calculate full reward + add to flags_raised1146. Update episode_reward1157. **Early termination**: If perfectly_completed_pairs == ground_truth_keys, set done=True1168. **Budget exhaustion**: If step_count >= max_steps, apply penalty and end1179. Return observation (what agent sees), step reward, done flag, internal state118 119**`_apply_termination_penalty() → float`**120Penalizes unidentified interactions by severity:121- Severe unidentified: -0.4 each122- Moderate unidentified: -0.3 each123- Mild unidentified: -0.2 each124 125---126 127### 2. **Drug Database** (`drug_database.py`)128 129Hardcoded dictionary of **35 real FDA drug interactions**:130 131```python132DRUG_INTERACTIONS: dict[tuple[str, str], dict] = {133 ("aspirin", "warfarin"): {134 "severity": "severe",135 "action": "replace_drug",136 "explanation": "Increased bleeding risk — dual antiplatelet + anticoagulant"137 },138 # ... 34 more entries139}140```141 142**Distribution**:143- 12 severe (high clinical risk)144- 13 moderate (medium risk, requires monitoring)145- 10 mild (low risk, mostly monitoring)146 147**Key Lookup** is always deterministic:148```python149key = tuple(sorted([drug_a.lower(), drug_b.lower()]))150interaction = DRUG_INTERACTIONS[key]151```152 153---154 155### 3. **Patient Scenarios** (`patients.py`)156 157Three difficulty levels with carefully curated medication lists:158 159| Level | Meds | Interactions | Structure |160|----|----|----|----|161| **easy** | 6 | 1 (severe) | Simple starting point |162| **medium** | 10 | 3 (1 severe, 1 moderate, 1 mild) | Curriculum step |163| **hard** | 15 | 5 (mixed severity, avoiding redundant pairs) | Full complexity |164 165**Critical Design**: Medications are chosen such that **only the intended interactions fire**. "Filler" drugs (losartan, gabapentin, acetaminophen, etc.) exist in the medication list but don't interact with anything else in that patient's list.166 167Example (easy):168```python169"easy": {170 "medications": ["warfarin", "aspirin", "losartan", "amlodipine", "gabapentin", "pantoprazole"],171 # Only interaction: (aspirin, warfarin) → severe/replace_drug172 # "losartan", "amlodipine" etc. = safe filler drugs for that combination173}174```175 176**Ground Truth Derivation** (runtime):177```python178meds = patient["medications"]179for a, b in combinations(meds, 2):180 key = tuple(sorted([a.lower(), b.lower()]))181 if key in DRUG_INTERACTIONS:182 ground_truth_keys.add(key) # This is the true positive set183```184 185---186 187### 4. **Pydantic Models** (`models.py`)188 189Structured data models for serialization:190 191**`DrugInteractionAction`** (agent output):192```python193action_type: "flag_interaction" | "DONE"194drug_a: str (optional, only for flag_interaction)195drug_b: str (optional, only for flag_interaction)196severity: "mild" | "moderate" | "severe"197suggested_action: "monitor" | "reduce_dose" | "replace_drug"198```199 200**`DrugInteractionObservation`** (what agent sees each step):201```python202patient_id: str203age: int204conditions: list[str]205medications: list[str]206flags_raised_so_far: list[FlagEntry] # History of all flags so far207steps_remaining: int208```209 210**`PredictionEntry`** (stored for each interaction flagged):211```python212key: str # "('drug_a', 'drug_b')"213predicted_severity: str214predicted_action: str215ground_truth_severity: str216ground_truth_action: str217reward_received: float218perfectly_completed: bool219```220 221**`DrugInteractionState`** (full internal state for grading):222```python223patient_id: str224step_count: int225task_level: str226attempted_keys: list[str]227identified_pairs: list[str]228perfectly_completed_pairs: list[str]229predictions: dict230done: bool231```232 233---234 235## 🎮 RL Environment Design: The Game Loop236 237### **Episode Lifecycle**238 239```240┌─── RESET ───────────────────────────────────────────┐241│ Client calls: POST /reset {"task_level": "easy"} │242│ Environment: │243│ • Loads patient scenario & meds │244│ • Derives ground_truth (C(n,2) vs DRUG_DATABASE) │245│ • Sets max_steps = |ground_truth| * 3 │246│ • Returns: initial observation │247└──────────────────────────────────────────────────────┘248 ▼249┌─── STEP LOOP ────────────────────────────────────────┐250│ Repeats until done=True: │251│ │252│ 1. Agent receives observation (patient + history) │253│ 2. Agent outputs JSON action via LLM │254│ 3. POST /step with action │255│ 4. Environment validates & scores action: │256│ - validate() → True/False/None │257│ - calculate_reward() if True │258│ - Check termination conditions │259│ 5. Return (obs, reward, done, state) │260│ │261│ Termination conditions: │262│ • Agent sends DONE action │263│ • All pairs perfectly identified │264│ • Step budget exhausted (step_count >= max) │265└──────────────────────────────────────────────────────┘266 ▼267┌─── GRADING ───────────────────────────────────────┐268│ Client calls: GET /state → final state dict │269│ Grader (grader.py) computes normalized score: │270│ │271│ total_reward = sum(all step rewards) │272│ - penalties for unidentified pairs │273│ max_possible = |ground_truth| * 0.8 │274│ final_score = clamp(total_reward / max_possible) │275│ Range: [0.0, 1.0] │276└────────────────────────────────────────────────────┘277```278 279### **Action Processing Flow**280 281```282Agent sends: {"action_type": "flag_interaction", "drug_a": "...", "drug_b": "...", "severity": "...", "suggested_action": "..."}283 ▼284 validate(action)285 / | \286 True / | \ False/None287 / | \288 Valid Duplicate(None) Invalid289 Pair (-0.05) (-0.3)290 |291 calculate_reward()292 |293 Base: +0.4294 + Severity match: +0.2 or -(0.05-0.2)295 + Action match: +0.2 or -0.1296 |297 Step reward ∈ [-0.3, +0.8]298 |299 Update:300 • episode_reward += step_reward301 • identified_pairs.add(key) if valid302 • perfectly_completed_pairs.add(key) if perfect303 • predictions[key] = {...}304 • flags_raised_so_far.append(...)305 |306 Check termination:307 • perfectly_completed? → done=True308 • max_steps reached? → apply penalty, done=True309 • else: done=False, continue loop310```311 312---313 314## 🏥 Core Assumptions315 316### **1. Environment Setup Assumptions**317 318| Assumption | Impact | Justification |319|----|----|---|320| **Medications are case-insensitive** | All lookups normalized to lowercase | Universal clinical practice |321| **Pair order doesn't matter** | Always sort to (min, max) tuple | Symmetric interaction property |322| **Only one patient per episode** | Simplified state management | Focus on single task |323| **No drug quantity/dosage in state** | Simplified action space | Agent learns interaction existence, not dosing |324| **Interactions are static** | DRUG_INTERACTIONS never changes | Real-world drugs have fixed interactions |325 326### **2. Reward Structure Assumptions**327 328| Assumption | Impact | Consequence |329|----|----|---|330| **Base reward (0.4) for identifying pair** | Incentivizes finding all pairs | Agent must explore medication combinations |331| **Severity scoring is asymmetric** | Severe misclassification: -0.2, mild: -0.05 | Penalizes underestimating serious interactions |332| **Action scoring is symmetric** | Right action: +0.2, wrong: -0.1 | Encourages correct clinical judgment |333| **Perfect completion ends episode immediately** | No further steps allowed | Avoids unnecessary exploration noise |334| **Step budget = 3x ground truth** | Allows ~3 wrong attempts per interaction | Reasonable exploration allowance |335 336### **3. Ground Truth Derivation Assumptions**337 338| Assumption | Implication |339|----|---|340| **Ground truth = all C(n,2) pairs in DRUG_INTERACTIONS** | Any pair not in DB is a phantom (invalid) |341| **Patient medication list is exhaustive** | Only medications in list can be checked |342| **Drug names match exactly** | "ibuprofen" ≠ "Ibuprofen" initially; case-normalized |343| **No active metabolite interactions** | Only direct drug-drug interactions considered |344| **Interactions are bidirectional** | (A, B) = (B, A) after sorting |345 346### **4. Task Design Assumptions**347 348| Level | # Pairs | Assumption | Implication |349|----|----|----|----|350| **easy** | 1 | Simple pattern recognition | Agent learns basic interaction flagging |351| **medium** | 3 | Curriculum learning | Diverse severity exposure (1 each level) |352| **hard** | 5 | Complex search | Real drug cocktail complexity |353 354**Assumption**: Medications in hard don't create phantom interactions. (e.g., hard patient excludes warfarin+amiodarone even though both exist, to avoid unintended pair count.)355 356### **5. Episode Termination Assumptions**357 358| Condition | Reward | Assumption |359|----|----|----|360| **Agent sends DONE** | Termination penalty only | Assumes agent honestly signals completion |361| **Perfect completion** | No additional penalty | Assumes agent explores optimally |362| **Max steps reached** | Termination penalty applied | Assumes agent ran out of budget |363 364**Penalty Logic** (severity-weighted):365- Severe missed: -0.4 per pair366- Moderate missed: -0.3 per pair367- Mild missed: -0.2 per pair368 369### **6. LLM Agent Assumptions** (`inference.py`)370 371| Assumption | Impact |372|----|---|373| **LLM can output valid JSON** | System prompt strongly constrains output format |374| **LLM respects medication names** | "Exactly as spelled in medications list" |375| **LLM has pharmacological knowledge** | Can predict severity + action correctly |376| **LLM follows one-pair-per-step rule** | Won't flag multiple pairs in single action |377| **LLM's reasoning aligns with clinical guidelines** | FDA + pharmacology references used |378 379---380 381## 📡 API Interface382 383### **FastAPI Endpoints** (`app.py`)384 385```python386POST /reset387├─ Request: {"task_level": "easy|medium|hard"}388└─ Response: {389 "patient_id": "P001",390 "age": 67,391 "conditions": ["hypertension", "type2_diabetes"],392 "medications": ["warfarin", "aspirin", ...],393 "flags_raised_so_far": [],394 "steps_remaining": 3395}396 397POST /step398├─ Request: {399│ "action_type": "flag_interaction|DONE",400│ "drug_a": "warfarin",401│ "drug_b": "aspirin",402│ "severity": "severe",403│ "suggested_action": "replace_drug"404}405└─ Response: {406 "observation": {...},407 "reward": 0.8,408 "done": false,409 "state": {...}410}411 412GET /state413├─ Returns internal state (for grading/debugging)414└─ Includes: episode_score, predictions, attempted_keys, etc.415 416GET /health417└─ Returns: {"status": "ok"}418```419 420---421 422## 🎯 Episode Grading423 424### **grader.py Logic**425 426```python427def grade_episode(task_level: str, final_state: dict) -> float:428 # 1. Get expected # of true pairs for this level429 n_true_pairs = TASK_TRUE_PAIRS[task_level] # 1, 3, or 5430 max_possible = n_true_pairs * 0.8 # Maximum achievable reward431 432 # 2. Sum all step rewards433 total_reward = sum(p["reward_received"] for p in predictions.values())434 435 # 3. Apply termination penalties for unidentified pairs436 unidentified = ground_truth_keys - identified_pairs437 for key in unidentified:438 severity = DRUG_INTERACTIONS[key]["severity"]439 if severity == "severe":440 total_reward -= 0.4441 elif severity == "moderate":442 total_reward -= 0.3443 else:444 total_reward -= 0.2445 446 # 4. Normalize to [0.0, 1.0]447 return max(0.0, min(1.0, total_reward / max_possible))448```449 450**Score Interpretation**:451- **1.0**: All interactions identified perfectly (severity + action both correct)452- **0.8**: All interactions identified but some severity/action misclassifications453- **0.0-0.4**: Some interactions missed or significant misclassifications454- **0.0**: All interactions missed OR severe penalties apply455 456---457 458## 🔄 Inference Loop459 460### **inference.py Execution Flow**461 462```python463for task_level in ["easy", "medium", "hard"]:464 # 1. Reset environment465 observation = env_reset(task_level)466 messages = [467 {"role": "system", "content": SYSTEM_PROMPT},468 {"role": "user", "content": build_user_message(observation)}469 ]470 471 # 2. Multi-turn conversation loop472 while not done:473 # LLM generates next action474 response = client.chat.completions.create(475 model=MODEL_NAME,476 messages=messages,477 temperature=0.1, # Low temperature for deterministic output478 max_tokens=256479 )480 481 # Parse JSON action from response482 action = parse_llm_action(response.text)483 484 # Submit to environment485 result = env_step(action)486 reward = result["reward"]487 done = result["done"]488 observation = result["observation"]489 490 # Log step491 print(f"[STEP] step={step_num} action={action_type} reward={reward}")492 493 # Update conversation for next turn494 messages.append({"role": "assistant", "content": json.dumps(action)})495 if not done:496 messages.append({497 "role": "user",498 "content": build_user_message(observation)499 })500 501 # 3. Grade episode502 state_data = env_state()503 episode_score = grade_episode(task_level, state_data)504 print(f"[END] task={task_level} patient_id={patient_id} episode_score={episode_score}")505```506 507**Key Points**:508- **System prompt** constrains LLM to JSON-only output509- **Temperature = 0.1** for deterministic behavior510- **Multi-turn**: Each LLM response gets added to message history511- **Observation updates**: After each step, agent sees updated flags_raised_so_far512- **Episode score** computed independently by grader module513 514---515 516## 🧪 Testing Strategy517 518### **Unit Tests** (`test_env.py`)519 520Tests environment mechanics:5211. **Validation tests**: Valid pairs, invalid drugs, phantom pairs, duplicates5222. **Reward calculation**: Perfect flag (0.8), severity mismatch, action mismatch5233. **Episode completion**: Easy, medium, hard scenarios5244. **Termination penalties**: DONE action with unidentified pairs525 526### **API Tests** (`test_app.py`)527 528Tests FastAPI endpoints:5291. Health check (GET /health)5302. Reset endpoint (POST /reset)5313. Step endpoint (POST /step with valid action)5324. State endpoint (GET /state)533 534---535 536## 📊 Example Episode Trace537 538```539[START] task=easy patient_id=P001540 541RESET542→ Medications: [warfarin, aspirin, losartan, amlodipine, gabapentin, pantoprazole]543→ Ground truth: [(aspirin, warfarin)]544→ Max steps: 3545→ Initial observation returned546 547STEP 1548Agent action: {"action_type": "flag_interaction", "drug_a": "warfarin", "drug_b": "aspirin", "severity": "severe", "suggested_action": "replace_drug"}549Validation: ✓ Valid pair in DRUG_INTERACTIONS550Reward: 0.4 (base) + 0.2 (severity match) + 0.2 (action match) = 0.8551Perfectly completed? Yes → done=True552[STEP] step=1 action=flag_interaction drug_a=warfarin drug_b=aspirin severity=severe suggested_action=replace_drug reward=0.8553 554[END] task=easy patient_id=P001 episode_score=1.0555 556GRADING557total_reward = 0.8558unidentified = {} (empty)559max_possible = 1 * 0.8 = 0.8560episode_score = 0.8 / 0.8 = 1.0 ✓561```562 563---564 565## 🎓 Key Takeaways566 567| Component | Role | Critical Assumption |568|----|----|---|569| **Environment** | Manages episode state, validates actions, scores | Interactions are deterministic, pair order irrelevant |570| **Database** | Source of truth for drug interactions | Real FDA interactions, static throughout episode |571| **Reward** | Learning signal | Asymmetric severity penalties, base score for identification |572| **Task design** | Curriculum learning | Carefully curated medication lists (no phantom pairs) |573| **Grading** | Evaluate agent performance | Penalty-weighted score normalized to [0, 1] |574| **LLM Interface** | Agent-environment interaction | JSON-constrained output, structured prompting |575 576---577 578## 📚 Dependencies579 580```python581fastapi >= 0.100.0 # API framework582uvicorn >= 0.23.0 # ASGI server583pydantic >= 2.0 # Data validation584openai >= 1.0.0 # LLM client585huggingface_hub # HuggingFace deployment586python >= 3.10 # Required for type hints587```588 589---590 591## 🚀 Deployment592 593- **Docker**: Self-contained image with FastAPI server594- **HuggingFace Spaces**: OpenEnv-compatible deployment via `hf_deploy.py`595- **Local dev**: `uvicorn drug_interaction_env.server.app:app --reload`596 597