Update app.py

app.py

@@ -7,31 +7,32 @@ from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.tools import tool
 from langgraph.prebuilt import ToolExecutor
 from langgraph.graph import StateGraph, END
-from langgraph.checkpoint.memory import MemorySaver # For state persistence (optional but good)
+# from langgraph.checkpoint.memory import MemorySaverInMemory # Optional for state persistence
 
 from typing import Optional, List, Dict, Any, TypedDict, Annotated
 import json
 import re
 import operator
+import traceback # For detailed error logging
 
-# --- Configuration & Constants --- (Keep previous ones like ClinicalAppSettings)
+# --- Configuration & Constants ---
 class ClinicalAppSettings:
     APP_TITLE = "SynapseAI: Interactive Clinical Decision Support"
     PAGE_LAYOUT = "wide"
-    MODEL_NAME = "llama3-70b-8192"
+    MODEL_NAME = "llama3-70b-8192" # Groq Llama3 70b
     TEMPERATURE = 0.1
     MAX_SEARCH_RESULTS = 3
 
 class ClinicalPrompts:
-    # UPDATED SYSTEM PROMPT FOR CONVERSATIONAL FLOW & GUIDELINES
+    # UPDATED SYSTEM PROMPT FOR CONVERSATIONAL FLOW, GUIDELINES & STRUCTURED OUTPUT FOCUS
     SYSTEM_PROMPT = """
     You are SynapseAI, an expert AI clinical assistant engaged in an interactive consultation.
     Your goal is to support healthcare professionals by analyzing patient data, providing differential diagnoses, suggesting evidence-based management plans, and identifying risks according to current standards of care.
 
     **Core Directives for this Conversation:**
-    1.  **Analyze Sequentially:** Process information turn-by-turn. You will receive initial patient data, and potentially follow-up messages or results from tools you requested. Base your responses on the *entire* conversation history.
+    1.  **Analyze Sequentially:** Process information turn-by-turn. Base your responses on the *entire* conversation history.
     2.  **Seek Clarity:** If the provided information is insufficient or ambiguous for a safe assessment, CLEARLY STATE what specific additional information or clarification is needed. Do NOT guess or make unsafe assumptions.
-    3.  **Structured Assessment (When Ready):** When you have sufficient information and have performed necessary checks (like interactions), provide a comprehensive assessment using the following JSON structure. Only output this structure when you believe you have a complete initial analysis or plan. Do NOT output incomplete JSON.
+    3.  **Structured Assessment (When Ready):** When you have sufficient information and have performed necessary checks (like interactions, guideline searches), provide a comprehensive assessment using the following JSON structure. Output this JSON structure as the primary content of your response when you are providing the full analysis. Do NOT output incomplete JSON. If you need to ask a question or perform a tool call first, do that instead of outputting this structure.
         ```json
         {
           "assessment": "Concise summary of the patient's presentation and key findings based on the conversation.",
@@ -41,137 +42,212 @@ class ClinicalPrompts:
             {"diagnosis": "Alternative Diagnosis 2", "likelihood": "Low", "rationale": "Why it's less likely but considered..."}
           ],
           "risk_assessment": {
-            "identified_red_flags": ["List any triggered red flags"],
-            "immediate_concerns": ["Specific urgent issues (e.g., sepsis risk, ACS rule-out)"],
-            "potential_complications": ["Possible future issues"]
+            "identified_red_flags": ["List any triggered red flags based on input and analysis"],
+            "immediate_concerns": ["Specific urgent issues requiring attention (e.g., sepsis risk, ACS rule-out)"],
+            "potential_complications": ["Possible future issues based on presentation"]
           },
           "recommended_plan": {
-            "investigations": ["List specific lab tests or imaging needed. Use 'order_lab_test' tool."],
-            "therapeutics": ["Suggest specific treatments/prescriptions. Use 'prescribe_medication' tool. MUST check interactions first."],
-            "consultations": ["Recommend specialist consultations."],
+            "investigations": ["List specific lab tests or imaging required. Use 'order_lab_test' tool."],
+            "therapeutics": ["Suggest specific treatments or prescriptions. Use 'prescribe_medication' tool. MUST check interactions first using 'check_drug_interactions'."],
+            "consultations": ["Recommend specialist consultations if needed."],
             "patient_education": ["Key points for patient communication."]
           },
-          "rationale_summary": "Justification for assessment/plan. **Crucially, if relevant (e.g., ACS, sepsis, common infections), use 'tavily_search_results' to find and cite current clinical practice guidelines (e.g., 'latest ACC/AHA chest pain guidelines 202X', 'Surviving Sepsis Campaign guidelines') supporting your recommendations.**",
+          "rationale_summary": "Justification for assessment/plan. **Crucially, if relevant (e.g., ACS, sepsis, common infections), use 'tavily_search_results' to find and cite current clinical practice guidelines (e.g., 'latest ACC/AHA chest pain guidelines 202X', 'Surviving Sepsis Campaign guidelines') supporting your recommendations.** Include summary of guideline findings here.",
           "interaction_check_summary": "Summary of findings from 'check_drug_interactions' if performed."
         }
         ```
-    4.  **Safety First - Interactions:** BEFORE suggesting a new prescription via `prescribe_medication`, you MUST FIRST use `check_drug_interactions`. Report the findings. If interactions exist, modify the plan or state the contraindication.
-    5.  **Safety First - Red Flags:** Use the `flag_risk` tool IMMEDIATELY if critical red flags requiring urgent action are identified at any point.
-    6.  **Tool Use:** Employ tools (`order_lab_test`, `prescribe_medication`, `check_drug_interactions`, `flag_risk`, `tavily_search_results`) logically within the conversational flow. Wait for tool results before proceeding if the result is needed for the next step (e.g., wait for interaction check before confirming prescription).
+    4.  **Safety First - Interactions:** BEFORE suggesting a new prescription via `prescribe_medication`, you MUST FIRST use `check_drug_interactions` in a preceding or concurrent tool call. Report the findings from the interaction check. If significant interactions exist, modify the plan or state the contraindication clearly.
+    5.  **Safety First - Red Flags:** Use the `flag_risk` tool IMMEDIATELY if critical red flags requiring urgent action are identified at any point in the conversation.
+    6.  **Tool Use:** Employ tools (`order_lab_test`, `prescribe_medication`, `check_drug_interactions`, `flag_risk`, `tavily_search_results`) logically within the conversational flow. Wait for tool results before proceeding if the result is needed for the next step (e.g., wait for interaction check before confirming prescription in the structured JSON).
     7.  **Evidence & Guidelines:** Actively use `tavily_search_results` not just for general knowledge, but specifically to query for and incorporate **current clinical practice guidelines** relevant to the patient's presentation (e.g., chest pain, shortness of breath, suspected infection). Summarize findings in the `rationale_summary` when providing the structured output.
-    8.  **Conciseness:** Be medically accurate and concise. Use standard terminology. Respond naturally in conversation until ready for the full structured JSON output.
+    8.  **Conciseness & Flow:** Be medically accurate and concise. Use standard terminology. Respond naturally in conversation (asking questions, acknowledging info) until ready for the full structured JSON output.
     """
 
-# --- Mock Data / Helpers --- (Keep previous ones like MOCK_INTERACTION_DB, ALLERGY_INTERACTIONS, parse_bp, check_red_flags)
-# (Include the helper functions from the previous response here)
+# --- Mock Data / Helpers ---
 MOCK_INTERACTION_DB = {
     ("lisinopril", "spironolactone"): "High risk of hyperkalemia. Monitor potassium closely.",
     ("warfarin", "amiodarone"): "Increased bleeding risk. Monitor INR frequently and adjust Warfarin dose.",
     ("simvastatin", "clarithromycin"): "Increased risk of myopathy/rhabdomyolysis. Avoid combination or use lower statin dose.",
-    ("aspirin", "ibuprofen"): "Concurrent use may decrease Aspirin's cardioprotective effect. Potential for increased GI bleeding."
+    ("aspirin", "ibuprofen"): "Concurrent use may decrease Aspirin's cardioprotective effect. Potential for increased GI bleeding.",
+    # Add lower case versions for easier lookup
+    ("amiodarone", "warfarin"): "Increased bleeding risk. Monitor INR frequently and adjust Warfarin dose.",
+    ("clarithromycin", "simvastatin"): "Increased risk of myopathy/rhabdomyolysis. Avoid combination or use lower statin dose.",
+    ("ibuprofen", "aspirin"): "Concurrent use may decrease Aspirin's cardioprotective effect. Potential for increased GI bleeding.",
+    ("spironolactone", "lisinopril"): "High risk of hyperkalemia. Monitor potassium closely.",
 }
 
 ALLERGY_INTERACTIONS = {
-    "penicillin": ["amoxicillin", "ampicillin", "piperacillin"],
-    "sulfa": ["sulfamethoxazole", "sulfasalazine"],
-    "aspirin": ["ibuprofen", "naproxen"] # Cross-reactivity example for NSAIDs
+    "penicillin": ["amoxicillin", "ampicillin", "piperacillin", "augmentin"],
+    "sulfa": ["sulfamethoxazole", "sulfasalazine", "bactrim"],
+    "aspirin": ["ibuprofen", "naproxen", "nsaid"] # Cross-reactivity example for NSAIDs
 }
 
 def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
-    match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string)
-    if match: return int(match.group(1)), int(match.group(2))
+    """Parses BP string like '120/80' into (systolic, diastolic) integers."""
+    if not isinstance(bp_string, str): return None
+    match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string.strip())
+    if match:
+        return int(match.group(1)), int(match.group(2))
     return None
 
 def check_red_flags(patient_data: dict) -> List[str]:
+    """Checks patient data against predefined red flags."""
     flags = []
+    if not patient_data: return flags
+
     symptoms = patient_data.get("hpi", {}).get("symptoms", [])
     vitals = patient_data.get("vitals", {})
     history = patient_data.get("pmh", {}).get("conditions", "")
-    symptoms_lower = [s.lower() for s in symptoms]
+    # Ensure symptoms are strings and lowercased
+    symptoms_lower = [str(s).lower() for s in symptoms if isinstance(s, str)]
 
+    # Symptom Flags
     if "chest pain" in symptoms_lower: flags.append("Red Flag: Chest Pain reported.")
     if "shortness of breath" in symptoms_lower: flags.append("Red Flag: Shortness of Breath reported.")
     if "severe headache" in symptoms_lower: flags.append("Red Flag: Severe Headache reported.")
-    # Add other symptom checks...
-
-    if "temp_c" in vitals and vitals["temp_c"] >= 38.5: flags.append(f"Red Flag: Fever ({vitals['temp_c']}°C).")
-    if "hr_bpm" in vitals and vitals["hr_bpm"] >= 120: flags.append(f"Red Flag: Tachycardia ({vitals['hr_bpm']} bpm).")
-    if "bp_mmhg" in vitals:
-        bp = parse_bp(vitals["bp_mmhg"])
-        if bp and (bp[0] >= 180 or bp[1] >= 110): flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {vitals['bp_mmhg']} mmHg).")
-        if bp and (bp[0] <= 90 or bp[1] <= 60): flags.append(f"Red Flag: Hypotension (BP: {vitals['bp_mmhg']} mmHg).")
-    # Add other vital checks...
+    if "sudden vision loss" in symptoms_lower: flags.append("Red Flag: Sudden Vision Loss reported.")
+    if "weakness on one side" in symptoms_lower: flags.append("Red Flag: Unilateral Weakness reported (potential stroke).")
+    if "hemoptysis" in symptoms_lower: flags.append("Red Flag: Hemoptysis (coughing up blood).")
+    if "syncope" in symptoms_lower: flags.append("Red Flag: Syncope (fainting).")
+
+    # Vital Sign Flags
+    if vitals:
+        temp = vitals.get("temp_c")
+        hr = vitals.get("hr_bpm")
+        rr = vitals.get("rr_rpm")
+        spo2 = vitals.get("spo2_percent")
+        bp_str = vitals.get("bp_mmhg")
+
+        if temp is not None and temp >= 38.5: flags.append(f"Red Flag: Fever (Temperature: {temp}°C).")
+        if hr is not None and hr >= 120: flags.append(f"Red Flag: Tachycardia (Heart Rate: {hr} bpm).")
+        if hr is not None and hr <= 50: flags.append(f"Red Flag: Bradycardia (Heart Rate: {hr} bpm).")
+        if rr is not None and rr >= 24: flags.append(f"Red Flag: Tachypnea (Respiratory Rate: {rr} rpm).")
+        if spo2 is not None and spo2 <= 92: flags.append(f"Red Flag: Hypoxia (SpO2: {spo2}%).")
+        if bp_str:
+            bp = parse_bp(bp_str)
+            if bp:
+                if bp[0] >= 180 or bp[1] >= 110: flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {bp_str} mmHg).")
+                if bp[0] <= 90 or bp[1] <= 60: flags.append(f"Red Flag: Hypotension (BP: {bp_str} mmHg).")
+
+    # History Flags (Simple examples)
+    if history and "history of mi" in history.lower() and "chest pain" in symptoms_lower:
+        flags.append("Red Flag: History of MI with current Chest Pain.")
+    if history and "history of dvt/pe" in history.lower() and "shortness of breath" in symptoms_lower:
+         flags.append("Red Flag: History of DVT/PE with current Shortness of Breath.")
+
+    # Remove duplicates
+    return list(set(flags))
 
-    if "history of mi" in history.lower() and "chest pain" in symptoms_lower: flags.append("Red Flag: History of MI with current Chest Pain.")
-    # Add other history checks...
-    return flags
+def format_patient_data_for_prompt(data: dict) -> str:
+    """Formats the patient dictionary into a readable string for the LLM."""
+    if not data: return "No patient data provided."
+    prompt_str = ""
+    for key, value in data.items():
+        section_title = key.replace('_', ' ').title()
+        if isinstance(value, dict) and value:
+            has_content = any(sub_value for sub_value in value.values())
+            if has_content:
+                prompt_str += f"**{section_title}:**\n"
+                for sub_key, sub_value in value.items():
+                     if sub_value:
+                        prompt_str += f"  - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
+        elif isinstance(value, list) and value:
+             prompt_str += f"**{section_title}:** {', '.join(map(str, value))}\n"
+        elif value and not isinstance(value, dict): # Check it's not an empty dict
+             prompt_str += f"**{section_title}:** {value}\n"
+    return prompt_str.strip()
 
 
-# --- Enhanced Tool Definitions --- (Keep previous Pydantic models and @tool functions)
-# (Include LabOrderInput, PrescriptionInput, InteractionCheckInput, FlagRiskInput
-# and the corresponding @tool functions: order_lab_test, prescribe_medication,
-# check_drug_interactions, flag_risk from the previous response here)
+# --- Tool Definitions ---
 
+# Pydantic models for robust argument validation
 class LabOrderInput(BaseModel):
-    test_name: str = Field(..., description="Specific name of the lab test or panel (e.g., 'CBC', 'BMP', 'Troponin I', 'Urinalysis').")
-    reason: str = Field(..., description="Clinical justification for ordering the test (e.g., 'Rule out infection', 'Assess renal function', 'Evaluate for ACS').")
+    test_name: str = Field(..., description="Specific name of the lab test or panel (e.g., 'CBC', 'BMP', 'Troponin I', 'Urinalysis', 'D-dimer').")
+    reason: str = Field(..., description="Clinical justification for ordering the test (e.g., 'Rule out infection', 'Assess renal function', 'Evaluate for ACS', 'Assess for PE').")
     priority: str = Field("Routine", description="Priority of the test (e.g., 'STAT', 'Routine').")
 
 @tool("order_lab_test", args_schema=LabOrderInput)
 def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
     """Orders a specific lab test with clinical justification and priority."""
-    return json.dumps({"status": "success", "message": f"Lab Ordered: {test_name} ({priority})", "details": f"Reason: {reason}"})
+    print(f"Executing order_lab_test: {test_name}, Reason: {reason}, Priority: {priority}")
+    # In a real system, this would integrate with an LIS/EMR API
+    return json.dumps({
+        "status": "success",
+        "message": f"Lab Ordered: {test_name} ({priority})",
+        "details": f"Reason: {reason}"
+    })
 
 class PrescriptionInput(BaseModel):
     medication_name: str = Field(..., description="Name of the medication.")
-    dosage: str = Field(..., description="Dosage amount and unit (e.g., '500 mg', '10 mg').")
-    route: str = Field(..., description="Route of administration (e.g., 'PO', 'IV', 'IM', 'Topical').")
-    frequency: str = Field(..., description="How often the medication should be taken (e.g., 'BID', 'QDaily', 'Q4-6H PRN').")
-    duration: str = Field("As directed", description="Duration of treatment (e.g., '7 days', '1 month', 'Until follow-up').")
+    dosage: str = Field(..., description="Dosage amount and unit (e.g., '500 mg', '10 mg', '81 mg').")
+    route: str = Field(..., description="Route of administration (e.g., 'PO', 'IV', 'IM', 'Topical', 'SL').")
+    frequency: str = Field(..., description="How often the medication should be taken (e.g., 'BID', 'QDaily', 'Q4-6H PRN', 'once').")
+    duration: str = Field("As directed", description="Duration of treatment (e.g., '7 days', '1 month', 'Ongoing', 'Until follow-up').")
     reason: str = Field(..., description="Clinical indication for the prescription.")
 
 @tool("prescribe_medication", args_schema=PrescriptionInput)
 def prescribe_medication(medication_name: str, dosage: str, route: str, frequency: str, duration: str, reason: str) -> str:
-    """Prescribes a medication with detailed instructions and clinical indication."""
-    # NOTE: Interaction check should have been done *before* calling this via a separate tool call
-    return json.dumps({"status": "success", "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}", "details": f"Duration: {duration}. Reason: {reason}"})
+    """Prescribes a medication with detailed instructions and clinical indication. IMPORTANT: Requires prior interaction check."""
+    print(f"Executing prescribe_medication: {medication_name} {dosage}...")
+    # NOTE: The safety check (ensuring interaction check was requested) happens in the tool_node *before* this function is called.
+    # In a real system, this would trigger an e-prescription workflow.
+    return json.dumps({
+        "status": "success",
+        "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}",
+        "details": f"Duration: {duration}. Reason: {reason}"
+    })
 
 class InteractionCheckInput(BaseModel):
-    potential_prescription: str = Field(..., description="The name of the NEW medication being considered.")
-    current_medications: List[str] = Field(..., description="List of the patient's CURRENT medication names.")
-    allergies: List[str] = Field(..., description="List of the patient's known allergies.")
+    potential_prescription: str = Field(..., description="The name of the NEW medication being considered for prescribing.")
+    # These next two args are now populated by the tool_node using AgentState
+    # current_medications: List[str] = Field(..., description="List of the patient's CURRENT medication names.")
+    # allergies: List[str] = Field(..., description="List of the patient's known allergies.")
+    # Make them optional in the schema, mandatory in the node logic
+    current_medications: Optional[List[str]] = Field(None, description="List of patient's current medication names (populated from state).")
+    allergies: Optional[List[str]] = Field(None, description="List of patient's known allergies (populated from state).")
+
 
 @tool("check_drug_interactions", args_schema=InteractionCheckInput)
-def check_drug_interactions(potential_prescription: str, current_medications: List[str], allergies: List[str]) -> str:
+def check_drug_interactions(potential_prescription: str, current_medications: Optional[List[str]] = None, allergies: Optional[List[str]] = None) -> str:
     """Checks for potential drug-drug and drug-allergy interactions BEFORE prescribing."""
+    print(f"Executing check_drug_interactions for: {potential_prescription}")
     warnings = []
     potential_med_lower = potential_prescription.lower()
-    current_meds_lower = [med.lower() for med in current_medications]
-    allergies_lower = [a.lower() for a in allergies]
 
+    # Use provided lists or default to empty if None (should be populated by tool_node)
+    current_meds_list = current_medications or []
+    allergies_list = allergies or []
+
+    current_meds_lower = [str(med).lower() for med in current_meds_list]
+    allergies_lower = [str(a).lower() for a in allergies_list]
+
+    print(f"  Checking against Meds: {current_meds_lower}")
+    print(f"  Checking against Allergies: {allergies_lower}")
+
+    # Check Allergies
     for allergy in allergies_lower:
+        # Direct match
         if allergy == potential_med_lower:
-            warnings.append(f"CRITICAL ALLERGY: Patient allergic to {allergy}. Cannot prescribe {potential_prescription}.")
-            continue
+            warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{allergy}'. Cannot prescribe '{potential_prescription}'.")
+            continue # Don't check cross-reactivity if direct match
+        # Check cross-reactivity
         if allergy in ALLERGY_INTERACTIONS:
             for cross_reactant in ALLERGY_INTERACTIONS[allergy]:
                 if cross_reactant.lower() == potential_med_lower:
-                    warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to {allergy}. High risk with {potential_prescription}.")
+                    warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to '{allergy}'. High risk with '{potential_prescription}'.")
 
+    # Check Drug-Drug Interactions
     for current_med in current_meds_lower:
+        # Check pairs in both orders using the mock DB
         pair1 = (current_med, potential_med_lower)
         pair2 = (potential_med_lower, current_med)
-        # Normalize keys for lookup if necessary (e.g., if DB keys are canonical names)
-        key1 = tuple(sorted(pair1))
-        key2 = tuple(sorted(pair2)) # Although redundant if always sorted
-
-        if pair1 in MOCK_INTERACTION_DB:
-            warnings.append(f"Interaction: {potential_prescription.capitalize()} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair1]}")
-        elif pair2 in MOCK_INTERACTION_DB:
-             warnings.append(f"Interaction: {potential_prescription.capitalize()} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair2]}")
+        interaction_msg = MOCK_INTERACTION_DB.get(pair1) or MOCK_INTERACTION_DB.get(pair2)
+        if interaction_msg:
+            warnings.append(f"Interaction: {potential_prescription.capitalize()} with {current_med.capitalize()} - {interaction_msg}")
 
     status = "warning" if warnings else "clear"
-    message = f"Interaction check for {potential_prescription}: {len(warnings)} potential issue(s) found." if warnings else f"No major interactions identified for {potential_prescription}."
+    message = f"Interaction check for '{potential_prescription}': {len(warnings)} potential issue(s) found." if warnings else f"No major interactions identified for '{potential_prescription}' based on provided lists."
+    print(f"  Interaction Check Result: {status}, Message: {message}, Warnings: {warnings}")
     return json.dumps({"status": status, "message": message, "warnings": warnings})
 
 
@@ -182,21 +258,26 @@ class FlagRiskInput(BaseModel):
 @tool("flag_risk", args_schema=FlagRiskInput)
 def flag_risk(risk_description: str, urgency: str) -> str:
     """Flags a critical risk identified during analysis for immediate attention."""
+    print(f"Executing flag_risk: {risk_description}, Urgency: {urgency}")
     # Display in Streamlit immediately
     st.error(f"🚨 **{urgency.upper()} RISK FLAGGED by AI:** {risk_description}", icon="🚨")
-    return json.dumps({"status": "flagged", "message": f"Risk '{risk_description}' flagged with {urgency} urgency."})
+    return json.dumps({
+        "status": "flagged",
+        "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
+    })
 
 # Initialize Search Tool
-search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS, name="tavily_search_results")
-
+search_tool = TavilySearchResults(
+    max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS,
+    name="tavily_search_results" # Explicitly name the tool
+    )
 
 # --- LangGraph Setup ---
 
 # Define the state structure
 class AgentState(TypedDict):
     messages: Annotated[list[Any], operator.add] # Accumulates messages (Human, AI, Tool)
-    patient_data: Optional[dict] # Holds the structured patient data (can be updated if needed)
-    # Potentially add other state elements like 'interaction_check_needed_for': Optional[str]
+    patient_data: Optional[dict] # Holds the structured patient data
 
 # Define Tools and Tool Executor
 tools = [
@@ -211,99 +292,212 @@ tool_executor = ToolExecutor(tools)
 # Define the Agent Model
 model = ChatGroq(
     temperature=ClinicalAppSettings.TEMPERATURE,
-    model=ClinicalAppSettings.MODEL_NAME
+    model=ClinicalAppSettings.MODEL_NAME,
+    # Increase max_tokens if needed for large JSON output + conversation history
+    # max_tokens=4096
 )
-model_with_tools = model.bind_tools(tools) # Bind tools for the LLM to know about them
+# Bind tools FOR the model to know their schemas and descriptions
+model_with_tools = model.bind_tools(tools)
 
 # --- Graph Nodes ---
 
 # 1. Agent Node: Calls the LLM
 def agent_node(state: AgentState):
     """Invokes the LLM to decide the next action or response."""
-    print("---AGENT NODE---")
-    # Make sure patient data is included in the first message if not already there
-    # This is a basic way; more robust would be merging patient_data into context
+    print("\n---AGENT NODE---")
     current_messages = state['messages']
-    if len(current_messages) == 1 and isinstance(current_messages[0], HumanMessage) and state.get('patient_data'):
-         # Augment the first human message with formatted patient data
-         formatted_data = format_patient_data_for_prompt(state['patient_data']) # Need this helper function
-         current_messages = [
-             SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT), # Ensure system prompt is first
-             HumanMessage(content=f"{current_messages[0].content}\n\n**Initial Patient Data:**\n{formatted_data}")
-         ]
-    elif not any(isinstance(m, SystemMessage) for m in current_messages):
-         # Add system prompt if missing
-         current_messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + current_messages
-
-
-    response = model_with_tools.invoke(current_messages)
-    print(f"Agent response: {response}")
+    # Ensure System Prompt is present
+    if not current_messages or not isinstance(current_messages[0], SystemMessage):
+        print("Prepending System Prompt.")
+        current_messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + current_messages
+
+    # Optional: Augment first human message with patient data if not already done explicitly
+    # This helps ensure the LLM sees it early, though it's also in the state.
+    # Be mindful of context window limits.
+    # if len(current_messages) > 1 and isinstance(current_messages[1], HumanMessage) and state.get('patient_data'):
+    #     if "**Initial Patient Data:**" not in current_messages[1].content:
+    #         print("Augmenting first HumanMessage with patient data summary.")
+    #         formatted_data = format_patient_data_for_prompt(state['patient_data'])
+    #         current_messages[1] = HumanMessage(content=f"{current_messages[1].content}\n\n**Initial Patient Data Summary:**\n{formatted_data}")
+
+    print(f"Invoking LLM with {len(current_messages)} messages.")
+    # print(f"Messages Sent: {[m.type for m in current_messages]}") # Log message types
+    try:
+        response = model_with_tools.invoke(current_messages)
+        print(f"Agent Raw Response Type: {type(response)}")
+        # print(f"Agent Raw Response Content: {response.content}")
+        if hasattr(response, 'tool_calls') and response.tool_calls:
+             print(f"Agent Response Tool Calls: {response.tool_calls}")
+        else:
+             print("Agent Response: No tool calls.")
+
+    except Exception as e:
+        print(f"ERROR in agent_node during LLM invocation: {type(e).__name__} - {e}")
+        traceback.print_exc() # Print full traceback for debugging
+        # Return an error message to the graph state
+        error_message = AIMessage(content=f"Sorry, an internal error occurred while processing the request: {type(e).__name__}")
+        return {"messages": [error_message]}
+
     return {"messages": [response]}
 
-# 2. Tool Node: Executes tools called by the Agent
+# 2. Tool Node: Executes tools called by the Agent (REVISED WITH ROBUST ERROR HANDLING)
 def tool_node(state: AgentState):
     """Executes tools called by the LLM and returns results."""
-    print("---TOOL NODE---")
+    print("\n---TOOL NODE---")
+    tool_messages = [] # Initialize list to store results or errors
     last_message = state['messages'][-1]
-    if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
-        print("No tool calls in last message.")
-        return {} # Should not happen if routing is correct, but safety check
+
+    # Ensure the last message is an AIMessage with tool calls
+    if not isinstance(last_message, AIMessage) or not getattr(last_message, 'tool_calls', None):
+        print("Warning: Tool node called unexpectedly without tool calls in the last AI message.")
+        # If this happens, it might indicate a routing issue or the LLM hallucinating flow.
+        # Returning empty list lets the agent proceed, potentially without needed info.
+        # Consider adding a ToolMessage indicating the issue if needed.
+        return {"messages": []}
 
     tool_calls = last_message.tool_calls
-    tool_messages = []
+    print(f"Tool calls received: {json.dumps(tool_calls, indent=2)}") # Log received calls
 
-    # Safety Check: Ensure interaction check happens *before* prescribing the *same* drug
-    prescribe_calls = {call['args'].get('medication_name'): call['id'] for call in tool_calls if call['name'] == 'prescribe_medication'}
-    interaction_check_calls = {call['args'].get('potential_prescription'): call['id'] for call in tool_calls if call['name'] == 'check_drug_interactions'}
+    # Safety Check: Identify required interaction checks before prescriptions
+    prescriptions_requested = {} # medication_name_lower -> tool_call
+    interaction_checks_requested = {} # medication_name_lower -> tool_call
 
-    for med_name, prescribe_call_id in prescribe_calls.items():
-        if med_name not in interaction_check_calls:
-            st.error(f"**Safety Violation:** AI attempted to prescribe '{med_name}' without requesting `check_drug_interactions` in the *same turn*. Prescription blocked for this turn.")
-            # Create an error ToolMessage to send back to the LLM
+    for call in tool_calls:
+        tool_name = call.get('name')
+        tool_args = call.get('args', {})
+        if tool_name == 'prescribe_medication':
+            med_name = tool_args.get('medication_name', '').lower()
+            if med_name:
+                prescriptions_requested[med_name] = call
+        elif tool_name == 'check_drug_interactions':
+            potential_med = tool_args.get('potential_prescription', '').lower()
+            if potential_med:
+                interaction_checks_requested[potential_med] = call
+
+    valid_tool_calls_for_execution = []
+
+    # Validate prescriptions against interaction checks
+    for med_name, prescribe_call in prescriptions_requested.items():
+        if med_name not in interaction_checks_requested:
+            st.error(f"**Safety Violation:** AI attempted to prescribe '{med_name}' without requesting `check_drug_interactions` in the *same turn*. Prescription blocked.")
             error_msg = ToolMessage(
-                content=json.dumps({"status": "error", "message": f"Interaction check for {med_name} must be requested *before or alongside* the prescription call."}),
-                tool_call_id=prescribe_call_id
+                content=json.dumps({"status": "error", "message": f"Interaction check for '{med_name}' must be requested *before or alongside* the prescription call."}),
+                tool_call_id=prescribe_call['id'],
+                name=prescribe_call['name'] # Include tool name in ToolMessage
             )
             tool_messages.append(error_msg)
-            # Remove the invalid prescribe call to prevent execution
-            tool_calls = [call for call in tool_calls if call['id'] != prescribe_call_id]
+        else:
+             # Interaction check IS requested, allow prescription call to proceed
+             pass # The call will be added below if it's in the original tool_calls list
 
+    # Prepare list of calls to execute (all non-blocked calls)
+    blocked_ids = {msg.tool_call_id for msg in tool_messages if msg.content and '"status": "error"' in msg.content}
+    valid_tool_calls_for_execution = [call for call in tool_calls if call['id'] not in blocked_ids]
 
-    # Add patient context to interaction checks if needed
+    # Augment interaction checks with patient data from state
     patient_meds = state.get("patient_data", {}).get("medications", {}).get("names_only", [])
     patient_allergies = state.get("patient_data", {}).get("allergies", [])
-    for call in tool_calls:
+
+    for call in valid_tool_calls_for_execution:
          if call['name'] == 'check_drug_interactions':
+             # Ensure args exist before modifying
+             if 'args' not in call: call['args'] = {}
              call['args']['current_medications'] = patient_meds
              call['args']['allergies'] = patient_allergies
-             print(f"Augmented interaction check args: {call['args']}")
-
-
-    # Execute remaining valid tool calls
-    if tool_calls:
-        responses = tool_executor.batch(tool_calls)
-        # Responses is a list of tool outputs corresponding to tool_calls
-        # We need to create ToolMessage objects
-        tool_messages.extend([
-            ToolMessage(content=str(resp), tool_call_id=call['id'])
-            for call, resp in zip(tool_calls, responses)
-        ])
-        print(f"Tool results: {tool_messages}")
-
+             print(f"Augmented interaction check args for call ID {call['id']}: {call['args']}")
+
+
+    # Execute valid tool calls using batch for efficiency, capturing exceptions
+    if valid_tool_calls_for_execution:
+        print(f"Attempting to execute {len(valid_tool_calls_for_execution)} tools: {[c['name'] for c in valid_tool_calls_for_execution]}")
+        try:
+            responses = tool_executor.batch(valid_tool_calls_for_execution, return_exceptions=True)
+
+            # Process responses, creating ToolMessage for each
+            for call, resp in zip(valid_tool_calls_for_execution, responses):
+                tool_call_id = call['id']
+                tool_name = call['name']
+
+                if isinstance(resp, Exception):
+                    # Handle exceptions returned by the batch call
+                    error_type = type(resp).__name__
+                    error_str = str(resp)
+                    print(f"ERROR executing tool '{tool_name}' (ID: {tool_call_id}): {error_type} - {error_str}")
+                    traceback.print_exc() # Log full traceback
+                    st.error(f"Error executing action '{tool_name}': {error_type}")
+                    error_content = json.dumps({
+                        "status": "error",
+                        "message": f"Failed to execute '{tool_name}': {error_type} - {error_str}"
+                    })
+                    tool_messages.append(ToolMessage(content=error_content, tool_call_id=tool_call_id, name=tool_name))
+
+                    # Specific check for the error mentioned by user
+                    if isinstance(resp, AttributeError) and "'dict' object has no attribute 'tool'" in error_str:
+                         print("\n *** DETECTED SPECIFIC ATTRIBUTE ERROR ('dict' object has no attribute 'tool') ***")
+                         print(f"   Tool Call causing error: {json.dumps(call, indent=2)}")
+                         print("   This likely indicates an internal issue within Langchain/LangGraph or ToolExecutor expecting a different object structure.")
+                         print("   Ensure tool definitions (@tool decorators) and Pydantic schemas are correct.\n")
+
+                else:
+                    # Process successful results
+                    print(f"Tool '{tool_name}' (ID: {tool_call_id}) executed successfully. Result type: {type(resp)}")
+                    # Ensure content is string for ToolMessage
+                    content_str = str(resp)
+                    tool_messages.append(ToolMessage(content=content_str, tool_call_id=tool_call_id, name=tool_name))
+
+                    # Display result in Streamlit right away for feedback (optional, but helpful)
+                    # This part might be better handled purely in the UI display loop later
+                    # try:
+                    #     result_data = json.loads(content_str)
+                    #     status = result_data.get("status", "info")
+                    #     message = result_data.get("message", content_str)
+                    #     if status in ["success", "clear", "flagged"]: st.success(f"Action `{tool_name}` completed: {message}", icon="✅" if status != "flagged" else "🚨")
+                    #     elif status == "warning": st.warning(f"Action `{tool_name}` completed: {message}", icon="⚠️")
+                    #     else: st.info(f"Action `{tool_name}` completed: {message}") # Info for other statuses
+                    # except json.JSONDecodeError:
+                    #     st.info(f"Action `{tool_name}` completed (non-JSON output).")
+
+        # Catch potential errors within the tool_node logic itself (e.g., preparing calls)
+        except Exception as e:
+            print(f"CRITICAL UNEXPECTED ERROR within tool_node logic: {type(e).__name__} - {e}")
+            traceback.print_exc()
+            st.error(f"Critical internal error processing actions: {e}")
+            # Create generic error messages for all calls that were intended
+            error_content = json.dumps({"status": "error", "message": f"Internal error processing tools: {e}"})
+            # Add error messages for calls that didn't get processed yet
+            processed_ids = {msg.tool_call_id for msg in tool_messages}
+            for call in valid_tool_calls_for_execution:
+                 if call['id'] not in processed_ids:
+                     tool_messages.append(ToolMessage(content=error_content, tool_call_id=call['id'], name=call['name']))
+
+    print(f"Returning {len(tool_messages)} tool messages.")
+    # print(f"Tool messages content snippets: {[m.content[:100] + '...' if len(m.content)>100 else m.content for m in tool_messages]}")
     return {"messages": tool_messages}
 
 
 # --- Graph Edges (Routing Logic) ---
 def should_continue(state: AgentState) -> str:
-    """Determines whether to continue the loop or end."""
-    last_message = state['messages'][-1]
-    # If the LLM made tool calls, we execute them
-    if isinstance(last_message, AIMessage) and last_message.tool_calls:
-        print("Routing: continue_tools")
+    """Determines whether to call tools, end the conversation turn, or handle errors."""
+    print("\n---ROUTING DECISION---")
+    last_message = state['messages'][-1] if state['messages'] else None
+
+    if not isinstance(last_message, AIMessage):
+        # This case might happen if the graph starts with a non-AI message or after an error
+        print("Routing: Last message not AI. Ending turn.")
+        return "end_conversation_turn"
+
+    # If the LLM produced an error message (e.g., during invocation)
+    if "Sorry, an internal error occurred" in last_message.content:
+         print("Routing: AI returned internal error. Ending turn.")
+         return "end_conversation_turn"
+
+    # If the LLM made tool calls, execute them
+    if getattr(last_message, 'tool_calls', None):
+        print("Routing: AI requested tool calls. Continue to tools node.")
         return "continue_tools"
-    # Otherwise, we end the loop (AI provided a direct answer or finished)
+    # Otherwise, the AI provided a response without tool calls, end the turn
     else:
-        print("Routing: end_conversation_turn")
+        print("Routing: AI provided final response or asked question. Ending turn.")
         return "end_conversation_turn"
 
 # --- Graph Definition ---
@@ -322,7 +516,7 @@ workflow.add_conditional_edges(
     should_continue,              # Function to decide the route
     {
         "continue_tools": "tools", # If tool calls exist, go to tools node
-        "end_conversation_turn": END # Otherwise, end the graph iteration
+        "end_conversation_turn": END # Otherwise, end the graph iteration for this turn
     }
 )
 
@@ -333,61 +527,48 @@ workflow.add_edge("tools", "agent")
 # memory = MemorySaverInMemory() # Optional: for persisting state across runs
 # app = workflow.compile(checkpointer=memory)
 app = workflow.compile()
+print("LangGraph compiled successfully.")
 
-# --- Helper Function to Format Patient Data ---
-def format_patient_data_for_prompt(data: dict) -> str:
-    """Formats the patient dictionary into a readable string for the LLM."""
-    prompt_str = ""
-    for key, value in data.items():
-        if isinstance(value, dict):
-            section_title = key.replace('_', ' ').title()
-            prompt_str += f"**{section_title}:**\n"
-            for sub_key, sub_value in value.items():
-                 if sub_value:
-                    prompt_str += f"  - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
-        elif isinstance(value, list) and value:
-             prompt_str += f"**{key.replace('_', ' ').title()}:** {', '.join(map(str, value))}\n"
-        elif value:
-             prompt_str += f"**{key.replace('_', ' ').title()}:** {value}\n"
-    return prompt_str.strip()
-
-# --- Streamlit UI (Modified for Conversation) ---
+# --- Streamlit UI ---
 def main():
     st.set_page_config(page_title=ClinicalAppSettings.APP_TITLE, layout=ClinicalAppSettings.PAGE_LAYOUT)
     st.title(f"🩺 {ClinicalAppSettings.APP_TITLE}")
     st.caption(f"Interactive Assistant | Powered by Langchain/LangGraph & Groq ({ClinicalAppSettings.MODEL_NAME})")
 
-    # Initialize session state for conversation
+    # Initialize session state
     if "messages" not in st.session_state:
-        st.session_state.messages = [] # Store entire conversation history (Human, AI, Tool)
+        st.session_state.messages = [] # Stores full conversation history
     if "patient_data" not in st.session_state:
         st.session_state.patient_data = None
-    if "initial_analysis_done" not in st.session_state:
-        st.session_state.initial_analysis_done = False
     if "graph_app" not in st.session_state:
-        st.session_state.graph_app = app # Store compiled graph
+        st.session_state.graph_app = app
 
-    # --- Patient Data Input Sidebar --- (Similar to before)
+    # --- Patient Data Input Sidebar ---
     with st.sidebar:
         st.header("📄 Patient Intake Form")
-        # ... (Keep the input fields exactly as in the previous example) ...
         # Demographics
+        st.subheader("Demographics")
         age = st.number_input("Age", min_value=0, max_value=120, value=55, key="age_input")
         sex = st.selectbox("Biological Sex", ["Male", "Female", "Other/Prefer not to say"], key="sex_input")
         # HPI
+        st.subheader("History of Present Illness (HPI)")
         chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="cc_input")
-        hpi_details = st.text_area("Detailed HPI", "55 y/o male presents with substernal chest pain started 2 hours ago...", key="hpi_input")
-        symptoms = st.multiselect("Associated Symptoms", ["Nausea", "Diaphoresis", "Shortness of Breath", "Dizziness", "Palpitations", "Fever", "Cough"], default=["Nausea", "Diaphoresis"], key="sym_input")
+        hpi_details = st.text_area("Detailed HPI", "55 y/o male presents with substernal chest pain started 2 hours ago, described as pressure, radiating to left arm. Associated with nausea and diaphoresis. Pain is 8/10 severity. No relief with rest.", key="hpi_input", height=150)
+        symptoms = st.multiselect("Associated Symptoms", ["Nausea", "Diaphoresis", "Shortness of Breath", "Dizziness", "Palpitations", "Fever", "Cough", "Severe Headache", "Syncope", "Hemoptysis"], default=["Nausea", "Diaphoresis"], key="sym_input")
         # History
-        pmh = st.text_area("Past Medical History (PMH)", "Hypertension (HTN), Hyperlipidemia (HLD), Type 2 Diabetes Mellitus (DM2)", key="pmh_input")
+        st.subheader("Past History")
+        pmh = st.text_area("Past Medical History (PMH)", "Hypertension (HTN), Hyperlipidemia (HLD), Type 2 Diabetes Mellitus (DM2), History of MI", key="pmh_input")
         psh = st.text_area("Past Surgical History (PSH)", "Appendectomy (2005)", key="psh_input")
         # Meds & Allergies
+        st.subheader("Medications & Allergies")
         current_meds_str = st.text_area("Current Medications (name, dose, freq)", "Lisinopril 10mg daily\nMetformin 1000mg BID\nAtorvastatin 40mg daily\nAspirin 81mg daily", key="meds_input")
-        allergies_str = st.text_area("Allergies (comma separated)", "Penicillin (rash)", key="allergy_input")
+        allergies_str = st.text_area("Allergies (comma separated, specify reaction if known)", "Penicillin (rash), Sulfa (hives)", key="allergy_input")
         # Social/Family
+        st.subheader("Social/Family History")
         social_history = st.text_area("Social History (SH)", "Smoker (1 ppd x 30 years), occasional alcohol.", key="sh_input")
         family_history = st.text_area("Family History (FHx)", "Father had MI at age 60. Mother has HTN.", key="fhx_input")
         # Vitals/Exam
+        st.subheader("Vitals & Exam Findings")
         col1, col2 = st.columns(2)
         with col1:
             temp_c = st.number_input("Temp (°C)", 35.0, 42.0, 36.8, format="%.1f", key="temp_input")
@@ -397,19 +578,28 @@ def main():
             bp_mmhg = st.text_input("BP (SYS/DIA)", "155/90", key="bp_input")
             spo2_percent = st.number_input("SpO2 (%)", 70, 100, 96, key="spo2_input")
             pain_scale = st.slider("Pain (0-10)", 0, 10, 8, key="pain_input")
-        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3...", key="exam_input")
+        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3. Mild distress. Lungs clear bilaterally. Cardiac exam: Regular rhythm, S1/S2 normal, no murmurs/gallops/rubs. Abdomen soft, non-tender. No lower extremity edema.", key="exam_input", height=100)
 
         # Compile Patient Data Dictionary on button press
         if st.button("Start/Update Consultation", key="start_button"):
             current_meds_list = [med.strip() for med in current_meds_str.split('\n') if med.strip()]
+            # Basic name extraction (first word, lowercased) for interaction check
             current_med_names = []
-            # Improved parsing for names (still basic, assumes name is first word)
             for med in current_meds_list:
                 match = re.match(r"^\s*([a-zA-Z\-]+)", med)
                 if match:
-                     current_med_names.append(match.group(1).lower()) # Use lower case for matching
-
-            allergies_list = [a.strip().lower() for a in allergies_str.split(',') if a.strip()] # Lowercase allergies
+                     current_med_names.append(match.group(1).lower())
+
+            # Basic allergy extraction (first word or phrase before parenthesis, lowercased)
+            allergies_list = []
+            for a in allergies_str.split(','):
+                cleaned_allergy = a.strip()
+                if cleaned_allergy:
+                     match = re.match(r"^\s*([a-zA-Z\-\s]+)(?:\s*\(.*\))?", cleaned_allergy)
+                     if match:
+                         allergies_list.append(match.group(1).strip().lower())
+                     else: # Fallback if no parenthesis
+                          allergies_list.append(cleaned_allergy.lower())
 
             st.session_state.patient_data = {
                 "demographics": {"age": age, "sex": sex},
@@ -424,75 +614,131 @@ def main():
 
             # Initial Red Flag Check (Client-side)
             red_flags = check_red_flags(st.session_state.patient_data)
+            st.sidebar.markdown("---")
             if red_flags:
-                st.warning("**Initial Red Flags Detected:**")
-                for flag in red_flags: st.warning(f"- {flag}")
+                st.sidebar.warning("**Initial Red Flags Detected:**")
+                for flag in red_flags: st.sidebar.warning(f"- {flag.replace('Red Flag: ','')}")
+            else:
+                 st.sidebar.success("No immediate red flags detected in initial data.")
 
             # Prepare initial message for the graph
-            initial_prompt = f"Analyze the following patient case:\nChief Complaint: {chief_complaint}\nSummary: {age} y/o {sex} presenting with..." # Keep it brief, full data is in state
+            initial_prompt = f"Initiate consultation for the patient described in the intake form. Start the analysis."
+            # Clear previous messages and start fresh
             st.session_state.messages = [HumanMessage(content=initial_prompt)]
-            st.session_state.initial_analysis_done = False # Reset analysis state
             st.success("Patient data loaded. Ready for analysis.")
-            st.rerun() # Refresh main area to show chat
-
+            # No rerun needed here, chat input will trigger the graph
 
     # --- Main Chat Interface Area ---
     st.header("💬 Clinical Consultation")
 
-    # Display chat messages
-    for msg in st.session_state.messages:
+    # Display chat messages from history
+    for msg_index, msg in enumerate(st.session_state.messages):
+        unique_key = f"msg_{msg_index}" # Basic unique key
         if isinstance(msg, HumanMessage):
-            with st.chat_message("user"):
+            with st.chat_message("user", key=f"{unique_key}_user"):
                 st.markdown(msg.content)
         elif isinstance(msg, AIMessage):
-            with st.chat_message("assistant"):
-                # Check for structured JSON output
+            with st.chat_message("assistant", key=f"{unique_key}_ai"):
+                # Display AI text content
+                ai_content = msg.content
                 structured_output = None
+
+                # Attempt to parse structured JSON if present
                 try:
-                    # Try to find JSON block first
-                    json_match = re.search(r"```json\n(\{.*?\})\n```", msg.content, re.DOTALL)
+                    # Look for ```json ... ``` block
+                    json_match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
                     if json_match:
-                        structured_output = json.loads(json_match.group(1))
-                        # Display non-JSON parts if any
-                        non_json_content = msg.content.replace(json_match.group(0), "").strip()
-                        if non_json_content:
-                            st.markdown(non_json_content)
-                            st.divider() # Separate text from structured output visually
-                    elif msg.content.strip().startswith("{") and msg.content.strip().endswith("}"):
-                         # Maybe the whole message is JSON
-                         structured_output = json.loads(msg.content)
+                        json_str = json_match.group(1)
+                        # Display content before/after the JSON block if any
+                        prefix = ai_content[:json_match.start()].strip()
+                        suffix = ai_content[json_match.end():].strip()
+                        if prefix: st.markdown(prefix)
+                        structured_output = json.loads(json_str)
+                        if suffix: st.markdown(suffix)
+                    # Check if the entire message might be JSON
+                    elif ai_content.strip().startswith("{") and ai_content.strip().endswith("}"):
+                         structured_output = json.loads(ai_content)
+                         ai_content = "" # Don't display raw JSON if parsed ok
                     else:
-                        # No JSON found, display raw content
-                        st.markdown(msg.content)
-
-                    if structured_output:
-                        # Display the structured data nicely (reuse parts of previous UI display logic)
-                        st.subheader("📊 AI Analysis & Recommendations")
-                        # ... (Add logic here to display assessment, ddx, plan etc. from structured_output)
-                        # Example:
-                        st.write(f"**Assessment:** {structured_output.get('assessment', 'N/A')}")
-                        # Display DDx, Plan etc. using expanders or tabs
-                        # ...
-                        # Display Rationale & Interaction Summary
-                        with st.expander("Rationale & Guideline Check"):
-                             st.write(structured_output.get("rationale_summary", "N/A"))
-                        if structured_output.get("interaction_check_summary"):
-                             with st.expander("Interaction Check"):
-                                 st.write(structured_output.get("interaction_check_summary"))
-
+                         # No JSON found, display content as is
+                         st.markdown(ai_content)
 
                 except json.JSONDecodeError:
-                    st.markdown(msg.content) # Display raw if JSON parsing fails
-
-                # Display tool calls if any were made in this AI turn
-                if msg.tool_calls:
+                    # Failed to parse, display raw content
+                    st.markdown(ai_content)
+                    st.warning("Note: Could not parse structured JSON in AI response.", icon="⚠️")
+                except Exception as e:
+                     st.markdown(ai_content) # Display raw on other errors
+                     st.error(f"Error processing AI message display: {e}", icon="❌")
+
+                # Display structured data nicely if parsed
+                if structured_output and isinstance(structured_output, dict):
+                    st.divider()
+                    st.subheader("📊 AI Analysis & Recommendations")
+                    cols = st.columns(2)
+                    with cols[0]:
+                        st.markdown(f"**Assessment:**")
+                        st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
+
+                        st.markdown(f"**Differential Diagnosis:**")
+                        ddx = structured_output.get('differential_diagnosis', [])
+                        if ddx:
+                            for item in ddx:
+                                likelihood = item.get('likelihood', 'Unknown').capitalize()
+                                icon = "🥇" if likelihood=="High" else ("🥈" if likelihood=="Medium" else "🥉")
+                                with st.expander(f"{icon} {item.get('diagnosis', 'Unknown')} ({likelihood})"):
+                                    st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
+                        else: st.info("No differential diagnosis provided.")
+
+                        st.markdown(f"**Risk Assessment:**")
+                        risk = structured_output.get('risk_assessment', {})
+                        flags = risk.get('identified_red_flags', [])
+                        if flags: st.warning(f"**Flags:** {', '.join(flags)}")
+                        if risk.get("immediate_concerns"): st.warning(f"**Concerns:** {', '.join(risk.get('immediate_concerns'))}")
+                        if risk.get("potential_complications"): st.info(f"**Potential Complications:** {', '.join(risk.get('potential_complications'))}")
+                        if not flags and not risk.get("immediate_concerns"): st.success("No major risks highlighted in this assessment.")
+
+                    with cols[1]:
+                        st.markdown(f"**Recommended Plan:**")
+                        plan = structured_output.get('recommended_plan', {})
+                        sub_sections = ["investigations", "therapeutics", "consultations", "patient_education"]
+                        for section in sub_sections:
+                             st.markdown(f"_{section.replace('_',' ').capitalize()}:_")
+                             items = plan.get(section)
+                             if items and isinstance(items, list):
+                                 for item in items: st.markdown(f"- {item}")
+                             elif items: # Handle if it's just a string
+                                  st.markdown(f"- {items}")
+                             else: st.markdown("_None suggested._")
+                             st.markdown("") # Add space
+
+                    # Display Rationale and Interaction Summary below columns
+                    st.markdown(f"**Rationale & Guideline Check:**")
+                    st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}")
+                    interaction_summary = structured_output.get("interaction_check_summary", "")
+                    if interaction_summary:
+                         st.markdown(f"**Interaction Check Summary:**")
+                         st.markdown(f"> {interaction_summary}")
+
+                    st.divider()
+
+                # Display tool calls requested in this AI turn
+                if getattr(msg, 'tool_calls', None):
                      with st.expander("🛠️ AI requested actions", expanded=False):
                          for tc in msg.tool_calls:
-                             st.code(f"{tc['name']}(args={tc['args']})", language="python")
+                             try:
+                                 # Safely display, default args to empty dict if missing
+                                 st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nArgs: {json.dumps(tc.get('args', {}), indent=2)}", language="json")
+                             except Exception as display_e:
+                                  st.error(f"Could not display tool call: {display_e}")
+                                  st.code(str(tc)) # Raw display as fallback
 
         elif isinstance(msg, ToolMessage):
-             with st.chat_message("tool", avatar="🛠️"):
+            # Safely get tool name
+            tool_name_display = getattr(msg, 'name', 'tool_execution') # Use 'name' attribute added in tool_node
+            with st.chat_message(tool_name_display, avatar="🛠️", key=f"{unique_key}_tool"):
                 try:
+                    # Attempt to parse content as JSON for structured display
                     tool_data = json.loads(msg.content)
                     status = tool_data.get("status", "info")
                     message = tool_data.get("message", msg.content)
@@ -500,97 +746,73 @@ def main():
                     warnings = tool_data.get("warnings")
 
                     if status == "success" or status == "clear" or status == "flagged":
-                        st.success(f"Tool Result ({msg.name}): {message}", icon="✅" if status != "flagged" else "🚨")
+                        st.success(f"{message}", icon="✅" if status != "flagged" else "🚨")
                     elif status == "warning":
-                         st.warning(f"Tool Result ({msg.name}): {message}", icon="⚠️")
-                         if warnings:
+                         st.warning(f"{message}", icon="⚠️")
+                         if warnings and isinstance(warnings, list):
+                             st.caption("Details:")
                              for warn in warnings: st.caption(f"- {warn}")
                     else: # Error or unknown status
-                         st.error(f"Tool Result ({msg.name}): {message}", icon="❌")
+                         st.error(f"{message}", icon="❌")
 
                     if details: st.caption(f"Details: {details}")
 
                 except json.JSONDecodeError:
-                    st.info(f"Tool Result ({msg.name}): {msg.content}") # Display raw if not JSON
-
+                    # If content is not JSON, display it plainly
+                    st.info(f"{msg.content}")
+                except Exception as e:
+                     st.error(f"Error displaying tool message: {e}", icon="❌")
+                     st.caption(f"Raw content: {msg.content}")
 
-    # Chat input for user
+    # --- Chat Input Logic ---
     if prompt := st.chat_input("Your message or follow-up query..."):
         if not st.session_state.patient_data:
             st.warning("Please load patient data using the sidebar first.")
-        else:
-            # Add user message to state
-            st.session_state.messages.append(HumanMessage(content=prompt))
-            with st.chat_message("user"):
-                st.markdown(prompt)
-
-            # Prepare state for graph invocation
-            current_state = AgentState(
-                messages=st.session_state.messages,
-                patient_data=st.session_state.patient_data
-            )
-
-            # Stream graph execution
-            with st.chat_message("assistant"):
-                message_placeholder = st.empty()
-                full_response = ""
-
-                # Use stream to get intermediate steps (optional but good for UX)
-                # This shows AI thinking and tool calls/results progressively
-                try:
-                    for event in st.session_state.graph_app.stream(current_state, {"recursion_limit": 15}):
-                        # event is a dictionary, keys are node names
-                        if "agent" in event:
-                            ai_msg = event["agent"]["messages"][-1] # Get the latest AI message
-                            if isinstance(ai_msg, AIMessage):
-                                full_response += ai_msg.content # Append content for final display
-                                message_placeholder.markdown(full_response + "▌") # Show typing indicator
-
-                                # Display tool calls as they happen (optional)
-                                # if ai_msg.tool_calls:
-                                #     st.info(f"Requesting tools: {[tc['name'] for tc in ai_msg.tool_calls]}")
-
-                        elif "tools" in event:
-                            # Display tool results as they come back (optional, already handled by message display loop)
-                             pass
-                             # tool_msgs = event["tools"]["messages"]
-                             # for tool_msg in tool_msgs:
-                             #      st.info(f"Tool {tool_msg.name} result received.")
-
-
-                    # Final display after streaming
-                    message_placeholder.markdown(full_response)
-
-
-                    # Update session state with the final messages from the graph run
-                    # The graph state itself isn't directly accessible after streaming finishes easily this way
-                    # We need to get the final state if we used invoke, or reconstruct from stream events
-                    # A simpler way for now: just append the *last* AI message and any Tool messages from the stream
-                    # This assumes the stream provides the final state implicitly. For robust state, use invoke or checkpointer.
-
-                    # A more robust way: invoke and get final state
-                    # final_state = st.session_state.graph_app.invoke(current_state, {"recursion_limit": 15})
-                    # st.session_state.messages = final_state['messages']
-                    # --- Let's stick to appending for simplicity in this example ---
-                    # Find the last AI message and tool messages from the stream (needs careful event parsing)
-                    # Or, re-run invoke non-streamed just to get final state (less efficient)
-                    final_state_capture = st.session_state.graph_app.invoke(current_state, {"recursion_limit": 15})
-                    st.session_state.messages = final_state_capture['messages']
-
-
-                except Exception as e:
-                    st.error(f"An error occurred during analysis: {e}")
-                    # Attempt to add the error message to the history
-                    st.session_state.messages.append(AIMessage(content=f"Sorry, an error occurred: {e}"))
-
-
-            # Rerun to display the updated chat history correctly
-            st.rerun()
-
-
-    # Disclaimer
+            st.stop() # Prevent execution if no patient data
+
+        # Add user message to state and display it immediately
+        user_message = HumanMessage(content=prompt)
+        st.session_state.messages.append(user_message)
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        # Prepare state for graph invocation
+        current_state = AgentState(
+            messages=st.session_state.messages,
+            patient_data=st.session_state.patient_data
+        )
+
+        # Invoke the graph
+        with st.spinner("SynapseAI is thinking..."):
+            try:
+                # Use invoke to run the graph until it ends for this turn
+                final_state = st.session_state.graph_app.invoke(
+                     current_state,
+                     {"recursion_limit": 15} # Add recursion limit for safety
+                )
+                # Update the session state messages with the final list from the graph run
+                st.session_state.messages = final_state['messages']
+
+            except Exception as e:
+                print(f"CRITICAL ERROR during graph invocation: {type(e).__name__} - {e}")
+                traceback.print_exc()
+                st.error(f"An error occurred during the conversation turn: {e}", icon="❌")
+                # Attempt to add an error message to the history for visibility
+                error_ai_msg = AIMessage(content=f"Sorry, a critical error occurred: {type(e).__name__}. Please check logs or try again.")
+                # Avoid modifying state directly during exception handling if possible,
+                # but appending might be okay for display purposes.
+                # st.session_state.messages.append(error_ai_msg) # Be cautious with state modification here
+
+        # Rerun the script to display the updated chat history, including AI response and tool results
+        st.rerun()
+
+    # Disclaimer at the bottom
     st.markdown("---")
-    st.warning("**Disclaimer:** SynapseAI is for clinical decision support...") # Keep disclaimer
+    st.warning(
+        """**Disclaimer:** SynapseAI is an AI assistant for clinical decision support and does not replace professional medical judgment.
+        All outputs must be critically reviewed and verified by a qualified healthcare provider before making any clinical decisions.
+        Validate all information, especially diagnoses, dosages, and interactions, independently using standard medical resources."""
+    )
 
 if __name__ == "__main__":
     main()