diff --git a/PanTS-Demo/src/components/viewer/MeshViewer.tsx b/PanTS-Demo/src/components/viewer/MeshViewer.tsx index a53354ff..a5bbc234 100644 --- a/PanTS-Demo/src/components/viewer/MeshViewer.tsx +++ b/PanTS-Demo/src/components/viewer/MeshViewer.tsx @@ -76,7 +76,22 @@ export function SegmentationMeshViewer({ caseId, checkState, loading, opacity, c return (
- + {/* + preserveDrawingBuffer is REQUIRED for the AI assistant's snapshots. + WebGL clears the drawing buffer as soon as the frame is composited, so + without it canvas.toDataURL() reads an already-cleared buffer and the + captured "3D view" is a black rectangle. data-bodymaps-3d marks the + canvas so the capture helper picks this one and never an unrelated + canvas that happens to sit in the same pane. + */} + { + gl.domElement.setAttribute("data-bodymaps-3d", "1"); + }} + > diff --git a/PanTS-Demo/src/routes/VisualizationPage.tsx b/PanTS-Demo/src/routes/VisualizationPage.tsx index eb7a724e..a5bdb967 100644 --- a/PanTS-Demo/src/routes/VisualizationPage.tsx +++ b/PanTS-Demo/src/routes/VisualizationPage.tsx @@ -1543,6 +1543,10 @@ function VisualizationPage() { c.height = Math.round(img.height * scale); const ctx = c.getContext("2d"); if (!ctx) return resolve(dataUrl); + // JPEG has no alpha: paint the CT viewer's black ground first so a + // source with transparent pixels does not decode as white fringing. + ctx.fillStyle = "#000"; + ctx.fillRect(0, 0, c.width, c.height); ctx.drawImage(img, 0, 0, c.width, c.height); resolve(c.toDataURL("image/jpeg", 0.85)); }; @@ -1555,14 +1559,72 @@ function VisualizationPage() { // grid. The segmentation masks are left VISIBLE so the model can identify // each organ by its color (paired with the mask legend). Images are // downscaled before returning so the vision model responds quickly. + // Wait for a frame that has actually been presented. rAF never fires in a + // background tab, so cap the wait rather than hanging the capture. + const nextPresentedFrame = () => + new Promise((resolve) => { + const done = () => resolve(); + const timer = window.setTimeout(done, 250); + requestAnimationFrame(() => + requestAnimationFrame(() => { + window.clearTimeout(timer); + done(); + }) + ); + }); + + // A WebGL canvas read back after its drawing buffer was cleared comes out as + // one flat color — the "black 3D screenshot". Sample a tiny copy so a dead + // capture is detected here instead of being sent to the vision model, which + // would then confidently describe an empty image. + const captureLooksBlank = (source: HTMLCanvasElement): boolean => { + try { + const probe = document.createElement("canvas"); + probe.width = 32; + probe.height = 32; + const ctx = probe.getContext("2d", { willReadFrequently: true }); + if (!ctx) return false; + ctx.drawImage(source, 0, 0, probe.width, probe.height); + const { data } = ctx.getImageData(0, 0, probe.width, probe.height); + let min = 255; + let max = 0; + for (let i = 0; i < data.length; i += 4) { + const luma = (data[i] * 299 + data[i + 1] * 587 + data[i + 2] * 114) / 1000; + if (luma < min) min = luma; + if (luma > max) max = luma; + } + return max - min < 4; + } catch { + return false; // unreadable canvas — assume the shot is usable + } + }; + const captureAllViews = useCallback(async () => { const shots: { name: string; dataUrl: string }[] = await captureViewportImages(); try { const pane = document.querySelector(".render"); - const canvas = pane?.querySelector("canvas"); - if (canvas && canvas.width && pane && pane.offsetParent !== null) { - const url = canvas.toDataURL("image/png"); - if (url && url.length > 128) shots.push({ name: "3d", dataUrl: url }); + // Prefer the canvas the mesh viewer tags on creation; the positional + // lookup is only a fallback for an older render tree. + const canvas = + document.querySelector("canvas[data-bodymaps-3d]") ?? + pane?.querySelector("canvas") ?? + null; + const paneVisible = !pane || pane.offsetParent !== null; + if (canvas && canvas.width && paneVisible) { + await nextPresentedFrame(); + let url = canvas.toDataURL("image/png"); + if (captureLooksBlank(canvas)) { + // One more frame: the pane may have only just become visible. + await nextPresentedFrame(); + url = canvas.toDataURL("image/png"); + } + if (captureLooksBlank(canvas)) { + console.warn( + "[BodyMaps AI] 3D pane captured blank — omitting it rather than sending a black image" + ); + } else if (url && url.length > 128) { + shots.push({ name: "3d", dataUrl: url }); + } } } catch (error) { console.warn("[BodyMaps AI] 3D capture skipped", error); diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index 54c5f684..54c9e1a9 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -10,12 +10,17 @@ DEFAULT_OLLAMA_MODEL, DEFAULT_OLLAMA_VISION_MODEL, OLLAMA_THINK, + OllamaModelMissing, OllamaUnavailable, chat_json, chat_stream, chat_with_tools, + is_reasoning_model, list_ollama_models, + resolve_text_model, + resolve_vision_model, ) +from services import ai_reasoning from services.segmentation_metrics import calculate_session_metrics from services.search_ranking import rank_quality_results, select_balanced_tumor_results from services.site_normalization import site_country_label, split_site_codes @@ -2810,8 +2815,11 @@ def _ai_strip_think(text, orphan_closer=False): "FINAL ANSWER RULES: start directly with the answer (no reasoning, no " "preamble); never mention tools, data blocks, or any internal machinery; " "quote measured values verbatim with their units and never invent one; " - "1-3 sentences for a simple question, one short readable paragraph for a " - "clinical one; optionally end with one short, natural follow-up question." + "answer the question that was asked and never substitute an unrelated " + "measurement for an answer; 1-3 sentences for a simple question, one short " + "readable paragraph for a clinical one; END with one short, specific " + "question — when the next useful step is something to LOOK at, ask for that " + "(a slice level, a plane, a window preset, a structure to isolate)." ) @@ -3677,8 +3685,17 @@ def _ai_grounded_reply( return reply - # Deterministically ground exact case measurements. - for action in actions: + # Deterministically ground exact case measurements — but ONLY when the user + # asked for a measurement. The rule parser emits get_organ_metric whenever an + # organ name and a word like "size" appear anywhere in the message, which in + # a long clinical question is a coincidence. Returning the volume here threw + # away the model's actual answer and replied with an unrelated number. + wants_measurement = ( + question_mode == "case_measurement" + or ai_reasoning.asks_for_measurement(message) + ) + + for action in (actions if wants_measurement else []): if action.get("type") != "get_organ_metric": continue @@ -4111,13 +4128,20 @@ def ai_models(): models = list_ollama_models() model_names = [model["name"] for model in models] default_model = DEFAULT_OLLAMA_MODEL if DEFAULT_OLLAMA_MODEL in model_names else (model_names[0] if model_names else DEFAULT_OLLAMA_MODEL) + # Report the vision model that will ACTUALLY be used, not the one that + # happens to be configured. When they differ, the configured model was + # never pulled — the single most common reason image messages fail, and + # something worth being able to see from a browser before a demo. + active_vision_model = resolve_vision_model() return jsonify({ "available": True, "models": models, "default_model": default_model, # The model automatically used when a message carries images, so # the UI can show the switch the moment snapshots are attached. - "vision_model": DEFAULT_OLLAMA_VISION_MODEL, + "vision_model": active_vision_model or DEFAULT_OLLAMA_VISION_MODEL, + "vision_available": bool(active_vision_model), + "configured_vision_model": DEFAULT_OLLAMA_VISION_MODEL, }) except OllamaUnavailable as error: return jsonify({ @@ -4125,6 +4149,8 @@ def ai_models(): "models": [], "default_model": DEFAULT_OLLAMA_MODEL, "vision_model": DEFAULT_OLLAMA_VISION_MODEL, + "vision_available": False, + "configured_vision_model": DEFAULT_OLLAMA_VISION_MODEL, "error": f"Ollama is not reachable at the configured local endpoint: {error}", }), 200 @@ -4139,6 +4165,128 @@ def _ai_gate(): return None +def _ai_normalize_conversation(raw, limit=12, chars=2000): + """Defensively normalize client-supplied chat history.""" + turns = [] + if not isinstance(raw, list): + return turns + for turn in raw[-limit:]: + if not isinstance(turn, dict): + continue + role = str(turn.get("role") or "").strip() + content = str(turn.get("content") or "").strip() + if role in {"user", "assistant"} and content: + turns.append({"role": role, "content": content[:chars]}) + return turns + + +def _ai_normalize_legend(raw): + """Normalize the color→organ legend for the attached screenshots.""" + legend = [] + if not isinstance(raw, list): + return legend + for item in raw: + if not isinstance(item, dict): + continue + organ = str(item.get("organ") or "").strip() + color = str(item.get("color") or "").strip() + if organ and color: + legend.append({"organ": organ, "color": color}) + return legend + + +def _ai_command_vision_reply(*, message, images, body, case_id, selected_model): + """Non-streaming answer for a message that carries CT screenshots. + + The browser falls back to /ai-command whenever the streaming endpoint fails, + and a vision question is exactly the kind most likely to hit that fallback + (bigger payload, slower model). Answering it blind is worse than not + answering, so this path uses the same vision model, the same prompt, and the + same follow-up guarantee as the streaming endpoint — it just collects the + tokens into one string instead of forwarding them. + """ + mask_legend = _ai_normalize_legend(body.get("mask_legend")) + conversation = _ai_normalize_conversation(body.get("conversation")) + + vision_model = resolve_vision_model(selected_model) + + if not vision_model: + return jsonify({ + "reply": ai_reasoning.model_offline_reply( + has_images=True, + vision_model_missing=True, + configured_vision_model=DEFAULT_OLLAMA_VISION_MODEL, + ), + "actions": [], + "source": "vision_model_unavailable", + "model": None, + "intent": "read_images", + }) + + user_prompt = ( + f"{len(images)} CT viewer screenshot(s) are attached to this message, " + "in the viewer's pane order (axial, sagittal, coronal, then the 3D " + "surface rendering when present). Look at them before answering.\n\n" + f"{message or 'Describe what is shown in the attached views.'}" + ) + + legend_fact = ai_reasoning.build_legend_fact(mask_legend) + if legend_fact: + user_prompt += f"\n\nFacts:\n- {legend_fact}" + + collected = "" + try: + for kind, text in chat_stream( + model=vision_model, + system_prompt=ai_reasoning.build_system_prompt( + has_images=True, + has_case=bool(case_id), + ), + user_prompt=user_prompt, + images=images, + history=conversation[-6:], + ): + if kind == "content" and text: + collected += text + except OllamaModelMissing as error: + print(f"[ai_command vision] model '{vision_model}' not installed: {error}") + collected = "" + except (OllamaUnavailable, Exception) as error: + print("[ai_command vision]", type(error).__name__, str(error)) + collected = "" + + reply = _ai_strip_think( + collected, + orphan_closer=is_reasoning_model(vision_model), + ).strip() + + if not reply: + return jsonify({ + "reply": ai_reasoning.model_offline_reply( + has_images=True, + configured_vision_model=DEFAULT_OLLAMA_VISION_MODEL, + ), + "actions": [], + "source": "rule_fallback", + "model": None, + "intent": "read_images", + }) + + return jsonify({ + "reply": ai_reasoning.ensure_followup( + reply, + message, + has_images=True, + has_case=bool(case_id), + force=True, + ), + "actions": [], + "source": "ollama", + "model": vision_model, + "intent": "read_images", + }) + + @api_blueprint.route("/ai-command", methods=["POST"]) def ai_command(): try: @@ -4210,6 +4358,31 @@ def ai_command(): else DEFAULT_OLLAMA_MODEL ) + # Attached CT screenshots. This endpoint is what the browser retries on + # when the streaming call fails, so it has to be able to read images + # too: previously it dropped them silently and answered from text alone, + # producing a confident description of views it had never looked at. + raw_images = body.get("images") if isinstance(body.get("images"), list) else [] + images = [] + for item in raw_images: + if not isinstance(item, str) or not item.strip(): + continue + value = item.strip() + if value.startswith("data:"): + comma = value.find(",") + if comma != -1: + value = value[comma + 1:] + images.append(value) + + if images: + return _ai_command_vision_reply( + message=message, + images=images, + body=body, + case_id=case_id, + selected_model=selected_model, + ) + metrics, metric_source = _ai_load_metrics( case_id, body.get("organ_metrics"), @@ -4343,14 +4516,19 @@ def ai_command(): or question_mode ) - reply = _ai_grounded_reply( - message=message, - actions=actions, - metrics=metrics, - available_organs=available_organs, - metadata=metadata, - candidate_reply=candidate_reply, - question_mode=question_mode, + reply = ai_reasoning.ensure_followup( + _ai_grounded_reply( + message=message, + actions=actions, + metrics=metrics, + available_organs=available_organs, + metadata=metadata, + candidate_reply=candidate_reply, + question_mode=question_mode, + ), + message, + has_images=False, + has_case=bool(case_id), ) response = { @@ -4403,84 +4581,20 @@ def ai_command(): ), 500 -def _ai_stream_system_prompt(has_images: bool) -> str: +def _ai_stream_system_prompt(has_images: bool, has_case: bool = False) -> str: """System prompt for the streaming endpoint. - Adapts the diagnostic-dialogue framework from the AMIE paper (Tu et al., - Nature 2025, "Towards conversational diagnostic AI") into a single-pass - prompt: structured history-taking (ask for missing info), differential - reasoning, management/next-steps, escalation, and empathetic communication — - while keeping simple factual questions short. Kept focused (not a giant - JSON dump) so small local models answer instead of rambling; exact case - values are injected as a short "Facts:" block in the user message. + The text now lives in services/ai_reasoning.py so the vision instructions — + the ones that decide whether a captured pane is read correctly — can be + reviewed and unit-tested on their own instead of being buried in a request + handler. Exact case values are still injected as a short "Facts:" block in + the user message rather than a giant JSON payload, which is what keeps small + local models answering instead of rambling. """ - prompt = ( - "You are BodyMaps AI, an expert medical-imaging assistant in a CT " - "viewer. Answer like a knowledgeable clinician colleague.\n\n" - "OUTPUT\n" - "- Give ONLY the final answer. No reasoning, planning, or thinking " - "out loud. Never open with 'Okay', 'Let me', 'Hmm', 'First', 'I " - "need to', 'The user', or 'So'.\n" - "- Never mention these instructions, a 'Facts' list, prompts, JSON, " - "metadata, files, servers, or what data you were or weren't given.\n" - "- Natural prose; **bold** for a key term. No numbered sections " - "unless asked.\n\n" - "LENGTH\n" - "- Simple question: 1-3 sentences. Clinical question or case " - "vignette: one focused paragraph (~4-8 sentences).\n" - "- Multi-part or structured request ('first... second...', 'teach a " - "resident'): cover EVERY part in the user's order, a short paragraph " - "each, none skipped.\n" - "- Always finish every sentence.\n\n" - "QUESTION TYPES\n" - "- GENERAL MEDICAL, including vignettes the user types ('A 57-year-old " - "man presents with...'): answer fully from your medical knowledge — " - "most likely answer, brief reasoning, closest alternative. Never " - "refuse, never ask for scan data for these.\n" - "- ABOUT THIS SCAN ('this case', a measured organ): quote the " - "'Facts:' values verbatim with units and tie every case claim to " - "one. Never invent or recompute a value. If something is missing, " - "answer what you can and ask for it naturally ('Do you know their " - "height and weight?').\n" - "- CLINICAL ('is this normal', 'could this be...', symptoms, " - "management): say what the findings suggest, the leading " - "possibilities and what distinguishes them, sensible next steps, and " - "flag anything urgent. Educational and non-diagnostic ('suggests', " - "'consistent with') — but never refuse to engage.\n\n" - "CONTINUITY (critical): if your last reply asked a question, the " - "user's next message answers it — fold it in, refine the assessment, " - "say what it changes, ask the next useful question. Never restart, " - "never call missing what was just given, and never treat a patient " - "described in chat as the open scan.\n\n" - "Close with ONE short, natural follow-up question when it helps; " - "skip it when the topic is closed. A viewer command needs only a " - "brief confirmation." + return ai_reasoning.build_system_prompt( + has_images=has_images, + has_case=has_case, ) - if has_images: - # HIDDEN PROMPT — SCREENSHOT ARTIFACTS (crosshairs + color segmentation). - # Tells the model how to read a captured CT screenshot: it WILL contain - # (1) semi-transparent colored segmentation masks and (2) thin crosshair - # reference lines. The model names organs by mask color (via the legend) - # and treats crosshairs as navigation, not anatomy. - prompt += ( - "\n\nATTACHED IMAGES\n" - "CT viewer screenshots are attached (axial, sagittal, coronal, " - "sometimes 3D). Answer any question about them as fully as you " - "can: identify organs, describe the anatomy and anything notable, " - "compare views. If you need a different slice, view, or window, " - "say what you can and then ask for it.\n" - "- Semi-transparent colored shapes are segmentation masks, one " - "color per organ. Name organs using the provided color list; " - "never contradict it.\n" - "- Thin crosshair lines are slice-position guides — navigation, " - "never anatomy, a wire, or a fracture.\n" - "- Corner letters are orientation (A/P/L/R/S/I); corner numbers " - "are window width/level; the 3D view shows the same organs as " - "colored surfaces.\n" - "Keep the anatomy separate from these overlays; stay " - "non-diagnostic." - ) - return prompt @api_blueprint.route("/ai-command-stream", methods=["POST"]) @@ -4663,14 +4777,33 @@ def generate(): if fallback_actions: yield sse({"type": "actions", "actions": fallback_actions}) - exact_facts = _ai_case_facts(message, metrics, metadata) if references_case else [] - # Exact measurement sentences that MUST appear when the user asked # for an organ metric (guarantees the volume is always answered). required_facts = _ai_required_metric_facts(fallback_actions, metrics, available_organs) - for fact in required_facts: - if fact not in exact_facts: - exact_facts.append(fact) + + # Measured values for the open case are only allowed into the answer + # when the question is actually about them. Previously every fact + # that had been computed was appended, so a message carrying a + # patient's bilirubin and MRCP report came back with the open scan's + # liver volume and the patient's age — correct numbers, and a + # complete non-answer to what was asked. + conversation_text = " ".join(turn["content"] for turn in conversation[-4:]) + candidate_facts = _ai_case_facts(message, metrics, metadata) if references_case else [] + forced_facts = ( + required_facts + if ( + question_mode == "case_measurement" + or ai_reasoning.asks_for_measurement(message) + ) + else [] + ) + exact_facts = ai_reasoning.relevant_facts( + candidate_facts + required_facts, + message, + conversation_text=conversation_text, + always_include=forced_facts, + ) + required_facts = forced_facts # A clean, buttonless confirmation of any viewer action (never the # rule parser's "click below ..." text, which has no button here). @@ -4824,23 +4957,30 @@ def generate(): # actually needs. (A giant JSON payload makes small models ramble.) facts_lines = list(dict.fromkeys([f for f in (exact_facts + required_facts) if f])) if images and mask_legend: - legend_str = ", ".join( - f"{str(e['organ']).replace('_', ' ')}: {e['color']}" for e in mask_legend - ) - facts_lines.append(f"Segmentation mask colors — {legend_str}.") - - user_prompt = message or "Describe what is shown." - if conversation: - # Enough turns and characters that a long clinical vignette - # from earlier in the chat survives intact — a 200-char cap - # decapitated the case story and broke follow-up questions. - recent = conversation[-4:] - convo_str = "\n".join( - f"{t['role']}: {t['content'][:800]}" for t in recent + legend_fact = ai_reasoning.build_legend_fact(mask_legend) + if legend_fact: + facts_lines.append(legend_fact) + + user_prompt = message or "Describe what is shown in the attached views." + if images: + # Name the attachments explicitly. Without this the model has to + # infer from the images alone which pane is which, and a request + # phrased "for each pane" gets answered as one blended image. + user_prompt = ( + f"{len(images)} CT viewer screenshot(s) are attached to this " + "message, in the viewer's pane order (axial, sagittal, " + "coronal, then the 3D surface rendering when present). Look " + "at them before answering.\n\n" + f"{user_prompt}" ) - user_prompt = f"Recent conversation:\n{convo_str}\n\nQuestion: {user_prompt}" if facts_lines: user_prompt += "\n\nFacts:\n" + "\n".join(f"- {f}" for f in facts_lines) + + # Prior turns are sent as real conversation messages instead of being + # flattened into the prompt text: a model that can see its own last + # reply continues the thread, which is what makes "here are the labs + # you asked for" attach to the question that asked for them. + history = conversation[-6:] except Exception as error: print("[ai_command_stream setup error]", type(error).__name__, str(error)) yield sse({"type": "error", "message": "An internal error occurred while preparing the answer."}) @@ -4848,17 +4988,43 @@ def generate(): return # Stream the model answer. - vision_model = None + # + # Vision is the point of this product, so the model that reads images is + # resolved against what is actually installed rather than assumed. The + # old code sent the configured name blindly; when that model had never + # been pulled (qwen3-vl needs Ollama 0.12.7+), every image message failed + # at the first byte and the user was told the assistant was unavailable. + vision_missing = False if images: - vision_model = ( - DEFAULT_OLLAMA_VISION_MODEL - if DEFAULT_OLLAMA_VISION_MODEL - else selected_model - ) + vision_model = resolve_vision_model(selected_model) + if not vision_model: + vision_missing = True + model_for_call = None + else: + model_for_call = vision_model + else: + model_for_call = resolve_text_model(selected_model) - model_for_call = vision_model or selected_model model_ok = False + if vision_missing: + offline = ai_reasoning.model_offline_reply( + has_images=True, + vision_model_missing=True, + configured_vision_model=DEFAULT_OLLAMA_VISION_MODEL, + ) + yield sse({"type": "reply", "delta": offline}) + yield sse({ + "type": "final", + "reply": offline, + "actions": fallback_actions, + "source": "vision_model_unavailable", + "model": None, + "intent": fallback.get("intent") or question_mode, + }) + yield sse({"type": "done"}) + return + yield sse({"type": "status", "text": "Composing the answer"}) raw_content = "" # full model content so far (may contain blocks) @@ -4868,8 +5034,9 @@ def generate(): # real answer while it streams. For those models, hold the text back # until a think tag proves where the reasoning ends (or the stream # finishes); non-reasoning models (llama3.1, qwen3-vl) stream live. - hold_for_think = bool(_AI_REASONING_MODEL_RE.search(model_for_call or "")) + hold_for_think = is_reasoning_model(model_for_call or "") stream_error = False + stream_error_missing_model = False if agent_final is not None: # The agent loop already wrote the answer (and it is already # think-stripped) — emit it directly instead of generating twice. @@ -4881,9 +5048,13 @@ def generate(): try: for kind, text in chat_stream( model=model_for_call, - system_prompt=_ai_stream_system_prompt(bool(images)), + system_prompt=_ai_stream_system_prompt( + bool(images), + has_case=bool(case_id), + ), user_prompt=user_prompt, images=images or None, + history=history, ): if not text: continue @@ -4904,6 +5075,13 @@ def generate(): emitted = cleaned model_ok = True yield sse({"type": "reply", "delta": delta}) + except OllamaModelMissing as error: + stream_error = True + stream_error_missing_model = True + print( + f"[ai_command_stream] model '{model_for_call}' is not " + f"installed on the Ollama host: {error}" + ) except OllamaUnavailable as error: stream_error = True print("[ai_command_stream] Ollama unavailable:", str(error)) @@ -4948,30 +5126,47 @@ def generate(): final_reply += "\n\n" + fact except Exception as error: print("[ai_command_stream fact-check]", type(error).__name__, str(error)) + + # Guarantee the closing question. The prompt asks for one, but small + # local models drop it exactly when it matters — right after saying + # something could not be determined. On vision turns it is always + # appended, because there is always a next thing worth looking at. + final_reply = ai_reasoning.ensure_followup( + final_reply, + message, + has_images=bool(images), + has_case=bool(case_id), + force=bool(images), + ) else: - # Model offline/empty -> build a clean deterministic reply. Order: - # image color legend, then viewer-action confirmation + measured - # value(s), then a short friendly message. Never the rule parser's - # "click below ..." text or a scary "verify Ollama" string. + # No model answer. Say so honestly. + # + # This branch used to concatenate whatever measurements had been + # computed for the open case and present them as the reply, which is + # how a question about a patient's bilirubin was answered with + # "The liver volume is 1209.11 cm³ ... Age: 52.0". A viewer action + # the user asked for is still worth confirming, but a bare list of + # unrelated numbers is not an answer to anything. parts = [] - if images and mask_legend: + if action_confirmation: + parts.append(action_confirmation) + if images and mask_legend and _ai_norm(message): legend_reply = _ai_legend_answer(message, mask_legend) if legend_reply: parts.append(legend_reply) - if action_confirmation: - parts.append(action_confirmation) - parts.extend(required_facts) - if parts: - final_reply = " ".join(parts) - else: - final_reply = ( - "I couldn't get an answer just now — please try again in a moment." + parts.append( + ai_reasoning.model_offline_reply( + has_images=bool(images), + vision_model_missing=stream_error_missing_model and bool(images), + configured_vision_model=DEFAULT_OLLAMA_VISION_MODEL, ) + ) + final_reply = " ".join(part for part in parts if part).strip() if not final_reply: - final_reply = ( - "I could not generate a response. The local model may be " - "offline — viewer controls still work from the left panel." + final_reply = ai_reasoning.model_offline_reply( + has_images=bool(images), + configured_vision_model=DEFAULT_OLLAMA_VISION_MODEL, ) yield sse({ diff --git a/flask-server/services/ai_reasoning.py b/flask-server/services/ai_reasoning.py new file mode 100644 index 00000000..881e5677 --- /dev/null +++ b/flask-server/services/ai_reasoning.py @@ -0,0 +1,506 @@ +"""Answer shaping for the BodyMaps AI assistant. + +Everything in here is about *what the assistant says*, kept out of the request +plumbing in api_blueprint.py so it can be read and tuned on its own: + + * the system prompts (text turns and, first-class, vision turns), + * the relevance gate that stops measured case facts from being pasted onto a + question that never asked for them, + * the guarantee that a reply which still needs information ends by asking for + it — phrased around what can be *seen* when screenshots are in play, + * the honest failure text used when no model could answer. + +The module is deliberately dependency-free (stdlib only) so it is importable +from a test, a script, or the blueprint without dragging in Flask. +""" + +from __future__ import annotations + +import re +from typing import Any, Iterable, Sequence + + +# --------------------------------------------------------------------------- +# Small shared helpers +# --------------------------------------------------------------------------- + +def normalize(value: Any) -> str: + """Lowercase, punctuation-flattened text used for all keyword matching.""" + text = str(value or "").lower() + text = text.replace(".nii.gz", " ").replace(".nii", " ") + text = re.sub(r"[_/]+", " ", text) + return " ".join(text.split()) + + +def _strip_markdown(value: str) -> str: + return re.sub(r"[*_`#]+", "", str(value or "")) + + +# Organs the viewer segments, plus the words a clinician actually uses for them. +# Used both to detect what a question is about and to pick a follow-up that is +# specific instead of generic. +_ORGAN_SYNONYMS: dict[str, tuple[str, ...]] = { + "pancreas": ("pancreas", "pancreatic", "uncinate", "ampulla", "whipple"), + "liver": ("liver", "hepatic", "hepato", "cirrhosis", "steatosis"), + "gallbladder": ("gallbladder", "gall bladder", "cholecyst", "biliary", "bile duct", "cbd"), + "spleen": ("spleen", "splenic"), + "stomach": ("stomach", "gastric"), + "kidney": ("kidney", "kidneys", "renal", "nephro"), + "adrenal gland": ("adrenal",), + "aorta": ("aorta", "aortic"), + "inferior vena cava": ("vena cava", "ivc"), + "portal vein": ("portal vein", "portal", "splenic vein", "smv", "mesenteric"), + "duodenum": ("duodenum", "duodenal"), + "colon": ("colon", "colonic", "bowel"), + "small bowel": ("small bowel", "jejunum", "ileum"), + "esophagus": ("esophagus", "oesophagus", "esophageal"), + "bladder": ("bladder", "vesical"), + "prostate": ("prostate", "prostatic"), + "lung": ("lung", "lungs", "pulmonary", "pleural"), + "vertebrae": ("vertebra", "vertebrae", "spine", "spinal", "vertebral"), + "rib": ("rib", "ribs"), + "femur": ("femur", "femoral"), +} + +_DEMOGRAPHIC_WORDS = ( + "age", "aged", "old", "sex", "male", "female", "man", "woman", + "bmi", "body mass", "height", "weight", "demographic", +) + +_MEASUREMENT_WORDS = ( + "volume", "cm3", "cm³", "size", "how big", "how large", "measure", + "measured", "measurement", "mean hu", "hounsfield", "attenuation", + "density", "percentile", "largest", "smallest", +) + + +def organs_mentioned(message: str) -> list[str]: + """Canonical organ names referenced anywhere in a message.""" + norm = normalize(message) + found: list[str] = [] + for organ, words in _ORGAN_SYNONYMS.items(): + if any(word in norm for word in words): + found.append(organ) + return found + + +def asks_for_measurement(message: str) -> bool: + norm = normalize(message) + return any(word in norm for word in _MEASUREMENT_WORDS) + + +def asks_about_demographics(message: str) -> bool: + norm = normalize(message) + return any(re.search(rf"\b{re.escape(word)}\b", norm) for word in _DEMOGRAPHIC_WORDS) + + +# --------------------------------------------------------------------------- +# Relevance gate for measured case facts +# --------------------------------------------------------------------------- + +def _fact_subject(fact: str) -> str: + """The thing a generated fact sentence is about ('Liver', 'Age', ...).""" + plain = _strip_markdown(fact).strip() + + labelled = re.match(r"^([A-Za-z][A-Za-z \-]{0,30}):", plain) + if labelled: + return labelled.group(1).strip() + + inline = re.search( + r"\bsegmented\s+([A-Za-z][A-Za-z \-]{0,30}?)\s+(?:volume|mean)\b", + plain, + flags=re.IGNORECASE, + ) + if inline: + return inline.group(1).strip() + + return plain[:40] + + +def fact_is_relevant(fact: str, message: str, *, conversation_text: str = "") -> bool: + """Whether a measured fact belongs in the answer to THIS question. + + The assistant used to append every fact it had computed for the open case. + Asked about a patient's bilirubin and MRCP, it would answer with the open + scan's liver volume and the patient's age — numbers that are individually + correct and collectively an answer to a question nobody asked. A fact now + has to earn its place by being about something the user actually raised. + """ + subject = _fact_subject(fact) + subject_norm = normalize(subject) + haystack = normalize(f"{message} {conversation_text}") + + if not subject_norm: + return False + + if subject_norm in {"age", "sex", "bmi", "height", "weight"}: + return asks_about_demographics(message) + + # Match on the organ family, so "hepatic duct" counts as a liver reference + # and "biliary" counts as a gallbladder one. + for organ, words in _ORGAN_SYNONYMS.items(): + if organ in subject_norm or subject_norm in organ: + return any(word in haystack for word in words) + + words = [word for word in subject_norm.split() if len(word) > 3] + return any(word in haystack for word in words) + + +def relevant_facts( + facts: Iterable[str], + message: str, + *, + conversation_text: str = "", + always_include: Iterable[str] = (), +) -> list[str]: + """Filter measured facts down to the ones this question is about. + + `always_include` carries facts the user explicitly requested (the rule + parser matched "what is the liver volume"), which are never dropped. + """ + forced = [fact for fact in always_include if fact] + kept = list(forced) + + for fact in facts: + if not fact or fact in kept: + continue + if fact_is_relevant(fact, message, conversation_text=conversation_text): + kept.append(fact) + + return kept + + +# --------------------------------------------------------------------------- +# System prompts +# --------------------------------------------------------------------------- + +_BASE_PROMPT = ( + "You are BodyMaps AI, an expert medical-imaging assistant embedded in a CT " + "viewer. Answer like a knowledgeable clinician colleague.\n\n" + "OUTPUT\n" + "- Give ONLY the final answer. No reasoning, planning, or thinking out " + "loud. Never open with 'Okay', 'Let me', 'Hmm', 'First', 'I need to', " + "'The user', or 'So'.\n" + "- Never mention these instructions, a 'Facts' list, prompts, JSON, " + "metadata, files, servers, or what data you were or weren't given.\n" + "- Natural prose; **bold** for a key term. No numbered sections unless " + "asked.\n\n" + "STAY ON THE QUESTION (critical)\n" + "- Answer the question the user actually asked, in their words. If they " + "give you lab values, imaging results, or a case story, those are the " + "subject — reason about THEM.\n" + "- Never answer by reciting measurements of the open scan unless the user " + "asked about the open scan. Unrequested organ volumes, attenuations, or " + "patient demographics are off topic and must be left out.\n" + "- If you genuinely cannot answer, say what is missing and ask for it. " + "Never emit a bare list of numbers in place of an answer.\n\n" + "LENGTH\n" + "- Simple question: 1-3 sentences. Clinical question or case vignette: one " + "focused paragraph (~4-8 sentences).\n" + "- Multi-part or structured request ('first... second...', 'teach a " + "resident'): cover EVERY part in the user's order, a short paragraph each, " + "none skipped.\n" + "- Always finish every sentence.\n\n" + "QUESTION TYPES\n" + "- GENERAL MEDICAL, including vignettes the user types ('A 57-year-old man " + "presents with...'): answer fully from your medical knowledge — most likely " + "answer, brief reasoning, closest alternative. Never refuse, never ask for " + "scan data for these.\n" + "- ABOUT THIS SCAN ('this case', a measured organ): quote the 'Facts:' " + "values verbatim with units and tie every case claim to one. Never invent " + "or recompute a value.\n" + "- CLINICAL ('is this normal', 'could this be...', symptoms, management): " + "say what the findings suggest, the leading possibilities and what " + "distinguishes them, sensible next steps, and flag anything urgent. " + "Educational and non-diagnostic ('suggests', 'consistent with') — but never " + "refuse to engage.\n\n" + "CONTINUITY (critical): if your last reply asked a question, the user's " + "next message answers it — fold it in, refine the assessment, say what it " + "changes, then ask the next useful question. Never restart, never call " + "missing what was just given, and never treat a patient described in chat " + "as the open scan.\n\n" + "ENDING\n" + "- If anything you would need to be more certain is missing, END with ONE " + "short, specific question asking for exactly that.\n" + "- If the answer is complete, you may still end with one short question " + "offering the natural next step.\n" + "- One question, never a list. A pure viewer command needs only a brief " + "confirmation." +) + +# HIDDEN VISION PROMPT. +# +# This is the instruction set that makes attached screenshots usable. It is the +# most important prompt in the product — BodyMaps is a vision-first application, +# and the model has to be told exactly what the artifacts in a captured CT pane +# mean, or it reads crosshairs as hardware and mask colors as pathology. +_VISION_PROMPT = ( + "\n\n=== ATTACHED CT VIEWER SCREENSHOTS — THIS IS THE PRIMARY EVIDENCE ===\n" + "Images from the CT viewer are attached. They are the subject of this turn: " + "look at them and describe what is actually there. Ground every visual " + "claim in something visible in a specific pane, and name the pane you saw " + "it in.\n\n" + "WHAT THE PANES ARE\n" + "- Up to four panes may be attached, in this order: axial (cross-section, " + "viewed from the feet — the patient's LEFT is on the RIGHT of the image), " + "sagittal (side view, anterior to one side), coronal (front view), and a 3D " + "surface rendering of the segmented organs.\n" + "- Name each pane you are describing so the reader can follow along.\n\n" + "OVERLAYS ARE NOT ANATOMY\n" + "- Semi-transparent colored regions are SEGMENTATION MASKS, one color per " + "organ. Identify organs using the supplied color list and never contradict " + "it; if a color is not in the list, say the region is unlabeled rather than " + "guessing an organ.\n" + "- Thin straight crosshair lines are slice-position guides. They are " + "navigation, never a wire, catheter, fracture, or vessel.\n" + "- Corner letters are orientation (A/P/L/R/S/I). Corner numbers are window " + "width/level and zoom, not measurements of the patient.\n\n" + "HOW TO READ\n" + "- Work through the user's request in the order they asked for it, covering " + "every part.\n" + "- Describe position relationally (anterior/posterior, medial/lateral, " + "cranial/caudal) and relative to neighboring structures.\n" + "- Compare paired structures (the two kidneys) for size, level, and " + "symmetry when relevant.\n\n" + "HONESTY ABOUT WHAT A SCREENSHOT CANNOT SHOW (required)\n" + "- A single captured slice cannot establish contrast phase, lesion " + "conspicuity below screen resolution, true HU values, or anything outside " + "the captured field of view. When the user asks what cannot be judged, say " + "so plainly and specifically.\n" + "- Never invent a finding to fill a gap. 'Not assessable from this capture' " + "is a correct and useful answer.\n" + "- If a pane is blank, black, or unreadable, say which pane and ask the " + "user to re-capture it. Do not describe an image you cannot see.\n\n" + "ENDING A VISION ANSWER\n" + "- Finish with ONE specific request for what you would need to SEE next — a " + "named slice level, a different plane, a window preset (soft tissue, liver, " + "bone, lung), a zoom on a named structure, or a re-capture of a pane. Make " + "it something the user can do in this viewer.\n" + "Stay educational and non-diagnostic." +) + + +def build_system_prompt( + *, + has_images: bool, + has_case: bool = False, +) -> str: + """System prompt for one streamed turn.""" + prompt = _BASE_PROMPT + if has_case: + prompt += ( + "\n\nA CT case is open in the viewer. Only bring it up if the user's " + "question is about it." + ) + if has_images: + prompt += _VISION_PROMPT + return prompt + + +def build_legend_fact(mask_legend: Sequence[dict[str, Any]]) -> str | None: + """One line mapping every visible mask color to its organ.""" + pairs = [ + f"{str(entry.get('organ') or '').replace('_', ' ')}: {entry.get('color')}" + for entry in mask_legend or [] + if entry.get("organ") and entry.get("color") + ] + if not pairs: + return None + return ( + "Segmentation mask colors in the attached screenshots — " + + ", ".join(pairs) + + "." + ) + + +# --------------------------------------------------------------------------- +# Follow-up question guarantee +# --------------------------------------------------------------------------- + +# Phrases that mean the model knows it is short of information. When one of +# these shows up without a question mark, the reply has stated a need and then +# failed to ask — exactly the behavior we are correcting. +_UNCERTAINTY_MARKERS = ( + "cannot be determined", "can't be determined", "not assessable", + "cannot be assessed", "can't be assessed", "would need", "i would need", + "not enough information", "insufficient information", "unclear from", + "not possible to tell", "cannot tell", "can't tell", "more information", + "additional information", "not visible", "not shown", "unable to", + "cannot confirm", "can't confirm", "further evaluation", "further imaging", +) + + +def _ends_with_question(reply: str) -> bool: + """Whether the reply closes by asking something. + + Only the tail counts: a rhetorical question in the middle of a teaching + paragraph is not the assistant asking the user for anything. + """ + text = _strip_markdown(reply).strip() + if not text: + return False + tail = text[-320:] + sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", tail) if s.strip()] + if not sentences: + return False + return sentences[-1].endswith("?") + + +def signals_missing_information(reply: str) -> bool: + lowered = _strip_markdown(reply).lower() + return any(marker in lowered for marker in _UNCERTAINTY_MARKERS) + + +_WINDOW_FOR_ORGAN = { + "liver": "liver", + "pancreas": "soft tissue", + "gallbladder": "soft tissue", + "spleen": "soft tissue", + "kidney": "soft tissue", + "lung": "lung", + "vertebrae": "bone", + "rib": "bone", + "femur": "bone", + "aorta": "soft tissue", +} + +_GENERIC_VISION_QUESTIONS = ( + "Which view should I look at next — a different slice level, another plane, " + "or the 3D surface rendering?", + "Would it help if you re-captured the panes at a different slice level, or " + "in a different window preset?", + "Is there a particular structure in these views you want me to zoom in on?", +) + +_GENERIC_CASE_QUESTIONS = ( + "Would you like me to isolate any of these structures in the viewer so we " + "can look at them directly?", + "Do you want me to pull up the measurements for a specific structure in " + "this case?", +) + +_GENERIC_CLINICAL_QUESTIONS = ( + "What other findings, labs, or history do you have for this patient?", + "Is there anything else from the workup you can share so I can narrow this " + "down?", +) + + +def _pick(options: Sequence[str], seed_text: str) -> str: + """Stable, varied choice — the same question never repeats twice in a row + for different messages, but one message always gets the same follow-up.""" + if not options: + return "" + return options[sum(ord(ch) for ch in seed_text[:64]) % len(options)] + + +def suggest_followup( + message: str, + *, + has_images: bool, + has_case: bool, +) -> str: + """One short, specific question to close a reply with. + + Vision-focused whenever screenshots are in play: BodyMaps is a viewer, so + the most useful next step is nearly always something to LOOK at. + """ + organs = organs_mentioned(message) + + if has_images: + if organs: + organ = organs[0] + window = _WINDOW_FOR_ORGAN.get(organ, "soft tissue") + return ( + f"Which {organ} slice would you like me to look at next — a " + f"different level, or the same one in the {window} window?" + ) + return _pick(_GENERIC_VISION_QUESTIONS, message) + + if has_case and organs: + organ = organs[0] + return ( + f"Would you like me to isolate the {organ} in the viewer and capture " + "the views so I can look at it directly?" + ) + + if has_case: + return _pick(_GENERIC_CASE_QUESTIONS, message) + + return _pick(_GENERIC_CLINICAL_QUESTIONS, message) + + +def ensure_followup( + reply: str, + message: str, + *, + has_images: bool, + has_case: bool = False, + force: bool = False, +) -> str: + """Guarantee the reply ends by asking for what it still needs. + + The model is instructed to do this, but small local models drop the closing + question exactly when it matters most — after admitting something could not + be determined. `force` appends one unconditionally (used for vision turns, + where there is always a next thing worth looking at). + """ + text = str(reply or "").rstrip() + if not text: + return text + + if _ends_with_question(text): + return text + + if not force and not signals_missing_information(text): + return text + + question = suggest_followup(message, has_images=has_images, has_case=has_case) + if not question: + return text + + separator = "\n\n" if "\n" in text else " " + return f"{text}{separator}{question}" + + +# --------------------------------------------------------------------------- +# Honest failure text +# --------------------------------------------------------------------------- + +def model_offline_reply( + *, + has_images: bool, + vision_model_missing: bool = False, + configured_vision_model: str = "", +) -> str: + """What to say when no model produced an answer. + + Never a dump of whatever numbers happened to be computed for the open case: + that reads as a confident non-sequitur. Say what failed and what fixes it. + """ + if vision_model_missing: + model = configured_vision_model or "qwen3-vl:4b" + return ( + "I can't read the attached views right now — no vision model is " + f"available on this server. Pulling one (`ollama pull {model}`) and " + "restarting the backend will enable image reading. In the meantime " + "I can still answer from text, and the viewer controls in the left " + "panel all work.\n\n" + "Would you like me to answer the anatomy question from the case " + "measurements instead, while the model is set up?" + ) + + if has_images: + return ( + "I couldn't finish reading the attached views — the local model " + "didn't return an answer. Please send them again in a moment.\n\n" + "If it keeps failing, would you re-capture the panes? A blank or " + "partially rendered pane can stall the read." + ) + + return ( + "I couldn't get an answer just now — the local model didn't respond. " + "Please try again in a moment; the viewer controls in the left panel " + "still work.\n\n" + "Would you like me to retry the same question?" + ) diff --git a/flask-server/services/ollama_client.py b/flask-server/services/ollama_client.py index c77b075a..74789134 100644 --- a/flask-server/services/ollama_client.py +++ b/flask-server/services/ollama_client.py @@ -3,7 +3,9 @@ import json import os import re -from typing import Any +import threading +import time +from typing import Any, Iterable import requests @@ -40,6 +42,21 @@ OLLAMA_NUM_CTX = int(os.getenv("OLLAMA_NUM_CTX", "8192")) +# Images are expensive in context: a single 768px CT pane costs roughly 700-1200 +# tokens for qwen3-vl, so four panes plus a long structured question overflow the +# 8k text window. When that happens Ollama silently drops the OLDEST tokens — +# which is the system prompt carrying every instruction about how to read the +# screenshots. The result looks exactly like "the hidden image prompt did +# nothing". Vision turns therefore get their own, larger window. +OLLAMA_VISION_NUM_CTX = int( + os.getenv("OLLAMA_VISION_NUM_CTX", "12288") +) + +# Hard ceiling on how many images are forwarded in one turn. The viewer sends at +# most four (axial/sagittal/coronal/3D); more than that is a client bug and only +# guarantees a context overflow. +OLLAMA_MAX_IMAGES = int(os.getenv("OLLAMA_MAX_IMAGES", "4")) + # Ceiling on generated tokens. -1 means "no cap": the model generates until it # naturally finishes, so answers are never cut off mid-sentence. Brevity for # simple questions comes from the prompt, not from a low cap. @@ -59,15 +76,68 @@ "on", } +# How long the installed-model list is trusted before it is re-fetched. Pulling a +# model mid-session should be picked up without a backend restart, but every +# assistant message must not pay for an extra HTTP round trip either. +_MODEL_LIST_TTL_SECONDS = float(os.getenv("OLLAMA_MODEL_LIST_TTL_SECONDS", "60")) + +# A failed probe is remembered only briefly, so the assistant comes back on its +# own within seconds of Ollama being restarted. +_MODEL_LIST_FAILURE_TTL_SECONDS = float( + os.getenv("OLLAMA_MODEL_LIST_FAILURE_TTL_SECONDS", "10") +) + class OllamaUnavailable(RuntimeError): """Raised when the local Ollama server cannot complete a request.""" +class OllamaModelMissing(OllamaUnavailable): + """Raised when Ollama is reachable but the requested model is not pulled. + + A subclass of OllamaUnavailable so existing `except OllamaUnavailable` + handlers keep working, while callers that can pick a different model (the + vision path) are able to tell "Ollama is down" apart from "this particular + model was never downloaded" — two failures that need very different fixes. + """ + + class OllamaInvalidResponse(RuntimeError): """Raised when Ollama returns a response that cannot be parsed.""" +# Model families that accept images. Ollama does not expose a "supports vision" +# flag on /api/tags in every version, so the family name is the reliable signal. +_VISION_MODEL_RE = re.compile( + r"(?:^|[/_-])(?:" + r"qwen[\d.]*-?vl|llava|llama[\d.]*-vision|vision|" + r"minicpm-?v|moondream|bakllava|granite[\d.]*-vision|" + r"gemma3|mistral-small3|pixtral|internvl|cogvlm|glm-4v" + r")", + re.IGNORECASE, +) + +# Reasoning families whose output can carry a chain-of-thought. qwen3-vl +# (instruct) is deliberately excluded: it does not reason, and sending it the +# "/no_think" directive just injects stray text into a vision prompt. +_REASONING_MODEL_RE = re.compile( + r"qwen3(?!-vl)|deepseek-r1|-r1\b|qwq|marco-o1|thinking", re.IGNORECASE +) + +_MODEL_CACHE: dict[str, Any] = {"names": [], "fetched_at": 0.0, "ok": False} +_MODEL_CACHE_LOCK = threading.Lock() + + +def is_vision_model(name: str) -> bool: + """Whether a model name belongs to a family that accepts images.""" + return bool(_VISION_MODEL_RE.search(str(name or ""))) + + +def is_reasoning_model(name: str) -> bool: + """Whether a model name belongs to a family that emits chain-of-thought.""" + return bool(_REASONING_MODEL_RE.search(str(name or ""))) + + def list_ollama_models( timeout: float | None = None, ) -> list[dict[str, Any]]: @@ -114,6 +184,124 @@ def list_ollama_models( return result +def installed_model_names(force: bool = False) -> tuple[list[str], bool]: + """Cached list of installed model names. + + Returns (names, listing_ok). listing_ok is False when Ollama could not be + reached — callers must then fall back to the configured model rather than + concluding that nothing is installed, because "the tag listing timed out" is + not evidence that a model is missing. + """ + now = time.monotonic() + + with _MODEL_CACHE_LOCK: + age = now - float(_MODEL_CACHE["fetched_at"]) + # A failed probe is cached too, on a shorter timer. Otherwise every + # message sent while Ollama is down pays the full list timeout again + # before the chat call fails for the same reason. + ttl = _MODEL_LIST_TTL_SECONDS if _MODEL_CACHE["ok"] else _MODEL_LIST_FAILURE_TTL_SECONDS + if not force and age < ttl: + return list(_MODEL_CACHE["names"]), bool(_MODEL_CACHE["ok"]) + + try: + names = [model["name"] for model in list_ollama_models()] + ok = True + except OllamaUnavailable: + names, ok = [], False + + with _MODEL_CACHE_LOCK: + _MODEL_CACHE["fetched_at"] = now + _MODEL_CACHE["ok"] = ok + if ok: + _MODEL_CACHE["names"] = list(names) + cached = list(_MODEL_CACHE["names"]) + + return (names if ok else cached), ok + + +def _name_matches(candidate: str, installed: str) -> bool: + """Whether `installed` satisfies a request for `candidate`. + + Ollama reports fully qualified tags ("llama3.1:latest"), while configuration + and the UI often use the bare name ("llama3.1"). Treat a missing tag as + ":latest" so the two forms resolve to each other. + """ + want = str(candidate or "").strip().lower() + have = str(installed or "").strip().lower() + if not want or not have: + return False + if want == have: + return True + if ":" not in want and have == f"{want}:latest": + return True + if ":" not in have and want == f"{have}:latest": + return True + return False + + +def model_is_installed(name: str, names: Iterable[str] | None = None) -> bool: + """Whether a model is present locally (best effort).""" + if names is None: + names, ok = installed_model_names() + if not ok: + # Unknown, not absent — assume present and let the chat call decide. + return True + return any(_name_matches(name, installed) for installed in names) + + +def resolve_vision_model(preferred: str | None = None) -> str | None: + """Pick a vision-capable model that is actually installed. + + Order: an explicitly requested vision model, then the configured default, + then any installed model from a known vision family. Returns None only when + Ollama is reachable and genuinely has no vision model — the one case where + the caller must tell the user to pull one instead of silently sending images + to a text-only model, which answers confidently about an image it never saw. + """ + names, listing_ok = installed_model_names() + + candidates = [ + preferred, + DEFAULT_OLLAMA_VISION_MODEL, + ] + for candidate in candidates: + candidate = str(candidate or "").strip() + if not candidate or not is_vision_model(candidate): + continue + if not listing_ok: + # Cannot verify; trust configuration rather than refusing to try. + return candidate + if model_is_installed(candidate, names): + return candidate + + if not listing_ok: + return DEFAULT_OLLAMA_VISION_MODEL or None + + for installed in names: + if is_vision_model(installed): + return installed + + return None + + +def resolve_text_model(preferred: str | None = None) -> str: + """Pick an installed text model, falling back to the configured default.""" + names, listing_ok = installed_model_names() + + for candidate in (preferred, DEFAULT_OLLAMA_MODEL): + candidate = str(candidate or "").strip() + if not candidate: + continue + if not listing_ok or model_is_installed(candidate, names): + return candidate + + for installed in names: + if not is_vision_model(installed): + return installed + + return str(preferred or DEFAULT_OLLAMA_MODEL).strip() + + def _extract_json_object(text: str) -> dict[str, Any]: value = (text or "").strip() @@ -152,6 +340,38 @@ def _extract_json_object(text: str) -> dict[str, Any]: return parsed +def _describes_missing_model(text: str) -> bool: + lowered = str(text or "").lower() + return ( + "not found" in lowered + or "no such model" in lowered + or "try pulling it first" in lowered + or "unknown model" in lowered + ) + + +def _raise_for_chat_error(exc: requests.RequestException) -> None: + """Translate a failed /api/chat call into the most specific error we can.""" + response = getattr(exc, "response", None) + detail = "" + status = None + + if response is not None: + status = response.status_code + try: + payload = response.json() + detail = str(payload.get("error") or payload) + except ValueError: + detail = (response.text or "")[:400] + + message = detail or str(exc) + + if status == 404 or _describes_missing_model(message): + raise OllamaModelMissing(message) from exc + + raise OllamaUnavailable(message) from exc + + def _post_chat( payload: dict[str, Any], timeout: float, @@ -168,7 +388,10 @@ def _post_chat( raise OllamaUnavailable( "Ollama timed out while generating a response." ) from exc - except (requests.RequestException, ValueError) as exc: + except requests.RequestException as exc: + _raise_for_chat_error(exc) + raise # unreachable; keeps type checkers happy + except ValueError as exc: raise OllamaUnavailable(str(exc)) from exc if not isinstance(data, dict): @@ -245,7 +468,7 @@ def chat_json( ) json_user_prompt = user_prompt - if not OLLAMA_THINK and "qwen3" in selected_model.lower(): + if not OLLAMA_THINK and is_reasoning_model(selected_model): json_user_prompt = f"{json_user_prompt}\n\n/no_think" payload: dict[str, Any] = { @@ -373,12 +596,58 @@ def chat_with_tools( return message +def _build_chat_messages( + *, + system_prompt: str, + user_prompt: str, + history: list[dict[str, Any]] | None, + images: list[str] | None, + selected_model: str, +) -> list[dict[str, Any]]: + """Assemble the /api/chat message list for one streamed turn. + + Prior turns are sent as real assistant/user messages instead of being + flattened into the prompt text: a model that can see its own last message + actually continues the conversation, which is what makes "you asked me for + the bilirubin, here it is" resolve to the right thread instead of restarting. + """ + content = user_prompt + + # "/no_think" is a qwen3 *reasoning* directive. Sending it to qwen3-vl (an + # instruct model) just prepends noise to a vision prompt, so it is gated on + # the reasoning family rather than on the substring "qwen3". + if not OLLAMA_THINK and is_reasoning_model(selected_model): + content = f"{content}\n\n/no_think" + + messages: list[dict[str, Any]] = [ + {"role": "system", "content": system_prompt}, + ] + + for turn in history or []: + if not isinstance(turn, dict): + continue + role = str(turn.get("role") or "").strip() + text = str(turn.get("content") or "").strip() + if role in {"user", "assistant"} and text: + messages.append({"role": role, "content": text}) + + user_message: dict[str, Any] = {"role": "user", "content": content} + + if images: + user_message["images"] = list(images)[:OLLAMA_MAX_IMAGES] + + messages.append(user_message) + return messages + + def chat_stream( *, model: str | None, system_prompt: str, user_prompt: str, images: list[str] | None = None, + history: list[dict[str, Any]] | None = None, + num_ctx: int | None = None, timeout: float | None = None, temperature: float = 0.3, ): @@ -391,7 +660,12 @@ def chat_stream( appears progressively instead of all at once. `images` is an optional list of base64-encoded images (no data: prefix); - when present the caller should pass a vision-capable model. + when present the caller should pass a vision-capable model, and the context + window is widened automatically so the system prompt is not evicted by the + image tokens. + + Raises OllamaModelMissing when the model is not pulled, so the caller can + say so plainly instead of reporting a generic outage. """ selected_model = str(model or DEFAULT_OLLAMA_MODEL).strip() @@ -399,6 +673,12 @@ def chat_stream( if not selected_model: raise ValueError("An Ollama model name is required.") + if images and not is_vision_model(selected_model): + raise OllamaModelMissing( + f"'{selected_model}' cannot read images. Pull a vision model " + f"(for example: ollama pull {DEFAULT_OLLAMA_VISION_MODEL or 'qwen3-vl:4b'})." + ) + # (connect, read) timeout. Read is None = unlimited, so a slow but still # streaming answer is NEVER disconnected before it finishes — only a failure # to connect (Ollama down) errors out quickly. `timeout` (if passed) only @@ -406,77 +686,95 @@ def chat_stream( connect_timeout = timeout if timeout is not None else OLLAMA_CONNECT_TIMEOUT request_timeout = (connect_timeout, None) - content = user_prompt - # qwen3 is a reasoning model: without this it emits a long chain - # before answering (slow, and shows "steps" the user doesn't want). The - # "/no_think" directive turns that off so it answers directly and fast. - if not OLLAMA_THINK and "qwen3" in selected_model.lower(): - content = f"{content}\n\n/no_think" - - user_message: dict[str, Any] = { - "role": "user", - "content": content, - } - - if images: - user_message["images"] = images + if num_ctx is not None: + window = int(num_ctx) + elif images: + window = max(OLLAMA_NUM_CTX, OLLAMA_VISION_NUM_CTX) + else: + window = OLLAMA_NUM_CTX + + messages = _build_chat_messages( + system_prompt=system_prompt, + user_prompt=user_prompt, + history=history, + images=images, + selected_model=selected_model, + ) payload: dict[str, Any] = { "model": selected_model, "stream": True, "keep_alive": OLLAMA_KEEP_ALIVE, - "messages": [ - {"role": "system", "content": system_prompt}, - user_message, - ], + "messages": messages, "options": { "temperature": temperature, - "num_ctx": OLLAMA_NUM_CTX, + "num_ctx": window, "num_predict": OLLAMA_NUM_PREDICT, }, "think": OLLAMA_THINK, } - try: - with requests.post( - f"{OLLAMA_BASE_URL}/api/chat", - json=payload, - timeout=request_timeout, - stream=True, - ) as response: - response.raise_for_status() - - for line in response.iter_lines(decode_unicode=True): - if not line: - continue - - try: - chunk = json.loads(line) - except ValueError: - continue - - if not isinstance(chunk, dict): - continue - - if chunk.get("error"): - raise OllamaUnavailable(str(chunk.get("error"))) - - message = chunk.get("message") or {} - - thinking = message.get("thinking") - if thinking: - yield ("thinking", thinking) - - content = message.get("content") - if content: - yield ("content", content) - - if chunk.get("done"): - break - - except requests.Timeout as exc: - raise OllamaUnavailable( - "Ollama timed out while generating a response." - ) from exc - except (requests.RequestException, ValueError) as exc: - raise OllamaUnavailable(str(exc)) from exc + produced_any = False + + # Older Ollama builds reject the "think" field outright. Retrying once + # without it keeps those servers working instead of surfacing a hard outage + # for what is only an optional feature flag. + for attempt, body in enumerate( + (payload, {k: v for k, v in payload.items() if k != "think"}) + ): + try: + with requests.post( + f"{OLLAMA_BASE_URL}/api/chat", + json=body, + timeout=request_timeout, + stream=True, + ) as response: + response.raise_for_status() + + for line in response.iter_lines(decode_unicode=True): + if not line: + continue + + try: + chunk = json.loads(line) + except ValueError: + continue + + if not isinstance(chunk, dict): + continue + + if chunk.get("error"): + detail = str(chunk.get("error")) + if _describes_missing_model(detail): + raise OllamaModelMissing(detail) + raise OllamaUnavailable(detail) + + message = chunk.get("message") or {} + + thinking = message.get("thinking") + if thinking: + produced_any = True + yield ("thinking", thinking) + + content = message.get("content") + if content: + produced_any = True + yield ("content", content) + + if chunk.get("done"): + return + return + + except requests.Timeout as exc: + raise OllamaUnavailable( + "Ollama timed out while generating a response." + ) from exc + except requests.RequestException as exc: + # Only the "think" retry is safe to attempt: once tokens have been + # yielded, restarting would duplicate text in the user's reply. + status = getattr(getattr(exc, "response", None), "status_code", None) + if attempt == 0 and not produced_any and status == 400: + continue + _raise_for_chat_error(exc) + except ValueError as exc: + raise OllamaUnavailable(str(exc)) from exc diff --git a/flask-server/tests/unit/test_ai_reasoning.py b/flask-server/tests/unit/test_ai_reasoning.py new file mode 100644 index 00000000..bd287148 --- /dev/null +++ b/flask-server/tests/unit/test_ai_reasoning.py @@ -0,0 +1,201 @@ +"""Answer-shaping rules for the BodyMaps AI assistant. + +Each test below is a regression: the assistant shipped the failing behavior at +some point, and the case is written the way it was actually reported. +""" + +import services.ai_reasoning as ai_reasoning + + +BILIRUBIN_MESSAGE = """Total Bilirubin: 4.2 mg/dL (Elevated) +Direct (Conjugated) Bilirubin: 3.1 mg/dL (Conjugated predominance) +Indirect (Unconjugated) Bilirubin: 1.1 mg/dL + +RUQ Abdominal Ultrasound: Cholelithiasis present without gallbladder wall +thickening. Common bile duct (CBD) is dilated at 9.5 mm. +MRCP: Confirms a 5 mm gallstone in the distal common bile duct with upstream +intrahepatic and extrahepatic biliary ductal dilation.""" + +CASE_FACTS = [ + "**Liver:** segmented volume 1209.11 cm³, mean attenuation 67.9 HU", + "**Age:** 52.0", + "**Sex:** M", + "**Spleen:** segmented volume 180.40 cm³", +] + + +# --------------------------------------------------------------------------- +# Relevance gate +# --------------------------------------------------------------------------- + +def test_demographics_are_not_appended_to_a_lab_result_question(): + # Reported failure: pasting bilirubin values and an MRCP report was answered + # with "The liver volume is 1209.11 cm3 ... Age: 52.0". + kept = ai_reasoning.relevant_facts(CASE_FACTS, BILIRUBIN_MESSAGE) + + assert not any("Age" in fact for fact in kept) + assert not any("Sex" in fact for fact in kept) + + +def test_an_unrelated_organ_is_not_appended(): + kept = ai_reasoning.relevant_facts(CASE_FACTS, BILIRUBIN_MESSAGE) + + assert not any("Spleen" in fact for fact in kept) + + +def test_a_hepatobiliary_question_keeps_the_liver_fact(): + # "intrahepatic" is a liver reference, so the liver measurement is on topic. + kept = ai_reasoning.relevant_facts(CASE_FACTS, BILIRUBIN_MESSAGE) + + assert any("Liver" in fact for fact in kept) + + +def test_an_explicit_measurement_question_keeps_its_fact(): + kept = ai_reasoning.relevant_facts(CASE_FACTS, "What is the liver volume in this case?") + + assert any("1209.11" in fact for fact in kept) + + +def test_a_demographics_question_keeps_the_age(): + kept = ai_reasoning.relevant_facts(CASE_FACTS, "How old is this patient?") + + assert any("Age" in fact for fact in kept) + + +def test_forced_facts_survive_the_gate(): + forced = ["The segmented Pancreas volume is **82.10 cm³**."] + kept = ai_reasoning.relevant_facts([], "anything at all", always_include=forced) + + assert kept == forced + + +# --------------------------------------------------------------------------- +# Follow-up question guarantee +# --------------------------------------------------------------------------- + +def test_a_vision_answer_always_ends_by_asking_what_to_look_at_next(): + reply = ( + "In the axial pane the pancreatic head sits medial to the duodenum. " + "The contrast phase cannot be determined from this capture." + ) + + out = ai_reasoning.ensure_followup( + reply, + "trace the pancreas across the panes", + has_images=True, + has_case=True, + force=True, + ) + + assert out.rstrip().endswith("?") + assert "pancreas" in out.lower() + + +def test_an_answer_that_admits_missing_information_asks_for_it(): + reply = "I would need more information to say whether this is normal." + + out = ai_reasoning.ensure_followup( + reply, "is the pancreas normal", has_images=False, has_case=True + ) + + assert out.rstrip().endswith("?") + + +def test_a_complete_text_answer_is_left_alone(): + reply = "Hemochromatosis is the most likely diagnosis; check ferritin and the HFE gene." + + out = ai_reasoning.ensure_followup( + reply, "52-year-old man with fatigue", has_images=False, has_case=False + ) + + assert out == reply + + +def test_an_existing_closing_question_is_not_duplicated(): + reply = "The kidneys look symmetric. Want me to check a lower slice?" + + out = ai_reasoning.ensure_followup( + reply, "compare the kidneys", has_images=True, has_case=True, force=True + ) + + assert out == reply + + +def test_a_mid_paragraph_question_does_not_count_as_a_closing_question(): + reply = ( + "Why does a head tumor jaundice earlier? Because it obstructs the bile " + "duct sooner than a tail tumor does." + ) + + out = ai_reasoning.ensure_followup( + reply, "explain pancreatic head versus tail", has_images=True, has_case=True, force=True + ) + + assert out != reply + assert out.rstrip().endswith("?") + + +# --------------------------------------------------------------------------- +# Failure text +# --------------------------------------------------------------------------- + +def test_a_missing_vision_model_says_what_to_pull(): + reply = ai_reasoning.model_offline_reply( + has_images=True, + vision_model_missing=True, + configured_vision_model="qwen3-vl:4b", + ) + + assert "ollama pull qwen3-vl:4b" in reply + assert reply.rstrip().endswith("?") + + +def test_the_offline_reply_never_recites_measurements(): + reply = ai_reasoning.model_offline_reply(has_images=False) + + assert "cm³" not in reply + assert "HU" not in reply + + +# --------------------------------------------------------------------------- +# Prompts +# --------------------------------------------------------------------------- + +def test_the_vision_prompt_is_only_attached_when_images_are(): + with_images = ai_reasoning.build_system_prompt(has_images=True) + without_images = ai_reasoning.build_system_prompt(has_images=False) + + assert "ATTACHED CT VIEWER SCREENSHOTS" in with_images + assert "ATTACHED CT VIEWER SCREENSHOTS" not in without_images + + +def test_the_vision_prompt_covers_the_screenshot_artifacts(): + prompt = ai_reasoning.build_system_prompt(has_images=True) + + # The artifacts a model misreads without being told: mask colors read as + # pathology, crosshairs read as hardware, corner numbers read as patient data. + assert "SEGMENTATION MASKS" in prompt + assert "crosshair" in prompt.lower() + assert "window width/level" in prompt + + +def test_every_prompt_forbids_substituting_measurements_for_an_answer(): + prompt = ai_reasoning.build_system_prompt(has_images=False) + + assert "STAY ON THE QUESTION" in prompt + + +def test_the_legend_fact_lists_every_visible_organ_color(): + legend = [ + {"organ": "liver", "color": "brownish red"}, + {"organ": "left_kidney", "color": "teal"}, + ] + + fact = ai_reasoning.build_legend_fact(legend) + + assert "liver: brownish red" in fact + assert "left kidney: teal" in fact + + +def test_no_legend_produces_no_fact(): + assert ai_reasoning.build_legend_fact([]) is None diff --git a/flask-server/tests/unit/test_ollama_client.py b/flask-server/tests/unit/test_ollama_client.py index 542943af..9bbfdb6f 100644 --- a/flask-server/tests/unit/test_ollama_client.py +++ b/flask-server/tests/unit/test_ollama_client.py @@ -52,3 +52,176 @@ def test_chat_structured_json_requires_object_schema(): messages=[{"role": "user", "content": "inspect"}], schema={"type": "array"}, ) + + +# --------------------------------------------------------------------------- +# Vision model resolution. +# +# Image messages used to be sent to whatever BODYMAPS_OLLAMA_VISION_MODEL named, +# installed or not. When it was not (qwen3-vl needs Ollama 0.12.7+), the very +# first byte failed and the user was told the whole assistant was unavailable. +# --------------------------------------------------------------------------- + +def _fake_tags(monkeypatch, names, ok=True): + def fake_installed(force=False): + return list(names), ok + + monkeypatch.setattr(ollama_client, "installed_model_names", fake_installed) + + +def test_the_configured_vision_model_wins_when_it_is_installed(monkeypatch): + _fake_tags(monkeypatch, ["llama3.1:latest", "qwen3-vl:4b"]) + + assert ollama_client.resolve_vision_model() == "qwen3-vl:4b" + + +def test_an_installed_vision_model_is_used_when_the_configured_one_is_missing(monkeypatch): + _fake_tags(monkeypatch, ["llama3.1:latest", "llava:13b"]) + + assert ollama_client.resolve_vision_model() == "llava:13b" + + +def test_no_vision_model_installed_resolves_to_none(monkeypatch): + _fake_tags(monkeypatch, ["llama3.1:latest", "qwen3:4b"]) + + assert ollama_client.resolve_vision_model() is None + + +def test_an_unreachable_tag_listing_is_not_treated_as_nothing_installed(monkeypatch): + # "The listing timed out" is not evidence a model is absent — try anyway. + _fake_tags(monkeypatch, [], ok=False) + + assert ollama_client.resolve_vision_model() == ollama_client.DEFAULT_OLLAMA_VISION_MODEL + + +def test_a_bare_name_matches_the_latest_tag(monkeypatch): + _fake_tags(monkeypatch, ["llama3.1:latest"]) + + assert ollama_client.model_is_installed("llama3.1") + + +def test_qwen3_vl_is_not_classified_as_a_reasoning_model(): + # It is an instruct model: sending it "/no_think" only injects stray text + # into a vision prompt. + assert not ollama_client.is_reasoning_model("qwen3-vl:4b") + assert ollama_client.is_reasoning_model("qwen3:4b") + + +def test_vision_families_are_recognized(): + for name in ["qwen3-vl:4b", "llava:13b", "llama3.2-vision:11b", "minicpm-v", "gemma3:4b"]: + assert ollama_client.is_vision_model(name), name + for name in ["llama3.1:latest", "qwen3:4b", "mistral:7b"]: + assert not ollama_client.is_vision_model(name), name + + +# --------------------------------------------------------------------------- +# chat_stream request shape +# --------------------------------------------------------------------------- + +def test_images_are_refused_by_a_text_only_model(): + with pytest.raises(ollama_client.OllamaModelMissing): + list( + ollama_client.chat_stream( + model="llama3.1:latest", + system_prompt="s", + user_prompt="u", + images=["base64"], + ) + ) + + +class _FakeResponse: + def __init__(self, lines): + self._lines = lines + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def raise_for_status(self): + return None + + def iter_lines(self, decode_unicode=False): + return iter(self._lines) + + +def _capture_stream(monkeypatch): + captured = {} + + def fake_post(url, json=None, timeout=None, stream=None): + captured["payload"] = json + return _FakeResponse(['{"message":{"content":"hi"},"done":true}']) + + monkeypatch.setattr(ollama_client.requests, "post", fake_post) + return captured + + +def test_an_image_turn_gets_a_wider_context_window(monkeypatch): + # Four 768px panes plus a long question overflow the 8k text window, and + # Ollama then evicts the OLDEST tokens — the system prompt carrying every + # instruction about how to read a screenshot. + captured = _capture_stream(monkeypatch) + + list( + ollama_client.chat_stream( + model="qwen3-vl:4b", + system_prompt="s", + user_prompt="u", + images=["a", "b", "c", "d"], + ) + ) + + assert captured["payload"]["options"]["num_ctx"] >= ollama_client.OLLAMA_VISION_NUM_CTX + assert captured["payload"]["messages"][-1]["images"] == ["a", "b", "c", "d"] + + +def test_no_think_is_not_injected_into_a_vision_prompt(monkeypatch): + captured = _capture_stream(monkeypatch) + + list( + ollama_client.chat_stream( + model="qwen3-vl:4b", + system_prompt="s", + user_prompt="read these views", + images=["a"], + ) + ) + + assert "/no_think" not in captured["payload"]["messages"][-1]["content"] + + +def test_prior_turns_are_sent_as_real_conversation_messages(monkeypatch): + captured = _capture_stream(monkeypatch) + + list( + ollama_client.chat_stream( + model="llama3.1:latest", + system_prompt="s", + user_prompt="here are the labs", + history=[ + {"role": "user", "content": "68-year-old woman with jaundice"}, + {"role": "assistant", "content": "What is her bilirubin?"}, + ], + ) + ) + + roles = [m["role"] for m in captured["payload"]["messages"]] + assert roles == ["system", "user", "assistant", "user"] + + +def test_too_many_images_are_capped(monkeypatch): + captured = _capture_stream(monkeypatch) + + list( + ollama_client.chat_stream( + model="qwen3-vl:4b", + system_prompt="s", + user_prompt="u", + images=[str(i) for i in range(12)], + ) + ) + + sent = captured["payload"]["messages"][-1]["images"] + assert len(sent) == ollama_client.OLLAMA_MAX_IMAGES