Update pages/2_Consult.py

pages/2_Consult.py

@@ -1,46 +1,75 @@
+# /home/user/app/pages/2_Consult.py
 import streamlit as st
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage
 from datetime import datetime
+from typing import List, Optional # Corrected import for List and Optional
 
 from config.settings import settings
 from agent import get_agent_executor
-from models import ChatMessage, ChatSession, User # Assuming User is in session_state
-from models.db import get_session_context
+from models import ChatMessage, ChatSession, User # User might not be needed directly if ID is used
+from models.db import get_session_context # Or from models import get_session_context
 from services.logger import app_logger
-from services.metrics import log_consultation_start
+from services.metrics import log_consultation_start # Assuming this function exists
 
-st.set_page_config(page_title=f"Consult - {settings.APP_TITLE}", layout="wide")
+# Page config typically in app.py
+# st.set_page_config(page_title=f"Consult - {settings.APP_TITLE}", layout="wide")
 
-if not st.session_state.get("authenticated_user"):
+# --- Authentication Check ---
+if not st.session_state.get("authenticated_user_id"):
     st.warning("Please log in to access the consultation page.")
-    st.switch_page("app.py") # Redirect to login
+    try:
+        st.switch_page("app.py")
+    except st.errors.StreamlitAPIException as e:
+        if "st.switch_page can only be called when running in MPA mode" in str(e):
+            app_logger.warning("Consult: Running in single-page mode or st.switch_page issue. Stopping script.")
+            st.info("Please navigate to the main login page.")
+        else:
+            app_logger.error(f"Consult: Error during st.switch_page: {e}")
+            st.error("Redirection error. Please go to the login page manually.")
+    st.stop()
+
+# Get authenticated user's ID and username
+authenticated_user_id = st.session_state.get("authenticated_user_id")
+authenticated_username = st.session_state.get("authenticated_username", "User")
+app_logger.info(f"User {authenticated_username} (ID: {authenticated_user_id}) accessed Consult page.")
 
 # --- Initialize Agent ---
 try:
     agent_executor = get_agent_executor()
-except ValueError as e: # Handles missing API key
+except ValueError as e: # Handles missing API key or other init issues
     st.error(f"Could not initialize AI Agent: {e}")
+    app_logger.critical(f"AI Agent initialization failed: {e}", exc_info=True)
+    st.stop()
+except Exception as e:
+    st.error(f"An unexpected error occurred while initializing the AI Agent: {e}")
+    app_logger.critical(f"Unexpected AI Agent initialization error: {e}", exc_info=True)
     st.stop()
 
 
 # --- Helper Functions ---
-def load_chat_history(session_id: int) -> list:
-    """Loads chat history from DB for the current session"""
+@st.cache_data(ttl=60) # Short cache for chat history to avoid constant DB hits on reruns
+def load_chat_history_for_agent(session_id: int) -> List: # Type hint for return
+    """Loads chat history from DB for the current session, formatted for LangChain agent."""
     messages = []
+    app_logger.debug(f"Loading agent chat history for session_id: {session_id}")
     with get_session_context() as db:
+        # If using SQLModel: from sqlmodel import select
+        # db_messages = db.exec(select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp)).all()
         db_messages = db.query(ChatMessage).filter(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp).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))
-            # Add tool message handling if you store them as distinct roles in DB
-            # elif msg.role == "tool":
-            #     messages.append(ToolMessage(content=msg.content, tool_call_id=msg.tool_call_id))
+            elif msg.role == "tool" and hasattr(msg, 'tool_call_id') and msg.tool_call_id: # Ensure tool_call_id exists
+                 messages.append(ToolMessage(content=msg.content, tool_call_id=str(msg.tool_call_id))) # Cast to str just in case
+            # Add other roles if necessary
+    app_logger.debug(f"Loaded {len(messages)} messages for agent history for session {session_id}.")
     return messages
 
-def save_chat_message(session_id: int, role: str, content: str, tool_call_id: Optional[str]=None, tool_name: Optional[str]=None):
+def save_chat_message_to_db(session_id: int, role: str, content: str, tool_call_id: Optional[str]=None, tool_name: Optional[str]=None):
     """Saves a chat message to the database."""
+    app_logger.debug(f"Saving message to DB for session {session_id}: Role={role}, Content snippet='{content[:50]}...'")
     with get_session_context() as db:
         chat_message = ChatMessage(
             session_id=session_id,
@@ -52,78 +81,85 @@ def save_chat_message(session_id: int, role: str, content: str, tool_call_id: Op
         )
         db.add(chat_message)
         db.commit()
+        app_logger.info(f"Message saved to DB for session {session_id}. Role: {role}.")
 
 # --- Page Logic ---
 st.title("AI Consultation Room")
-st.markdown("Interact with the Quantum Health Navigator AI.")
+st.markdown(f"Interacting as: **{authenticated_username}**")
 
-current_user: User = st.session_state.authenticated_user
 chat_session_id = st.session_state.get("current_chat_session_id")
 
 if not chat_session_id:
-    st.error("No active chat session. Please re-login or contact support.")
+    st.error("No active chat session ID found in session state. This might happen if you logged in before this feature was fully active. Please try logging out and logging back in.")
+    app_logger.error(f"User {authenticated_username} (ID: {authenticated_user_id}) on Consult page with no current_chat_session_id.")
     st.stop()
 
-# Load initial chat history for the agent (from Langchain Message objects)
-# For the agent, we need history in LangChain message format
-if "agent_chat_history" not in st.session_state:
-    st.session_state.agent_chat_history = load_chat_history(chat_session_id)
-    if not st.session_state.agent_chat_history: # If no history, maybe add a system greeting
-        log_consultation_start()
-        # You could add an initial AIMessage here if desired
-        # initial_ai_message = AIMessage(content="Hello! How can I assist you today?")
-        # st.session_state.agent_chat_history.append(initial_ai_message)
-        # save_chat_message(chat_session_id, "assistant", initial_ai_message.content)
-
-
-# Display chat messages from DB (for UI)
-with get_session_context() as db:
-    ui_messages = db.query(ChatMessage).filter(ChatMessage.session_id == chat_session_id).order_by(ChatMessage.timestamp).all()
-    for msg in ui_messages:
-        with st.chat_message(msg.role):
-            st.markdown(msg.content)
+# Initialize agent's chat history if not already present for this session_id
+# We use a more specific key for agent_chat_history to handle session changes
+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 no history, maybe add a system greeting
+        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}")
+        initial_ai_message_content = "Hello! I am your AI Health Navigator. How can I assist you today?"
+        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 for session {chat_session_id} with a greeting.")
+
+
+# Display chat messages for UI (always fetch fresh from DB for UI consistency)
+# This ensures UI reflects what's actually in the DB.
+with st.container(): # Use a container for the chat display area
+    with get_session_context() as db:
+        # If using SQLModel: from sqlmodel import select
+        # ui_messages = db.exec(select(ChatMessage).where(ChatMessage.session_id == chat_session_id).order_by(ChatMessage.timestamp)).all()
+        ui_messages = db.query(ChatMessage).filter(ChatMessage.session_id == chat_session_id).order_by(ChatMessage.timestamp).all()
+        for msg in ui_messages:
+            avatar = "🧑‍⚕️" if msg.role == "assistant" else "👤"
+            if msg.role == "tool": avatar = "🛠️"
+
+            with st.chat_message(msg.role, avatar=avatar):
+                st.markdown(msg.content)
 
 # Chat input
-if prompt := st.chat_input("Ask the AI... (e.g., 'What is hypertension?' or 'Optimize treatment for patient X with diabetes')"):
-    # Add user message to UI and save to DB
-    with st.chat_message("user"):
+if prompt := st.chat_input("Ask the AI... (e.g., 'What is hypertension?' or 'Suggest diagnostic tests for chest pain')"):
+    # Add user message to UI immediately (optimistic update)
+    with st.chat_message("user", avatar="👤"):
         st.markdown(prompt)
-    save_chat_message(chat_session_id, "user", prompt)
-    
+    # Save user message to DB
+    save_chat_message_to_db(chat_session_id, "user", prompt)
     # Add to agent's history (LangChain format)
-    st.session_state.agent_chat_history.append(HumanMessage(content=prompt))
+    st.session_state[agent_history_key].append(HumanMessage(content=prompt))
 
     # Get AI response
-    with st.spinner("AI is thinking..."):
-        try:
-            response = agent_executor.invoke({
-                "input": prompt,
-                "chat_history": st.session_state.agent_chat_history
-            })
-            ai_response_content = response['output']
-            
-            # Display AI response in UI and save to DB
-            with st.chat_message("assistant"):
-                st.markdown(ai_response_content)
-            save_chat_message(chat_session_id, "assistant", ai_response_content)
-            
-            # Add AI response to agent's history
-            st.session_state.agent_chat_history.append(AIMessage(content=ai_response_content))
-
-            # Note: The agent executor might make tool calls. The create_openai_functions_agent
-            # and AgentExecutor handle the tool invocation and adding ToolMessages to history internally
-            # before producing the final 'output'. If you need to log individual tool calls/results
-            # to your DB, you might need a more custom agent loop or callbacks.
-
-        except Exception as e:
-            app_logger.error(f"Error during agent invocation: {e}")
-            st.error(f"An error occurred: {e}")
-            # Save error message as AI response?
-            error_message = f"Sorry, I encountered an error: {str(e)[:200]}" # Truncate for DB
-            with st.chat_message("assistant"): # Or a custom error role
-                st.markdown(error_message)
-            save_chat_message(chat_session_id, "assistant", error_message) # Or "error" role
-            st.session_state.agent_chat_history.append(AIMessage(content=error_message))
-    
-    # Rerun to show the latest messages immediately (though Streamlit usually does this)
-    # st.rerun() # Usually not needed with st.chat_input and context managers
+    with st.chat_message("assistant", avatar="🧑‍⚕️"): # Prepare AI's chat message bubble
+        with st.spinner("AI is thinking..."):
+            try:
+                response = agent_executor.invoke({
+                    "input": prompt,
+                    "chat_history": st.session_state[agent_history_key] # Pass the current agent history
+                })
+                ai_response_content = response.get('output', "No output from AI.")
+                if not isinstance(ai_response_content, str): # Ensure it's a string
+                    ai_response_content = str(ai_response_content)
+
+                st.markdown(ai_response_content) # Display AI response in UI
+                save_chat_message_to_db(chat_session_id, "assistant", ai_response_content) # Save to DB
+                st.session_state[agent_history_key].append(AIMessage(content=ai_response_content)) # Add to agent's history
+
+            except Exception as e:
+                app_logger.error(f"Error during agent invocation for session {chat_session_id}: {e}", exc_info=True)
+                error_message_user = f"Sorry, I encountered an error and could not process your request. Please try again or rephrase. (Error: {type(e).__name__})"
+                st.error(error_message_user) # Display error in the AI's bubble
+                # Save a generic error message to DB for the assistant's turn
+                save_chat_message_to_db(chat_session_id, "assistant", f"Error processing request: {type(e).__name__}")
+                # Add error representation to agent history so it's aware
+                st.session_state[agent_history_key].append(AIMessage(content=f"Observed internal error: {type(e).__name__}"))
+
+    # A full st.rerun() can be a bit disruptive if not needed.
+    # Streamlit's chat_input and context managers usually handle updates well.
+    # If messages aren't appearing correctly, a targeted rerun might be useful.
+    # st.rerun()