Update pages/2_Consult.py

pages/2_Consult.py

@@ -1,23 +1,24 @@
 # /home/user/app/pages/2_Consult.py
 import streamlit as st
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage # <--- CORRECTED IMPORT
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage # Ensure SystemMessage is imported
 from datetime import datetime
 from typing import List, Optional, Dict, Any
 from sqlmodel import select
 
 from config.settings import settings
-from agent import get_agent_executor 
+from agent import get_agent_executor # This now returns the Gemini-based agent
 from models import ChatMessage, ChatSession
 from models.db import get_session_context
 from services.logger import app_logger
 from services.metrics import log_consultation_start
 
-
 # --- Authentication Check ---
 if not st.session_state.get("authenticated_user_id"):
     st.warning("Please log in to access the consultation page.")
-    try: st.switch_page("app.py")
-    except st.errors.StreamlitAPIException: st.info("Please navigate to the main login page.")
+    try:
+        st.switch_page("app.py")
+    except st.errors.StreamlitAPIException:
+        st.info("Please navigate to the main login page.")
     st.stop()
 
 authenticated_user_id = st.session_state.get("authenticated_user_id")
@@ -26,36 +27,44 @@ app_logger.info(f"User '{authenticated_username}' (ID: {authenticated_user_id})
 
 # --- Initialize Agent ---
 try:
-    agent_executor = get_agent_executor()
+    agent_executor = get_agent_executor() # Gets the Gemini agent executor
+    app_logger.info("Gemini-based agent executor initialized for Consult page.")
+except ValueError as e: # Catch specific error from get_agent_executor if API key is missing
+    st.error(f"AI Agent Initialization Error: {e}")
+    app_logger.critical(f"Fatal: AI Agent could not be initialized in Consult page: {e}", exc_info=True)
+    st.info("Please ensure the necessary API keys (e.g., Google API Key for Gemini) are configured in the application settings.")
+    st.stop()
 except Exception as e:
-    st.error(f"Fatal Error: Could not initialize AI Agent: {e}. Please check API keys and configurations.")
-    app_logger.critical(f"AI Agent initialization failed: {e}", exc_info=True)
+    st.error(f"An unexpected error occurred while initializing the AI Agent: {e}")
+    app_logger.critical(f"Fatal: Unexpected AI Agent initialization error: {e}", exc_info=True)
     st.stop()
 
+
 # --- Session State for Consult Page ---
 if 'current_consult_patient_context' not in st.session_state:
-    st.session_state.current_consult_patient_context = {} # Stores structured context for current consult
+    st.session_state.current_consult_patient_context = {}
 if 'consult_context_submitted' not in st.session_state:
     st.session_state.consult_context_submitted = False
 
 # --- Helper Functions ---
-@st.cache_data(ttl=30, show_spinner=False) # Short cache for agent history
-def load_chat_history_for_agent(session_id: int) -> List:
+@st.cache_data(ttl=30, show_spinner=False)
+def load_chat_history_for_agent(session_id: int) -> List: # List of LangChain messages
     messages = []
-    # ... (load_chat_history_for_agent from previous full rewrite of 2_Consult.py, using SQLModel select) ...
-    # This function should convert DB ChatMessages to LangChain HumanMessage/AIMessage
     app_logger.debug(f"Loading agent chat history for session_id: {session_id}")
     with get_session_context() as db:
         statement = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp)
         db_messages = db.exec(statement).all()
         for msg in db_messages:
-            if msg.role == "user": messages.append(HumanMessage(content=msg.content))
-            elif msg.role == "assistant": messages.append(AIMessage(content=msg.content))
+            if msg.role == "user":
+                messages.append(HumanMessage(content=msg.content))
+            elif msg.role == "assistant":
+                messages.append(AIMessage(content=msg.content))
+            elif msg.role == "system": # Include system messages in agent history if they were saved
+                messages.append(SystemMessage(content=msg.content))
+    app_logger.debug(f"Loaded {len(messages)} messages for agent history for session {session_id}.")
     return messages
 
-
 def save_chat_message_to_db(session_id: int, role: str, content: str, tool_call_id: Optional[str]=None, tool_name: Optional[str]=None):
-    # ... (save_chat_message_to_db from previous full rewrite of 2_Consult.py) ...
     app_logger.debug(f"Saving message to DB for session {session_id}: Role={role}")
     with get_session_context() as db:
         chat_message = ChatMessage(
@@ -65,8 +74,7 @@ def save_chat_message_to_db(session_id: int, role: str, content: str, tool_call_
         db.add(chat_message) # Commit handled by context manager
     app_logger.info(f"Message saved to DB for session {session_id}. Role: {role}.")
 
-
-def update_chat_session_with_context(session_id: int, context_summary: str):
+def update_chat_session_with_context_summary(session_id: int, context_summary: str):
     with get_session_context() as db:
         session_to_update = db.get(ChatSession, session_id)
         if session_to_update:
@@ -74,86 +82,95 @@ def update_chat_session_with_context(session_id: int, context_summary: str):
             db.add(session_to_update) # Stage for commit
             app_logger.info(f"Updated ChatSession {session_id} with patient context summary.")
         else:
-            app_logger.error(f"Could not find ChatSession {session_id} to update with context.")
+            app_logger.error(f"Could not find ChatSession {session_id} to update with context summary.")
 
 # --- Page Logic ---
 st.title("AI Consultation Room")
 st.markdown(f"Interacting as: **{authenticated_username}**")
-st.info(settings.MAIN_DISCLAIMER_SHORT + " Do not enter real PHI.")
+st.info(f"{settings.MAIN_DISCLAIMER_SHORT} Remember to use only anonymized, simulated data.")
 
 chat_session_id = st.session_state.get("current_chat_session_id")
 if not chat_session_id:
-    st.error("No active chat session. This may occur if you logged out and back in. A new session was created. If issues persist, please re-login fully or contact support.")
-    app_logger.error(f"User '{authenticated_username}' on Consult page with no current_chat_session_id.")
-    # Attempt to create a new one if truly missing, or guide to re-login
-    # For now, stopping is safer if app.py is supposed to always create one.
+    st.error("No active chat session ID. This can occur if a session wasn't properly created on login. Please try logging out and then logging back in. If the problem persists, contact support.")
+    app_logger.error(f"User '{authenticated_username}' (ID: {authenticated_user_id}) on Consult page with NO current_chat_session_id.")
     st.stop()
 
 # --- Patient Context Input Form ---
 if not st.session_state.consult_context_submitted:
-    st.subheader("Optional: Provide Patient Context (Simulated Data Only)")
-    with st.form(key="patient_context_form"):
+    st.subheader("Step 1: Provide Patient Context (Optional, Simulated Data Only)")
+    with st.form(key="patient_context_form_consult"):
         st.markdown("**Reminder: Use only anonymized, simulated data for this demonstration.**")
-        age = st.number_input("Patient Age (Simulated)", min_value=0, max_value=120, step=1)
-        gender = st.selectbox("Patient Gender (Simulated)", ["Not Specified", "Male", "Female", "Other"])
-        chief_complaint = st.text_area("Chief Complaint / Reason for Consult (Simulated)", height=100)
-        key_history = st.text_area("Key Medical History (Simulated - e.g., diabetes, hypertension)", height=100)
-        current_meds = st.text_area("Current Medications (Simulated - e.g., metformin, lisinopril)", height=100)
+        age = st.number_input("Patient Age (Simulated)", min_value=0, max_value=120, step=1, value=None) # Default to None
+        gender_options = ["Not Specified", "Male", "Female", "Other"]
+        gender = st.selectbox("Patient Gender (Simulated)", gender_options, index=0)
+        chief_complaint = st.text_area("Chief Complaint / Reason for Consult (Simulated)", height=100, placeholder="e.g., Persistent cough for 2 weeks")
+        key_history = st.text_area("Key Medical History (Simulated)", height=100, placeholder="e.g., Type 2 Diabetes, Hypertension, Asthma")
+        current_meds = st.text_area("Current Medications (Simulated)", height=100, placeholder="e.g., Metformin 500mg BID, Lisinopril 10mg OD")
         submit_context_button = st.form_submit_button("Start Consult with this Context")
 
         if submit_context_button:
-            context = {
-                "age": age if age > 0 else "Not Specified",
-                "gender": gender,
-                "chief_complaint": chief_complaint.strip() if chief_complaint.strip() else "Not Specified",
-                "key_medical_history": key_history.strip() if key_history.strip() else "Not Specified",
-                "current_medications": current_meds.strip() if current_meds.strip() else "Not Specified",
+            context_dict = {
+                "age": age if age is not None and age > 0 else None, # Store None if not specified
+                "gender": gender if gender != "Not Specified" else None,
+                "chief_complaint": chief_complaint.strip() or None,
+                "key_medical_history": key_history.strip() or None,
+                "current_medications": current_meds.strip() or None,
             }
-            st.session_state.current_consult_patient_context = context
-            st.session_state.consult_context_submitted = True
-
-            # Create a summary for the agent and DB
-            context_summary_parts = [f"{k.replace('_', ' ').title()}: {v}" for k, v in context.items() if v != "Not Specified" and v != ""]
-            context_summary_for_agent = "Patient Context: " + "; ".join(context_summary_parts) if context_summary_parts else "No specific patient context provided."
+            # Filter out None values for the summary string
+            valid_context_parts = {k: v for k, v in context_dict.items() if v is not None}
+            st.session_state.current_consult_patient_context = valid_context_parts # Store the filtered dict
+
+            if valid_context_parts:
+                context_summary_str_parts = [f"{k.replace('_', ' ').title()}: {v}" for k, v in valid_context_parts.items()]
+                context_summary_for_db_and_agent = "; ".join(context_summary_str_parts)
+            else:
+                context_summary_for_db_and_agent = "No specific patient context provided for this session."
+            
+            update_chat_session_with_context_summary(chat_session_id, context_summary_for_db_and_agent)
             
-            # Save context summary to ChatSession model
-            update_chat_session_with_context(chat_session_id, context_summary_for_agent)
-
-            # Prepend context to agent's chat history as a system message or initial user message
-            # For this example, let's add it as a system message to guide the AI
             agent_history_key = f"agent_chat_history_{chat_session_id}"
             if agent_history_key not in st.session_state: st.session_state[agent_history_key] = []
-            st.session_state[agent_history_key].insert(0, SystemMessage(content=context_summary_for_agent))
-            # Also save this "system" context message to DB for record keeping if desired
-            save_chat_message_to_db(chat_session_id, "system", context_summary_for_agent)
             
-            app_logger.info(f"Patient context submitted for session {chat_session_id}: {context_summary_for_agent}")
-            st.rerun() # Rerun to hide form and show chat
-    st.stop() # Don't proceed to chat until context is submitted or skipped
+            # Don't add patient context as a SystemMessage if it's passed as a variable to invoke
+            # The agent's main system prompt will now include a placeholder for it.
+            # However, we save it to DB for record keeping.
+            if valid_context_parts: # Save a system message indicating context was provided
+                 save_chat_message_to_db(chat_session_id, "system", f"Initial Patient Context Provided: {context_summary_for_db_and_agent}")
 
-# --- Chat Interface (Shown after context is submitted/skipped) ---
+            st.session_state.consult_context_submitted = True
+            app_logger.info(f"Patient context submitted for session {chat_session_id}: {context_summary_for_db_and_agent}")
+            st.rerun()
+    st.stop()
+
+# --- Chat Interface (Shown after context is submitted or if skipped by some other logic not yet present) ---
+st.subheader("Step 2: Interact with AI Health Navigator")
 agent_history_key = f"agent_chat_history_{chat_session_id}"
+
 if agent_history_key not in st.session_state:
     st.session_state[agent_history_key] = load_chat_history_for_agent(chat_session_id)
-    if not st.session_state[agent_history_key]: # If history is empty (even after context attempt)
+    if not st.session_state[agent_history_key]: # If history is empty
         try: log_consultation_start(user_id=authenticated_user_id, session_id=chat_session_id)
-        except Exception as e: app_logger.warning(f"Failed to log consultation start: {e}")
+        except Exception as e: app_logger.warning(f"Failed to log consultation start metric: {e}")
+        
         initial_ai_message_content = "Hello! I am your AI Health Navigator. How can I assist you today?"
+        if st.session_state.get('current_consult_patient_context'):
+            initial_ai_message_content += " I have noted the patient context you provided."
+        
         st.session_state[agent_history_key].append(AIMessage(content=initial_ai_message_content))
         save_chat_message_to_db(chat_session_id, "assistant", initial_ai_message_content)
+        app_logger.info(f"Initialized new consultation (session {chat_session_id}) with a greeting.")
 
-# Display chat messages from DB for UI
+# Display chat messages for UI
 with st.container():
     with get_session_context() as db:
         stmt = select(ChatMessage).where(ChatMessage.session_id == chat_session_id).order_by(ChatMessage.timestamp)
         ui_messages = db.exec(stmt).all()
         for msg in ui_messages:
-            if msg.role == "system": # Don't show system context messages directly in chat UI
-                continue
+            if msg.role == "system": continue # Don't display system context messages directly
             avatar = "🧑‍⚕️" if msg.role == "assistant" else "👤"
-            if msg.role == "tool": avatar = "🛠️"
+            if msg.role == "tool": avatar = "🛠️" # Assuming you might log tool calls this way
             with st.chat_message(msg.role, avatar=avatar):
-                st.markdown(msg.content) # Add source/confidence here if msg object supports it
+                st.markdown(msg.content) # Potentially enhance to show sources/confidence if agent provides
 
 if prompt := st.chat_input("Ask the AI..."):
     with st.chat_message("user", avatar="👤"): st.markdown(prompt)
@@ -163,24 +180,41 @@ if prompt := st.chat_input("Ask the AI..."):
     with st.chat_message("assistant", avatar="🧑‍⚕️"):
         with st.spinner("AI is thinking..."):
             try:
-                # Pass patient context if your agent is designed to use it explicitly
-                # current_context = st.session_state.get('current_consult_patient_context', {})
-                # context_str_for_invoke = "; ".join([f"{k}: {v}" for k,v in current_context.items() if v and v!="Not Specified"])
-
-                response = agent_executor.invoke({
+                # Prepare patient context string for the agent, if any was provided
+                patient_context_dict = st.session_state.get('current_consult_patient_context', {})
+                if patient_context_dict:
+                    context_parts_for_invoke = [f"{k.replace('_', ' ').title()}: {v}" for k, v in patient_context_dict.items()]
+                    patient_context_str_for_invoke = "; ".join(context_parts_for_invoke)
+                else:
+                    patient_context_str_for_invoke = "No specific patient context was provided for this interaction."
+
+                invoke_payload = {
                     "input": prompt,
                     "chat_history": st.session_state[agent_history_key],
-                    # "patient_context": context_str_for_invoke # If agent expects this
-                })
-                ai_response_content = response.get('output', "I could not generate a response.")
+                    "patient_context": patient_context_str_for_invoke # Pass to agent
+                }
+                app_logger.debug(f"Invoking agent with payload: {invoke_payload}")
+
+                response = agent_executor.invoke(invoke_payload)
+                
+                ai_response_content = response.get('output', "I could not generate a valid response.")
                 if not isinstance(ai_response_content, str): ai_response_content = str(ai_response_content)
                 
-                st.markdown(ai_response_content) # Display sources/confidence here if available in ai_response_content
+                app_logger.info(f"Agent response for session {chat_session_id}: '{ai_response_content[:100]}...'")
+                st.markdown(ai_response_content) # Display AI response
                 save_chat_message_to_db(chat_session_id, "assistant", ai_response_content)
                 st.session_state[agent_history_key].append(AIMessage(content=ai_response_content))
+
             except Exception as e:
                 app_logger.error(f"Error during agent invocation for session {chat_session_id}: {e}", exc_info=True)
-                error_msg_user = f"Sorry, an error occurred: {type(e).__name__}. Please try again."
-                st.error(error_msg_user)
-                save_chat_message_to_db(chat_session_id, "assistant", f"Internal error: {type(e).__name__}")
-                st.session_state[agent_history_key].append(AIMessage(content=f"Internal error: {type(e).__name__}"))
+                # The user-facing error was: "Sorry, an error occurred: ValidationError. Please try again."
+                # Let's try to be a bit more specific if we can, or keep it generic but log details.
+                error_type_name = type(e).__name__ # e.g., "ValidationError", "APIError"
+                user_friendly_error_message = f"Sorry, an error occurred ({error_type_name}). Please try rephrasing your query or contact support if the issue persists."
+                st.error(user_friendly_error_message)
+                
+                # Save a representation of the error to DB for the assistant's turn
+                db_error_message = f"System encountered an error: {error_type_name} while processing user query. Details logged."
+                save_chat_message_to_db(chat_session_id, "assistant", db_error_message)
+                # Add error representation to agent history so it's aware for next turn (optional)
+                st.session_state[agent_history_key].append(AIMessage(content=f"Note to self: Encountered an error ({error_type_name}) on the previous turn."))