Update app.py

app.py

@@ -1,116 +1,117 @@
 import streamlit as st
 from langchain_groq import ChatGroq
 from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
 from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.tools import tool
-from typing import Optional, List, Dict, Any
+from langgraph.prebuilt import ToolExecutor
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import MemorySaver # For state persistence (optional but good)
+
+from typing import Optional, List, Dict, Any, TypedDict, Annotated
 import json
-import re # For parsing vitals like BP
+import re
+import operator
 
-# --- Configuration & Constants ---
+# --- Configuration & Constants --- (Keep previous ones like ClinicalAppSettings)
 class ClinicalAppSettings:
-    APP_TITLE = "SynapseAI: Advanced Clinical Decision Support"
+    APP_TITLE = "SynapseAI: Interactive Clinical Decision Support"
     PAGE_LAYOUT = "wide"
-    MODEL_NAME = "llama3-70b-8192" # Use a powerful model like Groq's Llama3 70b
+    MODEL_NAME = "llama3-70b-8192"
     TEMPERATURE = 0.1
     MAX_SEARCH_RESULTS = 3
 
 class ClinicalPrompts:
+    # UPDATED SYSTEM PROMPT FOR CONVERSATIONAL FLOW & GUIDELINES
     SYSTEM_PROMPT = """
-    You are SynapseAI, an expert AI clinical assistant designed to support healthcare professionals.
-    Your primary function is to analyze patient data, provide differential diagnoses, suggest evidence-based management plans, and identify potential risks according to the latest medical guidelines and safety protocols.
+    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:**
-    1.  **Comprehensive Analysis:** Thoroughly analyze ALL provided patient data (demographics, HPI, PMH, PSH, Allergies, Meds, SH, FH, ROS, Vitals, Exam).
-    2.  **Structured Output:** ALWAYS format your response using the following JSON structure:
+    **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.
+    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.
         ```json
         {
-          "assessment": "Concise summary of the patient's presentation and key findings.",
+          "assessment": "Concise summary of the patient's presentation and key findings based on the conversation.",
           "differential_diagnosis": [
-            {"diagnosis": "Primary Diagnosis", "likelihood": "High/Medium/Low", "rationale": "Supporting evidence..."},
+            {"diagnosis": "Primary Diagnosis", "likelihood": "High/Medium/Low", "rationale": "Supporting evidence from conversation..."},
             {"diagnosis": "Alternative Diagnosis 1", "likelihood": "Medium/Low", "rationale": "Supporting/Refuting evidence..."},
             {"diagnosis": "Alternative Diagnosis 2", "likelihood": "Low", "rationale": "Why it's less likely but considered..."}
           ],
           "risk_assessment": {
-            "identified_red_flags": ["List any triggered red flags based on input"],
-            "immediate_concerns": ["Specific urgent issues requiring attention (e.g., sepsis risk, ACS rule-out)"],
-            "potential_complications": ["Possible future issues based on presentation"]
+            "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"]
           },
           "recommended_plan": {
-            "investigations": ["List specific lab tests or imaging required. Use 'order_lab_test' tool."],
-            "therapeutics": ["Suggest specific treatments or prescriptions. Use 'prescribe_medication' tool."],
-            "consultations": ["Recommend specialist consultations if needed."],
+            "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."],
             "patient_education": ["Key points for patient communication."]
           },
-          "rationale_summary": "Brief justification for the overall assessment and plan, referencing guidelines or evidence where possible. Use 'tavily_search_results' tool if needed to find supporting evidence/guidelines.",
-          "interaction_check_summary": "Summary of findings from the 'check_drug_interactions' tool IF a new medication was considered or prescribed."
+          "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.**",
+          "interaction_check_summary": "Summary of findings from 'check_drug_interactions' if performed."
         }
         ```
-    3.  **Safety First - Red Flags:** Immediately identify and report any conditions matching the defined RED_FLAGS. Use the `flag_risk` tool if critical.
-    4.  **Safety First - Drug Interactions:** BEFORE suggesting *any* new prescription, you MUST use the `check_drug_interactions` tool to verify against the patient's current medications and allergies. Mention the result in `interaction_check_summary`.
-    5.  **Tool Utilization:** Employ the provided tools (`order_lab_test`, `prescribe_medication`, `check_drug_interactions`, `flag_risk`, `tavily_search_results`) precisely when indicated by your plan. Adhere strictly to tool schemas. Do NOT hallucinate tool usage results; wait for actual tool output if required in a multi-turn scenario (though this implementation focuses on single-turn analysis with tool calls).
-    6.  **Evidence-Based:** Briefly cite reasoning, drawing on general medical knowledge. Use Tavily Search for specific guideline checks or novel information when necessary.
-    7.  **Clarity and Conciseness:** Be clear, avoiding ambiguity. Use standard medical terminology.
+    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).
+    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.
     """
 
-# --- Mock Data / Helpers ---
-# (In a real system, this would be a proper API/database)
+# --- 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_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."
+    ("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."
 }
 
 ALLERGY_INTERACTIONS = {
-    "Penicillin": ["Amoxicillin", "Ampicillin", "Piperacillin"],
-    "Sulfa": ["Sulfamethoxazole", "Sulfasalazine"],
-    "Aspirin": ["Ibuprofen", "Naproxen"] # Cross-reactivity example for NSAIDs
+    "penicillin": ["amoxicillin", "ampicillin", "piperacillin"],
+    "sulfa": ["sulfamethoxazole", "sulfasalazine"],
+    "aspirin": ["ibuprofen", "naproxen"] # Cross-reactivity example for NSAIDs
 }
 
 def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
-    """Parses BP string like '120/80' into (systolic, diastolic) integers."""
     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))
+    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 = []
     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]
+
+    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...
 
-    # Symptom Flags
-    if "chest pain" in [s.lower() for s in symptoms]: flags.append("Red Flag: Chest Pain reported.")
-    if "shortness of breath" in [s.lower() for s in symptoms]: flags.append("Red Flag: Shortness of Breath reported.")
-    if "severe headache" in [s.lower() for s in symptoms]: flags.append("Red Flag: Severe Headache reported.")
-    if "sudden vision loss" in [s.lower() for s in symptoms]: flags.append("Red Flag: Sudden Vision Loss reported.")
-    if "weakness on one side" in [s.lower() for s in symptoms]: flags.append("Red Flag: Unilateral Weakness reported (potential stroke).")
-
-    # Vital Sign Flags (add more checks as needed)
-    if "temp_c" in vitals and vitals["temp_c"] >= 38.5: flags.append(f"Red Flag: Fever (Temperature: {vitals['temp_c']}°C).")
-    if "hr_bpm" in vitals and vitals["hr_bpm"] >= 120: flags.append(f"Red Flag: Tachycardia (Heart Rate: {vitals['hr_bpm']} bpm).")
-    if "rr_rpm" in vitals and vitals["rr_rpm"] >= 24: flags.append(f"Red Flag: Tachypnea (Respiratory Rate: {vitals['rr_rpm']} rpm).")
-    if "spo2_percent" in vitals and vitals["spo2_percent"] <= 92: flags.append(f"Red Flag: Hypoxia (SpO2: {vitals['spo2_percent']}%).")
+    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:
-            if bp[0] >= 180 or bp[1] >= 110: flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {vitals['bp_mmhg']} mmHg).")
-            if bp[0] <= 90 and bp[1] <= 60: flags.append(f"Red Flag: Hypotension (BP: {vitals['bp_mmhg']} mmHg).")
-
-    # History Flags (Simple examples)
-    if "history of mi" in history.lower() and "chest pain" in [s.lower() for s in symptoms]: flags.append("Red Flag: History of MI with current Chest Pain.")
+        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 "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
 
-# --- Enhanced Tool Definitions ---
 
-# Use Pydantic models for robust argument validation
+# --- 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)
+
 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').")
@@ -119,11 +120,7 @@ class LabOrderInput(BaseModel):
 @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}"
-    })
+    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.")
@@ -136,12 +133,8 @@ class PrescriptionInput(BaseModel):
 @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."""
-    # In a real scenario, 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}"
-    })
+    # 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}"})
 
 class InteractionCheckInput(BaseModel):
     potential_prescription: str = Field(..., description="The name of the NEW medication being considered.")
@@ -153,35 +146,34 @@ def check_drug_interactions(potential_prescription: str, current_medications: Li
     """Checks for potential drug-drug and drug-allergy interactions BEFORE prescribing."""
     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]
 
-    # Check Allergies
-    for allergy in allergies:
-        allergy_lower = allergy.lower()
-        # Simple direct check
-        if allergy_lower == potential_med_lower:
+    for allergy in allergies_lower:
+        if allergy == potential_med_lower:
             warnings.append(f"CRITICAL ALLERGY: Patient allergic to {allergy}. Cannot prescribe {potential_prescription}.")
             continue
-        # Check cross-reactivity (using simplified mock data)
-        if allergy_lower in ALLERGY_INTERACTIONS:
-            for cross_reactant in ALLERGY_INTERACTIONS[allergy_lower]:
+        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}.")
 
-    # Check Drug-Drug Interactions (using simplified mock data)
-    current_meds_lower = [med.lower() for med in current_medications]
     for current_med in current_meds_lower:
-        # Check pairs in both orders
         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 Found: {potential_prescription} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair1]}")
+            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 Found: {potential_prescription} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair2]}")
+             warnings.append(f"Interaction: {potential_prescription.capitalize()} with {current_med.capitalize()} - {MOCK_INTERACTION_DB[pair2]}")
+
+    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}."
+    return json.dumps({"status": status, "message": message, "warnings": warnings})
 
-    if not warnings:
-        return json.dumps({"status": "clear", "message": f"No major interactions identified for {potential_prescription} with current meds/allergies.", "warnings": []})
-    else:
-        return json.dumps({"status": "warning", "message": f"Potential interactions identified for {potential_prescription}.", "warnings": warnings})
 
 class FlagRiskInput(BaseModel):
     risk_description: str = Field(..., description="Specific critical risk identified (e.g., 'Suspected Sepsis', 'Acute Coronary Syndrome', 'Stroke Alert').")
@@ -190,372 +182,415 @@ 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."""
-    st.error(f"🚨 **{urgency.upper()} RISK FLAGGED:** {risk_description}", icon="🚨")
-    return json.dumps({
-        "status": "flagged",
-        "message": f"Risk '{risk_description}' flagged with {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."})
 
 # Initialize Search Tool
-search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS)
-
-# --- Core Agent Logic ---
-class ClinicalAgent:
-    def __init__(self):
-        self.model = ChatGroq(
-            temperature=ClinicalAppSettings.TEMPERATURE,
-            model=ClinicalAppSettings.MODEL_NAME
-        )
-        # Combine all tools
-        self.tools = [
-            order_lab_test,
-            prescribe_medication,
-            check_drug_interactions,
-            flag_risk,
-            search_tool
-        ]
-        # Bind tools to the model
-        self.model_with_tools = self.model.bind_tools(self.tools)
-        # History for context (simple implementation)
-        self.history = []
-
-    def _format_patient_data_for_prompt(self, data: dict) -> str:
-        """Formats the patient dictionary into a readable string for the LLM."""
-        prompt_str = "Patient Data:\n"
-        for key, value in data.items():
-            if isinstance(value, dict):
-                prompt_str += f"  {key.replace('_', ' ').title()}:\n"
-                for sub_key, sub_value in value.items():
-                     if sub_value: # Only include if there's data
-                        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: # Only include non-empty fields
-                 prompt_str += f"  {key.replace('_', ' ').title()}: {value}\n"
-        return prompt_str.strip()
-
-
-    def analyze(self, patient_data: dict) -> tuple[Optional[dict], List[dict]]:
-        """Runs the analysis, handling tool calls and parsing the structured output."""
-        try:
-            # Add System Prompt and formatted Patient Data
-            # Simple history management: add previous messages if any
-            messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)]
-            # Include history if needed - consider token limits
-            # messages.extend(self.history)
-            formatted_data = self._format_patient_data_for_prompt(patient_data)
-            messages.append(HumanMessage(content=formatted_data))
-
-            # Invoke the model
-            ai_response = self.model_with_tools.invoke(messages)
-
-            # Store conversation turn
-            # self.history.append(HumanMessage(content=formatted_data))
-            # self.history.append(ai_response) # AIMessage includes tool calls
-
-            response_content = None
-            tool_calls = []
-
-            if isinstance(ai_response, AIMessage):
-                # Check if the response contains the structured JSON output
+search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS, name="tavily_search_results")
+
+
+# --- 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]
+
+# Define Tools and Tool Executor
+tools = [
+    order_lab_test,
+    prescribe_medication,
+    check_drug_interactions,
+    flag_risk,
+    search_tool
+]
+tool_executor = ToolExecutor(tools)
+
+# Define the Agent Model
+model = ChatGroq(
+    temperature=ClinicalAppSettings.TEMPERATURE,
+    model=ClinicalAppSettings.MODEL_NAME
+)
+model_with_tools = model.bind_tools(tools) # Bind tools for the LLM to know about them
+
+# --- 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
+    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}")
+    return {"messages": [response]}
+
+# 2. Tool Node: Executes tools called by the Agent
+def tool_node(state: AgentState):
+    """Executes tools called by the LLM and returns results."""
+    print("---TOOL NODE---")
+    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
+
+    tool_calls = last_message.tool_calls
+    tool_messages = []
+
+    # 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'}
+
+    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
+            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
+            )
+            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]
+
+
+    # Add patient context to interaction checks if needed
+    patient_meds = state.get("patient_data", {}).get("medications", {}).get("names_only", [])
+    patient_allergies = state.get("patient_data", {}).get("allergies", [])
+    for call in tool_calls:
+         if call['name'] == 'check_drug_interactions':
+             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}")
+
+    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")
+        return "continue_tools"
+    # Otherwise, we end the loop (AI provided a direct answer or finished)
+    else:
+        print("Routing: end_conversation_turn")
+        return "end_conversation_turn"
+
+# --- Graph Definition ---
+workflow = StateGraph(AgentState)
+
+# Add nodes
+workflow.add_node("agent", agent_node)
+workflow.add_node("tools", tool_node)
+
+# Define entry point
+workflow.set_entry_point("agent")
+
+# Add conditional edges
+workflow.add_conditional_edges(
+    "agent",                      # Source node
+    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
+    }
+)
+
+# Add edge from tools back to agent
+workflow.add_edge("tools", "agent")
+
+# Compile the graph
+# memory = MemorySaverInMemory() # Optional: for persisting state across runs
+# app = workflow.compile(checkpointer=memory)
+app = workflow.compile()
+
+# --- 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) ---
+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
+    if "messages" not in st.session_state:
+        st.session_state.messages = [] # Store entire conversation history (Human, AI, Tool)
+    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
+
+    # --- Patient Data Input Sidebar --- (Similar to before)
+    with st.sidebar:
+        st.header("📄 Patient Intake Form")
+        # ... (Keep the input fields exactly as in the previous example) ...
+        # 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
+        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")
+        # History
+        pmh = st.text_area("Past Medical History (PMH)", "Hypertension (HTN), Hyperlipidemia (HLD), Type 2 Diabetes Mellitus (DM2)", key="pmh_input")
+        psh = st.text_area("Past Surgical History (PSH)", "Appendectomy (2005)", key="psh_input")
+        # Meds & 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")
+        # Social/Family
+        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
+        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")
+            hr_bpm = st.number_input("HR (bpm)", 30, 250, 95, key="hr_input")
+            rr_rpm = st.number_input("RR (rpm)", 5, 50, 18, key="rr_input")
+        with col2:
+            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")
+
+        # 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()]
+            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
+
+            st.session_state.patient_data = {
+                "demographics": {"age": age, "sex": sex},
+                "hpi": {"chief_complaint": chief_complaint, "details": hpi_details, "symptoms": symptoms},
+                "pmh": {"conditions": pmh}, "psh": {"procedures": psh},
+                "medications": {"current": current_meds_list, "names_only": current_med_names},
+                "allergies": allergies_list,
+                "social_history": {"details": social_history}, "family_history": {"details": family_history},
+                "vitals": { "temp_c": temp_c, "hr_bpm": hr_bpm, "bp_mmhg": bp_mmhg, "rr_rpm": rr_rpm, "spo2_percent": spo2_percent, "pain_scale": pain_scale},
+                "exam_findings": {"notes": exam_notes}
+            }
+
+            # Initial Red Flag Check (Client-side)
+            red_flags = check_red_flags(st.session_state.patient_data)
+            if red_flags:
+                st.warning("**Initial Red Flags Detected:**")
+                for flag in red_flags: st.warning(f"- {flag}")
+
+            # 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
+            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
+
+
+    # --- Main Chat Interface Area ---
+    st.header("💬 Clinical Consultation")
+
+    # Display chat messages
+    for msg in st.session_state.messages:
+        if isinstance(msg, HumanMessage):
+            with st.chat_message("user"):
+                st.markdown(msg.content)
+        elif isinstance(msg, AIMessage):
+            with st.chat_message("assistant"):
+                # Check for structured JSON output
+                structured_output = None
                 try:
-                    # Sometimes the JSON is embedded in the content, sometimes it's the primary content
-                    # Look for ```json ... ``` block first
-                    json_match = re.search(r"```json\n(\{.*?\})\n```", ai_response.content, re.DOTALL)
+                    # Try to find JSON block first
+                    json_match = re.search(r"```json\n(\{.*?\})\n```", msg.content, re.DOTALL)
                     if json_match:
-                        response_content = json.loads(json_match.group(1))
+                        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)
                     else:
-                        # Try parsing the whole content as JSON
-                        response_content = json.loads(ai_response.content)
+                        # 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"))
+
+
+                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:
+                     with st.expander("🛠️ AI requested actions", expanded=False):
+                         for tc in msg.tool_calls:
+                             st.code(f"{tc['name']}(args={tc['args']})", language="python")
+
+        elif isinstance(msg, ToolMessage):
+             with st.chat_message("tool", avatar="🛠️"):
+                try:
+                    tool_data = json.loads(msg.content)
+                    status = tool_data.get("status", "info")
+                    message = tool_data.get("message", msg.content)
+                    details = tool_data.get("details")
+                    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 "🚨")
+                    elif status == "warning":
+                         st.warning(f"Tool Result ({msg.name}): {message}", icon="⚠️")
+                         if warnings:
+                             for warn in warnings: st.caption(f"- {warn}")
+                    else: # Error or unknown status
+                         st.error(f"Tool Result ({msg.name}): {message}", icon="❌")
+
+                    if details: st.caption(f"Details: {details}")
+
                 except json.JSONDecodeError:
-                    st.warning("AI did not return valid JSON in the expected format. Displaying raw content.")
-                    st.code(ai_response.content, language=None) # Display raw if not JSON
-                    response_content = {"assessment": ai_response.content, "error": "Output format incorrect"}
+                    st.info(f"Tool Result ({msg.name}): {msg.content}") # Display raw if not JSON
 
-                # Extract tool calls separately
-                if ai_response.tool_calls:
-                    tool_calls = ai_response.tool_calls
 
-            return response_content, tool_calls
+    # Chat input for user
+    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
 
-        except Exception as e:
-            st.error(f"Error during AI analysis: {str(e)}")
-            return None, []
+                                # 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]}")
 
-    def process_tool_call(self, tool_call: Dict[str, Any]) -> Any:
-        """Executes a single tool call."""
-        tool_name = tool_call.get("name")
-        tool_args = tool_call.get("args", {})
-        selected_tool = {t.name: t for t in self.tools}.get(tool_name)
+                        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.")
 
-        if not selected_tool:
-            return json.dumps({"status": "error", "message": f"Unknown tool: {tool_name}"})
 
-        try:
-            # Ensure args are correctly passed (Pydantic models handle validation)
-            return selected_tool.invoke(tool_args)
-        except Exception as e:
-            st.error(f"Error executing tool '{tool_name}': {str(e)}")
-            return json.dumps({"status": "error", "message": f"Failed to execute {tool_name}: {str(e)}"})
+                    # Final display after streaming
+                    message_placeholder.markdown(full_response)
 
-# --- 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"Powered by Langchain & Groq ({ClinicalAppSettings.MODEL_NAME})")
-
-    # Initialize Agent in session state
-    if 'agent' not in st.session_state:
-        st.session_state.agent = ClinicalAgent()
-    if 'analysis_complete' not in st.session_state:
-        st.session_state.analysis_complete = False
-    if 'analysis_result' not in st.session_state:
-        st.session_state.analysis_result = None
-    if 'tool_call_results' not in st.session_state:
-        st.session_state.tool_call_results = []
-    if 'red_flags' not in st.session_state:
-         st.session_state.red_flags = []
-
-    # --- Patient Data Input Sidebar ---
-    with st.sidebar:
-        st.header("📄 Patient Intake Form")
 
-        # Demographics
-        st.subheader("Demographics")
-        age = st.number_input("Age", min_value=0, max_value=120, value=55)
-        sex = st.selectbox("Biological Sex", ["Male", "Female", "Other/Prefer not to say"])
-
-        # History of Present Illness (HPI)
-        st.subheader("History of Present Illness (HPI)")
-        chief_complaint = st.text_input("Chief Complaint", "Chest pain")
-        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.")
-        symptoms = st.multiselect("Associated Symptoms", ["Nausea", "Diaphoresis", "Shortness of Breath", "Dizziness", "Palpitations", "Fever", "Cough"], default=["Nausea", "Diaphoresis"])
-
-        # Past Medical/Surgical History (PMH/PSH)
-        st.subheader("Past History")
-        pmh = st.text_area("Past Medical History (PMH)", "Hypertension (HTN), Hyperlipidemia (HLD), Type 2 Diabetes Mellitus (DM2)")
-        psh = st.text_area("Past Surgical History (PSH)", "Appendectomy (2005)")
-
-        # Medications & Allergies
-        st.subheader("Medications & Allergies")
-        current_meds = st.text_area("Current Medications (name, dose, freq)", "Lisinopril 10mg daily\nMetformin 1000mg BID\nAtorvastatin 40mg daily\nAspirin 81mg daily")
-        allergies = st.text_area("Allergies (comma separated)", "Penicillin (rash)")
-
-        # Social & Family History (SH/FH)
-        st.subheader("Social/Family History")
-        social_history = st.text_area("Social History (SH)", "Smoker (1 ppd x 30 years), occasional alcohol.")
-        family_history = st.text_area("Family History (FHx)", "Father had MI at age 60. Mother has HTN.")
-
-        # Review of Systems (ROS) - Simplified
-        # st.subheader("Review of Systems (ROS)") # Keep UI cleaner for now
-        # ros_constitutional = st.checkbox("ROS: Constitutional (Fever, Chills, Weight loss)")
-        # ros_cardiac = st.checkbox("ROS: Cardiac (Chest pain, Palpitations)", value=True) # Pre-check based on HPI
-
-        # Vitals & Basic Exam
-        st.subheader("Vitals & Exam Findings")
-        col1, col2 = st.columns(2)
-        with col1:
-            temp_c = st.number_input("Temperature (°C)", 35.0, 42.0, 36.8, format="%.1f")
-            hr_bpm = st.number_input("Heart Rate (bpm)", 30, 250, 95)
-            rr_rpm = st.number_input("Respiratory Rate (rpm)", 5, 50, 18)
-        with col2:
-            bp_mmhg = st.text_input("Blood Pressure (SYS/DIA)", "155/90")
-            spo2_percent = st.number_input("SpO2 (%)", 70, 100, 96)
-            pain_scale = st.slider("Pain (0-10)", 0, 10, 8)
-        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3. Mild distress. Lungs clear. Cardiac exam: Regular rhythm, no murmurs/gallops. Abdomen soft. No edema.")
-
-        # Clean medication list and allergies for processing
-        current_meds_list = [med.strip() for med in current_meds.split('\n') if med.strip()]
-        current_med_names = [med.split(' ')[0].strip() for med in current_meds_list] # Simplified name extraction
-        allergies_list = [a.strip() for a in allergies.split(',') if a.strip()]
-
-        # Compile Patient Data Dictionary
-        patient_data = {
-            "demographics": {"age": age, "sex": sex},
-            "hpi": {"chief_complaint": chief_complaint, "details": hpi_details, "symptoms": symptoms},
-            "pmh": {"conditions": pmh},
-            "psh": {"procedures": psh},
-            "medications": {"current": current_meds_list, "names_only": current_med_names},
-            "allergies": allergies_list,
-            "social_history": {"details": social_history},
-            "family_history": {"details": family_history},
-            # "ros": {"constitutional": ros_constitutional, "cardiac": ros_cardiac}, # Add if using ROS inputs
-            "vitals": {
-                "temp_c": temp_c, "hr_bpm": hr_bpm, "bp_mmhg": bp_mmhg,
-                "rr_rpm": rr_rpm, "spo2_percent": spo2_percent, "pain_scale": pain_scale
-            },
-            "exam_findings": {"notes": exam_notes}
-        }
+                    # 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()
 
-    # --- Main Analysis Area ---
-    st.header("🤖 AI Clinical Analysis")
-
-    # Action Button
-    if st.button("Analyze Patient Data", type="primary", use_container_width=True):
-        st.session_state.analysis_complete = False
-        st.session_state.analysis_result = None
-        st.session_state.tool_call_results = []
-        st.session_state.red_flags = []
-
-        # 1. Initial Red Flag Check (Client-side before LLM)
-        st.session_state.red_flags = check_red_flags(patient_data)
-        if st.session_state.red_flags:
-            st.warning("**Initial Red Flags Detected:**")
-            for flag in st.session_state.red_flags:
-                st.warning(f"- {flag}")
-            st.warning("Proceeding with AI analysis, but these require immediate attention.")
-
-        # 2. Call AI Agent
-        with st.spinner("SynapseAI is processing the case... Please wait."):
-            analysis_output, tool_calls = st.session_state.agent.analyze(patient_data)
-
-        if analysis_output:
-            st.session_state.analysis_result = analysis_output
-            st.session_state.analysis_complete = True
-
-            # 3. Process any Tool Calls requested by the AI
-            if tool_calls:
-                st.info(f"AI recommended {len(tool_calls)} action(s). Executing...")
-                tool_results = []
-                with st.spinner("Executing recommended actions..."):
-                    for call in tool_calls:
-                        st.write(f"⚙️ Requesting: `{call['name']}` with args `{call['args']}`")
-                        # Pass patient context if needed (e.g., for interaction check)
-                        if call['name'] == 'check_drug_interactions':
-                           call['args']['current_medications'] = patient_data['medications']['names_only']
-                           call['args']['allergies'] = patient_data['allergies']
-                        elif call['name'] == 'prescribe_medication':
-                             # Pre-flight check: Ensure interaction check was requested *before* this prescribe call
-                             interaction_check_requested = any(tc['name'] == 'check_drug_interactions' and tc['args'].get('potential_prescription') == call['args'].get('medication_name') for tc in tool_calls)
-                             if not interaction_check_requested:
-                                 st.error(f"**Safety Violation:** AI attempted to prescribe '{call['args'].get('medication_name')}' without requesting `check_drug_interactions` first. Prescription blocked.")
-                                 tool_results.append({"tool_call_id": call['id'], "name": call['name'], "output": json.dumps({"status":"error", "message": "Interaction check not performed prior to prescription attempt."})})
-                                 continue # Skip this tool call
-
-                        result = st.session_state.agent.process_tool_call(call)
-                        tool_results.append({"tool_call_id": call['id'], "name": call['name'], "output": result}) # Store result with ID
-
-                        # Display tool result immediately
-                        try:
-                            result_data = json.loads(result)
-                            if result_data.get("status") == "success" or result_data.get("status") == "clear" or result_data.get("status") == "flagged":
-                                st.success(f"✅ Action `{call['name']}`: {result_data.get('message')}", icon="✅")
-                                if result_data.get("details"): st.caption(f"Details: {result_data.get('details')}")
-                            elif result_data.get("status") == "warning":
-                                st.warning(f"⚠️ Action `{call['name']}`: {result_data.get('message')}", icon="⚠️")
-                                if result_data.get("warnings"):
-                                    for warn in result_data["warnings"]: st.caption(f"- {warn}")
-                            else:
-                                st.error(f"❌ Action `{call['name']}`: {result_data.get('message')}", icon="❌")
-                        except json.JSONDecodeError:
-                             st.error(f"Tool `{call['name']}` returned non-JSON: {result}") # Fallback for non-JSON results
-
-                st.session_state.tool_call_results = tool_results
-                # Optionally: Send results back to LLM for final summary (requires multi-turn agent)
-        else:
-            st.error("Analysis failed. Please check the input data or try again.")
-
-    # --- Display Analysis Results ---
-    if st.session_state.analysis_complete and st.session_state.analysis_result:
-        st.divider()
-        st.header("📊 Analysis & Recommendations")
-
-        res = st.session_state.analysis_result
-
-        # Layout columns for better readability
-        col_assessment, col_plan = st.columns(2)
-
-        with col_assessment:
-            st.subheader("📋 Assessment")
-            st.write(res.get("assessment", "N/A"))
-
-            st.subheader("🤔 Differential Diagnosis")
-            ddx = res.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 Diagnosis')} ({likelihood} Likelihood)", expanded=(likelihood=="High")):
-                        st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
-            else:
-                st.info("No differential diagnosis provided.")
-
-            st.subheader("🚨 Risk Assessment")
-            risk = res.get("risk_assessment", {})
-            flags = risk.get("identified_red_flags", []) + [f.replace("Red Flag: ", "") for f in st.session_state.red_flags] # Combine AI and initial flags
-            if flags:
-                st.warning(f"**Identified Red Flags:** {', '.join(flags)}")
-            else:
-                st.success("No immediate red flags identified by AI in this analysis.")
-
-            if risk.get("immediate_concerns"):
-                st.warning(f"**Immediate Concerns:** {', '.join(risk.get('immediate_concerns'))}")
-            if risk.get("potential_complications"):
-                 st.info(f"**Potential Complications:** {', '.join(risk.get('potential_complications'))}")
-
-
-        with col_plan:
-            st.subheader("📝 Recommended Plan")
-            plan = res.get("recommended_plan", {})
-
-            st.markdown("**Investigations:**")
-            if plan.get("investigations"):
-                 st.markdown("\n".join([f"- {inv}" for inv in plan.get("investigations")]))
-            else: st.markdown("_None suggested._")
-
-            st.markdown("**Therapeutics:**")
-            if plan.get("therapeutics"):
-                 st.markdown("\n".join([f"- {thx}" for thx in plan.get("therapeutics")]))
-            else: st.markdown("_None suggested._")
-
-            st.markdown("**Consultations:**")
-            if plan.get("consultations"):
-                 st.markdown("\n".join([f"- {con}" for con in plan.get("consultations")]))
-            else: st.markdown("_None suggested._")
-
-            st.markdown("**Patient Education:**")
-            if plan.get("patient_education"):
-                 st.markdown("\n".join([f"- {edu}" for edu in plan.get("patient_education")]))
-            else: st.markdown("_None specified._")
-
-        # Display Rationale and Interaction Summary below the columns
-        st.subheader("🧠 AI Rationale & Checks")
-        with st.expander("Show AI Reasoning Summary", expanded=False):
-            st.write(res.get("rationale_summary", "No rationale summary provided."))
-
-        interaction_summary = res.get("interaction_check_summary", "")
-        if interaction_summary: # Only show if interaction check was relevant/performed
-             with st.expander("Drug Interaction Check Summary", expanded=True):
-                 st.write(interaction_summary)
-                 # Also show detailed results from the tool call itself if available
-                 for tool_res in st.session_state.tool_call_results:
-                     if tool_res['name'] == 'check_drug_interactions':
-                         try:
-                            data = json.loads(tool_res['output'])
-                            if data.get('warnings'):
-                                st.warning("Interaction Details:")
-                                for warn in data['warnings']:
-                                    st.caption(f"- {warn}")
-                            else:
-                                st.success("Interaction Details: " + data.get('message', 'Check complete.'))
-                         except: pass # Ignore parsing errors here
-
-        # Display raw JSON if needed for debugging
-        with st.expander("Show Raw AI Output (JSON)"):
-            st.json(res)
-
-        st.divider()
-        st.success("Analysis Complete.")
 
     # Disclaimer
     st.markdown("---")
-    st.warning(
-        """**Disclaimer:** SynapseAI is an AI assistant for clinical decision support and does not replace professional medical judgment.
-        All outputs should be critically reviewed by a qualified healthcare provider before making any clinical decisions.
-        Verify all information, especially dosages and interactions, independently."""
-    )
-
+    st.warning("**Disclaimer:** SynapseAI is for clinical decision support...") # Keep disclaimer
 
 if __name__ == "__main__":
     main()