Update app.py

app.py

@@ -6,9 +6,9 @@ import os
 import traceback
 from dotenv import load_dotenv
 
-# Import agent logic and message types from agent.py
+# Import agent logic and message types
 try:
-    from agent import ClinicalAgent, AgentState, check_red_flags
+    from agent import ClinicalAgent, AgentState, check_red_flags # Import necessary components
     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.")
@@ -46,26 +46,22 @@ def main():
     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:
         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()
-
+            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)
+        # 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") # Added Warfarin/Aspirin for testing
+        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")
@@ -92,7 +88,6 @@ def main():
             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 ---
@@ -110,20 +105,44 @@ def main():
                     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
+                    else: st.markdown(ai_content)
                 except Exception as e: st.markdown(ai_content); print(f"Error parsing/displaying AI JSON: {e}")
+
                 if structured_output and isinstance(structured_output, dict): # 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: [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()
+                    with cols[0]: # Assessment, DDx, Risk
+                        st.markdown("**Assessment:**"); st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
+                        st.markdown("**Differential Diagnosis:**"); ddx = structured_output.get('differential_diagnosis', []);
+                        if ddx: [st.expander(f"{'🥇🥈🥉'[('High','Medium','Low').index(item.get('likelihood','Low')[0])] if item.get('likelihood','?')[0] in 'HML' else '?'} {item.get('diagnosis', 'Unknown')} ({item.get('likelihood','?')})").write(f"**Rationale:** {item.get('rationale', 'N/A')}") for item in ddx]
+                        else: st.info("No DDx provided.")
+
+                        # Risk Assessment Display (CORRECTED - Separate lines)
+                        st.markdown(f"**Risk Assessment:**")
+                        risk = structured_output.get('risk_assessment', {})
+                        flags = risk.get('identified_red_flags', [])
+                        concerns = risk.get("immediate_concerns", [])
+                        comps = risk.get("potential_complications", [])
+
+                        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)}")
+                        # Add a message if no risks were highlighted by the AI assessment
+                        if not flags and not concerns and not comps:
+                            st.success("No specific risks highlighted in this AI assessment.")
+
+                    with cols[1]: # Plan
+                         st.markdown("**Recommended Plan:**"); plan = structured_output.get('recommended_plan', {});
+                         for section in ["investigations","therapeutics","consultations","patient_education"]: st.markdown(f"_{section.replace('_',' ').capitalize()}:_"); items = plan.get(section); [st.markdown(f"- {item}") for item in items] if items and isinstance(items, list) else (st.markdown(f"- {items}") if items else st.markdown("_None_")); st.markdown("")
+
+                    # Rationale & Interaction Summary
+                    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):
@@ -138,9 +157,7 @@ def main():
             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");
-                    # 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="🚨") # Show flag in UI too
+                    if tool_name_display == "flag_risk" and status == "flagged": st.error(f"🚨 **RISK FLAGGED:** {message}", icon="🚨") # Show flag in UI too
                     elif status == "success" or status == "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]
@@ -154,37 +171,15 @@ 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)
+        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
+        current_state_dict = {"messages": st.session_state.messages, "patient_data": st.session_state.patient_data, "summary": st.session_state.get("summary"), "interaction_warnings": None}
         with st.spinner("SynapseAI is processing..."):
             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')
-
-            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="❌")
-                # 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
+            except Exception as e: print(f"CRITICAL ERROR: {e}"); traceback.print_exc(); st.error(f"Error: {e}"); st.session_state.messages.append(AIMessage(content=f"Error processing request: {e}"))
         st.rerun()
 
     # Disclaimer