Spaces:
Sleeping
Sleeping
| """ | |
| PharmaAgent — Clinical Decision RL Environment | |
| app.py: FastAPI server exposing OpenMV-compatible endpoints | |
| """ | |
| import json | |
| import uuid | |
| import time | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import HTMLResponse | |
| import os | |
| from .models import Action, Observation, State | |
| from .environment import PharmaAgentEnvironment | |
| app = FastAPI( | |
| title="PharmaAgent RL Environment", | |
| description="Clinical Decision Agent — Drug Intelligence RL Environment powered by DrugBank", | |
| version="1.0.0", | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| env = PharmaAgentEnvironment() | |
| # Fix #10: session store now tracks last-access timestamps so stale sessions | |
| # can be evicted, preventing unbounded memory growth under sustained load. | |
| SESSION_TTL_SECONDS = 3600 # evict sessions idle for more than 1 hour | |
| sessions: dict = {} # session_id -> state dict | |
| session_touched: dict = {} # session_id -> last access timestamp | |
| def _evict_stale_sessions(): | |
| """Remove sessions that have been idle longer than SESSION_TTL_SECONDS.""" | |
| now = time.time() | |
| stale = [sid for sid, t in session_touched.items() if now - t > SESSION_TTL_SECONDS] | |
| for sid in stale: | |
| sessions.pop(sid, None) | |
| session_touched.pop(sid, None) | |
| def _touch(session_id: str): | |
| session_touched[session_id] = time.time() | |
| # ── OpenMV-compatible endpoints ──────────────────────────────────────── | |
| def reset(session_id: str = None): | |
| """Reset the environment and return the initial observation.""" | |
| _evict_stale_sessions() | |
| if session_id is None: | |
| session_id = str(uuid.uuid4()) | |
| state, obs = env.reset() | |
| sessions[session_id] = state.model_dump() | |
| _touch(session_id) | |
| return { | |
| "session_id": session_id, | |
| "observation": obs.model_dump(), | |
| "state": state.model_dump(), | |
| } | |
| def step(action: Action, session_id: str = None): | |
| """Take one step in the environment.""" | |
| if session_id is None or session_id not in sessions: | |
| raise HTTPException(status_code=400, detail="Invalid or missing session_id. Call /reset first.") | |
| state = State(**sessions[session_id]) | |
| if state.done: | |
| raise HTTPException(status_code=400, detail="Episode is done. Call /reset to start a new episode.") | |
| new_state, obs, reward, done = env.step(state, action) | |
| sessions[session_id] = new_state.model_dump() | |
| if done: | |
| # Fix #10: eagerly evict completed sessions — no need to keep them around | |
| sessions.pop(session_id, None) | |
| session_touched.pop(session_id, None) | |
| else: | |
| _touch(session_id) | |
| return { | |
| "session_id": session_id, | |
| "observation": obs.model_dump(), | |
| "reward": reward, | |
| "done": done, | |
| "state": new_state.model_dump(), | |
| } | |
| def get_state(session_id: str): | |
| """Get the current state of a session.""" | |
| if session_id not in sessions: | |
| raise HTTPException(status_code=404, detail="Session not found.") | |
| _touch(session_id) | |
| return {"session_id": session_id, "state": sessions[session_id]} | |
| def health(): | |
| """Health check endpoint.""" | |
| return { | |
| "status": "ok", | |
| "environment": "PharmaAgent", | |
| "version": "1.0.0", | |
| "active_sessions": len(sessions), | |
| } | |
| def root(): | |
| return { | |
| "name": "PharmaAgent RL Environment", | |
| "description": "Clinical Decision Agent powered by DrugBank", | |
| "endpoints": ["/reset", "/step", "/state", "/health", "/web"], | |
| } | |
| # ── Optional Web UI ──────────────────────────────────────────── | |
| ENABLE_WEB = os.environ.get("ENABLE_WEB_INTERFACE", "true").lower() == "true" | |
| if ENABLE_WEB: | |
| def web_ui(): | |
| """Simple browser UI for manually testing the environment.""" | |
| return """ | |
| <!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"/> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"/> | |
| <title>PharmaAgent — RL Environment Tester</title> | |
| <style> | |
| * { box-sizing: border-box; margin: 0; padding: 0; } | |
| body { font-family: 'Segoe UI', sans-serif; background: #0f0f0f; color: #e0e0e0; min-height: 100vh; } | |
| .container { max-width: 860px; margin: 0 auto; padding: 32px 20px; } | |
| h1 { font-size: 28px; font-weight: 700; margin-bottom: 4px; color: #fff; } | |
| .subtitle { font-size: 13px; color: #666; margin-bottom: 32px; } | |
| .card { background: #1a1a1a; border: 1px solid #2a2a2a; border-radius: 8px; padding: 20px; margin-bottom: 16px; } | |
| .card-title { font-size: 11px; letter-spacing: 2px; text-transform: uppercase; color: #555; margin-bottom: 12px; } | |
| .btn { background: #fff; color: #000; border: none; padding: 10px 22px; font-size: 13px; font-weight: 600; cursor: pointer; border-radius: 4px; margin-right: 8px; } | |
| .btn:hover { background: #ddd; } | |
| .btn-outline { background: transparent; color: #fff; border: 1px solid #444; } | |
| .btn-outline:hover { border-color: #fff; } | |
| select, input { background: #111; border: 1px solid #333; color: #e0e0e0; padding: 8px 12px; font-size: 14px; border-radius: 4px; width: 100%; margin-bottom: 10px; } | |
| .output { background: #111; border: 1px solid #2a2a2a; border-radius: 6px; padding: 16px; font-family: monospace; font-size: 13px; white-space: pre-wrap; min-height: 80px; color: #a0e0a0; } | |
| .reward-bar { display: flex; align-items: center; gap: 12px; margin-top: 10px; } | |
| .reward-val { font-size: 24px; font-weight: 700; color: #4ade80; } | |
| .phase-pill { font-size: 10px; letter-spacing: 1.5px; text-transform: uppercase; padding: 3px 10px; border-radius: 20px; background: #2a2a2a; color: #aaa; } | |
| .phase-pill.triage { background: #1e3a5f; color: #60a5fa; } | |
| .phase-pill.selection { background: #1a3a1a; color: #4ade80; } | |
| .phase-pill.safety { background: #3a2a00; color: #fbbf24; } | |
| .phase-pill.finalize { background: #3a1a3a; color: #c084fc; } | |
| .phase-pill.done { background: #1a1a1a; color: #555; } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="container"> | |
| <h1>🏥 PharmaAgent</h1> | |
| <div class="subtitle">Clinical Decision RL Environment — Manual Tester</div> | |
| <div class="card"> | |
| <div class="card-title">Session Control</div> | |
| <button class="btn" onclick="resetEnv()">New Episode</button> | |
| <span id="session-id" style="font-family:monospace;font-size:11px;color:#555;margin-left:12px"></span> | |
| </div> | |
| <div class="card" id="patient-card" style="display:none"> | |
| <div class="card-title">Patient Case</div> | |
| <div id="patient-info" class="output"></div> | |
| </div> | |
| <div class="card" id="action-card" style="display:none"> | |
| <div class="card-title">Take Action</div> | |
| <select id="action-type"> | |
| <option value="diagnose">diagnose</option> | |
| <option value="select_drug">select_drug</option> | |
| <option value="check_ddi">check_ddi</option> | |
| <option value="finalize">finalize</option> | |
| </select> | |
| <input id="action-value" placeholder="e.g. Type 2 Diabetes / Metformin / Warfarin,Aspirin"/> | |
| <button class="btn" onclick="takeStep()">Submit Action</button> | |
| </div> | |
| <div class="card" id="feedback-card" style="display:none"> | |
| <div class="card-title">Environment Feedback</div> | |
| <div style="display:flex;align-items:center;gap:12px;margin-bottom:12px"> | |
| <span id="phase-pill" class="phase-pill">—</span> | |
| <span style="font-size:12px;color:#555">Step <span id="step-num">0</span></span> | |
| </div> | |
| <div id="feedback" class="output"></div> | |
| <div class="reward-bar"> | |
| <span style="font-size:12px;color:#555">Cumulative Reward:</span> | |
| <span class="reward-val" id="reward-val">0.0</span> | |
| </div> | |
| </div> | |
| </div> | |
| <script> | |
| let sessionId = null; | |
| async function resetEnv() { | |
| const r = await fetch('/reset', { method: 'POST' }); | |
| const d = await r.json(); | |
| sessionId = d.session_id; | |
| document.getElementById('session-id').textContent = 'Session: ' + sessionId.slice(0,8) + '…'; | |
| document.getElementById('patient-card').style.display = 'block'; | |
| document.getElementById('action-card').style.display = 'block'; | |
| document.getElementById('feedback-card').style.display = 'block'; | |
| const c = d.observation.patient_case; | |
| document.getElementById('patient-info').textContent = | |
| 'Symptoms: ' + (c.symptoms || []).join(', ') + '\\n' + | |
| 'Current Medications: ' + (c.existing_medications?.join(', ') || 'None'); | |
| document.getElementById('feedback').textContent = d.observation.feedback; | |
| document.getElementById('reward-val').textContent = '0.0'; | |
| document.getElementById('step-num').textContent = '0'; | |
| document.getElementById('action-card').style.display = 'block'; | |
| setPhase('triage'); | |
| } | |
| async function takeStep() { | |
| if (!sessionId) { alert('Click "New Episode" first.'); return; } | |
| const actionType = document.getElementById('action-type').value; | |
| const actionValue = document.getElementById('action-value').value.trim(); | |
| if (!actionValue && actionType !== 'finalize') { alert('Please enter a value.'); return; } | |
| const r = await fetch('/step?session_id=' + sessionId, { | |
| method: 'POST', | |
| headers: {'Content-Type': 'application/json'}, | |
| body: JSON.stringify({ action_type: actionType, value: actionValue || 'finalize' }) | |
| }); | |
| const d = await r.json(); | |
| document.getElementById('feedback').textContent = d.observation.feedback; | |
| document.getElementById('reward-val').textContent = d.observation.reward_so_far; | |
| document.getElementById('step-num').textContent = d.observation.step; | |
| setPhase(d.observation.phase); | |
| if (d.done) { | |
| document.getElementById('action-card').style.display = 'none'; | |
| } | |
| } | |
| function setPhase(phase) { | |
| const pill = document.getElementById('phase-pill'); | |
| pill.textContent = phase; | |
| pill.className = 'phase-pill ' + phase; | |
| } | |
| </script> | |
| </body> | |
| </html> | |
| """ | |