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" && (
)}
{/* 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}
))}
);
}