Update app.py

app.py

@@ -12,7 +12,6 @@ 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
@@ -38,16 +37,16 @@ class ClinicalPrompts: SYSTEM_PROMPT = """
     """
 
 # --- API Helper Functions (get_rxcui, get_openfda_label, search_text_list) ---
-# ... (Keep these functions exactly as they were in the previous 'full code' response) ...
+# ... (Keep these functions exactly as they were) ...
 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:
+    try: # Try direct lookup first
         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:
+        else: # Fallback to /drugs 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"]:
@@ -79,30 +78,71 @@ def search_text_list(text_list: Optional[List[str]], search_terms: List[str]) ->
         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
+                # Highlight first match for clarity
+                snippet = re.sub(f"({re.escape(term)})", r"**\1**", snippet, count=1, flags=re.IGNORECASE)
+                found_snippets.append(f"...{snippet}...")
+                break # Only report first match per text item
     return found_snippets
 
 
-# --- Other Helper Functions (parse_bp, check_red_flags, format_patient_data_for_prompt) ---
-# ... (Keep these functions exactly as they were) ...
+# --- Other Helper Functions ---
 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
+
+# CORRECTED check_red_flags function
 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))
+    """Checks patient data against predefined red flags."""
+    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)]
+
+    # Symptom Flags (CORRECTED - Separate lines)
+    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).")
+
+    # Vital Sign Flags
+    if vitals:
+        temp = vitals.get("temp_c"); hr = vitals.get("hr_bpm"); rr = vitals.get("rr_rpm")
+        spo2 = vitals.get("spo2_percent"); bp_str = vitals.get("bp_mmhg")
+        if temp is not None and temp >= 38.5: flags.append(f"Red Flag: Fever ({temp}°C).")
+        if hr is not None and hr >= 120: flags.append(f"Red Flag: Tachycardia ({hr} bpm).")
+        if hr is not None and hr <= 50: flags.append(f"Red Flag: Bradycardia ({hr} bpm).")
+        if rr is not None and rr >= 24: flags.append(f"Red Flag: Tachypnea ({rr} rpm).")
+        if spo2 is not None and spo2 <= 92: flags.append(f"Red Flag: Hypoxia ({spo2}%).")
+        if bp_str:
+            bp = parse_bp(bp_str)
+            if bp:
+                if bp[0] >= 180 or bp[1] >= 110: flags.append(f"Red Flag: Hypertensive Urgency/Emergency (BP: {bp_str} mmHg).")
+                if bp[0] <= 90 or bp[1] <= 60: flags.append(f"Red Flag: Hypotension (BP: {bp_str} mmHg).")
+
+    # History Flags
+    if history and isinstance(history, str):
+        history_lower = history.lower()
+        if "history of mi" in history_lower and "chest pain" in symptoms_lower:
+            flags.append("Red Flag: History of MI with current Chest Pain.")
+        if "history of dvt/pe" in history_lower and "shortness of breath" in symptoms_lower:
+             flags.append("Red Flag: History of DVT/PE with current Shortness of Breath.")
+
+    return list(set(flags)) # Unique flags
+
 def format_patient_data_for_prompt(data: dict) -> str:
+    # ... (Keep this function exactly as it was) ...
     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());
@@ -174,7 +214,7 @@ model = ChatGroq(temperature=ClinicalAppSettings.TEMPERATURE, model=ClinicalAppS
 model_with_tools = model.bind_tools(tools)
 
 # --- Graph Nodes (agent_node, tool_node) ---
-# ... (Keep agent_node and tool_node functions exactly as they were in the last 'full code' response) ...
+# ... (Keep these functions exactly as they were) ...
 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;
@@ -233,18 +273,18 @@ def main():
     # --- 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"])
-        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"])
-        st.subheader("History"); pmh = st.text_area("PMH", "HTN, HLD, DM2, History of 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"); allergies_str = st.text_area("Allergies", "Penicillin (rash)")
-        st.subheader("Social/Family"); social_history = st.text_area("SH", "Smoker"); family_history = st.text_area("FHx", "Father MI")
+        # Input fields...
+        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\nAtorvastatin 40mg daily", key="sb_meds"); allergies_str = st.text_area("Allergies", "Penicillin (rash), Sulfa", 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"); 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=50)
+        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=50, key="sb_exam")
 
-        if st.button("Start/Update Consultation"):
+        if st.button("Start/Update Consultation", key="sb_start"):
             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);
@@ -261,12 +301,12 @@ def main():
 
     # --- Main Chat Interface Area ---
     st.header("💬 Clinical Consultation")
-    # Display loop - SyntaxError Fixed
+    # Display loop - key= argument REMOVED, Tool Call Display Syntax FIXED
     for msg in st.session_state.messages:
         if isinstance(msg, HumanMessage):
-            with st.chat_message("user"): st.markdown(msg.content) # No key
+            with st.chat_message("user"): st.markdown(msg.content)
         elif isinstance(msg, AIMessage):
-            with st.chat_message("assistant"): # No key
+            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)
@@ -292,18 +332,15 @@ def main():
                 # CORRECTED Tool Call Display Block
                 if getattr(msg, 'tool_calls', None):
                      with st.expander("🛠️ AI requested actions", expanded=False):
-                         if msg.tool_calls: # Check if list is not empty
+                         if msg.tool_calls:
                             for tc in msg.tool_calls:
                                 try:
-                                    # Properly indented try block content
                                     st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nArgs: {json.dumps(tc.get('args', {}), indent=2)}", language="json")
                                 except Exception as display_e:
-                                    # Properly indented except block content
-                                    st.error(f"Could not display tool call arguments properly: {display_e}", icon="⚠️")
-                                    # Provide a fallback display
-                                    st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nRaw Args: {tc.get('args')}") # Show raw args if JSON fails
+                                    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="🛠️"): # No key
@@ -312,9 +349,9 @@ def main():
                     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:"); [st.caption(f"- {warn}") for warn in warnings]
-                    else: st.error(f"{message}", icon="❌")
+                    else: st.error(f"{message}", icon="❌") # Assume error if not success/clear/flagged/warning
                     if details: st.caption(f"Details: {details}")
-                except json.JSONDecodeError: st.info(f"{msg.content}")
+                except json.JSONDecodeError: st.info(f"{msg.content}") # Display raw if not JSON
                 except Exception as e: st.error(f"Error displaying tool message: {e}", icon="❌"); st.caption(f"Raw content: {msg.content}")
 
     # --- Chat Input Logic ---