Update pages/3_Reports.py

pages/3_Reports.py

@@ -1,98 +1,103 @@
 # /home/user/app/pages/3_Reports.py
 import streamlit as st
 from datetime import datetime
-from typing import List
-from sqlmodel import select # For SQLModel queries
+from typing import List, Dict, Any # For type hinting the dictionary
+from sqlmodel import select
 
 from config.settings import settings
-from models import ChatMessage, ChatSession # User model not directly needed here if ID is sufficient
-from models.db import get_session_context # SQLModel session context
+from models import ChatMessage, ChatSession # ChatSession model still needed 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 # Assuming this function exists and is set up
-
-# --- Page Configuration (Typically set in main app.py) ---
-# st.set_page_config(page_title=f"Reports - {settings.APP_TITLE}", layout="wide")
+from services.metrics import log_report_generated
 
 # --- Authentication Check ---
 if not st.session_state.get("authenticated_user_id"):
     st.warning("Please log in to access reports.")
     try:
-        st.switch_page("app.py") # Redirect to the main login page
-    except st.errors.StreamlitAPIException as e:
-        # Handle cases where switch_page might not work (e.g., not in MPA mode during dev)
-        app_logger.warning(f"Reports: Could not switch page (maybe not in MPA mode or other issue): {e}")
+        st.switch_page("app.py")
+    except st.errors.StreamlitAPIException:
         st.info("Please navigate to the main login page manually.")
-    st.stop() # Halt script execution for unauthenticated users
+    st.stop()
 
-# --- Get Authenticated User Info ---
 authenticated_user_id = st.session_state.get("authenticated_user_id")
-authenticated_username = st.session_state.get("authenticated_username", "User") # Default to "User"
+authenticated_username = st.session_state.get("authenticated_username", "User")
 app_logger.info(f"User '{authenticated_username}' (ID: {authenticated_user_id}) accessed Reports page.")
 
-# --- Page Title and Description ---
 st.title("Consultation Reports")
 st.markdown("View and download your past consultation sessions.")
 
-# --- Helper Function to Load User's Chat Sessions ---
-# @st.cache_data(ttl=300) # Cache for 5 minutes. Invalidate if session titles change or new sessions added frequently.
-# For caching to work correctly, arguments to the cached function must be hashable. user_id (int) is.
-def get_user_chat_sessions(user_id: int) -> List[ChatSession]:
-    """Fetches all chat sessions for a given user ID, ordered by start time descending."""
-    app_logger.debug(f"Fetching chat sessions for user_id: {user_id}")
-    sessions: List[ChatSession] = []
+# --- 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
+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: # db_session is a SQLModel Session
-            statement = select(ChatSession).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
-            results = db_session.exec(statement)
-            sessions = results.all()
-        app_logger.debug(f"Found {len(sessions)} sessions for user_id: {user_id}")
+        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
+
+            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
+                })
+        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 sessions for user_id {user_id}: {e}", exc_info=True)
+        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.")
-    return sessions
+    return session_data_list
 
-# Fetch the chat sessions for the current authenticated user
-chat_sessions = get_user_chat_sessions(authenticated_user_id)
+# Fetch session display data for the current authenticated user
+chat_session_display_items = get_user_chat_session_display_data(authenticated_user_id)
 
-# --- Display Logic ---
-if not chat_sessions:
+if not chat_session_display_items:
     st.info("You have no past consultation sessions recorded.")
-    st.stop() # No further processing if no sessions
+    st.stop()
 
+# --- Display Sessions and Download Option ---
 # Prepare options for the selectbox: (Display String, Session ID)
 session_options = []
-for s in chat_sessions:
-    start_time_display = s.start_time.strftime('%Y-%m-%d %H:%M') if s.start_time else "Date N/A"
-    # Use a more descriptive title if the session title is None
-    title_display = s.title if s.title else f"Consultation on {start_time_display}"
-    display_string = f"ID: {s.id} | Started: {start_time_display} | Title: {title_display}"
-    session_options.append((display_string, s.id)) # Store (display_text, actual_id)
-
-# Let the user select a session
+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}"
+    display_string = f"ID: {s_data['id']} | Started: {start_time_display} | Title: {title_display}"
+    session_options.append((display_string, s_data['id']))
+
 selected_option_tuple = st.selectbox(
     "Select a Consultation Session:",
     options=session_options,
-    format_func=lambda x: x[0], # Show only the display string in the selectbox
-    index=0 # Default to the first (most recent) session
+    format_func=lambda x: x[0],
+    index=0
 )
 
 if selected_option_tuple:
-    selected_session_id = selected_option_tuple[1] # Extract the session ID
+    selected_session_id = selected_option_tuple[1]
     app_logger.info(f"User '{authenticated_username}' selected session ID: {selected_session_id} for report viewing.")
 
-    # Find the full ChatSession object from the already fetched list
-    selected_session_object = next((s for s in chat_sessions if s.id == selected_session_id), None)
+    # Find the selected session's data from the list of dictionaries
+    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_object:
-        st.markdown("---") # Visual separator
-        st.subheader(f"Details for Session ID: {selected_session_object.id}")
+    if selected_session_data:
+        st.markdown("---")
+        st.subheader(f"Details for Session ID: {selected_session_data['id']}")
 
-        start_time_detail = selected_session_object.start_time.strftime('%Y-%m-%d %H:%M:%S UTC') if selected_session_object.start_time else "Not recorded"
-        st.write(f"**Started:** {start_time_detail}")
-        st.write(f"**Title:** {selected_session_object.title or 'Untitled Session'}")
+        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 for the selected session
+        # Fetch messages ONLY for the selected session when needed
         messages_for_report: List[ChatMessage] = []
         try:
             with get_session_context() as db_session:
@@ -106,35 +111,30 @@ if selected_option_tuple:
         if messages_for_report:
             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"
-
-                    # Displaying role and timestamp
                     st.markdown(f"**{icon} {msg.role.capitalize()}** ({timestamp_display})")
-                    # Displaying content in a blockquote or code block for better readability
-                    st.markdown(f"> ```\n{msg.content}\n```") # Markdown code block
+                    st.markdown(f"> ```\n{msg.content}\n```")
                     if msg_idx < len(messages_for_report) - 1:
-                        st.markdown("---") # Separator between messages
+                        st.markdown("---")
 
-            # PDF Download Button
-            st.markdown("---") # Visual separator before download button
+            st.markdown("---")
             try:
-                # Use the authenticated username as the patient name for the report
                 pdf_bytes = generate_pdf_report(messages_for_report, 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(
                     label="Download Report as PDF",
                     data=pdf_bytes,
                     file_name=pdf_file_name,
                     mime="application/pdf",
-                    key=f"download_pdf_btn_{selected_session_id}", # Unique key for the button
+                    key=f"download_pdf_btn_{selected_session_id}",
                     on_click=log_report_generated,
-                    # Pass necessary arguments to the on_click callback if it needs them
-                    args=(authenticated_user_id, selected_session_id), # Example arguments
+                    args=(authenticated_user_id, selected_session_id),
                     help="Click to download the consultation transcript as a PDF file."
                 )
             except Exception as e:
@@ -143,9 +143,7 @@ if selected_option_tuple:
         else:
             st.info("This session has no messages recorded.")
     else:
-        # This case should ideally not happen if selected_session_id comes from chat_sessions
-        app_logger.error(f"Selected session ID {selected_session_id} not found in fetched sessions for user '{authenticated_username}'.")
+        app_logger.error(f"Selected session ID {selected_session_id} not found in fetched display data for user '{authenticated_username}'.")
         st.error("Selected session could not be found. Please try again.")
 else:
-    # This case handles if session_options is empty, though the earlier check for `if not chat_sessions` should catch it.
     st.info("Please select a session from the dropdown menu to view details.")