Update services/pdf_report.py

services/pdf_report.py

@@ -5,7 +5,7 @@ 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, Optional # IMPORTED Optional
+from typing import List, Dict, Any, Optional
 from pathlib import Path
 from datetime import datetime
 
@@ -13,6 +13,185 @@ from config.settings import settings
 from assets.logo import get_logo_path # Assuming this function exists and works
 from services.logger import app_logger
 
+class MockChatMessage:
+    """
+    A simple data class to hold message attributes for PDF generation.
+    Ensures consistent attribute access and helps with type handling.
+    """
+    def __init__(self, role: str, content: str, timestamp: Optional[datetime], tool_name: Optional[str] = None, **kwargs):
+        self.role: str = str(role if role is not None else "unknown")
+        self.content: str = str(content if content is not None else "")
+        self.timestamp: Optional[datetime] = timestamp # Keep as datetime, format when used
+        self.tool_name: Optional[str] = str(tool_name) if tool_name is not None else None
+        
+        # Store any other potential attributes (e.g., source_references, confidence)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def get_formatted_timestamp(self, default_val: str = "Time N/A") -> str:
+        return self.timestamp.strftime('%Y-%m-%d %H:%M:%S UTC') if self.timestamp else default_val
+
+    def get_formatted_content_for_pdf(self) -> str:
+        # Ensure content is a string, handle newlines for PDF, and escape HTML-sensitive characters
+        text = self.content.replace('\n', '<br/>\n')
+        return text.replace("<", "<").replace(">", ">").replace("<br/>", "<br/>")
+
+
+def generate_pdf_report(report_data: Dict[str, Any]) -> BytesIO:
+    buffer = BytesIO()
+    doc = SimpleDocTemplate(buffer, pagesize=letter,
+                            leftMargin=0.75*inch, rightMargin=0.75*inch,
+                            topMargin=0.75*inch, bottomMargin=0.75*inch)
+    styles = getSampleStyleSheet()
+    story = []
+
+    # --- Custom Styles Definition ---
+    styles.add(ParagraphStyle(name='Justify', alignment=4, parent=styles['Normal']))
+    styles.add(ParagraphStyle(name='Disclaimer', parent=styles['Italic'], fontSize=8, leading=10, spaceBefore=6, spaceAfter=6, textColor=colors.dimgrey))
+    styles.add(ParagraphStyle(name='SectionHeader', parent=styles['h2'], spaceBefore=12, spaceAfter=6, keepWithNext=1, textColor=colors.darkblue))
+    styles.add(ParagraphStyle(name='SubHeader', parent=styles['h3'], spaceBefore=6, spaceAfter=3, keepWithNext=1, textColor=colors.cadetblue))
+    styles.add(ParagraphStyle(name='ListItem', parent=styles['Normal'], leftIndent=0.25*inch, bulletIndent=0.1*inch, spaceBefore=3))
+    
+    base_message_style = ParagraphStyle(name='BaseMessage', parent=styles['Normal'], spaceBefore=2, spaceAfter=2, leftIndent=0.1*inch, rightIndent=0.1*inch, leading=12, borderWidth=0.5, padding=(4, 4, 4, 4)) # top, right, bottom, left
+    user_msg_style = ParagraphStyle(name='UserMessage', parent=base_message_style, backColor=colors.whitesmoke, borderColor=colors.lightgrey)
+    ai_msg_style = ParagraphStyle(name='AIMessage', parent=base_message_style, backColor=colors.lightcyan, borderColor=colors.lightgrey)
+    tool_msg_style = ParagraphStyle(name='ToolMessage', parent=base_message_style, backColor=colors.lightgoldenrodyellow, textColor=colors.darkslategrey, borderColor=colors.darkgrey, fontName='Courier')
+    system_msg_style = ParagraphStyle(name='SystemMessage', parent=base_message_style, backColor=colors.beige, fontSize=9, textColor=colors.dimgrey, borderColor=colors.darkgrey, fontName='Helvetica-Oblique')
+
+    # --- Extract data safely from report_data dictionary ---
+    clinician_username = str(report_data.get("patient_name", "N/A Clinician")) # Actually the clinician's username
+    session_id_str = str(report_data.get("session_id", "N/A"))
+    session_title_str = str(report_data.get("session_title", "Untitled Consultation"))
+    session_start_time_obj = report_data.get("session_start_time") # This is a datetime object or None
+    patient_context_summary_str = str(report_data.get("patient_context_summary", "No specific patient context was provided."))
+    
+    # Messages are expected to be List[MockChatMessage] from the caller
+    messages: List[MockChatMessage] = report_data.get("messages", [])
+
+
+    # 1. Logo and Document Header
+    app_logger.debug("PDF Generation: Adding logo and header.")
+    logo_path_str = get_logo_path()
+    if logo_path_str and Path(logo_path_str).exists():
+        try:
+            img = Image(logo_path_str, width=0.75*inch, height=0.75*inch, preserveAspectRatio=True)
+            img.hAlign = 'LEFT'
+            story.append(img)
+            story.append(Spacer(1, 0.05*inch))
+        except Exception as e: app_logger.warning(f"PDF Report: Could not add logo image: {e}")
+    
+    story.append(Paragraph(str(settings.APP_TITLE), styles['h1']))
+    story.append(Paragraph("AI-Assisted Consultation Summary", styles['h2']))
+    story.append(Spacer(1, 0.2*inch))
+
+    # 2. Report Metadata Table
+    app_logger.debug("PDF Generation: Adding metadata table.")
+    report_date_str = datetime.now().strftime('%Y-%m-%d %H:%M UTC')
+    session_start_time_str_display = session_start_time_obj.strftime('%Y-%m-%d %H:%M UTC') if session_start_time_obj else "N/A"
+    
+    meta_data_content = [
+        [Paragraph("<b>Report Generated:</b>", styles['Normal']), Paragraph(report_date_str, styles['Normal'])],
+        [Paragraph("<b>Clinician:</b>", styles['Normal']), Paragraph(clinician_username, styles['Normal'])],
+        [Paragraph("<b>Consultation Session ID:</b>", styles['Normal']), Paragraph(session_id_str, styles['Normal'])],
+        [Paragraph("<b>Session Title:</b>", styles['Normal']), Paragraph(session_title_str, styles['Normal'])],
+        [Paragraph("<b>Session Start Time:</b>", styles['Normal']), Paragraph(session_start_time_str_display, styles['Normal'])],
+    ]
+    meta_table = Table(meta_data_content, colWidths=[2.0*inch, None])
+    meta_table.setStyle(TableStyle([
+        ('GRID', (0,0), (-1,-1), 0.5, colors.darkgrey), ('VALIGN', (0,0), (-1,-1), 'TOP'),
+        ('BACKGROUND', (0,0), (0,-1), colors.lightgrey), ('LEFTPADDING', (0,0), (-1,-1), 5),
+        ('RIGHTPADDING', (0,0), (-1,-1), 5), ('TOPPADDING', (0,0), (-1,-1), 3), ('BOTTOMPADDING', (0,0), (-1,-1), 3)
+    ]))
+    story.append(meta_table)
+    story.append(Spacer(1, 0.3*inch))
+
+    # 3. Disclaimer
+    app_logger.debug("PDF Generation: Adding disclaimers.")
+    story.append(Paragraph("<b>Important Disclaimer:</b>", styles['SubHeader']))
+    story.append(Paragraph(str(settings.MAIN_DISCLAIMER_LONG), styles['Disclaimer']))
+    if settings.SIMULATION_DISCLAIMER:
+        story.append(Paragraph(str(settings.SIMULATION_DISCLAIMER), styles['Disclaimer']))
+    story.append(Spacer(1, 0.3*inch))
+
+    # 4. Patient Context Provided (if any)
+    app_logger.debug("PDF Generation: Adding patient context summary.")
+    if patient_context_summary_str and patient_context_summary_str.lower() 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']))
+        # Split summary string into list items for better readability
+        context_items_raw = patient_context_summary_str.replace("Patient Context: ", "").replace("Initial Patient Context Set: ","").split(';')
+        for item_raw in context_items_raw:
+            item_clean = item_raw.strip()
+            if item_clean: story.append(Paragraph(f"• {item_clean}", styles['ListItem']))
+        story.append(Spacer(1, 0.2*inch))
+
+    # 5. Consultation Transcript
+    app_logger.debug(f"PDF Generation: Adding transcript with {len(messages)} messages.")
+    story.append(Paragraph("Consultation Transcript:", styles['SectionHeader']))
+    story.append(Spacer(1, 0.1*inch))
+
+    for msg_obj in messages: # msg_obj is now an instance of MockChatMessage
+        # Skip system messages that are just context logging, unless you want them
+        if msg_obj.role == 'system' and "Initial Patient Context Set:" in msg_obj.content:
+            app_logger.debug(f"Skipping system context message in PDF transcript: {msg_obj.content[:50]}...")
+            continue
+
+        formatted_timestamp = msg_obj.get_formatted_timestamp()
+        prefix_display = ""
+        active_message_style = styles['Normal']
+
+        if msg_obj.role == 'assistant':
+            prefix_display = f"AI Assistant ({formatted_timestamp}):"
+            active_message_style = ai_msg_style
+        elif msg_obj.role == 'user':
+            prefix_display = f"Clinician ({formatted_timestamp}):"
+            active_message_style = user_msg_style
+        elif msg_obj.role == 'tool':
+            tool_name_str = msg_obj.tool_name or "Tool"
+            prefix_display = f"{tool_name_str.capitalize()} Output ({formatted_timestamp}):"
+            active_message_style = tool_msg_style
+        elif msg_obj.role == 'system':
+            prefix_display = f"System Note ({formatted_timestamp}):"
+            active_message_style = system_msg_style
+        else: # Fallback for any other roles
+            prefix_display = f"{msg_obj.role.capitalize()} ({formatted_timestamp}):"
+        
+        formatted_content = msg_obj.get_formatted_content_for_pdf()
+        
+        story.append(Paragraph(f"<b>{prefix_display}</b>", styles['Normal']))
+        story.append(Paragraph(formatted_content, active_message_style))
+        # story.append(Spacer(1, 0.05*inch)) # Can make it too sparse
+
+    story.append(Spacer(1, 0.5*inch))
+    story.append(Paragraph("--- End of Report ---", ParagraphStyle(name='EndOfReport', parent=styles['Italic'], alignment=1, spaceBefore=12)))
+
+    # --- Build PDF ---
+    app_logger.debug("PDF Generation: Building document.")
+    try:
+        doc.build(story)
+        buffer.seek(0)
+        app_logger.info(f"PDF report generated successfully for session ID: {session_id_str}")
+    except Exception as e_build:
+        app_logger.error(f"Failed to build PDF document for session ID {session_id_str}: {e_build}", exc_info=True)
+        # Return an error PDF or an empty buffer
+        buffer = BytesIO()
+        error_styles_local = getSampleStyleSheet()
+        error_doc_local = SimpleDocTemplate(buffer, pagesize=letter)
+        error_story_local = [
+            Paragraph("Error: Could not generate PDF report.", error_styles_local['h1']),
+            Paragraph(f"An error occurred during PDF construction: {str(e_build)[:500]}", error_styles_local['Normal']),
+            Paragraph("Please check application logs for more details.", error_styles_local['Normal'])
+        ]
+        try:
+            error_doc_local.build(error_story_local)
+        except Exception as e_err_pdf:
+            app_logger.error(f"Failed to build even the error PDF: {e_err_pdf}", exc_info=True)
+            # If error PDF fails, buffer will be empty from the reset
+        buffer.seek(0)
+        
+    return buffer
+from assets.logo import get_logo_path # Assuming this function exists and works
+from services.logger import app_logger
+
 class MockChatMessage:
     """
     A simple class to hold message data for PDF generation if generate_pdf_report