Update pages/3_Reports.py

pages/3_Reports.py

@@ -1,17 +1,17 @@
 # /home/user/app/pages/3_Reports.py
 import streamlit as st
 from datetime import datetime
-from typing import List, Dict, Any # For type hinting the dictionary
+from typing import List, Dict, Any
 from sqlmodel import select
 
 from config.settings import settings
-from models import ChatMessage, ChatSession # ChatSession model still needed for query
+from models import ChatMessage, ChatSession # ChatMessage model for query
 from models.db import get_session_context
 from services.pdf_report import generate_pdf_report
 from services.logger import app_logger
 from services.metrics import log_report_generated
 
-# --- Authentication Check ---
+# --- Authentication Check (same as before) ---
 if not st.session_state.get("authenticated_user_id"):
     st.warning("Please log in to access reports.")
     try:
@@ -28,46 +28,36 @@ st.title("Consultation Reports")
 st.markdown("View and download your past consultation sessions.")
 
 # --- Helper Function to Load User's Chat Session Data (Primitives) ---
-# @st.cache_data(ttl=300) # If caching, ensure it works well with this primitive data structure
+# @st.cache_data(ttl=300)
 def get_user_chat_session_display_data(user_id: int) -> List[Dict[str, Any]]:
-    """
-    Fetches essential data for chat sessions for display, not full ORM objects.
-    Returns a list of dictionaries.
-    """
     app_logger.debug(f"Fetching chat session display data for user_id: {user_id}")
     session_data_list: List[Dict[str, Any]] = []
     try:
         with get_session_context() as db_session:
-            # Select specific columns needed for display
             statement = select(ChatSession.id, ChatSession.title, ChatSession.start_time)\
                         .where(ChatSession.user_id == user_id)\
                         .order_by(ChatSession.start_time.desc())
-            results = db_session.exec(statement).all() # results will be list of Rows/Tuples
-
+            results = db_session.exec(statement).all()
             for row in results:
                 session_data_list.append({
                     "id": row.id,
                     "title": row.title,
-                    "start_time": row.start_time # Store the datetime object, format it later
+                    "start_time": row.start_time
                 })
         app_logger.debug(f"Found {len(session_data_list)} session display entries for user_id: {user_id}")
     except Exception as e:
         app_logger.error(f"Error fetching chat session display data for user_id {user_id}: {e}", exc_info=True)
-        st.error("Could not load your chat sessions. Please try again later.")
+        st.error("Could not load your chat sessions.")
     return session_data_list
 
-# Fetch session display data for the current authenticated user
 chat_session_display_items = get_user_chat_session_display_data(authenticated_user_id)
 
 if not chat_session_display_items:
     st.info("You have no past consultation sessions recorded.")
     st.stop()
 
-# --- Display Sessions and Download Option ---
-# Prepare options for the selectbox: (Display String, Session ID)
 session_options = []
 for s_data in chat_session_display_items:
-    # Access data from the dictionary
     start_time_obj = s_data["start_time"]
     start_time_display = start_time_obj.strftime('%Y-%m-%d %H:%M') if start_time_obj else "Date N/A"
     title_display = s_data["title"] if s_data["title"] else f"Consultation on {start_time_display}"
@@ -81,50 +71,97 @@ selected_option_tuple = st.selectbox(
     index=0
 )
 
+# --- Store Fetched Message Data in Session State to Persist Across Reruns ---
+# This key will be specific to the selected session to avoid conflicts if the user selects another session.
 if selected_option_tuple:
     selected_session_id = selected_option_tuple[1]
-    app_logger.info(f"User '{authenticated_username}' selected session ID: {selected_session_id} for report viewing.")
+    MESSAGES_FOR_SESSION_KEY = f"messages_for_session_{selected_session_id}"
+
+    # Fetch messages only if not already in session state for this selected session, or if forced refresh
+    if MESSAGES_FOR_SESSION_KEY not in st.session_state:
+        app_logger.info(f"Fetching messages for session ID: {selected_session_id} as not in session_state.")
+        messages_data_for_display: List[Dict[str, Any]] = []
+        try:
+            with get_session_context() as db_session:
+                # Select specific columns needed
+                statement = select(
+                    ChatMessage.role,
+                    ChatMessage.content,
+                    ChatMessage.timestamp,
+                    ChatMessage.tool_name # If you use it
+                ).where(ChatMessage.session_id == selected_session_id).order_by(ChatMessage.timestamp)
+                
+                results = db_session.exec(statement).all() # List of Row objects
+
+                for msg_row in results:
+                    messages_data_for_display.append({
+                        "role": msg_row.role,
+                        "content": msg_row.content,
+                        "timestamp": msg_row.timestamp, # Keep as datetime for formatting
+                        "tool_name": msg_row.tool_name if hasattr(msg_row, 'tool_name') else None
+                    })
+            st.session_state[MESSAGES_FOR_SESSION_KEY] = messages_data_for_display
+            app_logger.info(f"Fetched and stored {len(messages_data_for_display)} messages for session {selected_session_id} in session_state.")
+        except Exception as e:
+            app_logger.error(f"Error fetching messages for session {selected_session_id}: {e}", exc_info=True)
+            st.error("Could not load messages for this session.")
+            st.session_state[MESSAGES_FOR_SESSION_KEY] = [] # Store empty list on error
+    else:
+        app_logger.debug(f"Using cached messages from session_state for session ID: {selected_session_id}")
+        messages_data_for_display = st.session_state[MESSAGES_FOR_SESSION_KEY]
+
 
-    # Find the selected session's data from the list of dictionaries
+    # --- Display selected session details and messages ---
     selected_session_data = next((s_data for s_data in chat_session_display_items if s_data['id'] == selected_session_id), None)
 
     if selected_session_data:
         st.markdown("---")
         st.subheader(f"Details for Session ID: {selected_session_data['id']}")
-
         start_time_obj_detail = selected_session_data["start_time"]
         start_time_detail_display = start_time_obj_detail.strftime('%Y-%m-%d %H:%M:%S UTC') if start_time_obj_detail else "Not recorded"
         st.write(f"**Started:** {start_time_detail_display}")
         st.write(f"**Title:** {selected_session_data['title'] or 'Untitled Session'}")
 
-        # Fetch messages ONLY for the selected session when needed
-        messages_for_report: List[ChatMessage] = []
-        try:
-            with get_session_context() as db_session:
-                statement = select(ChatMessage).where(ChatMessage.session_id == selected_session_id).order_by(ChatMessage.timestamp)
-                results = db_session.exec(statement)
-                messages_for_report = results.all()
-        except Exception as e:
-            app_logger.error(f"Error fetching messages for session {selected_session_id}: {e}", exc_info=True)
-            st.error("Could not load messages for this session.")
-
-        if messages_for_report:
+        if messages_data_for_display: # Use the data from session_state
             with st.expander("View Chat Transcript", expanded=False):
-                for msg_idx, msg in enumerate(messages_for_report):
-                    # Accessing msg.role, msg.timestamp, msg.content should be fine here
-                    # as `messages_for_report` are newly fetched ORM objects within an active (though soon to close) session.
-                    # The key is that their attributes are accessed while their session is implicitly active or just before it closes.
-                    icon = "🧑‍⚕️" if msg.role == "assistant" else "👤"
-                    if msg.role == "tool": icon = "🛠️"
-                    timestamp_display = msg.timestamp.strftime('%Y-%m-%d %H:%M:%S') if msg.timestamp else "Time N/A"
-                    st.markdown(f"**{icon} {msg.role.capitalize()}** ({timestamp_display})")
-                    st.markdown(f"> ```\n{msg.content}\n```")
-                    if msg_idx < len(messages_for_report) - 1:
+                for msg_idx, msg_data in enumerate(messages_data_for_display): # Iterate over list of dicts
+                    # Access data from the dictionary
+                    role = msg_data["role"]
+                    content = msg_data["content"]
+                    timestamp_obj = msg_data["timestamp"]
+
+                    icon = "🧑‍⚕️" if role == "assistant" else "👤"
+                    if role == "tool": icon = "🛠️"
+                    timestamp_display = timestamp_obj.strftime('%Y-%m-%d %H:%M:%S') if timestamp_obj else "Time N/A"
+                    
+                    st.markdown(f"**{icon} {role.capitalize()}** ({timestamp_display})")
+                    st.markdown(f"> ```\n{content}\n```")
+                    if msg_idx < len(messages_data_for_display) - 1:
                         st.markdown("---")
-
+            
             st.markdown("---")
             try:
-                pdf_bytes = generate_pdf_report(messages_for_report, patient_name=authenticated_username)
+                # For PDF generation, we might need to reconstruct ChatMessage-like objects
+                # if generate_pdf_report strictly expects List[ChatMessage].
+                # Or, modify generate_pdf_report to accept List[Dict].
+                # For now, assuming generate_pdf_report can handle List[Dict] or we adapt.
+                
+                # If generate_pdf_report strictly needs ChatMessage instances:
+                pdf_report_messages = []
+                for msg_d in messages_data_for_display:
+                    # This is a mock ChatMessage, not an ORM instance.
+                    # Ensure your generate_pdf_report can handle this structure.
+                    pdf_report_messages.append(
+                        ChatMessage( # Re-create a temporary (non-ORM) ChatMessage-like object
+                            role=msg_d["role"],
+                            content=msg_d["content"],
+                            timestamp=msg_d["timestamp"], # datetime object
+                            tool_name=msg_d.get("tool_name") # Get tool_name if exists
+                            # session_id and id are not strictly needed if generate_pdf_report doesn't use them.
+                        )
+                    )
+
+                pdf_bytes = generate_pdf_report(pdf_report_messages, patient_name=authenticated_username)
                 current_time_str = datetime.now().strftime('%Y%m%d_%H%M%S')
                 pdf_file_name = f"Consultation_Report_Session{selected_session_id}_{current_time_str}.pdf"
                 st.download_button(