feat: stage-specific prompts + contextual ROS

app/graph.py

@@ -119,12 +119,19 @@ def agent_node(state: IntakeState) -> dict:
     current_json = state.get("clinical_state") or CombinedOutput().model_dump_json()
     transcript = format_transcript(msgs)
 
+    # Compute the current stage BEFORE the LLM call so we can pick the right prompt
+    try:
+        pre_state = CombinedOutput.model_validate_json(current_json)
+        current_stage = compute_stage(pre_state)
+    except Exception:
+        current_stage = "intake"
+
     import time
     t_agent = time.time()
-    print(f"[{time.time():.3f}] [Graph Node] Requesting LLM inference...")
+    print(f"[{time.time():.3f}] [Graph Node] Requesting LLM inference (stage={current_stage})...")
 
     llm = get_llm()
-    result: CombinedOutput = llm.combined_call(transcript, current_json)
+    result: CombinedOutput = llm.combined_call(transcript, current_json, stage=current_stage)
 
     # ── Loop Guard: if LLM returned same reply as last turn, force-fill stuck field ──
     if _detect_repeat({"messages": msgs + [{"role": "assistant", "content": result.reply}]}):
@@ -141,11 +148,6 @@ def agent_node(state: IntakeState) -> dict:
                 break
 
     # ── ROS Hallucination Guard: LLM can only ADD one new ROS system per turn ──
-    ROS_QUESTIONS = {
-        "cardiac": "Have you experienced any palpitations, leg swelling, or dizziness?",
-        "respiratory": "Have you had any shortness of breath, coughing, or wheezing?",
-        "gi": "Have you had any nausea, vomiting, or heartburn?",
-    }
     try:
         prev_state = json.loads(current_json)
         prev_ros = prev_state.get("ros") or {}
@@ -153,7 +155,7 @@ def agent_node(state: IntakeState) -> dict:
         prev_ros = {}
     new_ros_keys = [k for k in result.ros if k not in prev_ros]
     if len(new_ros_keys) > 1:
-        print(f"[ROSGuard] LLM hallucinated {len(new_ros_keys)} new ROS systems in one turn: {new_ros_keys}. Keeping only first.")
+        print(f"[ROSGuard] LLM added {len(new_ros_keys)} new ROS systems in one turn: {new_ros_keys}. Keeping only first.")
         allowed_ros = dict(prev_ros)
         allowed_ros[new_ros_keys[0]] = result.ros[new_ros_keys[0]]
         object.__setattr__(result, "ros", allowed_ros)
@@ -162,19 +164,7 @@ def agent_node(state: IntakeState) -> dict:
 
     stage = compute_stage(result)
     missing = missing_from(result)
-
-    # ── ROS Question Forcing: if all HPI done but ROS incomplete, force a specific ROS question ──
-    if stage == "ros":
-        current_ros = result.ros or {}
-        for sys_name, question in ROS_QUESTIONS.items():
-            if sys_name not in current_ros:
-                print(f"[ROSForce] Forcing question for missing ROS system: {sys_name}")
-                reply = question
-                break
-        else:
-            reply = result.reply or "Could you tell me more?"
-    else:
-        reply = result.reply or "Could you tell me more?"
+    reply = result.reply or "Could you tell me more?"
 
     # All fields complete — build the brief inline so it's available this turn
     if stage == "done":

app/llm.py

@@ -3,46 +3,97 @@ import json
 import re
 from pydantic import BaseModel
 
-COMBINED_SYSTEM_PROMPT = """You are a clinical intake assistant AI. You have two jobs per turn:
+INTAKE_PROMPT = """You are a clinical intake assistant. The patient just arrived.
 
-JOB 1 (EXTRACT): Read the FULL conversation and update the clinical JSON state with any new information the patient provided.
-CRITICAL: If the patient denies a symptom, or replies with "none", "zero", "no", or "nothing", you MUST extract that exact word (e.g. "zero"). DO NOT leave it null if the patient has answered the question negatively.
+JOB: Extract the chief complaint from the conversation. Ask ONE simple question to identify their main symptom.
 
-JOB 2 (RESPOND): Based on what is STILL MISSING from the clinical state, ask the patient ONE natural, empathetic question. Do NOT ask about things already filled in.
-
-CRITICAL RULES:
-- Output ONLY valid JSON, nothing else.
+RULES:
+- Output ONLY valid JSON.
+- If you already know the chief complaint, ask about onset to move forward.
 - Do NOT diagnose or give medical advice.
-- Do NOT ask more than one question.
-- If all fields are complete, set reply to "Thank you — I have everything I need."
 
-OUTPUT FORMAT (strictly follow this, no extra text):
+OUTPUT FORMAT:
+{
+  "chief_complaint": "the main symptom" or null,
+  "onset": null, "location": null, "duration": null,
+  "character": null, "severity": null, "aggravating": null, "relieving": null,
+  "ros": {},
+  "reply": "Your question to the patient"
+}"""
+
+HPI_PROMPT = """You are a clinical intake assistant collecting History of Present Illness (HPI) using OLDCARTS.
+
+JOB 1 (EXTRACT): Read the conversation and update the JSON with any new patient info. If a patient denies something or says "none"/"zero"/"no", store that exact word — do NOT leave it null.
+
+JOB 2 (RESPOND): Ask ONE question about the FIRST missing field below. Do NOT re-ask fields already filled.
+
+FIELDS TO COLLECT (in order):
+- onset: when the symptom started
+- location: where in the body
+- duration: how long it has lasted
+- character: quality of pain (sharp, dull, pressure, burning, etc.)
+- severity: how bad on a scale of 1-10
+- aggravating: what makes it worse
+- relieving: what makes it better
+
+RULES:
+- Output ONLY valid JSON, no extra text.
+- Ask exactly ONE question per turn.
+- Keep existing values. Use null for unknowns.
+
+OUTPUT FORMAT:
+{
+  "chief_complaint": "...",
+  "onset": "..." or null,
+  "location": "..." or null,
+  "duration": "..." or null,
+  "character": "..." or null,
+  "severity": "..." or null,
+  "aggravating": "..." or null,
+  "relieving": "..." or null,
+  "ros": {},
+  "reply": "Your single question"
+}"""
+
+ROS_PROMPT = """You are a clinical intake assistant performing a Review of Systems (ROS).
+
+All HPI fields are already collected. Now you must screen for symptoms in OTHER body systems that are RELEVANT to the patient's chief complaint.
+
+JOB 1 (EXTRACT): The patient just answered a question about a body system. Extract their answer into the "ros" dict under the appropriate system key (e.g. "musculoskeletal": ["joint stiffness", "no swelling"]).
+
+JOB 2 (RESPOND): Ask about the NEXT relevant body system that is NOT yet in the "ros" dict.
+
+CHOOSING SYSTEMS: Pick 3 systems that are clinically relevant to the chief complaint. Examples:
+- Leg/knee/joint pain → musculoskeletal, neurological, vascular
+- Chest pain → cardiac, respiratory, gi
+- Headache → neurological, ophthalmologic, ent
+- Abdominal pain → gi, genitourinary, musculoskeletal
+- Back pain → musculoskeletal, neurological, genitourinary
+
+RULES:
+- Output ONLY valid JSON.
+- Ask about ONE system at a time.
+- If the patient denies symptoms, store as ["no X", "no Y"].
+- Once 3 systems are in "ros", set reply to "Thank you — I have everything I need."
+- Do NOT ask emotional, psychological, or off-topic questions.
+
+OUTPUT FORMAT:
 {
-  "chief_complaint": "the main symptom or reason for visit",
-  "onset": "when the symptom started",
-  "location": "where in the body",
-  "duration": "how long it has lasted, constant or intermittent",
-  "character": "quality of pain: sharp, dull, tightening, pressure, burning, squeezing, etc.",
-  "severity": "how bad, e.g. mild, moderate, severe, or a number out of 10",
-  "aggravating": "what makes it worse",
-  "relieving": "what makes it better or go away",
-  "ros": {"cardiac": ["findings"], "respiratory": ["findings"], "gi": ["findings"]},
-  "reply": "The single question to ask the patient next"
-}
-
-REVIEW OF SYSTEMS (ROS): Once all HPI fields above are filled, ask about these 3 systems ONE AT A TIME:
-1. Cardiac: palpitations, leg swelling, dizziness
-2. Respiratory: shortness of breath, cough, wheezing
-3. GI: nausea, vomiting, heartburn
-For each system the patient denies symptoms, store as ["no palpitations", "no leg swelling"]. Do NOT ask emotional or psychological questions — stick to the 3 systems above.
-
-Use null for any field not yet known. Keep existing values if the patient didn't add new info.
-
-IMPORTANT — ACCEPTING VAGUE ANSWERS:
-- If the patient gives ANY answer (even "none", "zero", "not sure", "it goes away", "very mild"), that IS a valid value. Store it as a string.
-- For relieving/aggravating: if patient implies rest helps (e.g. "very mild when not running", "zero at rest"), set relieving="rest" and aggravating="physical activity/running".
-- Do NOT ask the same question twice. If the patient has answered (even vaguely), move on to the next missing field.
-- "zero", "none", "not really", "it's fine otherwise" → treat as valid answer, fill the field."""
+  "chief_complaint": "...", "onset": "...", "location": "...", "duration": "...",
+  "character": "...", "severity": "...", "aggravating": "...", "relieving": "...",
+  "ros": {"system_name": ["findings"], ...},
+  "reply": "Your single ROS question"
+}"""
+
+
+def get_system_prompt(stage: str) -> str:
+    """Return the appropriate system prompt for the current clinical stage."""
+    if stage == "ros":
+        return ROS_PROMPT
+    elif stage == "hpi":
+        return HPI_PROMPT
+    else:
+        return INTAKE_PROMPT
 
 
 class CombinedOutput(BaseModel):
@@ -60,7 +111,7 @@ class CombinedOutput(BaseModel):
 
 
 class MockLLM:
-    def combined_call(self, transcript: str, current_json: str) -> CombinedOutput:
+    def combined_call(self, transcript: str, current_json: str, stage: str = "intake") -> CombinedOutput:
         """Single call: extract + generate reply. No real inference in mock mode."""
         t = transcript.lower()
         try:
@@ -163,7 +214,7 @@ class OllamaLLM:
         self.model_name = os.environ.get("MODEL_NAME", "qwen2.5:0.5b")
         self.api_url = "http://localhost:11434/api/chat"
 
-    def combined_call(self, transcript: str, current_json: str) -> CombinedOutput:
+    def combined_call(self, transcript: str, current_json: str, stage: str = "intake") -> CombinedOutput:
         """
         Calls the local Ollama instance using the /chat endpoint so system tags 
         are properly applied.
@@ -185,7 +236,7 @@ class OllamaLLM:
         payload = {
             "model": self.model_name,
             "messages": [
-                {"role": "system", "content": COMBINED_SYSTEM_PROMPT},
+                {"role": "system", "content": get_system_prompt(stage)},
                 {"role": "user", "content": prompt}
             ],
             "format": "json",