Update app.py

app.py

@@ -3,20 +3,36 @@ FastAPI Application for Multimodal RAG System
 US Army Medical Research Papers Q&A
 """
 
+"""
+FastAPI Application for Multimodal RAG System
+US Army Medical Research Papers Q&A
+"""
+
 import os
+import io
 import logging
 from typing import List, Dict, Optional, Union
 from contextlib import asynccontextmanager
+from datetime import datetime
 
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, Field
 
 # Import from query_index (standalone)
 from query_index import MultimodalRAGSystem
 
+# PDF generation libraries
+from reportlab.lib.pagesizes import letter
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
+from reportlab.lib.units import inch
+from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak, Image as RLImage, Table, TableStyle
+from reportlab.lib import colors
+import requests
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -87,7 +103,7 @@ class ImageSource(BaseModel):
     path: Optional[str]
     filename: Optional[str]
     score: Optional[float]
-    page: Optional[Union[str, int]] # could be int or str depending on metadata
+    page: Optional[Union[str, int]]
     file: Optional[str]
     link: Optional[str] = None
 
@@ -108,6 +124,16 @@ class HealthResponse(BaseModel):
     status: str
     rag_initialized: bool
 
+class ConversationItem(BaseModel):
+    question: str
+    answer: str
+    images: List[ImageSource]
+    texts: List[TextSource]
+    timestamp: str
+
+class ReportRequest(BaseModel):
+    conversations: List[ConversationItem]
+
 # API Endpoints
 
 @app.get("/", tags=["Root"])
@@ -132,7 +158,7 @@ async def query_rag(request: QueryRequest):
 
     try:
         result = rag_system.ask(
-            query_str=request.question,  # Fixed: changed from 'question=' to 'query_str='
+            query_str=request.question,
             chat_history=chat_history
         )
 
@@ -155,7 +181,172 @@ async def query_rag(request: QueryRequest):
         logger.error(f"Error processing query: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
+@app.post("/generate-report", tags=["Report"])
+async def generate_report(request: ReportRequest):
+    """Generate a PDF report from conversation data"""
+    try:
+        buffer = io.BytesIO()
+        
+        # Create PDF document
+        doc = SimpleDocTemplate(
+            buffer,
+            pagesize=letter,
+            rightMargin=0.75*inch,
+            leftMargin=0.75*inch,
+            topMargin=1*inch,
+            bottomMargin=0.75*inch
+        )
+        
+        # Container for the 'Flowable' objects
+        story = []
+        
+        # Define styles
+        styles = getSampleStyleSheet()
+        
+        # Custom styles
+        title_style = ParagraphStyle(
+            'CustomTitle',
+            parent=styles['Heading1'],
+            fontSize=24,
+            textColor=colors.HexColor('#3b82f6'),
+            spaceAfter=12,
+            alignment=TA_CENTER,
+            fontName='Helvetica-Bold'
+        )
+        
+        subtitle_style = ParagraphStyle(
+            'CustomSubtitle',
+            parent=styles['Normal'],
+            fontSize=12,
+            textColor=colors.HexColor('#64748b'),
+            spaceAfter=20,
+            alignment=TA_CENTER
+        )
+        
+        question_style = ParagraphStyle(
+            'QuestionStyle',
+            parent=styles['Heading2'],
+            fontSize=14,
+            textColor=colors.HexColor('#3b82f6'),
+            spaceAfter=10,
+            fontName='Helvetica-Bold'
+        )
+        
+        answer_style = ParagraphStyle(
+            'AnswerStyle',
+            parent=styles['Normal'],
+            fontSize=11,
+            alignment=TA_JUSTIFY,
+            spaceAfter=12
+        )
+        
+        source_title_style = ParagraphStyle(
+            'SourceTitle',
+            parent=styles['Heading3'],
+            fontSize=11,
+            textColor=colors.HexColor('#8b5cf6'),
+            spaceAfter=6,
+            fontName='Helvetica-Bold'
+        )
+        
+        # Add title
+        story.append(Paragraph("WHEC Research Assistant", title_style))
+        story.append(Paragraph("Conversation Report", title_style))
+        story.append(Spacer(1, 0.2*inch))
+        
+        # Add metadata
+        current_time = datetime.now().strftime("%B %d, %Y at %I:%M %p")
+        story.append(Paragraph(f"Generated: {current_time}", subtitle_style))
+        story.append(Paragraph("Source: WHEC (Warrior Heat- and Exertion-Related Events Collaborative)", subtitle_style))
+        story.append(Paragraph(f"Total Questions: {len(request.conversations)}", subtitle_style))
+        
+        story.append(Spacer(1, 0.3*inch))
+        
+        # Add horizontal line
+        story.append(Spacer(1, 0.1*inch))
+        
+        # Process each conversation
+        for idx, conv in enumerate(request.conversations, 1):
+            # Question
+            story.append(Paragraph(f"Question {idx}", question_style))
+            story.append(Paragraph(conv.question, answer_style))
+            story.append(Spacer(1, 0.15*inch))
+            
+            # Answer
+            story.append(Paragraph("Answer", source_title_style))
+            story.append(Paragraph(conv.answer, answer_style))
+            story.append(Spacer(1, 0.15*inch))
+            
+            # Text Sources
+            if conv.texts:
+                story.append(Paragraph("Referenced Text Sources", source_title_style))
+                for i, txt in enumerate(conv.texts, 1):
+                    source_text = f"[{i}] {txt.file or 'Unknown Document'} (Page {txt.page or 'N/A'}, {round((txt.score or 0) * 100)}% match)"
+                    story.append(Paragraph(source_text, styles['Normal']))
+                    
+                    excerpt = f'<i>"{txt.text[:300]}..."</i>'
+                    story.append(Paragraph(excerpt, styles['Normal']))
+                    story.append(Spacer(1, 0.1*inch))
+                
+                story.append(Spacer(1, 0.1*inch))
+            
+            # Image Sources
+            if conv.images:
+                relevant_images = [img for img in conv.images if (img.score or 0) >= 0.3]
+                if relevant_images:
+                    story.append(Paragraph("Referenced Images", source_title_style))
+                    
+                    for i, img in enumerate(relevant_images[:3], 1):  # Limit to 3 images
+                        # Add image metadata
+                        img_text = f"[{i}] {img.filename or 'Unknown'} from {img.file or 'Unknown Document'} (Page {img.page or 'N/A'}, {round((img.score or 0) * 100)}% match)"
+                        story.append(Paragraph(img_text, styles['Normal']))
+                        
+                        # Try to add the actual image
+                        if img.path:
+                            try:
+                                # Construct the full path to the image
+                                img_path = img.path.replace('/extracted_images/', '')
+                                full_img_path = os.path.join('extracted_images', img_path)
+                                
+                                if os.path.exists(full_img_path):
+                                    # Add image with max width of 4 inches
+                                    rl_img = RLImage(full_img_path, width=4*inch, height=3*inch, kind='proportional')
+                                    story.append(rl_img)
+                                    story.append(Spacer(1, 0.1*inch))
+                            except Exception as e:
+                                logger.warning(f"Could not add image to PDF: {e}")
+                    
+                    story.append(Spacer(1, 0.15*inch))
+            
+            # Add separator between conversations (except after last one)
+            if idx < len(request.conversations):
+                story.append(Spacer(1, 0.2*inch))
+                # Add a horizontal line as separator
+                line_table = Table([['']], colWidths=[6.5*inch])
+                line_table.setStyle(TableStyle([
+                    ('LINEABOVE', (0, 0), (-1, 0), 1, colors.HexColor('#334155')),
+                ]))
+                story.append(line_table)
+                story.append(Spacer(1, 0.2*inch))
+        
+        # Build PDF
+        doc.build(story)
+        
+        # Get PDF from buffer
+        buffer.seek(0)
+        
+        # Return PDF as download
+        return StreamingResponse(
+            buffer,
+            media_type="application/pdf",
+            headers={
+                "Content-Disposition": f"attachment; filename=WHEC_Report_{datetime.now().strftime('%Y-%m-%d')}.pdf"
+            }
+        )
+        
+    except Exception as e:
+        logger.error(f"Error generating PDF report: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error generating report: {str(e)}")
 
 if __name__ == "__main__":
     import uvicorn