Update app.py

app.py

@@ -1,121 +1,561 @@
 import streamlit as st
 from langchain_groq import ChatGroq
 from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+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
+from typing import Optional, List, Dict, Any
 import json
+import re # For parsing vitals like BP
 
-# Configuration
-class MedicalConfig:
-    SYSTEM_PROMPT = """You are an AI clinical assistant. Follow these rules:
-1. Analyze patient data using latest medical guidelines
-2. Always check for drug interactions
-3. Use structured actions when needed:
-   - lab_order: Order lab tests
-   - prescribe: Prescribe medication
-4. Flag high-risk conditions immediately"""
-
-    RED_FLAGS = {
-        "symptoms": ["chest pain", "shortness of breath", "severe headache"],
-        "vitals": {"temp": (38.5, "°C"), "hr": (120, "bpm"), "bp": ("180/120", "mmHg")}
-    }
-
-# Define tools with proper schemas
-@tool
-def order_lab_test(test_name: str, reason: str) -> str:
-    """Orders a lab test with specified parameters."""
-    return f"✅ Lab ordered: {test_name} ({reason})"
-
-@tool
-def prescribe_medication(name: str, dosage: str, frequency: str) -> str:
-    """Prescribes medication with specific dosage instructions."""
-    return f"✅ Prescribed: {name} {dosage} {frequency}"
-
-# Initialize tools and model
-tools = [order_lab_test, prescribe_medication, TavilySearchResults(max_results=3)]
-
-class MedicalAgent:
+# --- Configuration & Constants ---
+class ClinicalAppSettings:
+    APP_TITLE = "SynapseAI: Advanced Clinical Decision Support"
+    PAGE_LAYOUT = "wide"
+    MODEL_NAME = "llama3-70b-8192" # Use a powerful model like Groq's Llama3 70b
+    TEMPERATURE = 0.1
+    MAX_SEARCH_RESULTS = 3
+
+class ClinicalPrompts:
+    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.
+
+    **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:
+        ```json
+        {
+          "assessment": "Concise summary of the patient's presentation and key findings.",
+          "differential_diagnosis": [
+            {"diagnosis": "Primary Diagnosis", "likelihood": "High/Medium/Low", "rationale": "Supporting evidence..."},
+            {"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"]
+          },
+          "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."],
+            "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."
+        }
+        ```
+    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.
+    """
+
+# --- Mock Data / Helpers ---
+# (In a real system, this would be a proper API/database)
+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."
+}
+
+ALLERGY_INTERACTIONS = {
+    "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))
+    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", "")
+
+    # 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 "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.")
+
+    return flags
+
+# --- Enhanced Tool Definitions ---
+
+# Use 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').")
+    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}"
+    })
+
+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').")
+    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."""
+    # 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}"
+    })
+
+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.")
+
+@tool("check_drug_interactions", args_schema=InteractionCheckInput)
+def check_drug_interactions(potential_prescription: str, current_medications: List[str], allergies: List[str]) -> str:
+    """Checks for potential drug-drug and drug-allergy interactions BEFORE prescribing."""
+    warnings = []
+    potential_med_lower = potential_prescription.lower()
+
+    # Check Allergies
+    for allergy in allergies:
+        allergy_lower = allergy.lower()
+        # Simple direct check
+        if allergy_lower == 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 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)
+        if pair1 in MOCK_INTERACTION_DB:
+            warnings.append(f"Interaction Found: {potential_prescription} 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]}")
+
+    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').")
+    urgency: str = Field("High", description="Urgency level (e.g., 'Critical', 'High', 'Moderate').")
+
+@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."
+    })
+
+
+# Initialize Search Tool
+search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS)
+
+# --- Core Agent Logic ---
+class ClinicalAgent:
     def __init__(self):
-        self.model = ChatGroq(temperature=0.1, model="Llama-3.3-70b-Specdec")
-        self.model_with_tools = self.model.bind_tools(tools)
+        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 analyze_patient(self, patient_data: dict) -> Optional[dict]:
+    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:
-            response = self.model_with_tools.invoke([
-                SystemMessage(content=MedicalConfig.SYSTEM_PROMPT),
-                HumanMessage(content=f"Patient Data: {json.dumps(patient_data)}")
-            ])
-            return response
+            # 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
+                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)
+                    if json_match:
+                        response_content = json.loads(json_match.group(1))
+                    else:
+                        # Try parsing the whole content as JSON
+                        response_content = json.loads(ai_response.content)
+                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"}
+
+                # Extract tool calls separately
+                if ai_response.tool_calls:
+                    tool_calls = ai_response.tool_calls
+
+            return response_content, tool_calls
+
         except Exception as e:
-            st.error(f"Error analyzing patient data: {str(e)}")
-            return None
+            st.error(f"Error during AI analysis: {str(e)}")
+            return None, []
+
+    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)
+
+        if not selected_tool:
+            return json.dumps({"status": "error", "message": f"Unknown tool: {tool_name}"})
 
-    def process_action(self, action: dict) -> str:
         try:
-            tool_name = action['name']
-            args = action['args']
-            
-            if tool_name == "order_lab_test":
-                return order_lab_test.invoke(args)
-            elif tool_name == "prescribe_medication":
-                return prescribe_medication.invoke(args)
-            else:
-                return f"Unknown action: {tool_name}"
+            # Ensure args are correctly passed (Pydantic models handle validation)
+            return selected_tool.invoke(tool_args)
         except Exception as e:
-            return f"Error processing action: {str(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)}"})
 
-# Streamlit UI
+# --- Streamlit UI ---
 def main():
-    st.set_page_config(page_title="AI Clinic", layout="wide")
-    
-    # Initialize agent
+    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 = MedicalAgent()
-    
-    # Patient intake form
+        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")
-        symptoms = st.multiselect("Symptoms", ["Fever", "Cough", "Chest Pain", "Shortness of Breath"])
-        medical_history = st.text_input("Medical History")
-        current_meds = st.text_input("Current Medications")
-        temp = st.number_input("Temperature (°C)", 35.0, 42.0, 37.0)
-        bp = st.text_input("Blood Pressure (mmHg)", "120/80")
-        
+        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 = {
-            "symptoms": symptoms,
-            "history": {
-                "conditions": medical_history,
-                "medications": current_meds
-            },
+            "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": temp,
-                "bp": bp
-            }
+                "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}
         }
-    
-    # Main interface
-    st.title("AI-Powered Clinical Support System")
-    
-    # Analysis button
-    if st.button("Start Analysis", type="primary"):
-        with st.spinner("Analyzing patient data..."):
-            response = st.session_state.agent.analyze_patient(patient_data)
-            
-            if response:
-                st.subheader("Clinical Recommendations")
-                
-                # Check for tool calls
-                if hasattr(response, 'tool_calls') and response.tool_calls:
-                    for action in response.tool_calls:
-                        result = st.session_state.agent.process_action({
-                            "name": action['name'],
-                            "args": action['args']
-                        })
-                        st.success(result)
-                else:
-                    st.info("No specific actions recommended. Clinical Summary:")
-                    st.write(response.content)
+
+    # --- 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."""
+    )
+
 
 if __name__ == "__main__":
     main()