Update pages/3_Reports.py

pages/3_Reports.py

@@ -5,19 +5,17 @@ from typing import List, Dict, Any
 from sqlmodel import select
 
 from config.settings import settings
-from models import ChatMessage, ChatSession # ChatMessage model for query
+from models import ChatMessage, ChatSession # Model definitions
 from models.db import get_session_context
-from services.pdf_report import generate_pdf_report
+from services.pdf_report import generate_pdf_report # Needs to accept more data
 from services.logger import app_logger
 from services.metrics import log_report_generated
 
-# --- Authentication Check (same as before) ---
+# --- 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")
@@ -26,27 +24,25 @@ app_logger.info(f"User '{authenticated_username}' (ID: {authenticated_user_id})
 
 st.title("Consultation Reports")
 st.markdown("View and download your past consultation sessions.")
+st.info(settings.MAIN_DISCLAIMER_SHORT)
+
 
-# --- Helper Function to Load User's Chat Session Data (Primitives) ---
-# @st.cache_data(ttl=300)
 def get_user_chat_session_display_data(user_id: int) -> List[Dict[str, Any]]:
-    app_logger.debug(f"Fetching chat session display data for user_id: {user_id}")
+    app_logger.debug(f"Fetching session display data for user_id: {user_id}")
     session_data_list: List[Dict[str, Any]] = []
     try:
-        with get_session_context() as db_session:
-            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()
+        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:
                 session_data_list.append({
-                    "id": row.id,
-                    "title": row.title,
-                    "start_time": row.start_time
+                    "id": row.id, "title": row.title, "start_time": row.start_time,
+                    "patient_context_summary": row.patient_context_summary
                 })
-        app_logger.debug(f"Found {len(session_data_list)} session display entries for user_id: {user_id}")
+        app_logger.debug(f"Found {len(session_data_list)} session display entries for user {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)
+        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.")
     return session_data_list
 
@@ -59,128 +55,93 @@ if not chat_session_display_items:
 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 "Date N/A"
-    title_display = s_data["title"] if s_data["title"] else f"Consultation on {start_time_display}"
+    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']))
 
 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, format_func=lambda x: x[0], 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]
-    MESSAGES_FOR_SESSION_KEY = f"messages_for_session_{selected_session_id}"
+    MESSAGES_KEY = f"report_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]] = []
+    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_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
+            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_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.")
+            st.session_state[MESSAGES_KEY] = msg_data_list
         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
+            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 session ID: {selected_session_id}")
-        messages_data_for_display = st.session_state[MESSAGES_FOR_SESSION_KEY]
-
+        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]
 
-    # --- 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)
+    selected_session_data = next((s for s in chat_session_display_items if s['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.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 messages_data_for_display: # Use the data from session_state
+        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"])
+        
+        if msg_data_list:
             with st.expander("View Chat Transcript", expanded=False):
-                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("---")
+                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("---")
             
             st.markdown("---")
             try:
-                # 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.
+                # 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
+                }
+
+                pdf_bytes = generate_pdf_report(report_data_for_pdf) # Pass the whole dict
                 
-                # 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"
+                file_ts = datetime.now().strftime('%Y%m%d_%H%M%S')
+                pdf_fn = f"Consult_Report_S{selected_session_id}_{file_ts}.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}",
-                    on_click=log_report_generated,
-                    args=(authenticated_user_id, selected_session_id),
-                    help="Click to download the consultation transcript as a PDF file."
+                    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."
                 )
             except Exception as e:
-                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:
-        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:
-    st.info("Please select a session from the dropdown menu to view details.")
+                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.")