Update app.py

app.py

@@ -1,4 +1,13 @@
 import streamlit as st
+import requests
+import json
+import re
+import os
+import operator
+import traceback
+from functools import lru_cache
+from dotenv import load_dotenv
+
 from langchain_groq import ChatGroq
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
@@ -7,24 +16,37 @@ from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.tools import tool
 from langgraph.prebuilt import ToolExecutor
 from langgraph.graph import StateGraph, END
-# from langgraph.checkpoint.memory import MemorySaverInMemory # Optional for state persistence
 
 from typing import Optional, List, Dict, Any, TypedDict, Annotated
-import json
-import re
-import operator
-import traceback # For detailed error logging
+
+# --- Environment Variable Loading & Validation ---
+load_dotenv() # Load .env file if present (for local development)
+
+UMLS_API_KEY = os.environ.get("UMLS_API_KEY")
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
+TAVILY_API_KEY = os.environ.get("TAVILY_API_KEY")
+
+# Stop execution if essential keys are missing (crucial for HF Spaces)
+missing_keys = []
+if not UMLS_API_KEY: missing_keys.append("UMLS_API_KEY")
+if not GROQ_API_KEY: missing_keys.append("GROQ_API_KEY")
+if not TAVILY_API_KEY: missing_keys.append("TAVILY_API_KEY")
+
+if missing_keys:
+    st.error(f"Missing required API Key(s): {', '.join(missing_keys)}. Please set them in Hugging Face Space Secrets or your environment variables.")
+    st.stop()
 
 # --- Configuration & Constants ---
 class ClinicalAppSettings:
-    APP_TITLE = "SynapseAI: Interactive Clinical Decision Support"
+    APP_TITLE = "SynapseAI: Interactive Clinical Decision Support (UMLS/FDA Integrated)"
     PAGE_LAYOUT = "wide"
     MODEL_NAME = "llama3-70b-8192" # Groq Llama3 70b
     TEMPERATURE = 0.1
     MAX_SEARCH_RESULTS = 3
 
 class ClinicalPrompts:
-    # UPDATED SYSTEM PROMPT FOR CONVERSATIONAL FLOW, GUIDELINES & STRUCTURED OUTPUT FOCUS
+    # System prompt remains the same as the previous version, emphasizing structured output,
+    # safety checks, guideline search, and conversational flow.
     SYSTEM_PROMPT = """
     You are SynapseAI, an expert AI clinical assistant engaged in an interactive consultation.
     Your goal is to support healthcare professionals by analyzing patient data, providing differential diagnoses, suggesting evidence-based management plans, and identifying risks according to current standards of care.
@@ -63,45 +85,137 @@ class ClinicalPrompts:
     8.  **Conciseness & Flow:** Be medically accurate and concise. Use standard terminology. Respond naturally in conversation (asking questions, acknowledging info) until ready for the full structured JSON output.
     """
 
-# --- Mock Data / Helpers ---
-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.",
-    # Add lower case versions for easier lookup
-    ("amiodarone", "warfarin"): "Increased bleeding risk. Monitor INR frequently and adjust Warfarin dose.",
-    ("clarithromycin", "simvastatin"): "Increased risk of myopathy/rhabdomyolysis. Avoid combination or use lower statin dose.",
-    ("ibuprofen", "aspirin"): "Concurrent use may decrease Aspirin's cardioprotective effect. Potential for increased GI bleeding.",
-    ("spironolactone", "lisinopril"): "High risk of hyperkalemia. Monitor potassium closely.",
-}
-
-ALLERGY_INTERACTIONS = {
-    "penicillin": ["amoxicillin", "ampicillin", "piperacillin", "augmentin"],
-    "sulfa": ["sulfamethoxazole", "sulfasalazine", "bactrim"],
-    "aspirin": ["ibuprofen", "naproxen", "nsaid"] # Cross-reactivity example for NSAIDs
-}
+# --- UMLS/RxNorm & OpenFDA API Helper Functions ---
+UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key" # May not be needed if using apiKey directly
+RXNORM_API_BASE = "https://rxnav.nlm.nih.gov/REST"
+OPENFDA_API_BASE = "https://api.fda.gov/drug/label.json"
 
+@lru_cache(maxsize=256) # Cache RxCUI lookups
+def get_rxcui(drug_name: str) -> Optional[str]:
+    """Uses RxNorm API to find the RxCUI for a given drug name."""
+    if not drug_name or not isinstance(drug_name, str): return None
+    drug_name = drug_name.strip()
+    if not drug_name: return None
+
+    print(f"RxNorm Lookup for: '{drug_name}'")
+    try:
+        params = {"name": drug_name, "search": 1} # Search for concepts related to the name
+        response = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10)
+        response.raise_for_status()
+        data = response.json()
+        # Extract RxCUI - prioritize exact matches or common types
+        if data and "idGroup" in data and "rxnormId" in data["idGroup"]:
+            # Select the first one, assuming it's the most relevant by default.
+            # More sophisticated logic could check TTYs (Term Types) if needed.
+            rxcui = data["idGroup"]["rxnormId"][0]
+            print(f"  Found RxCUI: {rxcui} for '{drug_name}'")
+            return rxcui
+        else:
+            # Fallback: Search /drugs endpoint if direct rxcui lookup fails
+            params = {"name": drug_name}
+            response = requests.get(f"{RXNORM_API_BASE}/drugs.json", params=params, timeout=10)
+            response.raise_for_status()
+            data = response.json()
+            if data and "drugGroup" in data and "conceptGroup" in data["drugGroup"]:
+                for group in data["drugGroup"]["conceptGroup"]:
+                    # Prioritize Semantic Types like Brand/Clinical Drug/Ingredient
+                    if group.get("tty") in ["SBD", "SCD", "GPCK", "BPCK", "IN", "MIN", "PIN"]:
+                        if "conceptProperties" in group and group["conceptProperties"]:
+                            rxcui = group["conceptProperties"][0].get("rxcui")
+                            if rxcui:
+                                print(f"  Found RxCUI (via /drugs): {rxcui} for '{drug_name}'")
+                                return rxcui
+        print(f"  RxCUI not found for '{drug_name}'.")
+        return None
+    except requests.exceptions.RequestException as e:
+        print(f"  Error fetching RxCUI for '{drug_name}': {e}")
+        return None
+    except json.JSONDecodeError as e:
+        print(f"  Error decoding RxNorm JSON response for '{drug_name}': {e}")
+        return None
+    except Exception as e: # Catch any other unexpected error
+        print(f"  Unexpected error in get_rxcui for '{drug_name}': {e}")
+        return None
+
+@lru_cache(maxsize=128) # Cache OpenFDA lookups
+def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = None) -> Optional[dict]:
+    """Fetches drug label information from OpenFDA using RxCUI or drug name."""
+    if not rxcui and not drug_name: return None
+    print(f"OpenFDA Label Lookup for: RXCUI={rxcui}, Name={drug_name}")
+
+    search_terms = []
+    # Prioritize RxCUI lookup using multiple potential fields
+    if rxcui:
+        search_terms.append(f'spl_rxnorm_code:"{rxcui}"')
+        search_terms.append(f'openfda.rxcui:"{rxcui}"')
+    # Add name search as fallback or supplement
+    if drug_name:
+        search_terms.append(f'(openfda.brand_name:"{drug_name.lower()}" OR openfda.generic_name:"{drug_name.lower()}")')
+
+    search_query = " OR ".join(search_terms)
+    params = {"search": search_query, "limit": 1} # Get only the most relevant label
+
+    try:
+        response = requests.get(OPENFDA_API_BASE, params=params, timeout=15)
+        response.raise_for_status()
+        data = response.json()
+        if data and "results" in data and data["results"]:
+            print(f"  Found OpenFDA label for query: {search_query}")
+            return data["results"][0] # Return the first label found
+        print(f"  No OpenFDA label found for query: {search_query}")
+        return None
+    except requests.exceptions.RequestException as e:
+        print(f"  Error fetching OpenFDA label: {e}")
+        return None
+    except json.JSONDecodeError as e:
+        print(f"  Error decoding OpenFDA JSON response: {e}")
+        return None
+    except Exception as e:
+        print(f"  Unexpected error in get_openfda_label: {e}")
+        return None
+
+def search_text_list(text_list: Optional[List[str]], search_terms: List[str]) -> List[str]:
+    """ Case-insensitive search for any search_term within a list of text strings. Returns snippets. """
+    found_snippets = []
+    if not text_list or not search_terms: return found_snippets
+    # Ensure search terms are lowercased strings
+    search_terms_lower = [str(term).lower() for term in search_terms if term]
+
+    for text_item in text_list:
+        if not isinstance(text_item, str): continue # Skip non-string items
+        text_item_lower = text_item.lower()
+        for term in search_terms_lower:
+            if term in text_item_lower:
+                # Find the start index of the term
+                start_index = text_item_lower.find(term)
+                # Define snippet boundaries (e.g., 50 chars before, 100 after)
+                snippet_start = max(0, start_index - 50)
+                snippet_end = min(len(text_item), start_index + len(term) + 100)
+                snippet = text_item[snippet_start:snippet_end]
+                # Add indication of where the match is
+                snippet = snippet.replace(term, f"**{term}**", 1) # Highlight first match
+                found_snippets.append(f"...{snippet}...")
+                break # Move to the next text item once a match is found
+    return found_snippets
+
+# --- Other Helper Functions ---
 def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
     """Parses BP string like '120/80' into (systolic, diastolic) integers."""
     if not isinstance(bp_string, str): return None
     match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string.strip())
-    if match:
-        return int(match.group(1)), int(match.group(2))
+    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."""
+    # (Keep the implementation from the previous full code listing)
     flags = []
     if not patient_data: return flags
-
     symptoms = patient_data.get("hpi", {}).get("symptoms", [])
     vitals = patient_data.get("vitals", {})
     history = patient_data.get("pmh", {}).get("conditions", "")
-    # Ensure symptoms are strings and lowercased
     symptoms_lower = [str(s).lower() for s in symptoms if isinstance(s, str)]
 
-    # Symptom Flags
     if "chest pain" in symptoms_lower: flags.append("Red Flag: Chest Pain reported.")
     if "shortness of breath" in symptoms_lower: flags.append("Red Flag: Shortness of Breath reported.")
     if "severe headache" in symptoms_lower: flags.append("Red Flag: Severe Headache reported.")
@@ -110,36 +224,31 @@ def check_red_flags(patient_data: dict) -> List[str]:
     if "hemoptysis" in symptoms_lower: flags.append("Red Flag: Hemoptysis (coughing up blood).")
     if "syncope" in symptoms_lower: flags.append("Red Flag: Syncope (fainting).")
 
-    # Vital Sign Flags
     if vitals:
-        temp = vitals.get("temp_c")
-        hr = vitals.get("hr_bpm")
-        rr = vitals.get("rr_rpm")
-        spo2 = vitals.get("spo2_percent")
-        bp_str = vitals.get("bp_mmhg")
-
-        if temp is not None and temp >= 38.5: flags.append(f"Red Flag: Fever (Temperature: {temp}°C).")
-        if hr is not None and hr >= 120: flags.append(f"Red Flag: Tachycardia (Heart Rate: {hr} bpm).")
-        if hr is not None and hr <= 50: flags.append(f"Red Flag: Bradycardia (Heart Rate: {hr} bpm).")
-        if rr is not None and rr >= 24: flags.append(f"Red Flag: Tachypnea (Respiratory Rate: {rr} rpm).")
-        if spo2 is not None and spo2 <= 92: flags.append(f"Red Flag: Hypoxia (SpO2: {spo2}%).")
+        temp = vitals.get("temp_c"); hr = vitals.get("hr_bpm"); rr = vitals.get("rr_rpm")
+        spo2 = vitals.get("spo2_percent"); bp_str = vitals.get("bp_mmhg")
+        if temp is not None and temp >= 38.5: flags.append(f"Red Flag: Fever ({temp}°C).")
+        if hr is not None and hr >= 120: flags.append(f"Red Flag: Tachycardia ({hr} bpm).")
+        if hr is not None and hr <= 50: flags.append(f"Red Flag: Bradycardia ({hr} bpm).")
+        if rr is not None and rr >= 24: flags.append(f"Red Flag: Tachypnea ({rr} rpm).")
+        if spo2 is not None and spo2 <= 92: flags.append(f"Red Flag: Hypoxia ({spo2}%).")
         if bp_str:
             bp = parse_bp(bp_str)
             if bp:
                 if bp[0] >= 180 or bp[1] >= 110: flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {bp_str} mmHg).")
                 if bp[0] <= 90 or bp[1] <= 60: flags.append(f"Red Flag: Hypotension (BP: {bp_str} mmHg).")
 
-    # History Flags (Simple examples)
-    if history and "history of mi" in history.lower() and "chest pain" in symptoms_lower:
-        flags.append("Red Flag: History of MI with current Chest Pain.")
-    if history and "history of dvt/pe" in history.lower() and "shortness of breath" in symptoms_lower:
-         flags.append("Red Flag: History of DVT/PE with current Shortness of Breath.")
+    if history and isinstance(history, str):
+        history_lower = history.lower()
+        if "history of mi" in history_lower and "chest pain" in symptoms_lower: flags.append("Red Flag: History of MI with current Chest Pain.")
+        if "history of dvt/pe" in history_lower and "shortness of breath" in symptoms_lower: flags.append("Red Flag: History of DVT/PE with current Shortness of Breath.")
+
+    return list(set(flags)) # Unique flags
 
-    # Remove duplicates
-    return list(set(flags))
 
 def format_patient_data_for_prompt(data: dict) -> str:
     """Formats the patient dictionary into a readable string for the LLM."""
+    # (Keep the implementation from the previous full code listing)
     if not data: return "No patient data provided."
     prompt_str = ""
     for key, value in data.items():
@@ -149,34 +258,22 @@ def format_patient_data_for_prompt(data: dict) -> str:
             if has_content:
                 prompt_str += f"**{section_title}:**\n"
                 for sub_key, sub_value in value.items():
-                     if sub_value:
-                        prompt_str += f"  - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
+                     if sub_value: prompt_str += f"  - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
         elif isinstance(value, list) and value:
              prompt_str += f"**{section_title}:** {', '.join(map(str, value))}\n"
-        elif value and not isinstance(value, dict): # Check it's not an empty dict
+        elif value and not isinstance(value, dict):
              prompt_str += f"**{section_title}:** {value}\n"
     return prompt_str.strip()
 
 
 # --- Tool Definitions ---
 
-# Pydantic models for robust argument validation
+# Pydantic models for tool inputs
 class LabOrderInput(BaseModel):
     test_name: str = Field(..., description="Specific name of the lab test or panel (e.g., 'CBC', 'BMP', 'Troponin I', 'Urinalysis', 'D-dimer').")
     reason: str = Field(..., description="Clinical justification for ordering the test (e.g., 'Rule out infection', 'Assess renal function', 'Evaluate for ACS', 'Assess for PE').")
     priority: str = Field("Routine", description="Priority of the test (e.g., 'STAT', 'Routine').")
 
-@tool("order_lab_test", args_schema=LabOrderInput)
-def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
-    """Orders a specific lab test with clinical justification and priority."""
-    print(f"Executing order_lab_test: {test_name}, Reason: {reason}, Priority: {priority}")
-    # In a real system, this would integrate with an LIS/EMR API
-    return json.dumps({
-        "status": "success",
-        "message": f"Lab Ordered: {test_name} ({priority})",
-        "details": f"Reason: {reason}"
-    })
-
 class PrescriptionInput(BaseModel):
     medication_name: str = Field(..., description="Name of the medication.")
     dosage: str = Field(..., description="Dosage amount and unit (e.g., '500 mg', '10 mg', '81 mg').")
@@ -185,105 +282,165 @@ class PrescriptionInput(BaseModel):
     duration: str = Field("As directed", description="Duration of treatment (e.g., '7 days', '1 month', 'Ongoing', 'Until follow-up').")
     reason: str = Field(..., description="Clinical indication for the prescription.")
 
-@tool("prescribe_medication", args_schema=PrescriptionInput)
-def prescribe_medication(medication_name: str, dosage: str, route: str, frequency: str, duration: str, reason: str) -> str:
-    """Prescribes a medication with detailed instructions and clinical indication. IMPORTANT: Requires prior interaction check."""
-    print(f"Executing prescribe_medication: {medication_name} {dosage}...")
-    # NOTE: The safety check (ensuring interaction check was requested) happens in the tool_node *before* this function is called.
-    # In a real system, this would trigger an e-prescription workflow.
-    return json.dumps({
-        "status": "success",
-        "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}",
-        "details": f"Duration: {duration}. Reason: {reason}"
-    })
-
+# Updated InteractionCheckInput - Note: current_medications/allergies are Optional here
+# because they are populated by the tool_node from state *before* execution.
 class InteractionCheckInput(BaseModel):
     potential_prescription: str = Field(..., description="The name of the NEW medication being considered for prescribing.")
-    # These next two args are now populated by the tool_node using AgentState
-    # current_medications: List[str] = Field(..., description="List of the patient's CURRENT medication names.")
-    # allergies: List[str] = Field(..., description="List of the patient's known allergies.")
-    # Make them optional in the schema, mandatory in the node logic
     current_medications: Optional[List[str]] = Field(None, description="List of patient's current medication names (populated from state).")
     allergies: Optional[List[str]] = Field(None, description="List of patient's known allergies (populated from state).")
 
+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 functions
+@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."""
+    print(f"Executing order_lab_test: {test_name}, Reason: {reason}, Priority: {priority}")
+    return json.dumps({"status": "success", "message": f"Lab Ordered: {test_name} ({priority})", "details": f"Reason: {reason}"})
+
+@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. IMPORTANT: Requires prior interaction check."""
+    print(f"Executing prescribe_medication: {medication_name} {dosage}...")
+    # Safety check happens in tool_node *before* this is called.
+    return json.dumps({"status": "success", "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}", "details": f"Duration: {duration}. Reason: {reason}"})
 
+# --- NEW Interaction Check Tool using UMLS/RxNorm & OpenFDA ---
 @tool("check_drug_interactions", args_schema=InteractionCheckInput)
 def check_drug_interactions(potential_prescription: str, current_medications: Optional[List[str]] = None, allergies: Optional[List[str]] = None) -> str:
-    """Checks for potential drug-drug and drug-allergy interactions BEFORE prescribing."""
-    print(f"Executing check_drug_interactions for: {potential_prescription}")
+    """
+    Checks for potential drug-drug and drug-allergy interactions using RxNorm API for normalization
+    and OpenFDA drug labels for interaction/warning text. REQUIRES UMLS_API_KEY environment variable.
+    """
+    print(f"\n--- Executing REAL check_drug_interactions ---")
+    print(f"Checking potential prescription: '{potential_prescription}'")
     warnings = []
-    potential_med_lower = potential_prescription.lower()
+    potential_med_lower = potential_prescription.lower().strip()
 
-    # Use provided lists or default to empty if None (should be populated by tool_node)
+    # Use provided lists or default to empty
     current_meds_list = current_medications or []
     allergies_list = allergies or []
-
-    current_meds_lower = [str(med).lower() for med in current_meds_list]
-    allergies_lower = [str(a).lower() for a in allergies_list]
-
-    print(f"  Checking against Meds: {current_meds_lower}")
-    print(f"  Checking against Allergies: {allergies_lower}")
-
-    # Check Allergies
+    # Clean and lowercase current med names (basic extraction: first word)
+    current_med_names_lower = []
+    for med in current_meds_list:
+         match = re.match(r"^\s*([a-zA-Z\-]+)", str(med))
+         if match: current_med_names_lower.append(match.group(1).lower())
+    # Clean and lowercase allergies
+    allergies_lower = [str(a).lower().strip() for a in allergies_list if a]
+
+    print(f"  Against Current Meds (names): {current_med_names_lower}")
+    print(f"  Against Allergies: {allergies_lower}")
+
+    # --- Step 1: Normalize potential prescription ---
+    print(f"  Step 1: Normalizing '{potential_prescription}'...")
+    potential_rxcui = get_rxcui(potential_prescription)
+    potential_label = get_openfda_label(rxcui=potential_rxcui, drug_name=potential_prescription)
+    if not potential_rxcui and not potential_label:
+         print(f"  Warning: Could not find RxCUI or OpenFDA label for '{potential_prescription}'. Interaction check will be limited.")
+         warnings.append(f"INFO: Could not reliably identify '{potential_prescription}' in standard terminologies/databases. Checks may be incomplete.")
+
+    # --- Step 2: Allergy Check ---
+    print("  Step 2: Performing Allergy Check...")
+    # Direct name match against patient's allergy list
     for allergy in allergies_lower:
-        # Direct match
         if allergy == potential_med_lower:
-            warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{allergy}'. Cannot prescribe '{potential_prescription}'.")
-            continue # Don't check cross-reactivity if direct match
-        # Check cross-reactivity
-        if allergy in ALLERGY_INTERACTIONS:
-            for cross_reactant in ALLERGY_INTERACTIONS[allergy]:
-                if cross_reactant.lower() == potential_med_lower:
-                    warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to '{allergy}'. High risk with '{potential_prescription}'.")
-
-    # Check Drug-Drug Interactions
-    for current_med in current_meds_lower:
-        # Check pairs in both orders using the mock DB
-        pair1 = (current_med, potential_med_lower)
-        pair2 = (potential_med_lower, current_med)
-        interaction_msg = MOCK_INTERACTION_DB.get(pair1) or MOCK_INTERACTION_DB.get(pair2)
-        if interaction_msg:
-            warnings.append(f"Interaction: {potential_prescription.capitalize()} with {current_med.capitalize()} - {interaction_msg}")
-
-    status = "warning" if warnings else "clear"
-    message = f"Interaction check for '{potential_prescription}': {len(warnings)} potential issue(s) found." if warnings else f"No major interactions identified for '{potential_prescription}' based on provided lists."
-    print(f"  Interaction Check Result: {status}, Message: {message}, Warnings: {warnings}")
-    return json.dumps({"status": status, "message": message, "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').")
+            warnings.append(f"CRITICAL ALLERGY (Name Match): Patient allergic to '{allergy}'. Potential prescription is '{potential_prescription}'.")
+        # Basic cross-reactivity check (can be expanded)
+        elif allergy in ["penicillin", "pcns"] and potential_med_lower in ["amoxicillin", "ampicillin", "augmentin", "piperacillin"]:
+             warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to Penicillin. High risk with '{potential_prescription}'.")
+        elif allergy == "sulfa" and potential_med_lower in ["sulfamethoxazole", "bactrim", "sulfasalazine"]:
+             warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to Sulfa. High risk with '{potential_prescription}'.")
+        elif allergy in ["nsaids", "aspirin"] and potential_med_lower in ["ibuprofen", "naproxen", "ketorolac", "diclofenac"]:
+              warnings.append(f"POTENTIAL CROSS-ALLERGY: Patient allergic to NSAIDs/Aspirin. Risk with '{potential_prescription}'.")
+
+    # Check OpenFDA Label for Contraindications/Warnings related to ALLERGIES
+    if potential_label:
+        contraindications = potential_label.get("contraindications")
+        warnings_section = potential_label.get("warnings_and_cautions") or potential_label.get("warnings")
+
+        if contraindications:
+            allergy_mentions_ci = search_text_list(contraindications, allergies_lower)
+            if allergy_mentions_ci:
+                warnings.append(f"ALLERGY RISK (Contraindication Found): Label for '{potential_prescription}' mentions contraindication potentially related to patient allergies: {'; '.join(allergy_mentions_ci)}")
+
+        if warnings_section:
+            allergy_mentions_warn = search_text_list(warnings_section, allergies_lower)
+            if allergy_mentions_warn:
+                 warnings.append(f"ALLERGY RISK (Warning Found): Label for '{potential_prescription}' mentions warnings potentially related to patient allergies: {'; '.join(allergy_mentions_warn)}")
+
+    # --- Step 3: Drug-Drug Interaction Check ---
+    print("  Step 3: Performing Drug-Drug Interaction Check...")
+    if potential_rxcui or potential_label: # Proceed only if we have info on the potential drug
+        for current_med_name in current_med_names_lower:
+            if not current_med_name or current_med_name == potential_med_lower: continue # Skip empty or self-interaction
+
+            print(f"    Checking interaction between '{potential_prescription}' and '{current_med_name}'...")
+            current_rxcui = get_rxcui(current_med_name)
+            current_label = get_openfda_label(rxcui=current_rxcui, drug_name=current_med_name)
+
+            # Terms to search for in interaction text
+            search_terms_for_current = [current_med_name]
+            if current_rxcui: search_terms_for_current.append(current_rxcui) # Add RxCUI if found
+
+            search_terms_for_potential = [potential_med_lower]
+            if potential_rxcui: search_terms_for_potential.append(potential_rxcui) # Add RxCUI if found
+
+            interaction_found_flag = False
+            # Check Potential Drug's Label ('drug_interactions' section) for mentions of Current Drug
+            if potential_label and potential_label.get("drug_interactions"):
+                interaction_mentions = search_text_list(potential_label.get("drug_interactions"), search_terms_for_current)
+                if interaction_mentions:
+                    warnings.append(f"Potential Interaction ({potential_prescription.capitalize()} Label): Mentions '{current_med_name.capitalize()}'. Snippets: {'; '.join(interaction_mentions)}")
+                    interaction_found_flag = True
+
+            # Check Current Drug's Label ('drug_interactions' section) for mentions of Potential Drug
+            if current_label and current_label.get("drug_interactions") and not interaction_found_flag: # Avoid duplicate warnings if already found
+                interaction_mentions = search_text_list(current_label.get("drug_interactions"), search_terms_for_potential)
+                if interaction_mentions:
+                     warnings.append(f"Potential Interaction ({current_med_name.capitalize()} Label): Mentions '{potential_prescription.capitalize()}'. Snippets: {'; '.join(interaction_mentions)}")
+
+    else: # Case where potential drug wasn't identified
+         warnings.append(f"INFO: Drug-drug interaction check skipped for '{potential_prescription}' as it could not be identified via RxNorm/OpenFDA.")
+
+
+    # --- Step 4: Format Output ---
+    final_warnings = list(set(warnings)) # Remove duplicates
+    status = "warning" if any("CRITICAL" in w or "Interaction" in w or "RISK" in w for w in final_warnings) else "clear"
+    if not final_warnings: status = "clear" # Ensure clear if no warnings remain
+
+    message = f"Interaction/Allergy check for '{potential_prescription}': {len(final_warnings)} potential issue(s) identified using RxNorm/OpenFDA." if final_warnings else f"No major interactions or allergy issues identified for '{potential_prescription}' based on RxNorm/OpenFDA lookup."
+    print(f"--- Interaction Check Complete for '{potential_prescription}' ---")
+
+    return json.dumps({"status": status, "message": message, "warnings": final_warnings})
+# --- End of NEW Interaction Check Tool ---
 
 @tool("flag_risk", args_schema=FlagRiskInput)
 def flag_risk(risk_description: str, urgency: str) -> str:
     """Flags a critical risk identified during analysis for immediate attention."""
     print(f"Executing flag_risk: {risk_description}, Urgency: {urgency}")
-    # Display in Streamlit immediately
     st.error(f"🚨 **{urgency.upper()} RISK FLAGGED by AI:** {risk_description}", icon="🚨")
-    return json.dumps({
-        "status": "flagged",
-        "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
-    })
+    return json.dumps({"status": "flagged", "message": f"Risk '{risk_description}' flagged with {urgency} urgency."})
 
 # Initialize Search Tool
 search_tool = TavilySearchResults(
     max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS,
-    name="tavily_search_results" # Explicitly name the tool
+    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
+    messages: Annotated[list[Any], operator.add]
+    patient_data: Optional[dict]
 
 # Define Tools and Tool Executor
 tools = [
     order_lab_test,
     prescribe_medication,
-    check_drug_interactions,
+    check_drug_interactions, # Using the new implementation
     flag_risk,
     search_tool
 ]
@@ -293,239 +450,132 @@ tool_executor = ToolExecutor(tools)
 model = ChatGroq(
     temperature=ClinicalAppSettings.TEMPERATURE,
     model=ClinicalAppSettings.MODEL_NAME,
-    # Increase max_tokens if needed for large JSON output + conversation history
-    # max_tokens=4096
 )
-# Bind tools FOR the model to know their schemas and descriptions
 model_with_tools = model.bind_tools(tools)
 
-# --- Graph Nodes ---
+# --- Graph Nodes (agent_node, tool_node remain mostly the same structurally) ---
 
-# 1. Agent Node: Calls the LLM
+# 1. Agent Node: Calls the LLM (No change needed from previous version)
 def agent_node(state: AgentState):
     """Invokes the LLM to decide the next action or response."""
     print("\n---AGENT NODE---")
     current_messages = state['messages']
-    # Ensure System Prompt is present
     if not current_messages or not isinstance(current_messages[0], SystemMessage):
         print("Prepending System Prompt.")
         current_messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + current_messages
-
-    # Optional: Augment first human message with patient data if not already done explicitly
-    # This helps ensure the LLM sees it early, though it's also in the state.
-    # Be mindful of context window limits.
-    # if len(current_messages) > 1 and isinstance(current_messages[1], HumanMessage) and state.get('patient_data'):
-    #     if "**Initial Patient Data:**" not in current_messages[1].content:
-    #         print("Augmenting first HumanMessage with patient data summary.")
-    #         formatted_data = format_patient_data_for_prompt(state['patient_data'])
-    #         current_messages[1] = HumanMessage(content=f"{current_messages[1].content}\n\n**Initial Patient Data Summary:**\n{formatted_data}")
-
     print(f"Invoking LLM with {len(current_messages)} messages.")
-    # print(f"Messages Sent: {[m.type for m in current_messages]}") # Log message types
     try:
         response = model_with_tools.invoke(current_messages)
         print(f"Agent Raw Response Type: {type(response)}")
-        # print(f"Agent Raw Response Content: {response.content}")
-        if hasattr(response, 'tool_calls') and response.tool_calls:
-             print(f"Agent Response Tool Calls: {response.tool_calls}")
-        else:
-             print("Agent Response: No tool calls.")
-
+        if hasattr(response, 'tool_calls') and response.tool_calls: print(f"Agent Response Tool Calls: {response.tool_calls}")
+        else: print("Agent Response: No tool calls.")
     except Exception as e:
         print(f"ERROR in agent_node during LLM invocation: {type(e).__name__} - {e}")
-        traceback.print_exc() # Print full traceback for debugging
-        # Return an error message to the graph state
+        traceback.print_exc()
         error_message = AIMessage(content=f"Sorry, an internal error occurred while processing the request: {type(e).__name__}")
         return {"messages": [error_message]}
-
     return {"messages": [response]}
 
-# 2. Tool Node: Executes tools called by the Agent (REVISED WITH ROBUST ERROR HANDLING)
+# 2. Tool Node: Executes tools (Mostly the same, ensures context injection)
 def tool_node(state: AgentState):
     """Executes tools called by the LLM and returns results."""
     print("\n---TOOL NODE---")
-    tool_messages = [] # Initialize list to store results or errors
+    tool_messages = []
     last_message = state['messages'][-1]
 
-    # Ensure the last message is an AIMessage with tool calls
     if not isinstance(last_message, AIMessage) or not getattr(last_message, 'tool_calls', None):
-        print("Warning: Tool node called unexpectedly without tool calls in the last AI message.")
-        # If this happens, it might indicate a routing issue or the LLM hallucinating flow.
-        # Returning empty list lets the agent proceed, potentially without needed info.
-        # Consider adding a ToolMessage indicating the issue if needed.
+        print("Warning: Tool node called unexpectedly without tool calls.")
         return {"messages": []}
 
     tool_calls = last_message.tool_calls
-    print(f"Tool calls received: {json.dumps(tool_calls, indent=2)}") # Log received calls
-
-    # Safety Check: Identify required interaction checks before prescriptions
-    prescriptions_requested = {} # medication_name_lower -> tool_call
-    interaction_checks_requested = {} # medication_name_lower -> tool_call
+    print(f"Tool calls received: {json.dumps(tool_calls, indent=2)}")
 
+    # Safety Check Logic (No change needed from previous version)
+    prescriptions_requested = {}
+    interaction_checks_requested = {}
     for call in tool_calls:
-        tool_name = call.get('name')
-        tool_args = call.get('args', {})
+        tool_name = call.get('name'); tool_args = call.get('args', {})
         if tool_name == 'prescribe_medication':
-            med_name = tool_args.get('medication_name', '').lower()
-            if med_name:
-                prescriptions_requested[med_name] = call
+            med_name = tool_args.get('medication_name', '').lower();
+            if med_name: prescriptions_requested[med_name] = call
         elif tool_name == 'check_drug_interactions':
             potential_med = tool_args.get('potential_prescription', '').lower()
-            if potential_med:
-                interaction_checks_requested[potential_med] = call
+            if potential_med: interaction_checks_requested[potential_med] = call
 
     valid_tool_calls_for_execution = []
-
-    # Validate prescriptions against interaction checks
+    blocked_ids = set()
     for med_name, prescribe_call in prescriptions_requested.items():
         if med_name not in interaction_checks_requested:
             st.error(f"**Safety Violation:** AI attempted to prescribe '{med_name}' without requesting `check_drug_interactions` in the *same turn*. Prescription blocked.")
-            error_msg = ToolMessage(
-                content=json.dumps({"status": "error", "message": f"Interaction check for '{med_name}' must be requested *before or alongside* the prescription call."}),
-                tool_call_id=prescribe_call['id'],
-                name=prescribe_call['name'] # Include tool name in ToolMessage
-            )
+            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'], name=prescribe_call['name'])
             tool_messages.append(error_msg)
-        else:
-             # Interaction check IS requested, allow prescription call to proceed
-             pass # The call will be added below if it's in the original tool_calls list
+            blocked_ids.add(prescribe_call['id'])
 
-    # Prepare list of calls to execute (all non-blocked calls)
-    blocked_ids = {msg.tool_call_id for msg in tool_messages if msg.content and '"status": "error"' in msg.content}
     valid_tool_calls_for_execution = [call for call in tool_calls if call['id'] not in blocked_ids]
 
-    # Augment interaction checks with patient data from state
-    patient_meds = state.get("patient_data", {}).get("medications", {}).get("names_only", [])
-    patient_allergies = state.get("patient_data", {}).get("allergies", [])
+    # Augment interaction checks with patient data (Crucial part - no change needed here)
+    patient_data = state.get("patient_data", {})
+    patient_meds_full = patient_data.get("medications", {}).get("current", []) # Pass full med list if needed by tool
+    patient_allergies = patient_data.get("allergies", [])
 
     for call in valid_tool_calls_for_execution:
          if call['name'] == 'check_drug_interactions':
-             # Ensure args exist before modifying
              if 'args' not in call: call['args'] = {}
-             call['args']['current_medications'] = patient_meds
+             # Pass the necessary context from patient_data to the tool arguments
+             # The tool function expects 'current_medications' (list of names) and 'allergies'
+             call['args']['current_medications'] = patient_meds_full # Pass the full strings
              call['args']['allergies'] = patient_allergies
-             print(f"Augmented interaction check args for call ID {call['id']}: {call['args']}")
-
+             print(f"Augmented interaction check args for call ID {call['id']}") # Removed args content for brevity
 
-    # Execute valid tool calls using batch for efficiency, capturing exceptions
+    # Execute valid tool calls (No change needed from previous version)
     if valid_tool_calls_for_execution:
         print(f"Attempting to execute {len(valid_tool_calls_for_execution)} tools: {[c['name'] for c in valid_tool_calls_for_execution]}")
         try:
             responses = tool_executor.batch(valid_tool_calls_for_execution, return_exceptions=True)
-
-            # Process responses, creating ToolMessage for each
             for call, resp in zip(valid_tool_calls_for_execution, responses):
-                tool_call_id = call['id']
-                tool_name = call['name']
-
+                tool_call_id = call['id']; tool_name = call['name']
                 if isinstance(resp, Exception):
-                    # Handle exceptions returned by the batch call
-                    error_type = type(resp).__name__
-                    error_str = str(resp)
+                    error_type = type(resp).__name__; error_str = str(resp)
                     print(f"ERROR executing tool '{tool_name}' (ID: {tool_call_id}): {error_type} - {error_str}")
-                    traceback.print_exc() # Log full traceback
+                    traceback.print_exc()
                     st.error(f"Error executing action '{tool_name}': {error_type}")
-                    error_content = json.dumps({
-                        "status": "error",
-                        "message": f"Failed to execute '{tool_name}': {error_type} - {error_str}"
-                    })
+                    error_content = json.dumps({"status": "error", "message": f"Failed to execute '{tool_name}': {error_type} - {error_str}"})
                     tool_messages.append(ToolMessage(content=error_content, tool_call_id=tool_call_id, name=tool_name))
-
-                    # Specific check for the error mentioned by user
                     if isinstance(resp, AttributeError) and "'dict' object has no attribute 'tool'" in error_str:
-                         print("\n *** DETECTED SPECIFIC ATTRIBUTE ERROR ('dict' object has no attribute 'tool') ***")
-                         print(f"   Tool Call causing error: {json.dumps(call, indent=2)}")
-                         print("   This likely indicates an internal issue within Langchain/LangGraph or ToolExecutor expecting a different object structure.")
-                         print("   Ensure tool definitions (@tool decorators) and Pydantic schemas are correct.\n")
-
+                         print("\n *** DETECTED SPECIFIC ATTRIBUTE ERROR ('dict' object has no attribute 'tool') *** \n")
                 else:
-                    # Process successful results
-                    print(f"Tool '{tool_name}' (ID: {tool_call_id}) executed successfully. Result type: {type(resp)}")
-                    # Ensure content is string for ToolMessage
+                    print(f"Tool '{tool_name}' (ID: {tool_call_id}) executed successfully.")
                     content_str = str(resp)
                     tool_messages.append(ToolMessage(content=content_str, tool_call_id=tool_call_id, name=tool_name))
-
-                    # Display result in Streamlit right away for feedback (optional, but helpful)
-                    # This part might be better handled purely in the UI display loop later
-                    # try:
-                    #     result_data = json.loads(content_str)
-                    #     status = result_data.get("status", "info")
-                    #     message = result_data.get("message", content_str)
-                    #     if status in ["success", "clear", "flagged"]: st.success(f"Action `{tool_name}` completed: {message}", icon="✅" if status != "flagged" else "🚨")
-                    #     elif status == "warning": st.warning(f"Action `{tool_name}` completed: {message}", icon="⚠️")
-                    #     else: st.info(f"Action `{tool_name}` completed: {message}") # Info for other statuses
-                    # except json.JSONDecodeError:
-                    #     st.info(f"Action `{tool_name}` completed (non-JSON output).")
-
-        # Catch potential errors within the tool_node logic itself (e.g., preparing calls)
         except Exception as e:
             print(f"CRITICAL UNEXPECTED ERROR within tool_node logic: {type(e).__name__} - {e}")
-            traceback.print_exc()
-            st.error(f"Critical internal error processing actions: {e}")
-            # Create generic error messages for all calls that were intended
+            traceback.print_exc(); st.error(f"Critical internal error processing actions: {e}")
             error_content = json.dumps({"status": "error", "message": f"Internal error processing tools: {e}"})
-            # Add error messages for calls that didn't get processed yet
             processed_ids = {msg.tool_call_id for msg in tool_messages}
             for call in valid_tool_calls_for_execution:
-                 if call['id'] not in processed_ids:
-                     tool_messages.append(ToolMessage(content=error_content, tool_call_id=call['id'], name=call['name']))
+                 if call['id'] not in processed_ids: tool_messages.append(ToolMessage(content=error_content, tool_call_id=call['id'], name=call['name']))
 
     print(f"Returning {len(tool_messages)} tool messages.")
-    # print(f"Tool messages content snippets: {[m.content[:100] + '...' if len(m.content)>100 else m.content for m in tool_messages]}")
     return {"messages": tool_messages}
 
 
-# --- Graph Edges (Routing Logic) ---
+# --- Graph Edges (Routing Logic) --- (No change needed)
 def should_continue(state: AgentState) -> str:
     """Determines whether to call tools, end the conversation turn, or handle errors."""
     print("\n---ROUTING DECISION---")
     last_message = state['messages'][-1] if state['messages'] else None
+    if not isinstance(last_message, AIMessage): return "end_conversation_turn"
+    if "Sorry, an internal error occurred" in last_message.content: return "end_conversation_turn"
+    if getattr(last_message, 'tool_calls', None): return "continue_tools"
+    else: return "end_conversation_turn"
 
-    if not isinstance(last_message, AIMessage):
-        # This case might happen if the graph starts with a non-AI message or after an error
-        print("Routing: Last message not AI. Ending turn.")
-        return "end_conversation_turn"
-
-    # If the LLM produced an error message (e.g., during invocation)
-    if "Sorry, an internal error occurred" in last_message.content:
-         print("Routing: AI returned internal error. Ending turn.")
-         return "end_conversation_turn"
-
-    # If the LLM made tool calls, execute them
-    if getattr(last_message, 'tool_calls', None):
-        print("Routing: AI requested tool calls. Continue to tools node.")
-        return "continue_tools"
-    # Otherwise, the AI provided a response without tool calls, end the turn
-    else:
-        print("Routing: AI provided final response or asked question. Ending turn.")
-        return "end_conversation_turn"
-
-# --- Graph Definition ---
+# --- Graph Definition & Compilation --- (No change needed)
 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 for this turn
-    }
-)
-
-# Add edge from tools back to agent
+workflow.add_conditional_edges("agent", should_continue, {"continue_tools": "tools", "end_conversation_turn": END})
 workflow.add_edge("tools", "agent")
-
-# Compile the graph
-# memory = MemorySaverInMemory() # Optional: for persisting state across runs
-# app = workflow.compile(checkpointer=memory)
 app = workflow.compile()
 print("LangGraph compiled successfully.")
 
@@ -533,41 +583,34 @@ print("LangGraph compiled successfully.")
 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})")
+    st.caption(f"Interactive Assistant | LangGraph/Groq/Tavily/UMLS/OpenFDA | Model: {ClinicalAppSettings.MODEL_NAME}")
 
-    # Initialize session state
-    if "messages" not in st.session_state:
-        st.session_state.messages = [] # Stores full conversation history
-    if "patient_data" not in st.session_state:
-        st.session_state.patient_data = None
-    if "graph_app" not in st.session_state:
-        st.session_state.graph_app = app
+    # Initialize session state (No change needed)
+    if "messages" not in st.session_state: st.session_state.messages = []
+    if "patient_data" not in st.session_state: st.session_state.patient_data = None
+    if "graph_app" not in st.session_state: st.session_state.graph_app = app
 
-    # --- Patient Data Input Sidebar ---
+    # --- Patient Data Input Sidebar --- (Adjusted allergy/med extraction slightly)
     with st.sidebar:
         st.header("📄 Patient Intake Form")
-        # Demographics
+        # Demographics, HPI, History, Social/Family, Vitals/Exam sections remain the same input fields
+        # ... (Copy input fields from previous full code version) ...
         st.subheader("Demographics")
         age = st.number_input("Age", min_value=0, max_value=120, value=55, key="age_input")
         sex = st.selectbox("Biological Sex", ["Male", "Female", "Other/Prefer not to say"], key="sex_input")
-        # HPI
         st.subheader("History of Present Illness (HPI)")
         chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="cc_input")
-        hpi_details = st.text_area("Detailed HPI", "55 y/o male presents with substernal chest pain started 2 hours ago, described as pressure, radiating to left arm. Associated with nausea and diaphoresis. Pain is 8/10 severity. No relief with rest.", key="hpi_input", height=150)
+        hpi_details = st.text_area("Detailed HPI", "55 y/o male presents with substernal chest pain started 2 hours ago...", key="hpi_input", height=150)
         symptoms = st.multiselect("Associated Symptoms", ["Nausea", "Diaphoresis", "Shortness of Breath", "Dizziness", "Palpitations", "Fever", "Cough", "Severe Headache", "Syncope", "Hemoptysis"], default=["Nausea", "Diaphoresis"], key="sym_input")
-        # History
         st.subheader("Past History")
         pmh = st.text_area("Past Medical History (PMH)", "Hypertension (HTN), Hyperlipidemia (HLD), Type 2 Diabetes Mellitus (DM2), History of MI", key="pmh_input")
         psh = st.text_area("Past Surgical History (PSH)", "Appendectomy (2005)", key="psh_input")
-        # Meds & Allergies
         st.subheader("Medications & Allergies")
         current_meds_str = st.text_area("Current Medications (name, dose, freq)", "Lisinopril 10mg daily\nMetformin 1000mg BID\nAtorvastatin 40mg daily\nAspirin 81mg daily", key="meds_input")
         allergies_str = st.text_area("Allergies (comma separated, specify reaction if known)", "Penicillin (rash), Sulfa (hives)", key="allergy_input")
-        # Social/Family
         st.subheader("Social/Family History")
         social_history = st.text_area("Social History (SH)", "Smoker (1 ppd x 30 years), occasional alcohol.", key="sh_input")
         family_history = st.text_area("Family History (FHx)", "Father had MI at age 60. Mother has HTN.", key="fhx_input")
-        # Vitals/Exam
         st.subheader("Vitals & Exam Findings")
         col1, col2 = st.columns(2)
         with col1:
@@ -578,241 +621,161 @@ def main():
             bp_mmhg = st.text_input("BP (SYS/DIA)", "155/90", key="bp_input")
             spo2_percent = st.number_input("SpO2 (%)", 70, 100, 96, key="spo2_input")
             pain_scale = st.slider("Pain (0-10)", 0, 10, 8, key="pain_input")
-        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3. Mild distress. Lungs clear bilaterally. Cardiac exam: Regular rhythm, S1/S2 normal, no murmurs/gallops/rubs. Abdomen soft, non-tender. No lower extremity edema.", key="exam_input", height=100)
+        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3. Mild distress. Lungs clear bilaterally...", key="exam_input", height=100)
 
-        # Compile Patient Data Dictionary on button press
+
+        # Compile Patient Data Dictionary (Refined Extraction for Tool Use)
         if st.button("Start/Update Consultation", key="start_button"):
+            # Store full medication strings for display/context
             current_meds_list = [med.strip() for med in current_meds_str.split('\n') if med.strip()]
-            # Basic name extraction (first word, lowercased) for interaction check
-            current_med_names = []
+            # Extract just the names (simplified) for the interaction check tool's state population
+            current_med_names_only = []
             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())
+                if match: current_med_names_only.append(match.group(1).lower())
 
-            # Basic allergy extraction (first word or phrase before parenthesis, lowercased)
+            # Extract allergy names (simplified, before parenthesis)
             allergies_list = []
             for a in allergies_str.split(','):
                 cleaned_allergy = a.strip()
                 if cleaned_allergy:
-                     match = re.match(r"^\s*([a-zA-Z\-\s]+)(?:\s*\(.*\))?", cleaned_allergy)
-                     if match:
-                         allergies_list.append(match.group(1).strip().lower())
-                     else: # Fallback if no parenthesis
-                          allergies_list.append(cleaned_allergy.lower())
+                     match = re.match(r"^\s*([a-zA-Z\-\s/]+)(?:\s*\(.*\))?", cleaned_allergy)
+                     name_part = match.group(1).strip().lower() if match else cleaned_allergy.lower()
+                     allergies_list.append(name_part)
 
             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,
+                # Store both full list and names_only list
+                "medications": {"current": current_meds_list, "names_only": current_med_names_only},
+                "allergies": allergies_list, # Store cleaned 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)
+            # Initial Red Flag Check
             red_flags = check_red_flags(st.session_state.patient_data)
             st.sidebar.markdown("---")
             if red_flags:
                 st.sidebar.warning("**Initial Red Flags Detected:**")
                 for flag in red_flags: st.sidebar.warning(f"- {flag.replace('Red Flag: ','')}")
-            else:
-                 st.sidebar.success("No immediate red flags detected in initial data.")
+            else: st.sidebar.success("No immediate red flags detected.")
 
-            # Prepare initial message for the graph
-            initial_prompt = f"Initiate consultation for the patient described in the intake form. Start the analysis."
-            # Clear previous messages and start fresh
+            # Prepare initial message & reset history
+            initial_prompt = "Initiate consultation for the patient described in the intake form. Review the data and begin analysis."
             st.session_state.messages = [HumanMessage(content=initial_prompt)]
-            st.success("Patient data loaded. Ready for analysis.")
-            # No rerun needed here, chat input will trigger the graph
+            st.success("Patient data loaded/updated. Ready for analysis.")
+
 
-    # --- Main Chat Interface Area ---
+    # --- Main Chat Interface Area --- (No change needed in display logic)
     st.header("💬 Clinical Consultation")
 
     # Display chat messages from history
+    # (Copy the message display loop from the previous full code version)
     for msg_index, msg in enumerate(st.session_state.messages):
-        unique_key = f"msg_{msg_index}" # Basic unique key
+        unique_key = f"msg_{msg_index}"
         if isinstance(msg, HumanMessage):
-            with st.chat_message("user", key=f"{unique_key}_user"):
-                st.markdown(msg.content)
+            with st.chat_message("user", key=f"{unique_key}_user"): st.markdown(msg.content)
         elif isinstance(msg, AIMessage):
             with st.chat_message("assistant", key=f"{unique_key}_ai"):
-                # Display AI text content
-                ai_content = msg.content
-                structured_output = None
-
-                # Attempt to parse structured JSON if present
+                ai_content = msg.content; structured_output = None
                 try:
-                    # Look for ```json ... ``` block
                     json_match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
                     if json_match:
-                        json_str = json_match.group(1)
-                        # Display content before/after the JSON block if any
-                        prefix = ai_content[:json_match.start()].strip()
-                        suffix = ai_content[json_match.end():].strip()
+                        json_str = json_match.group(1); prefix = ai_content[:json_match.start()].strip(); suffix = ai_content[json_match.end():].strip()
                         if prefix: st.markdown(prefix)
                         structured_output = json.loads(json_str)
                         if suffix: st.markdown(suffix)
-                    # Check if the entire message might be JSON
                     elif ai_content.strip().startswith("{") and ai_content.strip().endswith("}"):
-                         structured_output = json.loads(ai_content)
-                         ai_content = "" # Don't display raw JSON if parsed ok
-                    else:
-                         # No JSON found, display content as is
-                         st.markdown(ai_content)
-
-                except json.JSONDecodeError:
-                    # Failed to parse, display raw content
-                    st.markdown(ai_content)
-                    st.warning("Note: Could not parse structured JSON in AI response.", icon="⚠️")
-                except Exception as e:
-                     st.markdown(ai_content) # Display raw on other errors
-                     st.error(f"Error processing AI message display: {e}", icon="❌")
-
-                # Display structured data nicely if parsed
+                         structured_output = json.loads(ai_content); ai_content = ""
+                    else: st.markdown(ai_content)
+                except Exception as e: st.markdown(ai_content); print(f"Error parsing/displaying AI JSON: {e}")
+
                 if structured_output and isinstance(structured_output, dict):
-                    st.divider()
-                    st.subheader("📊 AI Analysis & Recommendations")
+                    # (Copy the structured JSON display logic from previous full code)
+                    st.divider(); st.subheader("📊 AI Analysis & Recommendations")
                     cols = st.columns(2)
                     with cols[0]:
-                        st.markdown(f"**Assessment:**")
-                        st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
-
+                        st.markdown(f"**Assessment:**"); st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
                         st.markdown(f"**Differential Diagnosis:**")
-                        ddx = structured_output.get('differential_diagnosis', [])
+                        ddx = structured_output.get('differential_diagnosis', []);
                         if ddx:
                             for item in ddx:
-                                likelihood = item.get('likelihood', 'Unknown').capitalize()
-                                icon = "🥇" if likelihood=="High" else ("🥈" if likelihood=="Medium" else "🥉")
-                                with st.expander(f"{icon} {item.get('diagnosis', 'Unknown')} ({likelihood})"):
-                                    st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
-                        else: st.info("No differential diagnosis provided.")
-
-                        st.markdown(f"**Risk Assessment:**")
-                        risk = structured_output.get('risk_assessment', {})
-                        flags = risk.get('identified_red_flags', [])
+                                likelihood = item.get('likelihood', '?').capitalize(); icon = "🥇" if likelihood=="High" else ("🥈" if likelihood=="Medium" else "🥉")
+                                with st.expander(f"{icon} {item.get('diagnosis', 'Unknown')} ({likelihood})"): st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
+                        else: st.info("No DDx provided.")
+                        st.markdown(f"**Risk Assessment:**"); risk = structured_output.get('risk_assessment', {})
+                        flags = risk.get('identified_red_flags', []); concerns = risk.get("immediate_concerns", []); comps = risk.get("potential_complications", [])
                         if flags: st.warning(f"**Flags:** {', '.join(flags)}")
-                        if risk.get("immediate_concerns"): st.warning(f"**Concerns:** {', '.join(risk.get('immediate_concerns'))}")
-                        if risk.get("potential_complications"): st.info(f"**Potential Complications:** {', '.join(risk.get('potential_complications'))}")
-                        if not flags and not risk.get("immediate_concerns"): st.success("No major risks highlighted in this assessment.")
-
+                        if concerns: st.warning(f"**Concerns:** {', '.join(concerns)}")
+                        if comps: st.info(f"**Potential Complications:** {', '.join(comps)}")
+                        if not flags and not concerns: st.success("No major risks highlighted.")
                     with cols[1]:
-                        st.markdown(f"**Recommended Plan:**")
-                        plan = structured_output.get('recommended_plan', {})
-                        sub_sections = ["investigations", "therapeutics", "consultations", "patient_education"]
-                        for section in sub_sections:
-                             st.markdown(f"_{section.replace('_',' ').capitalize()}:_")
-                             items = plan.get(section)
-                             if items and isinstance(items, list):
-                                 for item in items: st.markdown(f"- {item}")
-                             elif items: # Handle if it's just a string
-                                  st.markdown(f"- {items}")
+                        st.markdown(f"**Recommended Plan:**"); plan = structured_output.get('recommended_plan', {})
+                        for section in ["investigations", "therapeutics", "consultations", "patient_education"]:
+                             st.markdown(f"_{section.replace('_',' ').capitalize()}:_"); items = plan.get(section)
+                             if items: [st.markdown(f"- {item}") for item in items] if isinstance(items, list) else st.markdown(f"- {items}")
                              else: st.markdown("_None suggested._")
-                             st.markdown("") # Add space
-
-                    # Display Rationale and Interaction Summary below columns
-                    st.markdown(f"**Rationale & Guideline Check:**")
-                    st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}")
+                             st.markdown("") # Space
+                    st.markdown(f"**Rationale & Guideline Check:**"); st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}")
                     interaction_summary = structured_output.get("interaction_check_summary", "")
-                    if interaction_summary:
-                         st.markdown(f"**Interaction Check Summary:**")
-                         st.markdown(f"> {interaction_summary}")
-
+                    if interaction_summary: st.markdown(f"**Interaction Check Summary:**"); st.markdown(f"> {interaction_summary}")
                     st.divider()
 
-                # Display tool calls requested in this AI turn
                 if getattr(msg, 'tool_calls', None):
                      with st.expander("🛠️ AI requested actions", expanded=False):
                          for tc in msg.tool_calls:
-                             try:
-                                 # Safely display, default args to empty dict if missing
-                                 st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nArgs: {json.dumps(tc.get('args', {}), indent=2)}", language="json")
-                             except Exception as display_e:
-                                  st.error(f"Could not display tool call: {display_e}")
-                                  st.code(str(tc)) # Raw display as fallback
+                             try: st.code(f"Action: {tc.get('name', 'Unknown')}\nArgs: {json.dumps(tc.get('args', {}), indent=2)}", language="json")
+                             except Exception as display_e: st.error(f"Could not display tool call: {display_e}"); st.code(str(tc))
 
         elif isinstance(msg, ToolMessage):
-            # Safely get tool name
-            tool_name_display = getattr(msg, 'name', 'tool_execution') # Use 'name' attribute added in tool_node
+            tool_name_display = getattr(msg, 'name', 'tool_execution')
             with st.chat_message(tool_name_display, avatar="🛠️", key=f"{unique_key}_tool"):
+                # (Copy the ToolMessage display logic from previous full code)
                 try:
-                    # Attempt to parse content as JSON for structured display
-                    tool_data = json.loads(msg.content)
-                    status = tool_data.get("status", "info")
-                    message = tool_data.get("message", msg.content)
-                    details = tool_data.get("details")
-                    warnings = tool_data.get("warnings")
-
-                    if status == "success" or status == "clear" or status == "flagged":
-                        st.success(f"{message}", icon="✅" if status != "flagged" else "🚨")
+                    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"{message}", icon="✅" if status != "flagged" else "🚨")
                     elif status == "warning":
                          st.warning(f"{message}", icon="⚠️")
                          if warnings and isinstance(warnings, list):
                              st.caption("Details:")
-                             for warn in warnings: st.caption(f"- {warn}")
-                    else: # Error or unknown status
-                         st.error(f"{message}", icon="❌")
-
+                             for warn in warnings: st.caption(f"- {warn}") # Display warnings from the tool output JSON
+                    else: st.error(f"{message}", icon="❌")
                     if details: st.caption(f"Details: {details}")
+                except json.JSONDecodeError: st.info(f"{msg.content}") # Display raw if not JSON
+                except Exception as e: st.error(f"Error displaying tool message: {e}", icon="❌"); st.caption(f"Raw content: {msg.content}")
 
-                except json.JSONDecodeError:
-                    # If content is not JSON, display it plainly
-                    st.info(f"{msg.content}")
-                except Exception as e:
-                     st.error(f"Error displaying tool message: {e}", icon="❌")
-                     st.caption(f"Raw content: {msg.content}")
 
-    # --- Chat Input Logic ---
+    # --- Chat Input Logic --- (No change needed)
     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.")
-            st.stop() # Prevent execution if no patient data
+            st.warning("Please load patient data using the sidebar first."); st.stop()
 
-        # Add user message to state and display it immediately
         user_message = HumanMessage(content=prompt)
         st.session_state.messages.append(user_message)
-        with st.chat_message("user"):
-            st.markdown(prompt)
+        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
-        )
+        current_state = AgentState(messages=st.session_state.messages, patient_data=st.session_state.patient_data)
 
-        # Invoke the graph
         with st.spinner("SynapseAI is thinking..."):
             try:
-                # Use invoke to run the graph until it ends for this turn
-                final_state = st.session_state.graph_app.invoke(
-                     current_state,
-                     {"recursion_limit": 15} # Add recursion limit for safety
-                )
-                # Update the session state messages with the final list from the graph run
+                final_state = st.session_state.graph_app.invoke(current_state, {"recursion_limit": 15})
                 st.session_state.messages = final_state['messages']
-
             except Exception as e:
-                print(f"CRITICAL ERROR during graph invocation: {type(e).__name__} - {e}")
-                traceback.print_exc()
+                print(f"CRITICAL ERROR during graph invocation: {type(e).__name__} - {e}"); traceback.print_exc()
                 st.error(f"An error occurred during the conversation turn: {e}", icon="❌")
-                # Attempt to add an error message to the history for visibility
-                error_ai_msg = AIMessage(content=f"Sorry, a critical error occurred: {type(e).__name__}. Please check logs or try again.")
-                # Avoid modifying state directly during exception handling if possible,
-                # but appending might be okay for display purposes.
-                # st.session_state.messages.append(error_ai_msg) # Be cautious with state modification here
+                # Optionally add error to history for user visibility
+                # error_ai_msg = AIMessage(content=f"Sorry, a critical error occurred: {type(e).__name__}. Please check logs or try again.")
+                # st.session_state.messages.append(error_ai_msg)
 
-        # Rerun the script to display the updated chat history, including AI response and tool results
-        st.rerun()
+        st.rerun() # Refresh display
 
-    # Disclaimer at the bottom
+    # Disclaimer (No change needed)
     st.markdown("---")
-    st.warning(
-        """**Disclaimer:** SynapseAI is an AI assistant for clinical decision support and does not replace professional medical judgment.
-        All outputs must be critically reviewed and verified by a qualified healthcare provider before making any clinical decisions.
-        Validate all information, especially diagnoses, dosages, and interactions, independently using standard medical resources."""
-    )
+    st.warning("""**Disclaimer:** SynapseAI is an AI assistant... (Verify all outputs)""")
 
 if __name__ == "__main__":
     main()