Update services/pdf_report.py

services/pdf_report.py

@@ -1,118 +1,168 @@
 # /home/user/app/services/pdf_report.py
 from reportlab.lib.pagesizes import letter
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
-from reportlab.lib.styles import getSampleStyleSheet
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle, PageBreak
+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, Optional # Optional for tool_name
+from typing import List, Dict, Any # For the input data structure
 from pathlib import Path
+from datetime import datetime
 
-# Assuming ChatMessage is defined in models.py or models/chat.py
-# Make sure ChatMessage has 'role', 'content', and optionally 'tool_name' attributes.
-# Example ChatMessage structure (Pydantic or dataclass):
-# class ChatMessage(BaseModel):
-#     role: str  # "user", "assistant", "tool"
-#     content: str
-#     tool_name: Optional[str] = None
-from models import ChatMessage # Or from models.chat import ChatMessage
-from config.settings import settings # For APP_TITLE
-from assets.logo import get_logo_path # Corrected import path
-from services.logger import app_logger # For logging issues
-
-def generate_pdf_report(chat_messages: List[ChatMessage], patient_name: str = "Patient") -> BytesIO:
+# 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 services.logger import app_logger
+
+class MockChatMessage: # Helper class if generate_pdf_report expects objects with attributes
+    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
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
     buffer = BytesIO()
-    # Reduce margins for more content space
     doc = SimpleDocTemplate(buffer, pagesize=letter,
                             leftMargin=0.75*inch, rightMargin=0.75*inch,
                             topMargin=0.75*inch, bottomMargin=0.75*inch)
     styles = getSampleStyleSheet()
     story = []
 
-    # Style adjustments
-    styles['h1'].alignment = 1 # Center align H1
-    styles['Normal'].spaceBefore = 6
-    styles['Normal'].spaceAfter = 6
-    code_style = styles['Code'] # For AI and Tool messages
-    code_style.spaceBefore = 6
-    code_style.spaceAfter = 6
-    code_style.leftIndent = 10 # Indent AI/Tool messages slightly
+    # 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='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))
+    
+    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)
+
 
-    # 1. Logo (optional)
+    # --- Extract data from report_data dictionary ---
+    patient_name = report_data.get("patient_name", "N/A")
+    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
+
+
+    # 1. Logo and Document Header
     logo_path_str = get_logo_path()
-    if logo_path_str:
-        logo_file = Path(logo_path_str) # Convert to Path object
-        if logo_file.exists():
-            try:
-                # Adjust width/height as needed, preserve aspect ratio if possible
-                img = Image(logo_file, width=1.0*inch, height=1.0*inch, preserveAspectRatio=True)
-                img.hAlign = 'LEFT' # Align logo to the left
-                story.append(img)
-                story.append(Spacer(1, 0.2*inch))
-            except Exception as e:
-                app_logger.warning(f"Could not add logo to PDF from path '{logo_path_str}': {e}")
-        else:
-            app_logger.warning(f"Logo path '{logo_path_str}' from get_logo_path() does not exist.")
-
-    # 2. Title
-    title_text = f"{settings.APP_TITLE} - Consultation Report"
-    title = Paragraph(title_text, styles['h1'])
-    story.append(title)
+    if logo_path_str and Path(logo_path_str).exists():
+        try:
+            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(Paragraph(f"{settings.APP_TITLE}", styles['h1']))
+    story.append(Paragraph("AI-Assisted Consultation Summary", styles['h2']))
     story.append(Spacer(1, 0.2*inch))
 
-    # 3. Patient Info
-    patient_info_text = f"<b>Patient:</b> {patient_name}" # Make "Patient:" bold
-    patient_info = Paragraph(patient_info_text, styles['h2'])
-    story.append(patient_info)
-    story.append(Spacer(1, 0.3*inch)) # More space after patient info
+    # 2. Report Metadata Table
+    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 = [
+        [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>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.setStyle(TableStyle([
+        ('GRID', (0,0), (-1,-1), 0.5, colors.grey),
+        ('VALIGN', (0,0), (-1,-1), 'TOP'),
+        ('BACKGROUND', (0,0), (0,-1), colors.whitesmoke)
+    ]))
+    story.append(meta_table)
+    story.append(Spacer(1, 0.3*inch))
 
-    # 4. Chat Transcript
-    story.append(Paragraph("<u>Consultation Transcript:</u>", styles['h3'])) # Underline transcript title
+    # 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']))
+    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']))
+        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 chat_messages:
+    for msg in messages:
+        timestamp_str = msg.timestamp.strftime('%Y-%m-%d %H:%M:%S UTC') if msg.timestamp else "N/A"
         prefix = ""
         current_style = styles['Normal']
 
         if msg.role == 'assistant':
-            prefix = "<b>AI Assistant:</b> "
-            current_style = code_style
+            prefix = f"AI Assistant ({timestamp_str}):"
+            current_style = ai_msg_style
         elif msg.role == 'user':
-            prefix = "<b>You:</b> "
-            current_style = styles['Normal']
+            prefix = f"Clinician ({timestamp_str}):"
+            current_style = user_msg_style
         elif msg.role == 'tool':
-            tool_name = getattr(msg, 'tool_name', 'UnknownTool') # Handle if tool_name is missing
-            prefix = f"<b>Tool ({tool_name}):</b> "
-            current_style = code_style
-        else:
-            prefix = f"<b>{msg.role.capitalize()}:</b> " # Fallback for other roles
-
-        # Sanitize content for ReportLab Paragraph (handles HTML-like tags)
-        # Replace newlines with <br/> for PDF line breaks
-        content = msg.content.replace('\n', '<br/>\n')
-        # Escape < and > that are not part of <br/>
-        content = content.replace("<", "<").replace(">", ">").replace("<br/>", "<br/>")
+            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']
+        
+        # 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/>")
 
-        try:
-            story.append(Paragraph(prefix + content, current_style))
-            story.append(Spacer(1, 0.05*inch)) # Smaller spacer between messages
-        except Exception as e:
-            app_logger.error(f"Error adding message to PDF: {prefix}{content}. Error: {e}")
-            story.append(Paragraph(f"<i>Error rendering message: {e}</i>", styles['Italic']))
+        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']))
 
     try:
         doc.build(story)
         buffer.seek(0)
-        app_logger.info(f"PDF report generated successfully for patient: {patient_name}")
+        app_logger.info(f"PDF report generated for session ID: {session_id}")
     except Exception as e:
-        app_logger.error(f"Failed to build PDF document: {e}", exc_info=True)
-        # Return an empty or error buffer if build fails
-        buffer = BytesIO() # Reset buffer
+        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 generating PDF report. Please check logs.", styles['h1'])]
+        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 doc fails, just return empty buffer
+        except: pass # If even error PDF fails, return empty
         buffer.seek(0)
-
+        
     return buffer