Update services/pdf_report.py

services/pdf_report.py

@@ -1,52 +1,118 @@
+# /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.lib.units import inch
 from io import BytesIO
-from typing import List
-from models.chat import ChatMessage
+from typing import List, Optional # Optional for tool_name
+from pathlib import Path
+
+# 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 # We'll create this
+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:
     buffer = BytesIO()
-    doc = SimpleDocTemplate(buffer, pagesize=letter)
+    # 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 = []
 
-    # Logo (optional)
-    logo_path = get_logo_path()
-    if logo_path:
-        try:
-            img = Image(logo_path, width=1*inch, height=1*inch)
-            img.hAlign = 'LEFT'
-            story.append(img)
-            story.append(Spacer(1, 0.2*inch))
-        except Exception:
-            pass # Handle if logo not found or invalid
-
-    # Title
-    title = Paragraph(f"{settings.APP_TITLE} - Consultation Report", styles['h1'])
+    # 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
+
+    # 1. Logo (optional)
+    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)
     story.append(Spacer(1, 0.2*inch))
 
-    # Patient Info
-    patient_info = Paragraph(f"Patient: {patient_name}", styles['h2'])
+    # 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.2*inch))
+    story.append(Spacer(1, 0.3*inch)) # More space after patient info
+
+    # 4. Chat Transcript
+    story.append(Paragraph("<u>Consultation Transcript:</u>", styles['h3'])) # Underline transcript title
+    story.append(Spacer(1, 0.1*inch))
 
-    # Chat Transcript
-    story.append(Paragraph("Consultation Transcript:", styles['h3']))
     for msg in chat_messages:
-        role_style = styles['Code'] if msg.role == 'assistant' else styles['Normal']
-        prefix = "AI Assistant: " if msg.role == 'assistant' else "You: "
-        if msg.role == 'tool':
-            prefix = f"Tool ({msg.tool_name}): "
-        
-        content = msg.content.replace('\n', '<br/>\n') # Handle newlines in PDF
-        story.append(Paragraph(f"<b>{prefix}</b>{content}", role_style))
-        story.append(Spacer(1, 0.1*inch))
-    
-    doc.build(story)
-    buffer.seek(0)
+        prefix = ""
+        current_style = styles['Normal']
+
+        if msg.role == 'assistant':
+            prefix = "<b>AI Assistant:</b> "
+            current_style = code_style
+        elif msg.role == 'user':
+            prefix = "<b>You:</b> "
+            current_style = styles['Normal']
+        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/>")
+
+        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']))
+
+
+    try:
+        doc.build(story)
+        buffer.seek(0)
+        app_logger.info(f"PDF report generated successfully for patient: {patient_name}")
+    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
+        error_doc = SimpleDocTemplate(buffer, pagesize=letter)
+        error_story = [Paragraph("Error generating PDF report. Please check logs.", styles['h1'])]
+        try:
+            error_doc.build(error_story)
+        except: pass # If even error doc fails, just return empty buffer
+        buffer.seek(0)
+
     return buffer