Update pages/2_Consult.py

pages/2_Consult.py

@@ -1,12 +1,11 @@
-# /home/user/app/pages/2_Consult.py
 import streamlit as st
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage # Ensure all are imported
+from langchain.schema import HumanMessage, AIMessage, SystemMessage
 from datetime import datetime
-from typing import List, Optional, Dict, Any
+from typing import List, Union, Optional, Dict, Any
 from sqlmodel import select
 
 from config.settings import settings
-from agent import get_agent_executor # This now returns the Gemini-based agent
+from agent import get_agent_executor  # Returns the configured Gemini-based agent executor
 from models import ChatMessage, ChatSession
 from models.db import get_session_context
 from services.logger import app_logger
@@ -21,181 +20,182 @@ if not st.session_state.get("authenticated_user_id"):
         st.info("Please navigate to the main login page.")
     st.stop()
 
-authenticated_user_id = st.session_state.get("authenticated_user_id")
+authenticated_user_id = st.session_state["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() # Gets the Gemini agent executor
+    agent_executor = get_agent_executor()
     app_logger.info("Gemini-based agent executor initialized for Consult page.")
 except ValueError as e:
     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)
+    app_logger.critical(f"Fatal: AI Agent could not be initialized: {e}", exc_info=True)
     st.info("Please ensure API keys (e.g., Google API Key for Gemini) are configured.")
     st.stop()
-except Exception as e: # Catch any other unexpected error during agent init
-    st.error(f"An unexpected error occurred while initializing the AI Agent: {e}")
+except Exception as e:
+    st.error(f"Unexpected error initializing 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_dict' not in st.session_state: # Renamed for clarity
+# --- Session State Setup ---
+if 'current_consult_patient_context_dict' not in st.session_state:
     st.session_state.current_consult_patient_context_dict = {}
 if 'consult_context_submitted' not in st.session_state:
     st.session_state.consult_context_submitted = False
 
-# --- Helper Functions ---
+# --- Caching helper to load history ---
 @st.cache_data(ttl=30, show_spinner=False)
-def load_chat_history_for_agent(session_id: int) -> List:
-    messages = []
-    app_logger.debug(f"Loading agent chat history for session_id: {session_id}")
+def load_chat_history_for_agent(session_id: int) -> List[Union[HumanMessage, AIMessage, SystemMessage]]:
+    messages: List[Union[HumanMessage, AIMessage, SystemMessage]] = []
+    app_logger.debug(f"Loading agent chat history for session {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))
-            elif msg.role == "system": messages.append(SystemMessage(content=msg.content))
-    app_logger.debug(f"Loaded {len(messages)} LangChain messages for agent history (session {session_id}).")
+        for msg in db.exec(statement).all():
+            if msg.role == "user":
+                messages.append(HumanMessage(content=msg.content))
+            elif msg.role == "assistant":
+                messages.append(AIMessage(content=msg.content))
+            elif msg.role == "system":
+                messages.append(SystemMessage(content=msg.content))
+    app_logger.debug(f"Loaded {len(messages)} messages 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):
+# --- Database write helpers ---
+def save_chat_message_to_db(
+    session_id: int,
+    role: str,
+    content: str,
+    tool_call_id: Optional[str] = None,
+    tool_name: Optional[str] = None,
+):
     app_logger.debug(f"Saving message to DB for session {session_id}: Role={role}")
     with get_session_context() as db:
-        chat_message_obj = ChatMessage(
-            session_id=session_id, role=role, content=content, timestamp=datetime.utcnow(),
-            tool_call_id=tool_call_id, tool_name=tool_name
+        obj = ChatMessage(
+            session_id=session_id,
+            role=role,
+            content=content,
+            timestamp=datetime.utcnow(),
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
         )
-        db.add(chat_message_obj)
+        db.add(obj)
     app_logger.info(f"Message saved to DB for session {session_id}. Role: {role}.")
 
-def update_chat_session_with_context_summary_in_db(session_id: int, context_summary: str):
+
+def update_chat_session_with_context_summary_in_db(session_id: int, summary: str):
     with get_session_context() as db:
-        session_to_update = db.get(ChatSession, session_id)
-        if session_to_update:
-            session_to_update.patient_context_summary = context_summary
-            db.add(session_to_update)
-            app_logger.info(f"Updated ChatSession {session_id} with patient context summary in DB.")
+        session = db.get(ChatSession, session_id)
+        if session:
+            session.patient_context_summary = summary
+            db.add(session)
+            app_logger.info(f"Updated ChatSession {session_id} with context summary.")
         else:
-            app_logger.error(f"Could not find ChatSession {session_id} to update context summary.")
+            app_logger.error(f"ChatSession {session_id} not found for updating context.")
 
-# --- Page Logic ---
+# --- Page Title & Disclaimer ---
 st.title("AI Consultation Room")
 st.markdown(f"Interacting as: **{authenticated_username}**")
-st.info(f"{settings.MAIN_DISCLAIMER_SHORT} Remember to use only anonymized, simulated data for this demonstration.")
+st.info(f"{settings.MAIN_DISCLAIMER_SHORT} Remember to use only anonymized, simulated data.")
 
+# --- Retrieve or Create Chat Session ---
 chat_session_id = st.session_state.get("current_chat_session_id")
 if not chat_session_id:
-    st.error("No active chat session ID found. Please try logging out and back in.")
-    app_logger.error(f"User '{authenticated_username}' on Consult page with NO current_chat_session_id.")
+    st.error("No active chat session found. Please logout and login again.")
+    app_logger.error("No current_chat_session_id in session_state.")
     st.stop()
 
-# --- Patient Context Input Form ---
+# --- Patient Context Form ---
 if not st.session_state.consult_context_submitted:
-    st.subheader("Step 1: Provide Patient Context (Optional, Simulated Data Only)")
-    with st.form(key="patient_context_form_consult_page"):
-        st.markdown("**Reminder: Use only anonymized, simulated data for this demonstration.**")
-        age_in = st.number_input("Patient Age (Simulated)", min_value=0, max_value=120, step=1, value=None)
-        gender_in = st.selectbox("Patient Gender (Simulated)", ["Not Specified", "Male", "Female", "Other"], index=0)
-        cc_in = st.text_area("Chief Complaint / Reason for Consult (Simulated)", height=100, placeholder="e.g., Persistent cough")
-        hist_in = st.text_area("Key Medical History (Simulated)", height=100, placeholder="e.g., Type 2 Diabetes")
-        meds_in = st.text_area("Current Medications (Simulated)", height=100, placeholder="e.g., Metformin")
-        submit_context_btn = st.form_submit_button("Start Consult with this Context")
-
-        if submit_context_btn:
-            raw_context = {
-                "age": age_in, "gender": gender_in, "chief_complaint": cc_in,
-                "key_medical_history": hist_in, "current_medications": meds_in,
-            }
-            # Filter out None/empty/"Not Specified" values for cleaner context dictionary
-            filtered_context_dict = {
-                k.replace('_', ' ').title(): v for k, v in raw_context.items()
-                if v is not None and str(v).strip() and str(v) != "Not Specified" and (isinstance(v, int) and v > 0 or isinstance(v, str)) # ensure age is >0
-            }
-            st.session_state.current_consult_patient_context_dict = filtered_context_dict
-
-            if filtered_context_dict:
-                context_summary_str = "; ".join([f"{k}: {v}" for k, v in filtered_context_dict.items()])
-            else:
-                context_summary_str = "No specific patient context was provided for this session."
-            
-            update_chat_session_with_context_summary_in_db(chat_session_id, context_summary_str)
-            save_chat_message_to_db(chat_session_id, "system", f"Initial Patient Context Provided: {context_summary_str}")
-            
-            st.session_state.consult_context_submitted = True
-            app_logger.info(f"Patient context submitted for session {chat_session_id}: {context_summary_str}")
-            st.rerun()
+    st.subheader("Step 1: Provide Patient Context (Optional)")
+    with st.form("patient_context_form"):
+        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"], index=0)
+        complaint = st.text_area("Chief Complaint", height=80, placeholder="e.g., Persistent cough")
+        history = st.text_area("Key Medical History", height=80, placeholder="e.g., Type 2 Diabetes")
+        meds = st.text_area("Current Medications", height=80, placeholder="e.g., Metformin")
+        submit = st.form_submit_button("Start Consult")
+
+    if submit:
+        raw = {"age": age, "gender": gender, "complaint": complaint, "history": history, "meds": meds}
+        filtered = {k.title(): v for k, v in raw.items() if v and not (isinstance(v, str) and v == "Not Specified")}
+        st.session_state.current_consult_patient_context_dict = filtered
+
+        summary = (
+            "; ".join(f"{k}: {v}" for k, v in filtered.items())
+            if filtered else
+            "No specific patient context provided."
+        )
+        update_chat_session_with_context_summary_in_db(chat_session_id, summary)
+        save_chat_message_to_db(chat_session_id, "system", f"Context: {summary}")
+        st.session_state.consult_context_submitted = True
+        app_logger.info(f"Patient context for session {chat_session_id}: {summary}")
+        st.rerun()
     st.stop()
 
-# --- Chat Interface ---
+# --- Chat Interaction ---
 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 after loading
-        try: log_consultation_start(user_id=authenticated_user_id, session_id=chat_session_id)
-        except Exception as e_metric: app_logger.warning(f"Failed to log consultation start metric: {e_metric}")
-        
-        initial_ai_msg = "Hello! I am your AI Health Navigator. How can I assist you today?"
-        if st.session_state.get('current_consult_patient_context_dict'): # Check the renamed key
-            initial_ai_msg += " I have noted the patient context you provided."
-        
-        st.session_state[agent_history_key].append(AIMessage(content=initial_ai_msg))
-        save_chat_message_to_db(chat_session_id, "assistant", initial_ai_msg)
-
-# Display chat messages from DB for UI
-with st.container(height=400): # Fixed height container for chat, makes it scrollable
+history_key = f"agent_history_{chat_session_id}"
+if history_key not in st.session_state:
+    st.session_state[history_key] = load_chat_history_for_agent(chat_session_id)
+    if not st.session_state[history_key]:
+        try:
+            log_consultation_start(authenticated_user_id, chat_session_id)
+        except Exception as e:
+            app_logger.warning(f"Metrics log failed: {e}")
+        greeting = "Hello! I am your AI Health Navigator. How can I help today?"
+        if st.session_state.current_consult_patient_context_dict:
+            greeting += " I have noted the patient context provided."
+        st.session_state[history_key].append(AIMessage(content=greeting))
+        save_chat_message_to_db(chat_session_id, "assistant", greeting)
+
+# Display past messages
+with st.container(height=400):
     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": continue
+        for msg in db.exec(stmt).all():
+            if msg.role == "system":
+                continue
             avatar = "🧑‍⚕️" if msg.role == "assistant" else "👤"
-            if msg.role == "tool": avatar = "🛠️"
+            if msg.role == "tool":
+                avatar = "🛠️"
             with st.chat_message(msg.role, avatar=avatar):
                 st.markdown(msg.content)
 
-# Chat input
-if user_prompt := st.chat_input("Ask the AI... (e.g., 'What is hypertension?')"):
-    with st.chat_message("user", avatar="👤"): st.markdown(user_prompt)
-    save_chat_message_to_db(chat_session_id, "user", user_prompt)
-    st.session_state[agent_history_key].append(HumanMessage(content=user_prompt))
+# Capture user input and invoke agent
+if prompt := st.chat_input("Ask the AI... (e.g., 'What is hypertension?')"):
+    with st.chat_message("user", avatar="👤"):
+        st.markdown(prompt)
+    save_chat_message_to_db(chat_session_id, "user", prompt)
+    st.session_state[history_key].append(HumanMessage(content=prompt))
 
     with st.chat_message("assistant", avatar="🧑‍⚕️"):
         with st.spinner("AI is thinking..."):
             try:
-                # Prepare patient context string for the agent
-                patient_context_dict = st.session_state.get('current_consult_patient_context_dict', {})
-                if patient_context_dict:
-                    context_parts_for_invoke = [f"{k}: {v}" for k, v in patient_context_dict.items()]
-                    patient_context_str_for_invoke = "; ".join(context_parts_for_invoke)
+                pctxt = st.session_state.current_consult_patient_context_dict
+                if pctxt:
+                    pc_str = "; ".join(f"{k}: {v}" for k, v in pctxt.items())
                 else:
-                    patient_context_str_for_invoke = "No specific patient context was provided for this interaction."
+                    pc_str = "No specific patient context provided."
 
-                invoke_payload = {
-                    "input": user_prompt,
-                    "chat_history": st.session_state[agent_history_key], # List of BaseMessage
-                    "patient_context": patient_context_str_for_invoke
+                payload = {
+                    "input": prompt,
+                    "chat_history": st.session_state[history_key],
+                    "patient_context": pc_str,
                 }
-                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 at this time.")
-                if not isinstance(ai_response_content, str): ai_response_content = str(ai_response_content)
-                
-                app_logger.info(f"Agent response for session {chat_session_id}: '{ai_response_content[:100]}...'")
-                st.markdown(ai_response_content)
-                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_type_name = type(e).__name__
-                user_friendly_error = 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)
-                db_error_msg = f"System encountered an error: {error_type_name}. Details logged."
-                save_chat_message_to_db(chat_session_id, "assistant", db_error_msg)
-                st.session_state[agent_history_key].append(AIMessage(content=f"Note: Encountered error ({error_type_name})."))
-    st.rerun() # Rerun to display the new messages immediately
+                app_logger.debug(f"Invoking agent with payload: {payload}")
+                resp = agent_executor.invoke(payload)
+                content = resp.get("output", "I could not generate a response.")
+                content = str(content)
+
+                app_logger.info(f"Agent response session {chat_session_id}: {content[:100]}")
+                st.markdown(content)
+                save_chat_message_to_db(chat_session_id, "assistant", content)
+                st.session_state[history_key].append(AIMessage(content=content))
+            except Exception as err:
+                app_logger.error(f"Error during agent invocation for session {chat_session_id}: {err}", exc_info=True)
+                st.error(f"Sorry, an error occurred: {type(err).__name__}. Please try again.")
+                err_msg = f"System encountered an error: {type(err).__name__}."
+                save_chat_message_to_db(chat_session_id, "assistant", err_msg)
+                st.session_state[history_key].append(AIMessage(content=err_msg))
+    st.rerun()