Update services/pdf_report.py

services/pdf_report.py

@@ -1,30 +1,32 @@
 # /home/user/app/services/pdf_report.py
 from reportlab.lib.pagesizes import letter
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle, PageBreak
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.lib.units import inch
 from reportlab.lib import colors
 from io import BytesIO
-from typing import List, Dict, Any # For the input data structure
+from typing import List, Dict, Any, Optional # IMPORTED Optional
 from pathlib import Path
 from datetime import datetime
 
-# from models import ChatMessage # Not strictly needed if we use dicts/mock objects
 from config.settings import settings
-from assets.logo import get_logo_path
+from assets.logo import get_logo_path # Assuming this function exists and works
 from services.logger import app_logger
 
-class MockChatMessage: # Helper class if generate_pdf_report expects objects with attributes
+class MockChatMessage:
+    """
+    A simple class to hold message data for PDF generation if generate_pdf_report
+    prefers object attribute access over dictionary key access for messages.
+    """
     def __init__(self, role: str, content: str, timestamp: datetime, tool_name: Optional[str] = None, **kwargs):
         self.role = role
         self.content = content
         self.timestamp = timestamp
         self.tool_name = tool_name
-        # Allow other attributes like source_references, confidence_score if added
+        # Capture any other attributes that might be passed (e.g., source_references)
         for key, value in kwargs.items():
             setattr(self, key, value)
 
-
 def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
     buffer = BytesIO()
     doc = SimpleDocTemplate(buffer, pagesize=letter,
@@ -35,28 +37,26 @@ def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
 
     # Custom Styles
     styles.add(ParagraphStyle(name='Justify', alignment=4, parent=styles['Normal']))
-    styles.add(ParagraphStyle(name='Disclaimer', parent=styles['Italic'], fontSize=8, leading=10))
+    styles.add(ParagraphStyle(name='Disclaimer', parent=styles['Italic'], fontSize=8, leading=10, spaceBefore=6, spaceAfter=6))
     styles.add(ParagraphStyle(name='SectionHeader', parent=styles['h2'], spaceBefore=12, spaceAfter=6, keepWithNext=1))
     styles.add(ParagraphStyle(name='SubHeader', parent=styles['h3'], spaceBefore=6, spaceAfter=3, keepWithNext=1))
-    styles.add(ParagraphStyle(name='ListItem', parent=styles['Normal'], leftIndent=0.25*inch, bulletIndent=0.1*inch))
+    styles.add(ParagraphStyle(name='ListItem', parent=styles['Normal'], leftIndent=0.25*inch, bulletIndent=0.1*inch, spaceBefore=3))
     
-    user_msg_style = ParagraphStyle(name='UserMessage', parent=styles['Normal'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.whitesmoke)
-    ai_msg_style = ParagraphStyle(name='AIMessage', parent=styles['Normal'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.lightcyan)
-    tool_msg_style = ParagraphStyle(name='ToolMessage', parent=styles['Code'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.lightgrey, textColor=colors.darkblue)
-    system_msg_style = ParagraphStyle(name='SystemMessage', parent=styles['Italic'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.beige, fontSize=9)
-
+    user_msg_style = ParagraphStyle(name='UserMessage', parent=styles['Normal'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.whitesmoke, borderColor=colors.lightgrey, borderWidth=0.5, padding=3)
+    ai_msg_style = ParagraphStyle(name='AIMessage', parent=styles['Normal'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.lightcyan, borderColor=colors.lightgrey, borderWidth=0.5, padding=3)
+    tool_msg_style = ParagraphStyle(name='ToolMessage', parent=styles['Code'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.lightgrey, textColor=colors.darkblue, borderColor=colors.grey, borderWidth=0.5, padding=3)
+    system_msg_style = ParagraphStyle(name='SystemMessage', parent=styles['Italic'], spaceBefore=3, spaceAfter=3, leftIndent=0.1*inch, backColor=colors.beige, fontSize=9, borderColor=colors.darkgrey, borderWidth=0.5, padding=3)
 
-    # --- Extract data from report_data dictionary ---
-    patient_name = report_data.get("patient_name", "N/A")
+    # --- Extract data ---
+    patient_name = report_data.get("patient_name", "N/A Clinician") # Clinician's username
     session_id = report_data.get("session_id", "N/A")
     session_title = report_data.get("session_title", "Untitled Consultation")
     session_start_time_obj = report_data.get("session_start_time")
     patient_context_summary = report_data.get("patient_context_summary", "No specific patient context was provided for this session.")
     
-    # Convert message dictionaries to MockChatMessage objects if needed, or use dicts directly
-    raw_messages = report_data.get("messages", [])
-    messages = [MockChatMessage(**msg_data) for msg_data in raw_messages] # If generate_pdf_report expects objects
-
+    raw_messages_data = report_data.get("messages", []) # Expects a list of dicts
+    # Convert message dictionaries to MockChatMessage objects for consistent attribute access
+    messages = [MockChatMessage(**msg_data) for msg_data in raw_messages_data]
 
     # 1. Logo and Document Header
     logo_path_str = get_logo_path()
@@ -65,11 +65,10 @@ def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
             img = Image(logo_path_str, width=0.8*inch, height=0.8*inch, preserveAspectRatio=True)
             img.hAlign = 'LEFT'
             story.append(img)
-            story.append(Spacer(1, 0.1*inch))
-        except Exception as e:
-            app_logger.warning(f"PDF Report: Could not add logo: {e}")
+            story.append(Spacer(1, 0.05*inch)) # Reduced spacer
+        except Exception as e: app_logger.warning(f"PDF Report: Could not add logo: {e}")
     
-    story.append(Paragraph(f"{settings.APP_TITLE}", styles['h1']))
+    story.append(Paragraph(settings.APP_TITLE, styles['h1']))
     story.append(Paragraph("AI-Assisted Consultation Summary", styles['h2']))
     story.append(Spacer(1, 0.2*inch))
 
@@ -77,18 +76,18 @@ def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
     report_date_str = datetime.now().strftime('%Y-%m-%d %H:%M UTC')
     session_start_time_str = session_start_time_obj.strftime('%Y-%m-%d %H:%M UTC') if session_start_time_obj else "N/A"
     
-    meta_data = [
+    meta_data_content = [
         [Paragraph("<b>Report Generated:</b>", styles['Normal']), Paragraph(report_date_str, styles['Normal'])],
-        [Paragraph("<b>Clinician:</b>", styles['Normal']), Paragraph(patient_name, styles['Normal'])], # "patient_name" is the clinician's username
+        [Paragraph("<b>Clinician:</b>", styles['Normal']), Paragraph(patient_name, styles['Normal'])],
         [Paragraph("<b>Consultation Session ID:</b>", styles['Normal']), Paragraph(str(session_id), styles['Normal'])],
         [Paragraph("<b>Session Title:</b>", styles['Normal']), Paragraph(session_title, styles['Normal'])],
         [Paragraph("<b>Session Start Time:</b>", styles['Normal']), Paragraph(session_start_time_str, styles['Normal'])],
     ]
-    meta_table = Table(meta_data, colWidths=[1.8*inch, None])
+    meta_table = Table(meta_data_content, colWidths=[2.0*inch, None]) # Adjusted colWidths
     meta_table.setStyle(TableStyle([
-        ('GRID', (0,0), (-1,-1), 0.5, colors.grey),
-        ('VALIGN', (0,0), (-1,-1), 'TOP'),
-        ('BACKGROUND', (0,0), (0,-1), colors.whitesmoke)
+        ('GRID', (0,0), (-1,-1), 0.5, colors.grey), ('VALIGN', (0,0), (-1,-1), 'TOP'),
+        ('BACKGROUND', (0,0), (0,-1), colors.lightgrey), ('LEFTPADDING', (0,0), (-1,-1), 3),
+        ('RIGHTPADDING', (0,0), (-1,-1), 3), ('TOPPADDING', (0,0), (-1,-1), 3), ('BOTTOMPADDING', (0,0), (-1,-1), 3)
     ]))
     story.append(meta_table)
     story.append(Spacer(1, 0.3*inch))
@@ -96,73 +95,69 @@ def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
     # 3. Disclaimer
     story.append(Paragraph("<b>Important Disclaimer:</b>", styles['SubHeader']))
     story.append(Paragraph(settings.MAIN_DISCLAIMER_LONG, styles['Disclaimer']))
-    story.append(Paragraph(settings.SIMULATION_DISCLAIMER, styles['Disclaimer']))
+    if settings.SIMULATION_DISCLAIMER: # Only add if defined
+        story.append(Paragraph(settings.SIMULATION_DISCLAIMER, styles['Disclaimer']))
     story.append(Spacer(1, 0.3*inch))
 
     # 4. Patient Context Provided (if any)
-    if patient_context_summary and patient_context_summary != "No specific patient context provided.":
-        story.append(Paragraph("Patient Context Provided (Simulated Data):", styles['SectionHeader']))
-        # Assuming patient_context_summary is a string. If it's structured, format it nicely.
-        # Example: if it's "Key1: Value1; Key2: Value2", split and list.
-        context_parts = patient_context_summary.replace("Patient Context: ", "").split(';')
-        for part in context_parts:
-            if part.strip():
-                 story.append(Paragraph(f"• {part.strip()}", styles['ListItem']))
+    if patient_context_summary and patient_context_summary not in ["No specific patient context was provided for this session.", "Not provided."]:
+        story.append(Paragraph("Patient Context Provided by Clinician (Simulated Data):", styles['SectionHeader']))
+        context_items = patient_context_summary.replace("Patient Context: ", "").split(';')
+        for item in context_items:
+            if item.strip(): story.append(Paragraph(f"• {item.strip()}", styles['ListItem']))
         story.append(Spacer(1, 0.2*inch))
 
     # 5. Consultation Transcript
     story.append(Paragraph("Consultation Transcript:", styles['SectionHeader']))
     story.append(Spacer(1, 0.1*inch))
 
-    for msg in messages:
+    for msg in messages: # Iterate over MockChatMessage objects
         timestamp_str = msg.timestamp.strftime('%Y-%m-%d %H:%M:%S UTC') if msg.timestamp else "N/A"
-        prefix = ""
-        current_style = styles['Normal']
+        prefix_text = ""
+        current_message_style = styles['Normal']
 
         if msg.role == 'assistant':
-            prefix = f"AI Assistant ({timestamp_str}):"
-            current_style = ai_msg_style
+            prefix_text = f"AI Assistant ({timestamp_str}):"
+            current_message_style = ai_msg_style
         elif msg.role == 'user':
-            prefix = f"Clinician ({timestamp_str}):"
-            current_style = user_msg_style
+            prefix_text = f"Clinician ({timestamp_str}):"
+            current_message_style = user_msg_style
         elif msg.role == 'tool':
-            tool_name = getattr(msg, 'tool_name', 'Tool')
-            prefix = f"{tool_name.capitalize()} ({timestamp_str}):"
-            current_style = tool_msg_style
-        elif msg.role == 'system': # Typically system context, might not always be included
-            prefix = f"System Context ({timestamp_str}):"
-            current_style = system_msg_style
-        else: # Fallback for any other roles
-            prefix = f"{msg.role.capitalize()} ({timestamp_str}):"
-            current_style = styles['Normal']
+            tool_name_display = getattr(msg, 'tool_name', 'Tool') or "Tool" # Handle None or empty tool_name
+            prefix_text = f"{tool_name_display.capitalize()} Output ({timestamp_str}):"
+            current_message_style = tool_msg_style
+        elif msg.role == 'system' and msg.content.startswith("Initial Patient Context Set:"): # Don't show this verbose system message
+            continue
+        elif msg.role == 'system':
+            prefix_text = f"System Note ({timestamp_str}):"
+            current_message_style = system_msg_style
+        else:
+            prefix_text = f"{msg.role.capitalize()} ({timestamp_str}):"
         
-        # Sanitize content for ReportLab Paragraph (handles HTML-like tags for <br/>)
         content_for_pdf = msg.content.replace('\n', '<br/>\n')
         content_for_pdf = content_for_pdf.replace("<", "<").replace(">", ">").replace("<br/>", "<br/>")
+        
+        story.append(Paragraph(f"<b>{prefix_text}</b>", styles['Normal'])) # Prefix bolded on its own line
+        story.append(Paragraph(content_for_pdf, current_message_style)) # Content in styled box
+        story.append(Spacer(1, 0.05*inch))
 
-        story.append(Paragraph(f"<b>{prefix}</b><br/>{content_for_pdf}", current_style))
-        # story.append(Spacer(1, 0.05*inch)) # Removed for tighter message packing
-
-    # Footer on each page (if needed, more complex - requires onLaterPages in SimpleDocTemplate)
-    # For simplicity, a final "End of Report"
     story.append(Spacer(1, 0.5*inch))
-    story.append(Paragraph("--- End of Report ---", styles['Italic']))
+    story.append(Paragraph("--- End of Report ---", ParagraphStyle(name='EndOfReport', parent=styles['Italic'], alignment=1)))
 
     try:
         doc.build(story)
         buffer.seek(0)
-        app_logger.info(f"PDF report generated for session ID: {session_id}")
+        app_logger.info(f"PDF report generated successfully for session ID: {session_id}")
     except Exception as e:
         app_logger.error(f"Failed to build PDF document for session ID {session_id}: {e}", exc_info=True)
-        buffer = BytesIO() # Return empty buffer on error
-        # Optionally, create a simple error PDF
-        error_styles = getSampleStyleSheet()
-        error_doc = SimpleDocTemplate(buffer, pagesize=letter)
-        error_story = [Paragraph("Error: Could not generate PDF report.", error_styles['h1']),
-                       Paragraph(f"Details: {str(e)[:500]}", error_styles['Normal'])]
-        try:
-            error_doc.build(error_story)
-        except: pass # If even error PDF fails, return empty
+        buffer = BytesIO() # Reset to empty buffer on error
+        # Create a simple error PDF
+        error_styles_local = getSampleStyleSheet() # Get fresh styles
+        error_doc_local = SimpleDocTemplate(buffer, pagesize=letter)
+        error_story_local = [Paragraph("Error: Could not generate PDF report.", error_styles_local['h1']),
+                             Paragraph(f"Details: {str(e)[:500]}", error_styles_local['Normal'])]
+        try: error_doc_local.build(error_story_local)
+        except: pass # If even error PDF fails, return empty buffer
         buffer.seek(0)
         
     return buffer