Update pages/2_Consult.py

pages/2_Consult.py

@@ -1,87 +1,61 @@
 # /home/user/app/pages/2_Consult.py
 import streamlit as st
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, ToolMessage
+from langchain_core.messages import HumanMessage, AIMessage # SystemMessage, ToolMessage removed if not used directly
 from datetime import datetime
-from typing import List, Optional # Corrected import for List and Optional
+from typing import List, Optional
+from sqlmodel import select # <--- IMPORT SELECT FOR SQLMODEL QUERIES
 
 from config.settings import settings
 from agent import get_agent_executor
-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 models import ChatMessage, ChatSession, User # User not directly used if ID is sufficient
+from models.db import get_session_context
 from services.logger import app_logger
-from services.metrics import log_consultation_start # Assuming this function exists
+from services.metrics import log_consultation_start
 
-# Page config typically in app.py
-# st.set_page_config(page_title=f"Consult - {settings.APP_TITLE}", layout="wide")
-
-# --- Authentication Check ---
+# --- Auth Check (same as before) ---
 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 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.")
+    except st.errors.StreamlitAPIException: # Catch specific error if needed
+        st.info("Please navigate to the main login page.")
     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 ---
+# --- Initialize Agent (same as before) ---
 try:
     agent_executor = get_agent_executor()
-except ValueError as e: # Handles missing API key or other init issues
+except Exception as e:
     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 ---
-@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."""
+@st.cache_data(ttl=60)
+def load_chat_history_for_agent(session_id: int) -> List:
     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()
+    with get_session_context() as db: # db is a SQLModel Session
+        # --- SQLMODEL QUERY ---
+        statement = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.timestamp)
+        db_messages_results = db.exec(statement)
+        db_messages = db_messages_results.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 == "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
+            # Add ToolMessage handling if you store and use them
+            # elif msg.role == "tool" and hasattr(msg, 'tool_call_id') and msg.tool_call_id:
+            #     messages.append(ToolMessage(content=msg.content, tool_call_id=str(msg.tool_call_id)))
     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):
-    """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,
-            role=role,
-            content=content,
-            timestamp=datetime.utcnow(),
-            tool_call_id=tool_call_id,
-            tool_name=tool_name
-        )
-        db.add(chat_message)
-        db.commit()
-        app_logger.info(f"Message saved to DB for session {session_id}. Role: {role}.")
+# save_chat_message_to_db remains the same as it's performing an insert, not a query.
 
 # --- Page Logic ---
 st.title("AI Consultation Room")
@@ -90,76 +64,84 @@ st.markdown(f"Interacting as: **{authenticated_username}**")
 chat_session_id = st.session_state.get("current_chat_session_id")
 
 if not chat_session_id:
-    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.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.stop()
 
-# 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
+    if not st.session_state[agent_history_key]:
         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)
+        # save_chat_message_to_db(chat_session_id, "assistant", initial_ai_message_content) # save_chat_message_to_db defined elsewhere
+        # Ensure save_chat_message_to_db is defined or called correctly. For this example, it's:
+        with get_session_context() as db:
+            chat_msg_obj = ChatMessage(session_id=chat_session_id, role="assistant", content=initial_ai_message_content)
+            db.add(chat_msg_obj) # commit handled by context manager
         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()
+# Display chat messages for UI
+with st.container():
+    with get_session_context() as db: # db is a SQLModel Session
+        # --- SQLMODEL QUERY ---
+        statement = select(ChatMessage).where(ChatMessage.session_id == chat_session_id).order_by(ChatMessage.timestamp)
+        ui_messages_results = db.exec(statement)
+        ui_messages = ui_messages_results.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 'Suggest diagnostic tests for chest pain')"):
-    # Add user message to UI immediately (optimistic update)
+# Chat input and AI response logic (remains largely the same as it calls agent_executor and save_chat_message_to_db)
+# ... (rest of 2_Consult.py from the previous good version)
+# For brevity, I'm omitting the chat input and AI response handling section as it primarily
+# uses `agent_executor.invoke` and `save_chat_message_to_db`, which itself doesn't change syntax for inserts.
+
+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}, Content snippet='{content[:50]}...'")
+    with get_session_context() as db:
+        chat_message = 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)
+        # db.commit() # Handled by context manager
+        app_logger.info(f"Message saved to DB for session {session_id}. Role: {role}.")
+
+if prompt := st.chat_input("Ask the AI..."):
     with st.chat_message("user", avatar="👤"):
         st.markdown(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_history_key].append(HumanMessage(content=prompt))
 
-    # Get AI response
-    with st.chat_message("assistant", avatar="🧑‍⚕️"): # Prepare AI's chat message bubble
+    with st.chat_message("assistant", avatar="🧑‍⚕️"):
         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
+                    "chat_history": st.session_state[agent_history_key]
                 })
                 ai_response_content = response.get('output', "No output from AI.")
-                if not isinstance(ai_response_content, str): # Ensure it's a string
+                if not isinstance(ai_response_content, str):
                     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
-
+                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_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
+                error_message_user = f"Sorry, I encountered an error: {type(e).__name__}"
+                st.error(error_message_user)
                 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()
+                st.session_state[agent_history_key].append(AIMessage(content=f"Observed internal error: {type(e).__name__}"))