Update pages/2_Consult.py

pages/2_Consult.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
 from datetime import datetime
-from typing import List, Union, Optional, Dict, Any
+from typing import List, Union, Optional
 from sqlmodel import select
 
 from config.settings import settings
@@ -20,56 +20,51 @@ if not st.session_state.get("authenticated_user_id"):
         st.info("Please navigate to the main login page.")
     st.stop()
 
+# Retrieve authenticated user info
 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 ---
+# --- Initialize AI Agent ---
 try:
     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: {e}", exc_info=True)
-    st.info("Please ensure API keys (e.g., Google API Key for Gemini) are configured.")
+    app_logger.critical(f"Fatal: AI Agent init failed: {e}", exc_info=True)
+    st.info("Ensure API keys (e.g., Gemini) are configured in settings or environment.")
     st.stop()
 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)
+    app_logger.critical(f"Fatal: Unexpected AI Agent init error: {e}", exc_info=True)
     st.stop()
 
 # --- 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
+st.session_state.setdefault('current_consult_patient_context', {})
+st.session_state.setdefault('consult_context_submitted', False)
 
-# --- Caching helper to load history ---
 @st.cache_data(ttl=30, show_spinner=False)
-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}")
+def load_chat_history(session_id: int) -> List[Union[HumanMessage, AIMessage, SystemMessage]]:
+    """
+    Load chat history from the database and convert to LangChain messages.
+    """
+    history: List[Union[HumanMessage, AIMessage, SystemMessage]] = []
     with get_session_context() as db:
-        statement = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp)
-        for msg in db.exec(statement).all():
+        stmt = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp)
+        for msg in db.exec(stmt).all():
             if msg.role == "user":
-                messages.append(HumanMessage(content=msg.content))
+                history.append(HumanMessage(content=msg.content))
             elif msg.role == "assistant":
-                messages.append(AIMessage(content=msg.content))
+                history.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
-
-# --- 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}")
+                history.append(SystemMessage(content=msg.content))
+    return history
+
+
+def save_message(session_id: int, role: str, content: str, tool_call_id: Optional[str] = None, tool_name: Optional[str] = None):
+    """
+    Persist a chat message to the database.
+    """
     with get_session_context() as db:
         obj = ChatMessage(
             session_id=session_id,
@@ -80,80 +75,79 @@ def save_chat_message_to_db(
             tool_name=tool_name,
         )
         db.add(obj)
-    app_logger.info(f"Message saved to DB for session {session_id}. Role: {role}.")
+    app_logger.info(f"Saved {role} message to DB for session {session_id}.")
 
 
-def update_chat_session_with_context_summary_in_db(session_id: int, summary: str):
+def update_session_context(session_id: int, summary: str):
+    """
+    Update the ChatSession record with a summary of patient context.
+    """
     with get_session_context() as 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.")
+            app_logger.info(f"Updated session {session_id} with context summary.")
         else:
-            app_logger.error(f"ChatSession {session_id} not found for updating context.")
+            app_logger.error(f"ChatSession {session_id} not found.")
 
-# --- Page Title & Disclaimer ---
+# --- Page Header ---
 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.")
+st.markdown(f"**User:** {authenticated_username}")
+st.info(f"{settings.MAIN_DISCLAIMER_SHORT} Use only anonymized, simulated data.")
 
-# --- Retrieve or Create Chat Session ---
+# --- Ensure Chat Session Exists ---
 chat_session_id = st.session_state.get("current_chat_session_id")
 if not 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.error("No active chat session. Please log out and log back in.")
+    app_logger.error("Missing current_chat_session_id in session_state.")
     st.stop()
 
-# --- Patient Context Form ---
+# --- Step 1: Patient Context Input ---
 if not st.session_state.consult_context_submitted:
     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")
+    with st.form("context_form"):
+        age = st.number_input("Patient Age (Sim)", min_value=0, max_value=120, step=1)
+        gender = st.selectbox("Patient Gender (Sim)", ["Not Specified", "Male", "Female", "Other"], index=0)
+        complaint = st.text_area("Chief Complaint (Sim)", height=80)
+        history = st.text_area("Key Medical History (Sim)", height=80)
+        meds = st.text_area("Current Medications (Sim)", height=80)
         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}")
+        raw = {"Age": age, "Gender": gender, "Complaint": complaint.strip(), "History": history.strip(), "Meds": meds.strip()}
+        filtered = {k: 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 = filtered
+
+        summary = "; ".join(f"{k}: {v}" for k, v in filtered.items()) if filtered else "No patient context provided."
+        update_session_context(chat_session_id, summary)
+        save_message(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()
+        app_logger.info(f"Context for session {chat_session_id}: {summary}")
+        st.experimental_rerun()
     st.stop()
 
-# --- Chat Interaction ---
-st.subheader("Step 2: Interact with AI Health Navigator")
-history_key = f"agent_history_{chat_session_id}"
+# --- Step 2: Chat Interface ---
+st.subheader("Step 2: Interact with the AI")
+history_key = f"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)
+    st.session_state[history_key] = load_chat_history(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."
+        except:
+            pass
+        greeting = "Hello! I'm your AI Health Navigator. How can I assist today?"
+        if st.session_state.current_consult_patient_context:
+            greeting += " (Patient context noted.)"
         st.session_state[history_key].append(AIMessage(content=greeting))
-        save_chat_message_to_db(chat_session_id, "assistant", greeting)
+        save_message(chat_session_id, "assistant", greeting)
 
-# Display past messages
+# Display chat 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)
-        for msg in db.exec(stmt).all():
+        for msg in db.exec(stmt):
             if msg.role == "system":
                 continue
             avatar = "🧑‍⚕️" if msg.role == "assistant" else "👤"
@@ -162,40 +156,38 @@ with st.container(height=400):
             with st.chat_message(msg.role, avatar=avatar):
                 st.markdown(msg.content)
 
-# Capture user input and invoke agent
-if prompt := st.chat_input("Ask the AI... (e.g., 'What is hypertension?')"):
+# Capture user input
+if user_input := st.chat_input("Ask the AI..."):
     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))
+        st.markdown(user_input)
+    save_message(chat_session_id, "user", user_input)
+    st.session_state[history_key].append(HumanMessage(content=user_input))
 
     with st.chat_message("assistant", avatar="🧑‍⚕️"):
         with st.spinner("AI is thinking..."):
             try:
-                pctxt = st.session_state.current_consult_patient_context_dict
-                if pctxt:
-                    pc_str = "; ".join(f"{k}: {v}" for k, v in pctxt.items())
-                else:
-                    pc_str = "No specific patient context provided."
+                # Prepare invocation payload
+                context_str = "; ".join(f"{k}: {v}" for k, v in st.session_state.current_consult_patient_context.items())
+                if not context_str:
+                    context_str = "No patient context provided."
 
                 payload = {
-                    "input": prompt,
+                    "input": user_input,
                     "chat_history": st.session_state[history_key],
-                    "patient_context": pc_str,
+                    "patient_context": context_str,
                 }
-                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()
+                app_logger.debug(f"Invoking agent: {payload}")
+                result = agent_executor.invoke(payload)
+                output = str(result.get("output", "No response generated."))
+
+                app_logger.info(f"Agent output: {output[:100]}")
+                st.markdown(output)
+                save_message(chat_session_id, "assistant", output)
+                st.session_state[history_key].append(AIMessage(content=output))
+            except Exception as e:
+                app_logger.error(f"Error during agent invocation: {e}", exc_info=True)
+                error_msg = f"Sorry, an error occurred ({type(e).__name__})."
+                st.error(error_msg)
+                save_message(chat_session_id, "assistant", error_msg)
+                st.session_state[history_key].append(AIMessage(content=error_msg))
+    st.experimental_rerun()