Update agent.py

agent.py

@@ -1,308 +1,439 @@
-# app.py
-import streamlit as st
-import json
-import re
+# agent.py
+
 import os
+import re
+import json
+import logging
 import traceback
-from dotenv import load_dotenv
-
-# 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")
-if not GROQ_API_KEY: missing_keys.append("GROQ_API_KEY")
-if not TAVILY_API_KEY: missing_keys.append("TAVILY_API_KEY")
-if missing_keys:
-    st.error(f"Missing required API Key(s): {', '.join(missing_keys)}. Please set them in Hugging Face Space Secrets or environment variables.")
-    st.stop()
-
-# --- App Configuration ---
-class ClinicalAppSettings:
-    APP_TITLE = "SynapseAI (UMLS/FDA Integrated)"
-    PAGE_LAYOUT = "wide"
-    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_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 "summary" not in st.session_state:
-        st.session_state.summary = None
-
-    # Initialize the agent instance only once
-    if "agent" not in st.session_state:
+from functools import lru_cache
+from typing import List, Dict, Any, Optional, TypedDict
+
+import requests
+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.pydantic_v1 import BaseModel, Field
+from langchain_core.tools import tool
+from langgraph.prebuilt import ToolExecutor
+from langgraph.graph import StateGraph, END
+
+# ── Logging Configuration ─────────────────────────────────────────────────────
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# ── Environment Variables ─────────────────────────────────────────────────────
+UMLS_API_KEY = os.getenv("UMLS_API_KEY")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+
+if not all([UMLS_API_KEY, GROQ_API_KEY, TAVILY_API_KEY]):
+    logger.error("Missing one or more required API keys: UMLS_API_KEY, GROQ_API_KEY, TAVILY_API_KEY")
+    raise RuntimeError("Missing required API keys")
+
+# ── Agent Configuration ───────────────────────────────────────────────────────
+AGENT_MODEL_NAME = "llama3-70b-8192"
+AGENT_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 CONTENT HERE]
+    """
+
+# ── Helper Functions ──────────────────────────────────────────────────────────
+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]:
+    """Lookup RxNorm CUI for a drug name."""
+    drug_name = (drug_name or "").strip()
+    if not drug_name:
+        return None
+    logger.info(f"Looking up RxCUI for '{drug_name}'")
+    try:
+        params = {"name": drug_name, "search": 1}
+        r = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10)
+        r.raise_for_status()
+        data = r.json().get("idGroup", {})
+        if ids := data.get("rxnormId"):
+            logger.info(f"Found RxCUI {ids[0]} for '{drug_name}'")
+            return ids[0]
+        # fallback to broader search
+        r = requests.get(f"{RXNORM_API_BASE}/drugs.json", params={"name": drug_name}, timeout=10)
+        r.raise_for_status()
+        for group in r.json().get("drugGroup", {}).get("conceptGroup", []):
+            if props := group.get("conceptProperties"):
+                logger.info(f"Found RxCUI {props[0]['rxcui']} via /drugs for '{drug_name}'")
+                return props[0]["rxcui"]
+    except Exception:
+        logger.exception(f"Error fetching RxCUI for '{drug_name}'")
+    return None
+
+@lru_cache(maxsize=128)
+def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
+    """Fetch label data from OpenFDA by RxCUI or drug name."""
+    if not (rxcui or drug_name):
+        return None
+    terms = []
+    if rxcui:
+        terms.append(f'spl_rxnorm_code:"{rxcui}" OR openfda.rxcui:"{rxcui}"')
+    if drug_name:
+        dn = drug_name.lower()
+        terms.append(f'(openfda.brand_name:"{dn}" OR openfda.generic_name:"{dn}")')
+    query = " OR ".join(terms)
+    logger.info(f"Looking up OpenFDA label with query: {query}")
+    try:
+        r = requests.get(OPENFDA_API_BASE, params={"search": query, "limit": 1}, timeout=15)
+        r.raise_for_status()
+        results = r.json().get("results", [])
+        if results:
+            return results[0]
+    except Exception:
+        logger.exception("Error fetching OpenFDA label")
+    return None
+
+def search_text_list(texts: List[str], terms: List[str]) -> List[str]:
+    """Return snippets from texts containing any of the search terms."""
+    snippets = []
+    lowers = [t.lower() for t in terms if t]
+    for text in texts or []:
+        tl = text.lower()
+        for term in lowers:
+            if term in tl:
+                i = tl.find(term)
+                start = max(0, i - 50)
+                end = min(len(text), i + len(term) + 100)
+                snippet = text[start:end]
+                snippet = re.sub(f"({re.escape(term)})", r"**\1**", snippet, flags=re.IGNORECASE)
+                snippets.append(f"...{snippet}...")
+                break
+    return snippets
+
+def parse_bp(bp: str) -> Optional[tuple[int, int]]:
+    """Parse a 'SYS/DIA' blood pressure string into a tuple."""
+    if m := re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", (bp or "").strip()):
+        return int(m.group(1)), int(m.group(2))
+    return None
+
+def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
+    """Identify immediate red flags in patient data."""
+    flags: List[str] = []
+    hpi = patient_data.get("hpi", {})
+    vitals = patient_data.get("vitals", {})
+    symptoms = [s.lower() for s in hpi.get("symptoms", []) if isinstance(s, str)]
+    # Symptom-based
+    mapping = {
+        "chest pain": "Chest pain reported",
+        "shortness of breath": "Shortness of breath reported",
+        "severe headache": "Severe headache reported",
+        "syncope": "Syncope (fainting) reported",
+        "hemoptysis": "Hemoptysis (coughing blood) reported"
+    }
+    for term, desc in mapping.items():
+        if term in symptoms:
+            flags.append(f"Red Flag: {desc}.")
+    # Vitals-based
+    temp = vitals.get("temp_c")
+    hr   = vitals.get("hr_bpm")
+    rr   = vitals.get("rr_rpm")
+    spo2 = vitals.get("spo2_percent")
+    bp   = parse_bp(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:
+        if hr >= 120:
+            flags.append(f"Red Flag: Tachycardia ({hr} bpm).")
+        if 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:
+        sys, dia = bp
+        if sys >= 180 or dia >= 110:
+            flags.append(f"Red Flag: Hypertensive urgency/emergency ({sys}/{dia} mmHg).")
+        if sys <= 90 or dia <= 60:
+            flags.append(f"Red Flag: Hypotension ({sys}/{dia} mmHg).")
+    return list(dict.fromkeys(flags))  # preserve order, dedupe
+
+def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
+    """Convert patient_data dict into a markdown-like prompt string."""
+    if not data:
+        return "No patient data provided."
+    lines: List[str] = []
+    for section, value in data.items():
+        title = section.replace("_", " ").title()
+        if isinstance(value, dict) and any(value.values()):
+            lines.append(f"**{title}:**")
+            for k, v in value.items():
+                if v:
+                    lines.append(f"- {k.replace('_',' ').title()}: {v}")
+        elif isinstance(value, list) and value:
+            lines.append(f"**{title}:** {', '.join(map(str, value))}")
+        elif value:
+            lines.append(f"**{title}:** {value}")
+    return "\n".join(lines)
+
+# ── Tool Input Schemas ────────────────────────────────────────────────────────
+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 Implementations ──────────────────────────────────────────────────────
+@tool("order_lab_test", args_schema=LabOrderInput)
+def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
+    logger.info(f"Ordering 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:
+    logger.info(f"Preparing prescription: {medication_name} {dosage}, route: {route}, freq: {frequency}")
+    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:
+    logger.info(f"Checking interactions for: {potential_prescription}")
+    warnings: List[str] = []
+    pm = [m.lower().strip() for m in (current_medications or []) if m]
+    al = [a.lower().strip() for a in (allergies or []) if a]
+
+    # Allergy checks
+    if potential_prescription.lower().strip() in al:
+        warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{potential_prescription}'.")
+    # RxNorm/OpenFDA lookups
+    rxcui = get_rxcui(potential_prescription)
+    label = get_openfda_label(rxcui=rxcui, drug_name=potential_prescription)
+    if not (rxcui or label):
+        warnings.append(f"INFO: Could not identify '{potential_prescription}'. Checks may be incomplete.")
+
+    # Contraindications & warnings sections
+    for section in ("contraindications", "warnings_and_cautions", "warnings"):
+        items = label.get(section) if label else None
+        if isinstance(items, list):
+            snippets = search_text_list(items, al)
+            if snippets:
+                warnings.append(f"ALLERGY RISK ({section}): {'; '.join(snippets)}")
+
+    # Drug–drug interactions
+    for med in pm:
+        mrxcui = get_rxcui(med)
+        mlabel = get_openfda_label(rxcui=mrxcui, drug_name=med)
+        # check in both labels
+        for sec in ("drug_interactions",):
+            for src_label, src_name in ((label, potential_prescription), (mlabel, med)):
+                items = src_label.get(sec) if src_label else None
+                if isinstance(items, list):
+                    snippets = search_text_list(items, [med if src_name==potential_prescription else potential_prescription])
+                    if snippets:
+                        warnings.append(f"Interaction ({src_name} label): {'; '.join(snippets)}")
+
+    status = "warning" if warnings else "clear"
+    message = (
+        f"{len(warnings)} issue(s) found for '{potential_prescription}'."
+        if warnings else
+        f"No major interactions or allergy issues identified for '{potential_prescription}'."
+    )
+    return json.dumps({"status": status, "message": message, "warnings": warnings})
+
+@tool("flag_risk", args_schema=FlagRiskInput)
+def flag_risk(risk_description: str, urgency: str = "High") -> str:
+    logger.info(f"Flagging risk: {risk_description} (urgency={urgency})")
+    return json.dumps({
+        "status": "flagged",
+        "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
+    })
+
+search_tool = TavilySearchResults(max_results=MAX_SEARCH_RESULTS, name="tavily_search_results")
+all_tools = [order_lab_test, prescribe_medication, check_drug_interactions, flag_risk, search_tool]
+
+# ── LLM & Tool Executor ──────────────────────────────────────────────────────
+llm = ChatGroq(temperature=AGENT_TEMPERATURE, model=AGENT_MODEL_NAME)
+model_with_tools = llm.bind_tools(all_tools)
+tool_executor = ToolExecutor(all_tools)
+
+# ── State Definition ──────────────────────────────────────────────────────────
+class AgentState(TypedDict):
+    messages: List[Any]
+    patient_data: Optional[Dict[str, Any]]
+    summary: Optional[str]
+    interaction_warnings: Optional[List[str]]
+
+# ── Graph Nodes ───────────────────────────────────────────────────────────────
+def agent_node(state: AgentState) -> Dict[str, Any]:
+    msgs = state["messages"]
+    if not msgs or not isinstance(msgs[0], SystemMessage):
+        msgs = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + msgs
+    logger.info(f"Invoking LLM with {len(msgs)} messages")
+    try:
+        response = model_with_tools.invoke(msgs)
+        return {"messages": [response]}
+    except Exception as e:
+        logger.exception("Error in agent_node")
+        return {"messages": [AIMessage(content=f"Error: {e}")]} 
+
+def tool_node(state: AgentState) -> Dict[str, Any]:
+    last = state["messages"][-1]
+    if not isinstance(last, AIMessage) or not getattr(last, "tool_calls", None):
+        logger.warning("tool_node invoked without pending tool_calls")
+        return {"messages": [], "interaction_warnings": None}
+
+    calls = last.tool_calls
+    # Safety: require interaction check before prescribing
+    blocked_ids = set()
+    for call in calls:
+        if call["name"] == "prescribe_medication":
+            med = call["args"].get("medication_name", "").lower()
+            if not any(c["name"] == "check_drug_interactions" and c["args"].get("potential_prescription","").lower() == med for c in calls):
+                logger.warning(f"Blocking prescribe_medication for '{med}' without interaction check")
+                blocked_ids.add(call["id"])
+
+    to_execute = [c for c in calls if c["id"] not in blocked_ids]
+    # Augment interaction checks with patient data
+    pd = state.get("patient_data", {})
+    for call in to_execute:
+        if call["name"] == "check_drug_interactions":
+            call["args"].setdefault("current_medications", pd.get("medications", {}).get("current", []))
+            call["args"].setdefault("allergies", pd.get("allergies", []))
+
+    messages: List[ToolMessage] = []
+    warnings: List[str] = []
+    try:
+        responses = tool_executor.batch(to_execute, return_exceptions=True)
+        for call, resp in zip(to_execute, responses):
+            if isinstance(resp, Exception):
+                logger.exception(f"Error executing tool {call['name']}")
+                content = json.dumps({"status": "error", "message": str(resp)})
+            else:
+                content = str(resp)
+                if call["name"] == "check_drug_interactions":
+                    data = json.loads(content)
+                    if data.get("status") == "warning":
+                        warnings.extend(data.get("warnings", []))
+            messages.append(ToolMessage(content=content, tool_call_id=call["id"], name=call["name"]))
+    except Exception as e:
+        logger.exception("Critical error in tool_node")
+        # return an error message for each pending call
+        for call in to_execute:
+            messages.append(ToolMessage(
+                content=json.dumps({"status": "error", "message": str(e)}),
+                tool_call_id=call["id"],
+                name=call["name"]
+            ))
+    return {"messages": messages, "interaction_warnings": warnings or None}
+
+def reflection_node(state: AgentState) -> Dict[str, Any]:
+    warns = state.get("interaction_warnings")
+    if not warns:
+        logger.warning("reflection_node called without warnings")
+        return {"messages": [], "interaction_warnings": None}
+
+    # Find the AIMessage that triggered the warnings
+    triggering: Optional[AIMessage] = None
+    for msg in reversed(state["messages"]):
+        if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
+            triggering = msg
+            break
+
+    if not triggering:
+        return {"messages": [AIMessage(content="Internal Error: reflection context missing.")], "interaction_warnings": None}
+
+    prompt = (
+        "You are SynapseAI, performing a focused safety review of the following plan:\n\n"
+        f"{triggering.content}\n\n"
+        "Highlight any issues based on these warnings:\n" + "\n".join(f"- {w}" for w in warns)
+    )
+    try:
+        resp = llm.invoke([SystemMessage(content="Safety reflection"), HumanMessage(content=prompt)])
+        return {"messages": [AIMessage(content=resp.content)], "interaction_warnings": None}
+    except Exception as e:
+        logger.exception("Error during reflection")
+        return {"messages": [AIMessage(content=f"Error during reflection: {e}")], "interaction_warnings": None}
+
+# ── Routing Functions ────────────────────────────────────────────────────────
+def should_continue(state: AgentState) -> str:
+    last = state["messages"][-1] if state["messages"] else None
+    if not isinstance(last, AIMessage):
+        return "end_conversation_turn"
+    if getattr(last, "tool_calls", None):
+        return "continue_tools"
+    return "end_conversation_turn"
+
+def after_tools_router(state: AgentState) -> str:
+    return "reflection" if state.get("interaction_warnings") else "agent"
+
+# ── ClinicalAgent ────────────────────────────────────────────────────────────
+class ClinicalAgent:
+    def __init__(self):
+        logger.info("Building ClinicalAgent workflow")
+        wf = StateGraph(AgentState)
+        wf.add_node("agent", agent_node)
+        wf.add_node("tools", tool_node)
+        wf.add_node("reflection", reflection_node)
+        wf.set_entry_point("agent")
+        wf.add_conditional_edges("agent", should_continue, {
+            "continue_tools": "tools",
+            "end_conversation_turn": END
+        })
+        wf.add_conditional_edges("tools", after_tools_router, {
+            "reflection": "reflection",
+            "agent": "agent"
+        })
+        wf.add_edge("reflection", "agent")
+        self.graph_app = wf.compile()
+        logger.info("ClinicalAgent ready")
+
+    def invoke_turn(self, state: Dict[str, Any]) -> Dict[str, Any]:
         try:
-            st.session_state.agent = ClinicalAgent()
-            print("ClinicalAgent successfully initialized in Streamlit session state.")
+            result = self.graph_app.invoke(state, {"recursion_limit": 15})
+            result.setdefault("summary", state.get("summary"))
+            result.setdefault("interaction_warnings", None)
+            return result
         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, 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", 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", 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\nWarfarin 5mg daily",
-            key="sb_meds"
-        )
-        allergies_str = st.text_area("Allergies", "Penicillin (rash), Aspirin", key="sb_allergies")
-
-        st.subheader("Social/Family")
-        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", 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", 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=68, key="sb_exam")
-
-        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()
-                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)
-
-            # 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},
-                "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}
+            logger.exception("Error during graph invocation")
+            return {
+                "messages": state.get("messages", []) + [AIMessage(content=f"Error: {e}")],
+                "patient_data": state.get("patient_data"),
+                "summary": state.get("summary"),
+                "interaction_warnings": None
             }
-
-            # 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:**")
-                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")
-    # Display loop
-    for msg in st.session_state.messages:
-        if isinstance(msg, HumanMessage):
-            with st.chat_message("user"):
-                st.markdown(msg.content)
-        elif isinstance(msg, AIMessage):
-            with st.chat_message("assistant"):
-                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)
-                        prefix = ai_content[:json_match.start()].strip()
-                        suffix = ai_content[json_match.end():].strip()
-                        if prefix:
-                            st.markdown(prefix)
-                        structured_output = json.loads(json_str)
-                        if suffix:
-                            st.markdown(suffix)
-                    elif ai_content.strip().startswith("{") and ai_content.strip().endswith("}"):
-                        structured_output = json.loads(ai_content)
-                        ai_content = ""
-                    else:
-                        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)
-                    with cols[0]:
-                        st.markdown("**Assessment:**")
-                        st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
-                        st.markdown("**Differential Diagnosis:**")
-                        ddx = structured_output.get('differential_diagnosis', [])
-                        if ddx:
-                            for item in ddx:
-                                likelihood = item.get('likelihood', 'Low')
-                                medal = ('🥇' if likelihood.startswith('H') else '🥈' if likelihood.startswith('M') else '🥉')
-                                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:**")
-                        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("**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 and isinstance(items, list):
-                                for it in items:
-                                    st.markdown(f"- {it}")
-                            elif items:
-                                st.markdown(f"- {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("**Interaction Check Summary:**")
-                            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"
-                                    )
-                                except Exception as display_e:
-                                    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._")
-    
-    # --- 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 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)
-
-        # 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.agent.invoke_turn(current_state_dict)
-                st.session_state.messages = final_state.get('messages', [])
-                st.session_state.summary = final_state.get('summary')
-            except Exception as e:
-                print(f"CRITICAL ERROR during agent invocation: {type(e).__name__} - {e}")
-                traceback.print_exc()
-                st.error(f"An error occurred during processing: {e}", icon="❌")
-                st.session_state.messages.append(AIMessage(content=f"Error during processing: {e}"))
-
-        st.rerun()
-
-    # Disclaimer
-    st.markdown("---")
-    st.warning("**Disclaimer:** SynapseAI is for demonstration...")
-
-if __name__ == "__main__":
-    main()