Create agent.py

agent.py

@@ -0,0 +1,191 @@
+# app.py
+import streamlit as st
+import json
+import re
+import os
+from dotenv import load_dotenv
+
+# Import agent logic and message types
+from agent import ClinicalAgent, AgentState, check_red_flags # Import necessary components
+from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
+
+# --- 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 is now primarily defined in agent.py, but can keep for display
+    MODEL_NAME_DISPLAY = "Llama3-70b (via Groq)"
+
+
+# --- 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
+    # Summary state for future memory enhancement
+    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... (Assume full fields as before)
+        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=50, 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} }
+            # 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]
+            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 - Uses messages from st.session_state
+    for msg in st.session_state.messages: # Removed enumerate and key
+        if isinstance(msg, HumanMessage):
+            with st.chat_message("user"): st.markdown(msg.content)
+        elif isinstance(msg, AIMessage):
+            with st.chat_message("assistant"):
+                # ... (Keep the detailed AI message display logic, including JSON parsing) ...
+                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")
+                    # ... (Keep detailed JSON display logic for assessment, ddx, plan, etc.) ...
+                    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: [st.expander(f"{'🥇🥈🥉'[('High','Medium','Low').index(item.get('likelihood','Low')[0])] if item.get('likelihood','?')[0] in 'HML' else '?'} {item.get('diagnosis', 'Unknown')} ({item.get('likelihood','?')})").write(f"**Rationale:** {item.get('rationale', 'N/A')}") for item in ddx]
+                    else: st.info("No DDx provided."); st.markdown("**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); [st.markdown(f"- {item}") for item in items] if items and isinstance(items, list) else (st.markdown(f"- {items}") if items else st.markdown("_None_")); st.markdown("")
+                    st.markdown("**Rationale & Guideline Check:**"); st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}"); interaction_summary = structured_output.get("interaction_check_summary", "");
+                    if interaction_summary: st.markdown("**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._")
+        elif isinstance(msg, ToolMessage):
+            tool_name_display = getattr(msg, 'name', 'tool_execution')
+            with st.chat_message(tool_name_display, avatar="🛠️"):
+                # ... (Keep ToolMessage display logic) ...
+                try:
+                    tool_data = json.loads(msg.content); status = tool_data.get("status", "info"); message = tool_data.get("message", msg.content); details = tool_data.get("details"); warnings = tool_data.get("warnings");
+                    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="❌") # Assume error if not known status
+                    if details: st.caption(f"Details: {details}")
+                except json.JSONDecodeError: st.info(f"{msg.content}") # Display raw if not JSON
+                except Exception as e: st.error(f"Error displaying tool message: {e}", icon="❌"); st.caption(f"Raw content: {msg.content}")
+
+
+    # --- 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 not st.session_state.agent: st.error("Agent not initialized. Check logs."); st.stop() # Add check for agent
+
+        # 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, including existing messages and patient data
+        current_state_dict = {
+            "messages": st.session_state.messages,
+            "patient_data": st.session_state.patient_data,
+            "summary": st.session_state.get("summary"), # Include summary if implemented
+            "interaction_warnings": None # Always start turn with no pending warnings
+        }
+
+        # Invoke the agent's graph for one turn
+        with st.spinner("SynapseAI is thinking..."):
+            try:
+                # 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') # Update summary if implemented
+
+            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="❌")
+                # Optionally append an error message to the chat display
+                # st.session_state.messages.append(AIMessage(content=f"Error processing request: {e}"))
+
+        # Rerun Streamlit script to update the chat display
+        st.rerun()
+
+    # Disclaimer
+    st.markdown("---"); st.warning("**Disclaimer:** SynapseAI is for demonstration...")
+
+if __name__ == "__main__":
+    main()