Update app.py

app.py

@@ -1,30 +1,26 @@
-# -*- coding: utf-8 -*-
+# app.py
 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
-# from langchain_core.prompts import ChatPromptTemplate # Not explicitly used
-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 typing import Optional, List, Dict, Any, TypedDict, Annotated
+# Import agent logic and message types from agent.py
+try:
+    from agent import ClinicalAgent, AgentState, check_red_flags
+    from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+except ImportError as e:
+    st.error(f"Failed to import from agent.py: {e}. Make sure agent.py is in the same directory.")
+    st.stop()
 
 # --- Environment Variable Loading & Validation ---
 load_dotenv()
+# Check keys required by agent.py are present before initializing the agent
 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")
+
 missing_keys = []
 if not UMLS_API_KEY:
     missing_keys.append("UMLS_API_KEY")
@@ -32,547 +28,97 @@ 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 API Key(s): {', '.join(missing_keys)}.")
+    st.error(f"Missing required API Key(s): {', '.join(missing_keys)}. Please set them in Hugging Face Space Secrets or environment variables.")
     st.stop()
 
-# --- Configuration & Constants ---
+# --- App Configuration ---
 class ClinicalAppSettings:
     APP_TITLE = "SynapseAI (UMLS/FDA Integrated)"
     PAGE_LAYOUT = "wide"
-    MODEL_NAME = "llama3-70b-8192"
-    TEMPERATURE = 0.1
-    MAX_SEARCH_RESULTS = 3
-
-class ClinicalPrompts:
-    SYSTEM_PROMPT = """
-    You are SynapseAI, an expert AI clinical assistant engaged in an interactive consultation... [SYSTEM PROMPT REMAINS THE SAME - OMITTED FOR BREVITY]
-    """
-
-# --- API Helper Functions (get_rxcui, get_openfda_label, search_text_list) ---
-UMLS_AUTH_ENDPOINT = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-RXNORM_API_BASE = "https://rxnav.nlm.nih.gov/REST"
-OPENFDA_API_BASE = "https://api.fda.gov/drug/label.json"
-
-@lru_cache(maxsize=256)
-def get_rxcui(drug_name: str) -> Optional[str]:
-    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}
-        response = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10)
-        response.raise_for_status()
-        data = response.json()
-        if data and "idGroup" in data and "rxnormId" in data["idGroup"]:
-            rxcui = data["idGroup"]["rxnormId"][0]
-            print(f"  Found RxCUI: {rxcui} for '{drug_name}'")
-            return rxcui
-        else:
-            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"]:
-                    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:
-        print(f"  Unexpected error in get_rxcui for '{drug_name}': {e}")
-        return None
-
-@lru_cache(maxsize=128)
-def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = None) -> Optional[dict]:
-    if not rxcui and not drug_name:
-        return None
-    print(f"OpenFDA Label Lookup for: RXCUI={rxcui}, Name={drug_name}")
-    search_terms = []
-    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]
-        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]:
-    found_snippets = []
-    if not text_list or not search_terms:
-        return found_snippets
-    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
-        text_item_lower = text_item.lower()
-        for term in search_terms_lower:
-            if term in text_item_lower:
-                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 (parse_bp, check_red_flags, format_patient_data_for_prompt) ---
-def parse_bp(bp_string: str) -> Optional[tuple[int, int]]:
-    if not isinstance(bp_string, str):
-        return None
-    match = re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", bp_string.strip())
-    if match:
-        return int(match.group(1)), int(match.group(2))
-    return None
-
-def check_red_flags(patient_data: dict) -> List[str]:
-    flags = []
-    if not patient_data:
-        return flags
-    symptoms = patient_data.get("hpi", {}).get("symptoms", [])
-    vitals = patient_data.get("vitals", {})
-    history = patient_data.get("pmh", {}).get("conditions", "")
-    symptoms_lower = [str(s).lower() for s in symptoms if isinstance(s, str)]
-    
-    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.")
-    if "sudden vision loss" in symptoms_lower:
-        flags.append("Red Flag: Sudden Vision Loss reported.")
-    if "weakness on one side" in symptoms_lower:
-        flags.append("Red Flag: Unilateral Weakness reported (potential stroke).")
-    if "hemoptysis" in symptoms_lower:
-        flags.append("Red Flag: Hemoptysis (coughing up blood).")
-    if "syncope" in symptoms_lower:
-        flags.append("Red Flag: Syncope (fainting).")
-    
-    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 ({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).")
-    
-    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))
-
-def format_patient_data_for_prompt(data: dict) -> str:
-    if not data:
-        return "No patient data provided."
-    prompt_str = ""
-    for key, value in data.items():
-        section_title = key.replace('_', ' ').title()
-        if isinstance(value, dict) and value:
-            has_content = any(sub_value for sub_value in value.values())
-            if has_content:
-                prompt_str += f"**{section_title}:**\n"
-                for sub_key, sub_value in value.items():
-                    if sub_value:
-                        prompt_str += f"  - {sub_key.replace('_', ' ').title()}: {sub_value}\n"
-        elif isinstance(value, list) and value:
-            prompt_str += f"**{section_title}:** {', '.join(map(str, value))}\n"
-        elif value and not isinstance(value, dict):
-            prompt_str += f"**{section_title}:** {value}\n"
-    return prompt_str.strip()
-
-# --- Tool Definitions ---
-class LabOrderInput(BaseModel):
-    test_name: str = Field(...)
-    reason: str = Field(...)
-    priority: str = Field("Routine")
-
-class PrescriptionInput(BaseModel):
-    medication_name: str = Field(...)
-    dosage: str = Field(...)
-    route: str = Field(...)
-    frequency: str = Field(...)
-    duration: str = Field("As directed")
-    reason: str = Field(...)
-
-class InteractionCheckInput(BaseModel):
-    potential_prescription: str = Field(...)
-    current_medications: Optional[List[str]] = Field(None)
-    allergies: Optional[List[str]] = Field(None)
-
-class FlagRiskInput(BaseModel):
-    risk_description: str = Field(...)
-    urgency: str = Field("High")
-
-@tool("order_lab_test", args_schema=LabOrderInput)
-def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
-    """
-    Orders a lab test with the specified test name, reason, 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:
-    """
-    Prepares a prescription for the specified medication including dosage, route, frequency, and duration.
-    """
-    print(f"Executing prescribe_medication: {medication_name} {dosage}...")
-    return json.dumps({
-        "status": "success",
-        "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}",
-        "details": f"Duration: {duration}. Reason: {reason}"
-    })
-
-@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 interactions and allergy risks for the given prescription.
-    """
-    print(f"\n--- Executing REAL check_drug_interactions ---")
-    print(f"Checking potential prescription: '{potential_prescription}'")
-    warnings = []
-    potential_med_lower = potential_prescription.lower().strip()
-    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))
-        if match:
-            current_med_names_lower.append(match.group(1).lower())
-    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"  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:
-        warnings.append(f"INFO: Could not reliably identify '{potential_prescription}'. Checks may be incomplete.")
-    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}'.")
-        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")
-        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)}")
-    print("  Step 3: Performing Drug-Drug Interaction Check...")
-    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
-            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)
-            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
-            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
-            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:
-        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
-    })
-
-@tool("flag_risk", args_schema=FlagRiskInput)
-def flag_risk(risk_description: str, urgency: str) -> str:
-    """
-    Flags a clinical risk with the provided description and urgency.
-    """
-    print(f"Executing flag_risk: {risk_description}, Urgency: {urgency}")
-    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."
-    })
-
-search_tool = TavilySearchResults(max_results=ClinicalAppSettings.MAX_SEARCH_RESULTS, name="tavily_search_results")
-
-# --- LangGraph Setup ---
-class AgentState(TypedDict):
-    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)
-model = ChatGroq(temperature=ClinicalAppSettings.TEMPERATURE, model=ClinicalAppSettings.MODEL_NAME)
-model_with_tools = model.bind_tools(tools)
-
-# --- Graph Nodes (agent_node, tool_node) ---
-def agent_node(state: AgentState):
-    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(f"Invoking LLM with {len(current_messages)} messages.")
-    try:
-        response = model_with_tools.invoke(current_messages)
-        print(f"Agent Raw Response Type: {type(response)}")
-        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: {e}")
-        traceback.print_exc()
-        error_message = AIMessage(content=f"Error: {e}")
-        return {"messages": [error_message]}
-    return {"messages": [response]}
-
-def tool_node(state: AgentState):
-    print("\n---TOOL NODE---")
-    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.")
-        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()
-    for med_name, prescribe_call in prescriptions_requested.items():
-        if med_name not in interaction_checks_requested:
-            st.error(f"**Safety Violation:** AI tried to prescribe '{med_name}' without check.")
-            error_msg = ToolMessage(
-                content=json.dumps({
-                    "status": "error",
-                    "message": f"Interaction check needed for '{med_name}'."
-                }),
-                tool_call_id=prescribe_call['id'],
-                name=prescribe_call['name']
-            )
-            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]
-    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'] = {}
-            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 execution: {[c['name'] for c in valid_tool_calls_for_execution]}")
-    try:
-        responses = tool_executor.batch(valid_tool_calls_for_execution, return_exceptions=True)
-        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}': {error_type} - {error_str}")
-                traceback.print_exc()
-                st.error(f"Error: {error_type}")
-                error_content = json.dumps({
-                    "status": "error",
-                    "message": f"Failed: {error_type} - {error_str}"
-                })
-                tool_messages.append(ToolMessage(content=error_content, tool_call_id=tool_call_id, name=tool_name))
-            else:
-                print(f"Tool '{tool_name}' executed.")
-                # If response is a dict, convert it to a JSON string.
-                if isinstance(resp, dict):
-                    content_str = json.dumps(resp)
-                else:
-                    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 TOOL NODE ERROR: {e}")
-        traceback.print_exc()
-        st.error(f"Critical error: {e}")
-        error_content = json.dumps({
-            "status": "error",
-            "message": f"Internal error: {e}"
-        })
-        processed_ids = {msg.tool_call_id for msg in tool_messages}
-        [tool_messages.append(ToolMessage(content=error_content, tool_call_id=call['id'], name=call['name']))
-         for call in valid_tool_calls_for_execution if call['id'] not in processed_ids]
-    print(f"Returning {len(tool_messages)} tool messages.")
-    return {"messages": tool_messages}
-
-# --- Graph Edges (Routing Logic) ---
-def should_continue(state: AgentState) -> str:
-    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 ---
-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.")
+    MODEL_NAME_DISPLAY = "Llama3-70b (via Groq)"  # Defined in agent.py
 
 # --- 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}")
+    st.caption(f"Interactive Assistant | LangGraph/Groq/Tavily/UMLS/OpenFDA | Model: {ClinicalAppSettings.MODEL_NAME_DISPLAY}")
+
+    # Initialize session state
     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
+    if "summary" not in st.session_state:
+        st.session_state.summary = None
+
+    # Initialize the agent instance only once
+    if "agent" not in st.session_state:
+        try:
+            st.session_state.agent = ClinicalAgent()
+            print("ClinicalAgent successfully initialized in Streamlit session state.")
+        except Exception as e:
+            st.error(f"Failed to initialize Clinical Agent: {e}. Check API keys and dependencies.")
+            print(f"ERROR Initializing ClinicalAgent: {e}")
+            traceback.print_exc()
+            st.stop()
 
     # --- Patient Data Input Sidebar ---
     with st.sidebar:
         st.header("📄 Patient Intake Form")
         # Input fields... (Using shorter versions for brevity, assume full fields are here)
         st.subheader("Demographics")
-        age = st.number_input("Age", 0, 120, 55)
-        sex = st.selectbox("Sex", ["Male", "Female", "Other"])
+        age = st.number_input("Age", 0, 120, 55, key="sb_age")
+        sex = st.selectbox("Sex", ["Male", "Female", "Other"], key="sb_sex")
+
         st.subheader("HPI")
-        chief_complaint = st.text_input("Chief Complaint", "Chest pain")
-        hpi_details = st.text_area("HPI Details", "55 y/o male...", height=100)
-        symptoms = st.multiselect("Symptoms", ["Nausea", "Diaphoresis", "SOB", "Dizziness"], default=["Nausea", "Diaphoresis"])
+        chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="sb_cc")
+        hpi_details = st.text_area("HPI Details", "55 y/o male...", height=100, key="sb_hpi")
+        symptoms = st.multiselect(
+            "Symptoms",
+            ["Nausea", "Diaphoresis", "SOB", "Dizziness", "Severe Headache", "Syncope", "Hemoptysis"],
+            default=["Nausea", "Diaphoresis"],
+            key="sb_sym"
+        )
+
         st.subheader("History")
-        pmh = st.text_area("PMH", "HTN, HLD, DM2, History of MI")
-        psh = st.text_area("PSH", "Appendectomy")
+        pmh = st.text_area("PMH", "HTN, HLD, DM2, History of MI", key="sb_pmh")
+        psh = st.text_area("PSH", "Appendectomy", key="sb_psh")
+
         st.subheader("Meds & Allergies")
-        current_meds_str = st.text_area("Current Meds", "Lisinopril 10mg daily\nMetformin 1000mg BID")
-        allergies_str = st.text_area("Allergies", "Penicillin (rash)")
+        current_meds_str = st.text_area(
+            "Current Meds",
+            "Lisinopril 10mg daily\nMetformin 1000mg BID\nWarfarin 5mg daily",
+            key="sb_meds"
+        )
+        allergies_str = st.text_area("Allergies", "Penicillin (rash), Aspirin", key="sb_allergies")  # Added Warfarin/Aspirin for testing
+
         st.subheader("Social/Family")
-        social_history = st.text_area("SH", "Smoker")
-        family_history = st.text_area("FHx", "Father MI")
+        social_history = st.text_area("SH", "Smoker", key="sb_sh")
+        family_history = st.text_area("FHx", "Father MI", key="sb_fhx")
+
         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)
+            temp_c = st.number_input("Temp C", 35.0, 42.0, 36.8, format="%.1f", key="sb_temp")
+            hr_bpm = st.number_input("HR", 30, 250, 95, key="sb_hr")
+            rr_rpm = st.number_input("RR", 5, 50, 18, key="sb_rr")
         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)
-        # Updated height to at least 68 pixels as required by Streamlit
-        exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=68)
+            bp_mmhg = st.text_input("BP", "155/90", key="sb_bp")
+            spo2_percent = st.number_input("SpO2", 70, 100, 96, key="sb_spo2")
+            pain_scale = st.slider("Pain", 0, 10, 8, key="sb_pain")
+        exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=50, key="sb_exam")
 
-        if st.button("Start/Update Consultation"):
+        if st.button("Start/Update Consultation", key="sb_start"):
+            # Compile data...
             current_meds_list = [med.strip() for med in current_meds_str.split('\n') if med.strip()]
             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()
@@ -580,6 +126,8 @@ def main():
                     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)
+
+            # Update patient data in session state
             st.session_state.patient_data = {
                 "demographics": {"age": age, "sex": sex},
                 "hpi": {"chief_complaint": chief_complaint, "details": hpi_details, "symptoms": symptoms},
@@ -599,16 +147,24 @@ def main():
                 },
                 "exam_findings": {"notes": exam_notes}
             }
+
+            # Call check_red_flags from agent module
             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]
+                for flag in red_flags:
+                    st.sidebar.warning(f"- {flag.replace('Red Flag: ', '')}")
             else:
                 st.sidebar.success("No immediate red flags.")
+
+            # Reset conversation and summary on new intake
             initial_prompt = "Initiate consultation. Review patient data and begin analysis."
             st.session_state.messages = [HumanMessage(content=initial_prompt)]
+            st.session_state.summary = None  # Reset summary
             st.success("Patient data loaded/updated.")
+            # Rerun might be needed if the main area should clear or update based on new data
+            st.rerun()
 
     # --- Main Chat Interface Area ---
     st.header("💬 Clinical Consultation")
@@ -616,12 +172,13 @@ def main():
     for msg in st.session_state.messages:
         if isinstance(msg, HumanMessage):
             with st.chat_message("user"):
-                st.markdown(msg.content)  # No key
+                st.markdown(msg.content)
         elif isinstance(msg, AIMessage):
             with st.chat_message("assistant"):
                 ai_content = msg.content
                 structured_output = None
                 try:
+                    # JSON Parsing logic...
                     json_match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
                     if json_match:
                         json_str = json_match.group(1)
@@ -636,11 +193,13 @@ def main():
                         structured_output = json.loads(ai_content)
                         ai_content = ""
                     else:
-                        st.markdown(ai_content)
+                        st.markdown(ai_content)  # Display non-JSON 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):
+                    # Structured JSON display logic...
                     st.divider()
                     st.subheader("📊 AI Analysis & Recommendations")
                     cols = st.columns(2)
@@ -650,7 +209,15 @@ def main():
                         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]
+                            for item in ddx:
+                                likelihood = item.get('likelihood', 'Low')
+                                if likelihood and likelihood[0] in 'HML':
+                                    medal = "🥇" if likelihood[0] == 'H' else "🥈" if likelihood[0] == 'M' else "🥉"
+                                else:
+                                    medal = "?"
+                                expander_title = f"{medal} {item.get('diagnosis', 'Unknown')} ({likelihood})"
+                                with st.expander(expander_title):
+                                    st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
                         else:
                             st.info("No DDx provided.")
                         st.markdown("**Risk Assessment:**")
@@ -669,11 +236,12 @@ def main():
                     with cols[1]:
                         st.markdown("**Recommended Plan:**")
                         plan = structured_output.get('recommended_plan', {})
-                        for section in ["investigations","therapeutics","consultations","patient_education"]:
-                            st.markdown(f"_{section.replace('_',' ').capitalize()}:_")
+                        for section in ["investigations", "therapeutics", "consultations", "patient_education"]:
+                            st.markdown(f"_{section.replace('_', ' ').capitalize()}:_")
                             items = plan.get(section)
                             if items and isinstance(items, list):
-                                [st.markdown(f"- {item}") for item in items]
+                                for item in items:
+                                    st.markdown(f"- {item}")
                             elif items:
                                 st.markdown(f"- {items}")
                             else:
@@ -687,35 +255,44 @@ def main():
                             st.markdown(f"> {interaction_summary}")
                         st.divider()
 
+                # Tool Call Display
                 if getattr(msg, 'tool_calls', None):
                     with st.expander("🛠️ AI requested actions", expanded=False):
                         if msg.tool_calls:
                             for tc in msg.tool_calls:
                                 try:
-                                    st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nArgs: {json.dumps(tc.get('args', {}), indent=2)}", language="json")
+                                    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 arguments properly: {display_e}", icon="⚠️")
+                                    st.error(f"Could not display tool call args: {display_e}", icon="⚠️")
                                     st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nRaw Args: {tc.get('args')}")
                         else:
-                            st.caption("_No actions requested in this turn._")
+                            st.caption("_No actions requested._")
         elif isinstance(msg, ToolMessage):
             tool_name_display = getattr(msg, 'name', 'tool_execution')
             with st.chat_message(tool_name_display, avatar="🛠️"):
                 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 in ["success", "clear", "flagged"]:
-                        st.success(f"{message}", icon="✅" if status != "flagged" else "🚨")
+                    # Display flagged risks immediately if the tool signals it
+                    if tool_name_display == "flag_risk" and status == "flagged":
+                        st.error(f"🚨 **RISK FLAGGED:** {message}", icon="🚨")
+                    elif status in ["success", "clear"]:
+                        st.success(f"{message}", icon="✅")
                     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 warnings and isinstance(warnings, list):
+                        st.caption("Details:")
+                        for warn in warnings:
+                            st.caption(f"- {warn}")
                     if details:
                         st.caption(f"Details: {details}")
                 except json.JSONDecodeError:
@@ -729,19 +306,42 @@ def main():
         if not st.session_state.patient_data:
             st.warning("Please load patient data first.")
             st.stop()
+        if 'agent' not in st.session_state or not st.session_state.agent:
+            st.error("Agent not initialized. Check logs.")
+            st.stop()
+
+        # Append user message and display immediately
         user_message = HumanMessage(content=prompt)
         st.session_state.messages.append(user_message)
         with st.chat_message("user"):
             st.markdown(prompt)
-        current_state = AgentState(messages=st.session_state.messages, patient_data=st.session_state.patient_data)
-        with st.spinner("SynapseAI is thinking..."):
+
+        # Prepare state for the agent
+        current_state_dict = {
+            "messages": st.session_state.messages,
+            "patient_data": st.session_state.patient_data,
+            "summary": st.session_state.get("summary"),
+            "interaction_warnings": None  # Start clean
+        }
+
+        # Invoke the agent's graph for one turn
+        with st.spinner("SynapseAI is processing..."):
             try:
-                final_state = st.session_state.graph_app.invoke(current_state, {"recursion_limit": 15})
-                st.session_state.messages = final_state['messages']
+                # Call the agent instance's method
+                final_state = st.session_state.agent.invoke_turn(current_state_dict)
+
+                # Update Streamlit session state from the returned agent state
+                st.session_state.messages = final_state.get('messages', [])
+                st.session_state.summary = final_state.get('summary')
+
             except Exception as e:
-                print(f"CRITICAL ERROR: {e}")
+                print(f"CRITICAL ERROR during agent invocation: {type(e).__name__} - {e}")
                 traceback.print_exc()
-                st.error(f"Error: {e}")
+                st.error(f"An error occurred during processing: {e}", icon="❌")
+                # Append error to messages for user visibility
+                st.session_state.messages.append(AIMessage(content=f"Error during processing: {e}"))
+
+        # Rerun Streamlit script to update the chat display
         st.rerun()
 
     # Disclaimer