Update pages/3_Reports.py

pages/3_Reports.py

@@ -5,143 +5,211 @@ from typing import List, Dict, Any
 from sqlmodel import select
 
 from config.settings import settings
-from models import ChatMessage, ChatSession # Model definitions
+from models import ChatMessage, ChatSession # SQLModel classes for DB query
 from models.db import get_session_context
-from services.pdf_report import generate_pdf_report # Needs to accept more data
+# Import generate_pdf_report AND MockChatMessage from services.pdf_report
+from services.pdf_report import generate_pdf_report, MockChatMessage
 from services.logger import app_logger
-from services.metrics import log_report_generated
+from services.metrics import log_report_generated # Assuming this function exists
 
 # --- 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 manually.")
+    try:
+        st.switch_page("app.py")
+    except st.errors.StreamlitAPIException:
+        st.info("Please navigate to the main login page manually.")
     st.stop()
 
 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.")
 
+# --- Page Title and Info ---
 st.title("Consultation Reports")
 st.markdown("View and download your past consultation sessions.")
-st.info(settings.MAIN_DISCLAIMER_SHORT)
-
+st.info(f"{settings.MAIN_DISCLAIMER_SHORT} PDFs are summaries for review.")
 
+# --- Helper Function to Load User's Chat Session Display Data ---
+# @st.cache_data(ttl=300, show_spinner=False) # Consider caching if fetching is slow and data doesn't change too rapidly
 def get_user_chat_session_display_data(user_id: int) -> List[Dict[str, Any]]:
-    app_logger.debug(f"Fetching session display data for user_id: {user_id}")
+    """Fetches essential data for chat sessions (ID, title, start_time, context) as dictionaries."""
+    app_logger.debug(f"ReportsPage: Fetching session display data for user_id: {user_id}")
     session_data_list: List[Dict[str, Any]] = []
     try:
-        with get_session_context() as db:
-            stmt = select(ChatSession.id, ChatSession.title, ChatSession.start_time, ChatSession.patient_context_summary)\
-                   .where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
-            results = db.exec(stmt).all()
-            for row in results:
+        with get_session_context() as db_session:
+            statement = select(
+                ChatSession.id,
+                ChatSession.title,
+                ChatSession.start_time,
+                ChatSession.patient_context_summary # Fetch this as well
+            ).where(ChatSession.user_id == user_id).order_by(ChatSession.start_time.desc())
+            
+            results = db_session.exec(statement).all() # List of Row objects
+            for row_item in results:
                 session_data_list.append({
-                    "id": row.id, "title": row.title, "start_time": row.start_time,
-                    "patient_context_summary": row.patient_context_summary
+                    "id": row_item.id,
+                    "title": row_item.title,
+                    "start_time": row_item.start_time, # datetime object
+                    "patient_context_summary": row_item.patient_context_summary
                 })
-        app_logger.debug(f"Found {len(session_data_list)} session display entries for user {user_id}")
+        app_logger.debug(f"ReportsPage: Found {len(session_data_list)} session display entries for user {user_id}")
     except Exception as e:
-        app_logger.error(f"Error fetching session display data for user {user_id}: {e}", exc_info=True)
-        st.error("Could not load your chat sessions.")
+        app_logger.error(f"ReportsPage: Error fetching session display data for user {user_id}: {e}", exc_info=True)
+        st.error("Could not load your chat sessions. Please try again later.")
     return session_data_list
 
+# Fetch session display data
 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.info("You have no past consultation sessions recorded to display.")
     st.stop()
 
-session_options = []
-for s_data in chat_session_display_items:
-    start_time_obj = s_data["start_time"]
-    start_time_display = start_time_obj.strftime('%Y-%m-%d %H:%M') if start_time_obj else "N/A"
-    title_display = s_data["title"] if s_data["title"] else f"Consult 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']))
+# --- UI: Select a Session ---
+session_options_for_selectbox = []
+for s_data_item in chat_session_display_items:
+    start_time = s_data_item.get("start_time")
+    start_time_str = start_time.strftime('%Y-%m-%d %H:%M') if start_time else "Date N/A"
+    title_str = s_data_item.get("title") or f"Consultation on {start_time_str}"
+    display_text = f"ID: {s_data_item.get('id')} | Started: {start_time_str} | Title: {title_str}"
+    session_options_for_selectbox.append((display_text, s_data_item.get('id'))) # (Display, Value)
 
 selected_option_tuple = st.selectbox(
-    "Select a Consultation Session:", options=session_options, format_func=lambda x: x[0], index=0
+    "Select a Consultation Session:",
+    options=session_options_for_selectbox,
+    format_func=lambda x_tuple: x_tuple[0], # Show only the display string
+    index=0, # Default to the first (most recent) session
+    key="report_session_selector"
 )
 
+# --- Display Details and PDF Download for Selected Session ---
 if selected_option_tuple:
-    selected_session_id = selected_option_tuple[1]
-    MESSAGES_KEY = f"report_messages_for_session_{selected_session_id}"
-
-    if MESSAGES_KEY not in st.session_state:
-        app_logger.info(f"Fetching messages for report, session ID: {selected_session_id}")
-        msg_data_list: List[Dict[str, Any]] = []
-        try:
-            with get_session_context() as db:
-                stmt = select(ChatMessage.role, ChatMessage.content, ChatMessage.timestamp, ChatMessage.tool_name)\
-                       .where(ChatMessage.session_id == selected_session_id).order_by(ChatMessage.timestamp)
-                results = db.exec(stmt).all()
-                for row in results:
-                    msg_data_list.append({
-                        "role": row.role, "content": row.content, "timestamp": row.timestamp,
-                        "tool_name": getattr(row, 'tool_name', None)
-                    })
-            st.session_state[MESSAGES_KEY] = msg_data_list
-        except Exception as e:
-            app_logger.error(f"Error fetching messages for report (session {selected_session_id}): {e}", exc_info=True)
-            st.error("Could not load messages for this report.")
-            st.session_state[MESSAGES_KEY] = []
-    else:
-        app_logger.debug(f"Using cached messages from session_state for report, session ID: {selected_session_id}")
-        msg_data_list = st.session_state[MESSAGES_KEY]
+    selected_session_id_val = selected_option_tuple[1] # The actual session ID
+    app_logger.info(f"ReportsPage: User '{authenticated_username}' selected session ID: {selected_session_id_val} for viewing.")
 
-    selected_session_data = next((s for s in chat_session_display_items if s['id'] == selected_session_id), None)
+    # Find the full data for the selected session from our fetched list of dicts
+    selected_session_dict_data = next(
+        (s_item for s_item in chat_session_display_items if s_item['id'] == selected_session_id_val), None
+    )
 
-    if selected_session_data:
+    if selected_session_dict_data:
         st.markdown("---")
-        st.subheader(f"Report Preview for Session ID: {selected_session_data['id']}")
-        start_t = selected_session_data["start_time"]
-        st.write(f"**Started:** {start_t.strftime('%Y-%m-%d %H:%M:%S UTC') if start_t else 'N/A'}")
-        st.write(f"**Title:** {selected_session_data['title'] or 'Untitled Session'}")
-        if selected_session_data.get("patient_context_summary"):
-            with st.expander("View Patient Context Provided", expanded=False):
-                st.markdown(selected_session_data["patient_context_summary"])
+        st.subheader(f"Report Preview for Session ID: {selected_session_dict_data['id']}")
+
+        session_start_time = selected_session_dict_data.get("start_time")
+        start_time_display_detail = session_start_time.strftime('%Y-%m-%d %H:%M:%S UTC') if session_start_time else "Not recorded"
+        st.write(f"**Started:** {start_time_display_detail}")
+        st.write(f"**Title:** {selected_session_dict_data.get('title') or 'Untitled Session'}")
+
+        patient_context_from_session = selected_session_dict_data.get("patient_context_summary")
+        if patient_context_from_session:
+            with st.expander("View Patient Context Provided for this Session", expanded=False):
+                st.markdown(patient_context_from_session)
         
-        if msg_data_list:
+        # --- Fetch and Display Messages for the Selected Session ---
+        # Store fetched messages in session_state to avoid re-fetching on every interaction within the page
+        MESSAGES_DATA_SESS_KEY = f"report_page_messages_for_session_{selected_session_id_val}"
+
+        if MESSAGES_DATA_SESS_KEY not in st.session_state:
+            app_logger.info(f"ReportsPage: Fetching DB messages for session ID: {selected_session_id_val}")
+            db_messages_as_dicts: List[Dict[str, Any]] = []
+            try:
+                with get_session_context() as db:
+                    # Fetch specific columns needed for display and PDF
+                    stmt = select(
+                        ChatMessage.role, ChatMessage.content, ChatMessage.timestamp, ChatMessage.tool_name
+                    ).where(ChatMessage.session_id == selected_session_id_val).order_by(ChatMessage.timestamp)
+                    
+                    message_results = db.exec(stmt).all() # List of Row objects
+                    for msg_row in message_results:
+                        db_messages_as_dicts.append({
+                            "role": msg_row.role, "content": msg_row.content,
+                            "timestamp": msg_row.timestamp, # datetime object
+                            "tool_name": getattr(msg_row, 'tool_name', None)
+                        })
+                st.session_state[MESSAGES_DATA_SESS_KEY] = db_messages_as_dicts
+                app_logger.info(f"ReportsPage: Fetched and stored {len(db_messages_as_dicts)} messages for session {selected_session_id_val} in st.session_state.")
+            except Exception as e_msg_fetch:
+                app_logger.error(f"ReportsPage: Error fetching DB messages for session {selected_session_id_val}: {e_msg_fetch}", exc_info=True)
+                st.error("Could not load messages for this session's transcript.")
+                st.session_state[MESSAGES_DATA_SESS_KEY] = [] # Store empty list on error
+        else:
+            app_logger.debug(f"ReportsPage: Using cached messages from st.session_state for session ID: {selected_session_id_val}")
+        
+        messages_to_display_and_pdf = st.session_state[MESSAGES_DATA_SESS_KEY]
+
+        if messages_to_display_and_pdf:
             with st.expander("View Chat Transcript", expanded=False):
-                for idx, msg_d in enumerate(msg_data_list):
-                    if msg_d["role"] == "system": continue # Don't show system context messages here
-                    icon = "🧑‍⚕️" if msg_d["role"] == "assistant" else "👤"
-                    if msg_d["role"] == "tool": icon = "🛠️"
-                    ts_display = msg_d["timestamp"].strftime('%Y-%m-%d %H:%M:%S') if msg_d["timestamp"] else "N/A"
-                    st.markdown(f"**{icon} {msg_d['role'].capitalize()}** ({ts_display})")
-                    st.markdown(f"> ```\n{msg_d['content']}\n```")
-                    if idx < len(msg_data_list) - 1: st.markdown("---")
+                for msg_idx, msg_dict_item in enumerate(messages_to_display_and_pdf):
+                    msg_role = str(msg_dict_item.get("role", "N/A"))
+                    msg_content = str(msg_dict_item.get("content", ""))
+                    msg_timestamp_obj = msg_dict_item.get("timestamp")
+                    
+                    # Skip verbose system messages about context in transcript view
+                    if msg_role == "system" and "Initial Patient Context Set:" in msg_content:
+                        continue
+
+                    icon_char = "🧑‍⚕️" if msg_role == "assistant" else "👤"
+                    if msg_role == "tool": icon_char = "🛠️"
+                    if msg_role == "system": icon_char = "⚙️"
+                    
+                    timestamp_str_display = msg_timestamp_obj.strftime('%Y-%m-%d %H:%M:%S') if msg_timestamp_obj else "N/A"
+                    
+                    st.markdown(f"**{icon_char} {msg_role.capitalize()}** ({timestamp_str_display})")
+                    st.markdown(f"> ```\n{msg_content}\n```") # Using markdown code block for content
+                    if msg_idx < len(messages_to_display_and_pdf) - 1:
+                        st.markdown("---")
             
+            # --- PDF Download Button ---
             st.markdown("---")
             try:
-                # Prepare data for PDF report
-                # This creates temporary ChatMessage-like objects.
-                # Ensure generate_pdf_report can handle these non-ORM objects.
-                pdf_report_messages = [
-                    ChatMessage(**msg_d) for msg_d in msg_data_list # Spread dict into ChatMessage constructor
-                ]
-                report_data_for_pdf = {
-                    "patient_name": authenticated_username,
-                    "session_id": selected_session_data['id'],
-                    "session_title": selected_session_data['title'] or 'Untitled Session',
-                    "session_start_time": selected_session_data['start_time'],
-                    "patient_context_summary": selected_session_data.get('patient_context_summary', "Not provided."),
-                    "messages": pdf_report_messages
+                # Prepare list of MockChatMessage instances for the PDF generator
+                pdf_report_mock_messages: List[MockChatMessage] = []
+                for msg_d_item in messages_to_display_and_pdf:
+                    pdf_report_mock_messages.append(
+                        MockChatMessage( # Use the imported MockChatMessage class
+                            role=msg_d_item.get("role"),
+                            content=msg_d_item.get("content"),
+                            timestamp=msg_d_item.get("timestamp"),
+                            tool_name=msg_d_item.get("tool_name")
+                        )
+                    )
+
+                # Prepare the consolidated data dictionary for generate_pdf_report
+                report_data_for_pdf_generation = {
+                    "patient_name": authenticated_username, # This is the clinician's username
+                    "session_id": selected_session_dict_data['id'],
+                    "session_title": selected_session_dict_data.get('title') or 'Untitled Session',
+                    "session_start_time": selected_session_dict_data.get('start_time'), # datetime object
+                    "patient_context_summary": selected_session_dict_data.get('patient_context_summary', "Not provided."),
+                    "messages": pdf_report_mock_messages # List of MockChatMessage instances
                 }
+                
+                app_logger.debug(f"ReportsPage: Data for PDF generation (session {selected_session_id_val}): Keys={list(report_data_for_pdf_generation.keys())}, NumMessages={len(pdf_report_mock_messages)}")
 
-                pdf_bytes = generate_pdf_report(report_data_for_pdf) # Pass the whole dict
+                pdf_bytes_output = generate_pdf_report(report_data_for_pdf_generation)
                 
-                file_ts = datetime.now().strftime('%Y%m%d_%H%M%S')
-                pdf_fn = f"Consult_Report_S{selected_session_id}_{file_ts}.pdf"
+                current_time_filename_str = datetime.now().strftime('%Y%m%d_%H%M%S')
+                pdf_report_filename = f"Consultation_Report_S{selected_session_id_val}_{current_time_filename_str}.pdf"
+
                 st.download_button(
-                    label="Download Report as PDF", data=pdf_bytes, file_name=pdf_fn, mime="application/pdf",
-                    key=f"dl_pdf_s{selected_session_id}", on_click=log_report_generated,
-                    args=(authenticated_user_id, selected_session_id), help="Download consultation summary."
+                    label="Download Report as PDF",
+                    data=pdf_bytes_output,
+                    file_name=pdf_report_filename,
+                    mime="application/pdf",
+                    key=f"download_pdf_button_session_{selected_session_id_val}", # Unique key
+                    on_click=log_report_generated,
+                    args=(authenticated_user_id, selected_session_id_val), # Example args for metrics
+                    help="Click to download a PDF summary of this consultation session."
                 )
-            except Exception as e:
-                app_logger.error(f"Error generating PDF for session {selected_session_id}: {e}", exc_info=True)
-                st.error(f"Could not generate PDF report. Error: {type(e).__name__}")
-        else: st.info("This session has no messages for the report.")
-    else: st.error("Selected session data not found.")
-else: st.info("Select a session to view details.")
+            except Exception as e_pdf_gen:
+                app_logger.error(f"ReportsPage: Error generating PDF for session {selected_session_id_val} (User: '{authenticated_username}'): {e_pdf_gen}", exc_info=True)
+                st.error(f"Could not generate PDF report at this time. Error: {type(e_pdf_gen).__name__}")
+        else:
+            st.info("This session has no messages recorded to include in the report.")
+    else:
+        app_logger.error(f"ReportsPage: Selected session ID {selected_session_id_val} not found in fetched display data for user '{authenticated_username}'. This is unexpected.")
+        st.error("An error occurred: The selected session data could not be retrieved. Please try again or select another session.")
+else:
+    st.info("Please select a session from the dropdown menu above to view details and generate a report.")