Update app.py

app.py

@@ -11,7 +11,7 @@ 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
-from langchain_core.prompts import ChatPromptTemplate
+# from langchain_core.prompts import ChatPromptTemplate # Not explicitly used in this version
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain_core.tools import tool
 from langgraph.prebuilt import ToolExecutor
@@ -33,9 +33,12 @@ 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:
+    # Use st.error which stops execution in recent Streamlit versions
     st.error(f"Missing required API Key(s): {', '.join(missing_keys)}. Please set them in Hugging Face Space Secrets or your environment variables.")
+    # Ensure execution stops if st.error doesn't automatically do it in the environment
     st.stop()
 
+
 # --- Configuration & Constants ---
 class ClinicalAppSettings:
     APP_TITLE = "SynapseAI: Interactive Clinical Decision Support (UMLS/FDA Integrated)"
@@ -103,99 +106,57 @@ def get_rxcui(drug_name: str) -> Optional[str]:
         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()
+        else: # Fallback search
+            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
+                            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
+    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: 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
-
+    if rxcui: search_terms.append(f'spl_rxnorm_code:"{rxcui}" OR openfda.rxcui:"{rxcui}"')
+    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}
     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
+        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]
+        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
+        if not isinstance(text_item, str): continue
         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
+                start_index = text_item_lower.find(term); 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]
+                snippet = snippet.replace(term, f"**{term}**", 1); found_snippets.append(f"...{snippet}...")
+                break
     return found_snippets
 
 # --- Other Helper Functions ---
@@ -208,7 +169,6 @@ def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
 
 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", [])
@@ -243,12 +203,10 @@ def check_red_flags(patient_data: dict) -> List[str]:
         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
-
+    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():
@@ -268,7 +226,7 @@ def format_patient_data_for_prompt(data: dict) -> str:
 
 # --- Tool Definitions ---
 
-# Pydantic models for tool inputs
+# Pydantic models
 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').")
@@ -282,8 +240,6 @@ 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.")
 
-# 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.")
     current_medications: Optional[List[str]] = Field(None, description="List of patient's current medication names (populated from state).")
@@ -304,10 +260,8 @@ def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> st
 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:
     """
@@ -318,103 +272,57 @@ def check_drug_interactions(potential_prescription: str, current_medications: Op
     print(f"Checking potential prescription: '{potential_prescription}'")
     warnings = []
     potential_med_lower = potential_prescription.lower().strip()
-
-    # Use provided lists or default to empty
-    current_meds_list = current_medications or []
-    allergies_list = allergies or []
-    # Clean and lowercase current med names (basic extraction: first word)
+    current_meds_list = current_medications or []; allergies_list = allergies or []
     current_med_names_lower = []
     for med in current_meds_list:
-         match = re.match(r"^\s*([a-zA-Z\-]+)", str(med))
+         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}")
 
-    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)
+    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.")
+    if not potential_rxcui and not potential_label: warnings.append(f"INFO: Could not reliably identify '{potential_prescription}'. Checks may be incomplete.")
 
-    # --- Step 2: Allergy Check ---
-    print("  Step 2: Performing Allergy Check...")
-    # Direct name match against patient's allergy list
+    print("  Step 2: Performing Allergy Check...");
     for allergy in allergies_lower:
-        if allergy == potential_med_lower:
-            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 allergy == potential_med_lower: warnings.append(f"CRITICAL ALLERGY (Name Match): Patient allergic to '{allergy}'. Potential prescription is '{potential_prescription}'.")
+        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}'.")
     if potential_label:
-        contraindications = potential_label.get("contraindications")
-        warnings_section = potential_label.get("warnings_and_cautions") or potential_label.get("warnings")
-
+        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 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)}")
+            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
+    if potential_rxcui or potential_label:
         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
-
+            if not current_med_name or current_med_name == potential_med_lower: continue
             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
-
+            current_rxcui = get_rxcui(current_med_name); current_label = get_openfda_label(rxcui=current_rxcui, drug_name=current_med_name)
+            search_terms_for_current = [current_med_name];
+            if current_rxcui: search_terms_for_current.append(current_rxcui)
+            search_terms_for_potential = [potential_med_lower];
+            if potential_rxcui: search_terms_for_potential.append(potential_rxcui)
             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
+                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
+            if current_label and current_label.get("drug_interactions") and not interaction_found_flag:
                 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
+                if interaction_mentions: warnings.append(f"Potential Interaction ({current_med_name.capitalize()} Label): Mentions '{potential_prescription.capitalize()}'. Snippets: {'; '.join(interaction_mentions)}")
+    else: warnings.append(f"INFO: Drug-drug interaction check skipped for '{potential_prescription}' as it could not be identified via RxNorm/OpenFDA.")
 
+    final_warnings = list(set(warnings)); 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"
     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:
@@ -424,45 +332,22 @@ def flag_risk(risk_description: str, urgency: str) -> str:
     return json.dumps({"status": "flagged", "message": f"Risk '{risk_description}' flagged with {urgency} urgency."})
 
 # Initialize Search Tool
-search_tool = TavilySearchResults(
-    max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS,
-    name="tavily_search_results"
-    )
+search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS, name="tavily_search_results")
 
 # --- LangGraph Setup ---
-
-# Define the state structure
 class AgentState(TypedDict):
-    messages: Annotated[list[Any], operator.add]
-    patient_data: Optional[dict]
-
-# Define Tools and Tool Executor
-tools = [
-    order_lab_test,
-    prescribe_medication,
-    check_drug_interactions, # Using the new implementation
-    flag_risk,
-    search_tool
-]
+    messages: Annotated[list[Any], operator.add]; patient_data: Optional[dict]
+tools = [order_lab_test, prescribe_medication, check_drug_interactions, flag_risk, search_tool]
 tool_executor = ToolExecutor(tools)
-
-# Define the Agent Model
-model = ChatGroq(
-    temperature=ClinicalAppSettings.TEMPERATURE,
-    model=ClinicalAppSettings.MODEL_NAME,
-)
+model = ChatGroq(temperature=ClinicalAppSettings.TEMPERATURE, model=ClinicalAppSettings.MODEL_NAME)
 model_with_tools = model.bind_tools(tools)
 
-# --- Graph Nodes (agent_node, tool_node remain mostly the same structurally) ---
-
-# 1. Agent Node: Calls the LLM (No change needed from previous version)
+# --- Graph Nodes ---
 def agent_node(state: AgentState):
-    """Invokes the LLM to decide the next action or response."""
     print("\n---AGENT NODE---")
     current_messages = state['messages']
     if not current_messages or not isinstance(current_messages[0], SystemMessage):
-        print("Prepending System Prompt.")
-        current_messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + current_messages
+        print("Prepending System Prompt."); current_messages = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + current_messages
     print(f"Invoking LLM with {len(current_messages)} messages.")
     try:
         response = model_with_tools.invoke(current_messages)
@@ -470,64 +355,36 @@ def agent_node(state: AgentState):
         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(f"ERROR in agent_node during LLM invocation: {type(e).__name__} - {e}"); 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 (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 = []
-    last_message = state['messages'][-1]
-
+    tool_messages = []; last_message = state['messages'][-1]
     if not isinstance(last_message, AIMessage) or not getattr(last_message, 'tool_calls', None):
-        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)}")
-
-    # Safety Check Logic (No change needed from previous version)
-    prescriptions_requested = {}
-    interaction_checks_requested = {}
+        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)}")
+    prescriptions_requested = {}; interaction_checks_requested = {}
     for call in tool_calls:
         tool_name = call.get('name'); tool_args = call.get('args', {})
-        if tool_name == 'prescribe_medication':
-            med_name = tool_args.get('medication_name', '').lower();
-            if med_name: prescriptions_requested[med_name] = call
-        elif tool_name == 'check_drug_interactions':
-            potential_med = tool_args.get('potential_prescription', '').lower()
-            if potential_med: interaction_checks_requested[potential_med] = call
-
-    valid_tool_calls_for_execution = []
-    blocked_ids = set()
+        if tool_name == 'prescribe_medication': med_name = tool_args.get('medication_name', '').lower();
+        if med_name: prescriptions_requested[med_name] = call
+        elif tool_name == 'check_drug_interactions': potential_med = tool_args.get('potential_prescription', '').lower()
+        if potential_med: interaction_checks_requested[potential_med] = call
+    valid_tool_calls_for_execution = []; 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'])
-            tool_messages.append(error_msg)
-            blocked_ids.add(prescribe_call['id'])
-
+            tool_messages.append(error_msg); blocked_ids.add(prescribe_call['id'])
     valid_tool_calls_for_execution = [call for call in tool_calls if call['id'] not in blocked_ids]
-
-    # 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", [])
-
+    patient_data = state.get("patient_data", {}); patient_meds_full = patient_data.get("medications", {}).get("current", []); patient_allergies = patient_data.get("allergies", [])
     for call in valid_tool_calls_for_execution:
          if call['name'] == 'check_drug_interactions':
              if 'args' not in call: call['args'] = {}
-             # 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']}") # Removed args content for brevity
-
-    # Execute valid tool calls (No change needed from previous version)
+             call['args']['current_medications'] = patient_meds_full; call['args']['allergies'] = patient_allergies; print(f"Augmented interaction check args for call ID {call['id']}")
     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:
@@ -535,151 +392,79 @@ def tool_node(state: AgentState):
             for call, resp in zip(valid_tool_calls_for_execution, responses):
                 tool_call_id = call['id']; tool_name = call['name']
                 if isinstance(resp, Exception):
-                    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()
-                    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_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()
+                    st.error(f"Error executing action '{tool_name}': {error_type}"); error_content = json.dumps({"status": "error", "message": f"Failed to execute '{tool_name}': {error_type} - {error_str}"})
                     tool_messages.append(ToolMessage(content=error_content, tool_call_id=tool_call_id, name=tool_name))
-                    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') *** \n")
+                    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') *** \n")
                 else:
-                    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))
+                    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))
         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}")
-            error_content = json.dumps({"status": "error", "message": f"Internal error processing tools: {e}"})
-            processed_ids = {msg.tool_call_id for msg in tool_messages}
+            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}")
+            error_content = json.dumps({"status": "error", "message": f"Internal error processing tools: {e}"}); processed_ids = {msg.tool_call_id for msg in tool_messages}
             for call in valid_tool_calls_for_execution:
                  if call['id'] not in processed_ids: tool_messages.append(ToolMessage(content=error_content, tool_call_id=call['id'], name=call['name']))
+    print(f"Returning {len(tool_messages)} tool messages."); return {"messages": tool_messages}
 
-    print(f"Returning {len(tool_messages)} tool messages.")
-    return {"messages": tool_messages}
-
-
-# --- Graph Edges (Routing Logic) --- (No change needed)
+# --- Graph Edges (Routing Logic) ---
 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
+    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"
 
-# --- Graph Definition & Compilation --- (No change needed)
-workflow = StateGraph(AgentState)
-workflow.add_node("agent", agent_node)
-workflow.add_node("tools", tool_node)
-workflow.set_entry_point("agent")
-workflow.add_conditional_edges("agent", should_continue, {"continue_tools": "tools", "end_conversation_turn": END})
-workflow.add_edge("tools", "agent")
-app = workflow.compile()
-print("LangGraph compiled successfully.")
+# --- Graph Definition & Compilation ---
+workflow = StateGraph(AgentState); workflow.add_node("agent", agent_node); workflow.add_node("tools", tool_node)
+workflow.set_entry_point("agent"); workflow.add_conditional_edges("agent", should_continue, {"continue_tools": "tools", "end_conversation_turn": END})
+workflow.add_edge("tools", "agent"); app = workflow.compile(); print("LangGraph compiled successfully.")
 
 # --- Streamlit UI ---
 def main():
     st.set_page_config(page_title=ClinicalAppSettings.APP_TITLE, layout=ClinicalAppSettings.PAGE_LAYOUT)
     st.title(f"🩺 {ClinicalAppSettings.APP_TITLE}")
     st.caption(f"Interactive Assistant | LangGraph/Groq/Tavily/UMLS/OpenFDA | Model: {ClinicalAppSettings.MODEL_NAME}")
-
-    # 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 --- (Adjusted allergy/med extraction slightly)
+    # --- Patient Data Input Sidebar ---
     with st.sidebar:
         st.header("📄 Patient Intake Form")
-        # 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")
-        st.subheader("History of Present Illness (HPI)")
-        chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="cc_input")
-        hpi_details = st.text_area("Detailed HPI", "55 y/o male presents with substernal chest pain started 2 hours ago...", key="hpi_input", 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")
-        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")
-        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")
-        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")
-        st.subheader("Vitals & Exam Findings")
-        col1, col2 = st.columns(2)
-        with col1:
-            temp_c = st.number_input("Temp (°C)", 35.0, 42.0, 36.8, format="%.1f", key="temp_input")
-            hr_bpm = st.number_input("HR (bpm)", 30, 250, 95, key="hr_input")
-            rr_rpm = st.number_input("RR (rpm)", 5, 50, 18, key="rr_input")
-        with col2:
-            bp_mmhg = st.text_input("BP (SYS/DIA)", "155/90", key="bp_input")
-            spo2_percent = st.number_input("SpO2 (%)", 70, 100, 96, key="spo2_input")
-            pain_scale = st.slider("Pain (0-10)", 0, 10, 8, key="pain_input")
-        exam_notes = st.text_area("Brief Physical Exam Notes", "Awake, alert, oriented x3. Mild distress. Lungs clear bilaterally...", key="exam_input", height=100)
-
-
-        # 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
+        # Input fields (Demographics, HPI, History, Meds/Allergies, Social/Family, Vitals/Exam)
+        st.subheader("Demographics"); age = st.number_input("Age", 0, 120, 55); sex = st.selectbox("Sex", ["Male", "Female", "Other"])
+        st.subheader("HPI"); chief_complaint = st.text_input("Chief Complaint", "Chest pain"); hpi_details = st.text_area("HPI Details", "55 y/o male...", height=150); symptoms = st.multiselect("Symptoms", ["Nausea", "Diaphoresis", "SOB", "Dizziness"], default=["Nausea", "Diaphoresis"])
+        st.subheader("History"); pmh = st.text_area("PMH", "HTN, HLD, DM2, MI"); psh = st.text_area("PSH", "Appendectomy")
+        st.subheader("Meds & Allergies"); current_meds_str = st.text_area("Current Meds", "Lisinopril 10mg daily\nMetformin 1000mg BID\nAtorvastatin 40mg daily"); allergies_str = st.text_area("Allergies", "Penicillin (rash), Sulfa")
+        st.subheader("Social/Family"); social_history = st.text_area("SH", "Smoker"); family_history = st.text_area("FHx", "Father MI")
+        st.subheader("Vitals & Exam"); col1, col2 = st.columns(2);
+        with col1: temp_c = st.number_input("Temp C", 35.0, 42.0, 36.8, format="%.1f"); hr_bpm = st.number_input("HR", 30, 250, 95); rr_rpm = st.number_input("RR", 5, 50, 18)
+        with col2: bp_mmhg = st.text_input("BP", "155/90"); spo2_percent = st.number_input("SpO2", 70, 100, 96); pain_scale = st.slider("Pain", 0, 10, 8)
+        exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=100)
+
+        if st.button("Start/Update Consultation"):
             current_meds_list = [med.strip() for med in current_meds_str.split('\n') if med.strip()]
-            # 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_only.append(match.group(1).lower())
-
-            # Extract allergy names (simplified, before parenthesis)
+            current_med_names_only = [];
+            for med in current_meds_list: match = re.match(r"^\s*([a-zA-Z\-]+)", med);
+            if match: current_med_names_only.append(match.group(1).lower())
             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)
-                     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},
-                # 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
-            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.")
-
-            # 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/updated. Ready for analysis.")
-
-
-    # --- Main Chat Interface Area --- (No change needed in display logic)
+                cleaned_allergy = a.strip();
+                if cleaned_allergy: 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_only}, "allergies": allergies_list, "social_history": {"details": social_history}, "family_history": {"details": family_history}, "vitals": { "temp_c": temp_c, "hr_bpm": hr_bpm, "bp_mmhg": bp_mmhg, "rr_rpm": rr_rpm, "spo2_percent": spo2_percent, "pain_scale": pain_scale}, "exam_findings": {"notes": exam_notes} }
+            red_flags = check_red_flags(st.session_state.patient_data); st.sidebar.markdown("---")
+            if red_flags: st.sidebar.warning("**Initial Red Flags:**"); [st.sidebar.warning(f"- {flag.replace('Red Flag: ','')}") for flag in red_flags]
+            else: st.sidebar.success("No immediate red flags.")
+            initial_prompt = "Initiate consultation for the patient described in the intake form. Review data and begin analysis."
+            st.session_state.messages = [HumanMessage(content=initial_prompt)]; st.success("Patient data loaded/updated.")
+
+    # --- Main Chat Interface Area ---
     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}"
+    # Display loop - REMOVED key= ARGUMENT
+    for msg in st.session_state.messages:
         if isinstance(msg, HumanMessage):
-            with st.chat_message("user", key=f"{unique_key}_user"): st.markdown(msg.content)
+            with st.chat_message("user"): st.markdown(msg.content) # No key
         elif isinstance(msg, AIMessage):
-            with st.chat_message("assistant", key=f"{unique_key}_ai"):
+            with st.chat_message("assistant"): # No key
                 ai_content = msg.content; structured_output = None
                 try:
                     json_match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
@@ -692,90 +477,58 @@ def main():
                          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):
-                    # (Copy the structured JSON display logic from previous full code)
-                    st.divider(); st.subheader("📊 AI Analysis & Recommendations")
+                    st.divider(); st.subheader("📊 AI Analysis & Recommendations") # Display logic for JSON...
                     cols = st.columns(2)
-                    with cols[0]:
-                        st.markdown(f"**Assessment:**"); st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
-                        st.markdown(f"**Differential Diagnosis:**")
-                        ddx = structured_output.get('differential_diagnosis', []);
-                        if ddx:
-                            for item in ddx:
-                                likelihood = item.get('likelihood', '?').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')}")
+                    with cols[0]: # Assessment, DDx, Risk
+                        st.markdown("**Assessment:**"); st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
+                        st.markdown("**Differential Diagnosis:**"); ddx = structured_output.get('differential_diagnosis', []);
+                        if ddx: [st.expander(f"{'🥇🥈🥉'[('High','Medium','Low').index(item.get('likelihood','Low')[0])] if item.get('likelihood','?')[0] in 'HML' else '?'} {item.get('diagnosis', 'Unknown')} ({item.get('likelihood','?')})").write(f"**Rationale:** {item.get('rationale', 'N/A')}") for item in ddx]
                         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", [])
+                        st.markdown("**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 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', {})
-                        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("") # Space
-                    st.markdown(f"**Rationale & Guideline Check:**"); st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}")
+                    with cols[1]: # Plan
+                         st.markdown("**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); [st.markdown(f"- {item}") for item in items] if items and isinstance(items, list) else (st.markdown(f"- {items}") if items else st.markdown("_None_")); st.markdown("")
+                    st.markdown("**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("**Interaction Check Summary:**"); st.markdown(f"> {interaction_summary}")
                     st.divider()
-
                 if getattr(msg, 'tool_calls', None):
-                     with st.expander("🛠️ AI requested actions", expanded=False):
-                         for tc in msg.tool_calls:
-                             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))
-
+                     with st.expander("🛠️ AI requested actions", expanded=False): # Tool call display logic...
+                        for tc in msg.tool_calls: 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):
             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:
-                    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")
+            with st.chat_message(tool_name_display, avatar="🛠️"): # No key
+                try: # Tool message display logic...
+                    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}") # Display warnings from the tool output JSON
+                    elif status == "warning": st.warning(f"{message}", icon="⚠️");
+                    if warnings and isinstance(warnings, list): st.caption("Details:"); [st.caption(f"- {warn}") for warn in warnings]
                     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 json.JSONDecodeError: 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 --- (No change needed)
+    # --- Chat Input Logic ---
     if prompt := st.chat_input("Your message or follow-up query..."):
-        if not st.session_state.patient_data:
-            st.warning("Please load patient data using the sidebar first."); st.stop()
-
-        user_message = HumanMessage(content=prompt)
-        st.session_state.messages.append(user_message)
-        with st.chat_message("user"): st.markdown(prompt)
-
+        if not st.session_state.patient_data: st.warning("Please load patient data first."); st.stop()
+        user_message = HumanMessage(content=prompt); st.session_state.messages.append(user_message)
+        with st.chat_message("user"): st.markdown(prompt) # Display user msg immediately
         current_state = AgentState(messages=st.session_state.messages, patient_data=st.session_state.patient_data)
-
         with st.spinner("SynapseAI is thinking..."):
             try:
                 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()
-                st.error(f"An error occurred during the conversation turn: {e}", icon="❌")
-                # 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)
-
+                st.session_state.messages = final_state['messages'] # Update state with results
+            except Exception as e: print(f"CRITICAL ERROR: {e}"); traceback.print_exc(); st.error(f"Error: {e}")
         st.rerun() # Refresh display
 
-    # Disclaimer (No change needed)
-    st.markdown("---")
-    st.warning("""**Disclaimer:** SynapseAI is an AI assistant... (Verify all outputs)""")
+    # Disclaimer
+    st.markdown("---"); st.warning("**Disclaimer:** SynapseAI is for demonstration...")
 
 if __name__ == "__main__":
     main()