import { useEffect, useMemo, useRef, useState, useCallback } from "react"; function resolveApiBase() { const explicitBase = import.meta.env.VITE_API_BASE; if (explicitBase) return explicitBase.replace(/\/$/, ""); const host = window.location.hostname; const isLocal = host === "localhost" || host === "127.0.0.1" || host === "0.0.0.0"; if (isLocal && window.location.port === "5173") { return "http://localhost:7860"; } return window.location.origin.replace(/\/$/, ""); } const API_BASE = resolveApiBase(); const WS_URL = `${API_BASE.replace(/^http/, "ws")}/ws`; const TASKS = [ { id: "easy_screening", label: "Easy Screening" }, { id: "budgeted_screening", label: "Budgeted Screening" }, { id: "complex_tradeoff", label: "Complex Tradeoff" }, ]; const TASK_LABEL_MAP = Object.fromEntries(TASKS.map((t) => [t.id, t.label])); const ACTION_LABELS = { query_ddi: "Check Drug Interaction", propose_intervention: "Propose Change", finish_review: "Finish Review", }; const INTERVENTION_LABELS = { stop: "Stop Medication", dose_reduce: "Reduce Dose", substitute: "Substitute with Safer Drug", add_monitoring: "Add Monitoring", }; // ── Contextual guide steps: each targets a specific UI section ────────────── const GUIDE_STEPS = [ { target: "topbar", position: "below", title: "Welcome to PolypharmacyEnv", body: `This tool helps review elderly patients' medication regimens for safety. You'll act as a pharmacist assistant: check pairs of drugs for harmful interactions, propose changes to reduce risk, and get scored on how well you protect the patient — all under limited budgets. Behind the scenes, an AI agent (Neural Bandit) learns which drug combinations to investigate first, getting smarter with each review.`, }, { target: "task-selector", position: "below", title: "Choose a Scenario", body: `Pick a difficulty level: • Easy Screening — 3–5 drugs, 1 known dangerous interaction. Great for getting started. • Budgeted Screening — 6–10 drugs, multiple problems to find, tighter budgets. • Complex Tradeoff — 10–15 drugs including critical ones (blood thinners, insulin). Removing critical drugs without a replacement is penalized. Click "Reset Episode" to load a new patient case.`, }, { target: "episode-panel", position: "below", title: "Patient Overview", body: `After resetting, this panel shows the patient's details: • Demographics (age, sex, medical conditions) • Your remaining query and intervention budgets • A risk bar comparing starting risk vs. current risk • How many review steps you've taken Each check and intervention uses up budget — use them wisely to get the best outcome.`, }, { target: "action-console", position: "right", title: "Check Drug Interactions", body: `Select "Check Drug Interaction" and pick two drugs from the patient's list: Example dangerous combinations: • Warfarin + Naproxen → severe bleeding risk • Diazepam + Tramadol → dangerous sedation • Apixaban + Naproxen → severe bleeding risk Each check costs a small amount of budget. Finding a serious interaction earns a bonus. A smart strategy checks high-risk pairs first.`, }, { target: "action-console", position: "right", title: "Propose Changes", body: `After finding a dangerous interaction, switch to "Propose Change": • Stop Medication — Remove the drug entirely • Reduce Dose — Lower the dose to reduce risk • Substitute Drug — Automatically finds a safer alternative in the same drug class • Add Monitoring — Flag for closer clinical monitoring Example: After finding warfarin + naproxen interaction, select Naproxen → "Substitute". The system finds a safer pain reliever.`, }, { target: "medications-panel", position: "left", title: "Current Medications", body: `This grid shows the patient's active medications. Each card shows: • Drug name and dose • Drug class (e.g., pain reliever, blood thinner) • "High Risk" badge for drugs that need extra caution in elderly patients • Safety flags (avoid, caution, adjust dose) Cards marked "avoid" or "High Risk" are prime candidates for a closer look. The list updates live as you make changes.`, }, { target: "event-log", position: "above", title: "Activity Log & Score", body: `The log tracks every action you take and its impact. When you click "Finish Review", you get a final score (0–100%): • Easy: Based on risk reduction + targeting the right dangerous drugs • Medium: Risk reduction + precision of your interventions + how well you used your budget • Hard: Risk reduction minus penalties for disrupting the patient's treatment plan The "Ask AI" button lets an AI agent make decisions using the same tools you have.`, }, ]; async function apiPost(path, body) { const res = await fetch(`${API_BASE}${path}`, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body), }); if (!res.ok) { const msg = await res.text(); throw new Error(msg || `HTTP ${res.status}`); } return res.json(); } // ── Spotlight Guide Component ─────────────────────────────────────────────── function SpotlightGuide({ step, steps, onNext, onPrev, onClose }) { const [rect, setRect] = useState(null); const tooltipRef = useRef(null); const updateRect = useCallback(() => { const target = steps[step]?.target; if (!target) return; const el = document.querySelector(`[data-guide="${target}"]`); if (el) { const r = el.getBoundingClientRect(); setRect({ top: r.top, left: r.left, width: r.width, height: r.height }); // scroll into view el.scrollIntoView({ behavior: "smooth", block: "nearest" }); } }, [step, steps]); useEffect(() => { updateRect(); window.addEventListener("resize", updateRect); window.addEventListener("scroll", updateRect, true); return () => { window.removeEventListener("resize", updateRect); window.removeEventListener("scroll", updateRect, true); }; }, [updateRect]); if (!rect) return null; const pad = 8; const current = steps[step]; // Calculate tooltip position const getTooltipStyle = () => { const pos = current.position || "below"; const base = {}; if (pos === "below") { base.top = rect.top + rect.height + pad + 12; base.left = rect.left; base.maxWidth = Math.min(440, window.innerWidth - 40); } else if (pos === "above") { base.bottom = window.innerHeight - rect.top + pad + 12; base.left = rect.left; base.maxWidth = Math.min(440, window.innerWidth - 40); } else if (pos === "right") { base.top = rect.top; base.left = rect.left + rect.width + pad + 12; base.maxWidth = Math.min(380, window.innerWidth - rect.left - rect.width - 40); } else if (pos === "left") { base.top = rect.top; base.right = window.innerWidth - rect.left + pad + 12; base.maxWidth = Math.min(380, rect.left - 40); } return base; }; return (
{/* Dark overlay with cutout */} {/* Highlight border around target */}
{/* Tooltip */}

{current.title}

{step + 1} / {steps.length}
{current.body.split("\n").map((line, i) => (

{line}

))}
{step < steps.length - 1 ? ( ) : ( )}
{steps.map((_, i) => ( ))}
); } // ── Main App ──────────────────────────────────────────────────────────────── export default function App() { const [taskId, setTaskId] = useState("budgeted_screening"); const [obs, setObs] = useState(null); const [log, setLog] = useState([]); const [loading, setLoading] = useState(false); const [guideStep, setGuideStep] = useState(0); const [showGuide, setShowGuide] = useState(true); const [action, setAction] = useState({ action_type: "query_ddi", drug_id_1: "", drug_id_2: "", target_drug_id: "", intervention_type: "stop", proposed_new_drug_id: "", rationale: "", }); const medIds = useMemo( () => (obs?.current_medications || []).map((m) => m.drug_id), [obs] ); const hasValidEpisode = Boolean(obs?.episode_id) && (obs?.current_medications?.length || 0) > 0; const isDone = Boolean(obs?.done); const finalScore = typeof obs?.metadata?.grader_score === "number" ? obs.metadata.grader_score : null; const noBudgetsLeft = hasValidEpisode && (obs?.remaining_query_budget ?? 0) <= 0 && (obs?.remaining_intervention_budget ?? 0) <= 0; const wsRef = useRef(null); const pendingRef = useRef([]); const wsEnsure = async () => { if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) return wsRef.current; if (wsRef.current && wsRef.current.readyState === WebSocket.CONNECTING) { await new Promise((r) => setTimeout(r, 80)); return wsEnsure(); } const ws = new WebSocket(WS_URL); wsRef.current = ws; ws.onmessage = (evt) => { try { const msg = JSON.parse(evt.data); const pending = pendingRef.current.shift(); if (pending) pending.resolve(msg); } catch (e) { const pending = pendingRef.current.shift(); if (pending) pending.reject(e); } }; ws.onerror = (err) => { const pending = pendingRef.current.shift(); if (pending) pending.reject(err); }; ws.onclose = () => { wsRef.current = null; }; await new Promise((resolve, reject) => { const t = setTimeout( () => reject(new Error("WebSocket connect timeout")), 2500 ); ws.onopen = () => { clearTimeout(t); resolve(); }; }); return ws; }; const wsSend = async (type, data) => { const ws = await wsEnsure(); return await new Promise((resolve, reject) => { pendingRef.current.push({ resolve, reject }); ws.send(JSON.stringify({ type, data })); }); }; useEffect(() => { return () => { try { wsRef.current?.close(); } catch { /* ignore */ } }; }, []); const appendLog = (text) => { setLog((prev) => [`${new Date().toLocaleTimeString()} ${text}`, ...prev].slice(0, 30) ); }; const normalizeObsFromWs = (packetData) => { const observation = packetData?.observation || {}; const mergedMetadata = { ...(observation?.metadata || {}), ...(packetData?.info || {}), }; return { ...observation, done: Boolean(packetData?.done ?? observation?.done ?? false), reward: packetData?.reward ?? observation?.reward ?? null, metadata: mergedMetadata, }; }; const handleReset = async () => { setLoading(true); try { const msg = await wsSend("reset", { task_id: taskId }); const data = msg?.data || {}; const normalized = normalizeObsFromWs(data); setObs(normalized); const ids = (normalized?.current_medications || []).map((m) => m.drug_id); setAction((prev) => ({ ...prev, drug_id_1: ids[0] || "", drug_id_2: ids[1] || "", target_drug_id: ids[0] || "", })); appendLog(`Reset — ${TASK_LABEL_MAP[taskId] || taskId}`); } catch (err) { appendLog(`Reset failed: ${err.message}`); } finally { setLoading(false); } }; const buildActionPayload = () => { if (noBudgetsLeft) { return { action_type: "finish_review" }; } if (action.action_type === "query_ddi") { return { action_type: "query_ddi", drug_id_1: action.drug_id_1, drug_id_2: action.drug_id_2, }; } if (action.action_type === "propose_intervention") { return { action_type: "propose_intervention", target_drug_id: action.target_drug_id, intervention_type: action.intervention_type, proposed_new_drug_id: action.proposed_new_drug_id || undefined, rationale: action.rationale || undefined, }; } return { action_type: "finish_review" }; }; const isActionValid = () => { if (!hasValidEpisode) return false; if (isDone) return false; if (noBudgetsLeft) return true; if (action.action_type === "query_ddi") { return Boolean(action.drug_id_1 && action.drug_id_2); } if (action.action_type === "propose_intervention") { return Boolean(action.target_drug_id && action.intervention_type); } return true; }; const handleStep = async (overrideAction = null) => { if (!hasValidEpisode) { appendLog("Run Reset Episode before stepping."); return; } setLoading(true); try { const payload = overrideAction || buildActionPayload(); const msg = await wsSend("step", payload); const data = msg?.data || {}; const normalized = normalizeObsFromWs(data); setObs(normalized); const label = ACTION_LABELS[payload.action_type] || payload.action_type; const rwd = data.reward ?? 0; appendLog(`${label} → reward: ${Number(rwd).toFixed(3)}`); } catch (err) { appendLog(`Step failed: ${err.message}`); } finally { setLoading(false); } }; const askAi = async () => { if (!hasValidEpisode) { appendLog("Run Reset Episode before asking AI."); return; } setLoading(true); try { const data = await apiPost("/agent/suggest", { observation: obs }); const label = ACTION_LABELS[data.action.action_type] || data.action.action_type; appendLog(`AI suggests: ${label}`); await handleStep(data.action); } catch (err) { appendLog(`AI suggestion failed: ${err.message}`); } finally { setLoading(false); } }; const formatDrugName = (drugId) => { if (!drugId) return ""; return drugId .replace(/^DRUG_/, "") .replace(/_/g, " ") .replace(/\b\w/g, (c) => c.toUpperCase()); }; const currentRisk = obs?.metadata?.current_risk; const baselineRisk = obs?.metadata?.baseline_risk; return (
{/* Spotlight Guide */} {showGuide && ( setGuideStep((s) => Math.min(s + 1, GUIDE_STEPS.length - 1))} onPrev={() => setGuideStep((s) => Math.max(0, s - 1))} onClose={() => setShowGuide(false)} /> )}

PolypharmacyEnv

Elderly Medication Safety — Powered by Neural Bandits

{hasValidEpisode ? isDone ? "Episode Complete" : "Session Live" : "Ready"}
{/* Episode Info */}

Episode Overview

{hasValidEpisode ? ( <>
Episode {obs.episode_id}
Task {TASK_LABEL_MAP[obs.task_id] || obs.task_id}
Patient Age {obs.age}, {obs.sex === "M" ? "Male" : "Female"}
Step {obs.step_index}
Query Budget {obs.remaining_query_budget} remaining
Intervention Budget {obs.remaining_intervention_budget} remaining
{currentRisk !== undefined && baselineRisk !== undefined && (
Baseline Risk:{" "} {Number(baselineRisk).toFixed(3)} Current Risk:{" "} {Number(currentRisk).toFixed(3)}
)} {obs.conditions && obs.conditions.length > 0 && (
Conditions: {obs.conditions.map((c) => ( {c.replace(/_/g, " ")} ))}
)} ) : (

Select a task difficulty and click Reset Episode{" "} to begin a patient case.

)} {noBudgetsLeft && !isDone && (
All budgets exhausted. Click Finish Review to receive your final score.
)} {isDone && (
Episode complete {finalScore !== null ? ` — Final score: ${(finalScore * 100).toFixed(1)}%` : ""} . Click Reset Episode to start a new case.
)}
{/* Action Console */}

Action Console

{action.action_type === "query_ddi" && (
)} {action.action_type === "propose_intervention" && (
setAction((a) => ({ ...a, proposed_new_drug_id: e.target.value, })) } />
setAction((a) => ({ ...a, rationale: e.target.value })) } />
)}
{/* Current Medications */}

Current Medications {obs?.current_medications?.length ? ` (${obs.current_medications.length})` : ""}

{(obs?.current_medications || []).map((m) => (
{formatDrugName(m.drug_id)} {m.is_high_risk_elderly && ( High Risk )}

{m.generic_name}

{m.dose_mg} mg {m.atc_class}
{m.beers_flags && m.beers_flags.length > 0 && (
{m.beers_flags.map((f, i) => ( {f} ))}
)}
))}
{(!obs?.current_medications || obs.current_medications.length === 0) && (

No medications loaded. Reset an episode to begin.

)}
{/* Interaction Queries & Interventions */} {hasValidEpisode && (

Drug Interaction Checks ({obs?.interaction_queries?.length || 0})

{(obs?.interaction_queries || []).map((q, i) => (
{formatDrugName(q.drug_id_1)} +{" "} {formatDrugName(q.drug_id_2)} {q.severity} {q.recommendation && (

{q.recommendation.replace(/_/g, " ")}

)}
))} {(!obs?.interaction_queries || obs.interaction_queries.length === 0) && (

No queries yet.

)}

Proposed Changes ({obs?.interventions?.length || 0})

{(obs?.interventions || []).map((iv, i) => (
{formatDrugName(iv.target_drug_id)} {INTERVENTION_LABELS[iv.action_type] || iv.action_type} {iv.proposed_new_drug_id && (

Replaced with: {formatDrugName(iv.proposed_new_drug_id)}

)} {iv.rationale && (

{iv.rationale}

)}
))} {(!obs?.interventions || obs.interventions.length === 0) && (

No interventions yet.

)}
)} {/* Event Log */}

Event Log

{log.length === 0 && (
Events will appear here as you interact with the environment.
)} {log.map((line, idx) => (
{line}
))}
); }