Update pages/3_Reports.py

pages/3_Reports.py

@@ -2,120 +2,150 @@
 import streamlit as st
 from datetime import datetime
 from typing import List
-from sqlmodel import select # <--- IMPORT SELECT FOR SQLMODEL QUERIES
+from sqlmodel import select # For SQLModel queries
 
 from config.settings import settings
-from models import ChatMessage, ChatSession, User # User not directly used if ID is sufficient
-from models.db import get_session_context
+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 services.pdf_report import generate_pdf_report
 from services.logger import app_logger
-from services.metrics import log_report_generated
+from services.metrics import log_report_generated # Assuming this function exists and is set up
 
-# --- Auth Check (same as before) ---
+# --- Page Configuration (Typically set in main app.py) ---
+# st.set_page_config(page_title=f"Reports - {settings.APP_TITLE}", layout="wide")
+
+# --- 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")
-    except st.errors.StreamlitAPIException:
-        st.info("Please navigate to the main login page.")
-    st.stop()
+        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.info("Please navigate to the main login page manually.")
+    st.stop() # Halt script execution for unauthenticated users
 
+# --- Get Authenticated User Info ---
 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 Reports page.")
+authenticated_username = st.session_state.get("authenticated_username", "User") # Default to "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.")
 
-# --- Load User's Chat Sessions ---
-# @st.cache_data(ttl=300) # Consider re-enabling with SQLModel if args are hashable
+# --- 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}")
-    with get_session_context() as db: # db is a SQLModel Session
-        # --- SQLMODEL QUERY ---
-        statement = select(ChatSession).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
-        results = db.exec(statement)
-        sessions = results.all()
-        # --------------------
+    sessions: List[ChatSession] = []
+    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}")
-        return sessions
+    except Exception as e:
+        app_logger.error(f"Error fetching chat sessions for user_id {user_id}: {e}", exc_info=True)
+        st.error("Could not load your chat sessions. Please try again later.")
+    return sessions
 
+# Fetch the chat sessions for the current authenticated user
 chat_sessions = get_user_chat_sessions(authenticated_user_id)
 
+# --- Display Logic ---
 if not chat_sessions:
-    st.info("You have no past consultation sessions to display.")
-    st.stop()
+    st.info("You have no past consultation sessions recorded.")
+    st.stop() # No further processing if no sessions
 
-# --- Display Sessions and Download Option (largely same, but ensure messages query is updated) ---
+# 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 "N/A"
-    title_display = s.title or f"Session on {start_time_display}"
+    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))
+    session_options.append((display_string, s.id)) # Store (display_text, actual_id)
 
+# Let the user select a session
 selected_option_tuple = st.selectbox(
     "Select a Consultation Session:",
     options=session_options,
-    format_func=lambda x: x[0]
+    format_func=lambda x: x[0], # Show only the display string in the selectbox
+    index=0 # Default to the first (most recent) 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.")
-    selected_session = next((s for s in chat_sessions if s.id == selected_session_id), None)
-
-    if selected_session:
-        st.subheader(f"Details for Session ID: {selected_session.id}")
-        # ... (display session details - same as before)
-
-        with get_session_context() as db: # db is a SQLModel Session
-            # --- SQLMODEL QUERY ---
-            statement = select(ChatMessage).where(ChatMessage.session_id == selected_session.id).order_by(ChatMessage.timestamp)
-            results = db.exec(statement)
-            messages: List[ChatMessage] = results.all()
-            # --------------------
-
-        if messages:
-            # ... (display transcript and PDF download - same as before)
-            # The generate_pdf_report(messages, ...) part doesn't change.
+    selected_session_id = selected_option_tuple[1] # Extract the session ID
+    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)
+
+    if selected_session_object:
+        st.markdown("---") # Visual separator
+        st.subheader(f"Details for Session ID: {selected_session_object.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'}")
+
+        # Fetch messages for the selected session
+        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:
             with st.expander("View Chat Transcript", expanded=False):
-                for msg_idx, msg in enumerate(messages):
+                for msg_idx, msg in enumerate(messages_for_report):
                     icon = "🧑‍⚕️" if msg.role == "assistant" else "👤"
                     if msg.role == "tool": icon = "🛠️"
-                    timestamp_display = msg.timestamp.strftime('%H:%M:%S') if msg.timestamp else "N/A"
-                    st.markdown(f"**{icon} {msg.role.capitalize()} ({timestamp_display}):**")
-                    st.markdown(f"> ```\n{msg.content}\n```") # Using markdown code block for content
-                    if msg_idx < len(messages) - 1:
-                        st.markdown("---")
+                    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
+                    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
             try:
-                pdf_bytes = generate_pdf_report(messages, patient_name=authenticated_username)
-                file_name_timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-                pdf_file_name = f"Consultation_Report_Session{selected_session.id}_{file_name_timestamp}.pdf"
+                # 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_{selected_session.id}",
+                    key=f"download_pdf_btn_{selected_session_id}", # Unique key for the button
                     on_click=log_report_generated,
-                    args=(authenticated_user_id, selected_session_id)
+                    # Pass necessary arguments to the on_click callback if it needs them
+                    args=(authenticated_user_id, selected_session_id), # Example arguments
+                    help="Click to download the consultation transcript as a PDF file."
                 )
             except Exception as e:
-                app_logger.error(f"Error generating PDF for session {selected_session.id} for user {authenticated_username}: {e}", exc_info=True)
-                st.error(f"Could not generate PDF report: {type(e).__name__}")
+                app_logger.error(f"Error generating PDF for session {selected_session_id} (User: '{authenticated_username}'): {e}", exc_info=True)
+                st.error(f"Could not generate PDF report for this session. Error: {type(e).__name__}")
         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}'.")
+        st.error("Selected session could not be found. Please try again.")
 else:
-    st.info("Select a session from the dropdown to view details and download the report.")
-
-# (Ensure the rest of 3_Reports.py from the previous good version is here if anything was omitted for brevity)
-# Example for start_time display in selected_session details:
-if selected_session:
-    # ... other details
-    start_time_detail = selected_session.start_time.strftime('%Y-%m-%d %H:%M:%S UTC') if selected_session.start_time else "Not recorded"
-    st.write(f"**Started:** {start_time_detail}")
-    st.write(f"**Title:** {selected_session.title or 'Untitled Session'}")
-    # ... rest of message display and PDF generation ...
+    # 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.")