Update app.py

app.py

@@ -12,11 +12,9 @@ from io import BytesIO
 import uuid
 import logging
 import textwrap
-from reportlab.lib.pagesizes import letter
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image
-from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
-from reportlab.lib.enums import TA_CENTER
-from reportlab.lib import colors
+from fpdf import FPDF
+from datetime import datetime
+from bs4 import BeautifulSoup
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -341,6 +339,9 @@ mark {
     margin-top: 20px !important;
     background: linear-gradient(135deg, #4CAF50, #2E7D32) !important;
 }
+.download-btn:hover {
+    background: linear-gradient(135deg, #2E7D32, #4CAF50) !important;
+}
 /* Hide elements */
 footer, .gradio-footer, .hide, [data-testid="Use via API"], [data-testid="mmsettings"], 
 .sentiment-analysis, .risk-visualization {
@@ -501,13 +502,7 @@ def generate_heatmap(risk_level: str):
        
         ax.set_axis_off()
         plt.tight_layout()
-        
-        # Save the figure to a BytesIO object
-        buf = BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight')
-        plt.close(fig)
-        buf.seek(0)
-        return buf
+        return fig
     except Exception as e:
         logger.error(f"Heatmap generation failed: {str(e)}")
         raise Exception(f"Heatmap generation failed: {str(e)}")
@@ -585,80 +580,112 @@ def format_clause_example(example: str, index: int) -> str:
     </div>
     """
 
-def generate_pdf_report(analysis_results: dict, heatmap_buffer: BytesIO) -> BytesIO:
-    """Generate a PDF report from the analysis results"""
-    buffer = BytesIO()
-    doc = SimpleDocTemplate(buffer, pagesize=letter)
-    
-    styles = getSampleStyleSheet()
-    styles.add(ParagraphStyle(name='Center', alignment=TA_CENTER))
-    styles.add(ParagraphStyle(name='RiskHigh', textColor=colors.red, fontSize=14))
-    styles.add(ParagraphStyle(name='RiskMedium', textColor=colors.orange, fontSize=14))
-    styles.add(ParagraphStyle(name='RiskLow', textColor=colors.green, fontSize=14))
-    styles.add(ParagraphStyle(name='NormalBold', fontName='Helvetica-Bold', parent=styles['Normal']))
-    
-    story = []
-    
-    # Title
-    story.append(Paragraph("Contract Risk Analysis Report", styles['Title']))
-    story.append(Spacer(1, 12))
-    
-    # Risk Summary
-    story.append(Paragraph("<b>Risk Summary</b>", styles['Heading2']))
-    risk_style = 'RiskHigh' if analysis_results['risk_level'] == 'High' else 'RiskMedium' if analysis_results['risk_level'] == 'Medium' else 'RiskLow'
-    story.append(Paragraph(f"Risk Score: {analysis_results['risk_score']:.1f}/100 - <b>{analysis_results['risk_level']} Risk</b>", styles[risk_style]))
-    story.append(Spacer(1, 12))
-    
-    # Add heatmap to PDF
-    if heatmap_buffer:
-        story.append(Image(heatmap_buffer, width=400, height=100))
-        story.append(Spacer(1, 24))
-    
-    # Key Findings
-    story.append(Paragraph("<b>Key Findings</b>", styles['Heading2']))
-    story.append(Paragraph(f"<b>Penalty Clauses:</b> {analysis_results['penalty_count']}", styles['Normal']))
-    story.append(Paragraph(f"<b>Obligation Clauses:</b> {analysis_results['obligation_count']}", styles['Normal']))
-    story.append(Paragraph(f"<b>Delay Clauses:</b> {analysis_results['delay_count']}", styles['Normal']))
-    story.append(Spacer(1, 12))
-    
-    if analysis_results['penalty_values']:
-        story.append(Paragraph("<b>Penalty Amounts Found</b>", styles['Heading3']))
-        for amount in analysis_results['penalty_values'][:5]:
-            story.append(Paragraph(f"${amount:,.2f}", styles['Normal']))
-        story.append(Spacer(1, 12))
-    
-    # Detailed Findings
-    story.append(Paragraph("<b>Detailed Findings</b>", styles['Heading2']))
-    
-    # Penalty Details
-    if analysis_results['penalty_details']:
-        story.append(Paragraph("<b>Penalty Clauses</b>", styles['Heading3']))
-        story.append(Paragraph(analysis_results['penalty_details'], styles['Normal']))
-        story.append(Spacer(1, 12))
+def generate_full_pdf_report(
+    risk_summary: str,
+    penalty_details: str,
+    penalty_amounts: str,
+    obligation_details: str,
+    delay_details: str,
+    penalty_examples: str,
+    keyword_matches: str,
+    file_name: str
+) -> BytesIO:
+    """Generate a comprehensive PDF report of all analysis results"""
+    try:
+        # Remove HTML tags for PDF content
+        def clean_html(html):
+            soup = BeautifulSoup(html, 'html.parser')
+            return soup.get_text(separator='\n')
+        
+        # Create PDF object
+        pdf = FPDF()
+        pdf.add_page()
+        pdf.set_auto_page_break(auto=True, margin=15)
+        
+        # Set font and colors
+        pdf.set_font("Arial", size=12)
+        
+        # Add title
+        pdf.set_fill_color(30, 144, 255)  # DodgerBlue
+        pdf.set_text_color(255, 255, 255)
+        pdf.cell(0, 10, f"Contract Risk Analysis Report - {file_name}", ln=1, fill=True, align='C')
+        pdf.ln(10)
+        
+        # Add report date
+        pdf.set_text_color(0, 0, 0)
+        pdf.set_font("Arial", 'I', 10)
+        pdf.cell(0, 10, f"Report generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=1)
+        pdf.ln(5)
+        
+        # Add sections
+        sections = [
+            ("Risk Summary", risk_summary),
+            ("Penalty Analysis", penalty_details),
+            ("Penalty Amounts", penalty_amounts),
+            ("Obligation Analysis", obligation_details),
+            ("Delay Analysis", delay_details),
+            ("Sample Clauses", penalty_examples),
+            ("Keyword Matches", keyword_matches)
+        ]
+        
+        for title, content in sections:
+            # Add section header
+            pdf.set_font("Arial", 'B', 14)
+            pdf.set_text_color(30, 144, 255)  # DodgerBlue
+            pdf.cell(0, 10, title, ln=1)
+            pdf.ln(2)
+            
+            # Add section content
+            pdf.set_font("Arial", size=10)
+            pdf.set_text_color(0, 0, 0)
+            
+            # Clean HTML and add multi_cell for proper wrapping
+            cleaned_content = clean_html(content)
+            pdf.multi_cell(0, 8, cleaned_content)
+            pdf.ln(5)
+            
+            # Add page break if needed
+            if pdf.get_y() > 250:
+                pdf.add_page()
+        
+        # Save to bytes buffer
+        pdf_bytes = BytesIO()
+        pdf.output(pdf_bytes)
+        pdf_bytes.seek(0)
+        return pdf_bytes
+        
+    except Exception as e:
+        logger.error(f"PDF generation failed: {str(e)}")
+        raise Exception(f"Failed to generate PDF report: {str(e)}")
+
+def generate_error_pdf(error_msg: str) -> BytesIO:
+    """Generate a PDF for error cases"""
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
     
-    # Obligation Details
-    if analysis_results['obligation_details']:
-        story.append(Paragraph("<b>Obligation Clauses</b>", styles['Heading3']))
-        story.append(Paragraph(analysis_results['obligation_details'], styles['Normal']))
-        story.append(Spacer(1, 12))
+    # Add title
+    pdf.set_fill_color(255, 0, 0)  # Red
+    pdf.set_text_color(255, 255, 255)
+    pdf.cell(0, 10, "Contract Analysis Error Report", ln=1, fill=True, align='C')
+    pdf.ln(20)
     
-    # Delay Details
-    if analysis_results['delay_details']:
-        story.append(Paragraph("<b>Delay Clauses</b>", styles['Heading3']))
-        story.append(Paragraph(analysis_results['delay_details'], styles['Normal']))
-        story.append(Spacer(1, 12))
+    # Add error message
+    pdf.set_text_color(0, 0, 0)
+    pdf.set_font("Arial", 'B', 14)
+    pdf.cell(0, 10, "Analysis Failed", ln=1)
+    pdf.ln(5)
     
-    # Sentiment Analysis
-    story.append(Paragraph("<b>Sentiment Analysis</b>", styles['Heading2']))
-    sentiment_text = "Positive" if analysis_results['sentiment_score'] > 0.6 else "Negative" if analysis_results['sentiment_score'] < 0.4 else "Neutral"
-    story.append(Paragraph(f"Sentiment Score: {analysis_results['sentiment_score']:.2f} ({sentiment_text})", styles['Normal']))
+    pdf.set_font("Arial", size=10)
+    pdf.multi_cell(0, 8, error_msg)
     
-    doc.build(story)
-    buffer.seek(0)
-    return buffer
+    pdf_bytes = BytesIO()
+    pdf.output(pdf_bytes)
+    pdf_bytes.seek(0)
+    return pdf_bytes
 
-def analyze_pdf(file_obj) -> List:
-    """Main analysis function for Gradio interface"""
+def analyze_pdf(file_obj) -> Tuple[List, BytesIO]:
+    """Main analysis function that now also returns PDF bytes"""
     try:
         if not file_obj:
             raise Exception("No PDF file uploaded. Please upload a valid PDF file.")
@@ -757,7 +784,7 @@ def analyze_pdf(file_obj) -> List:
             raise Exception(f"Risk score calculation failed: {str(e)}")
 
         try:
-            heatmap_buffer = generate_heatmap(risk_level)
+            heatmap = generate_heatmap(risk_level)
             risk_meter = generate_risk_meter(risk_score)
             sentiment_meter = generate_sentiment_meter(sentiment_score)
         except Exception as e:
@@ -779,7 +806,7 @@ def analyze_pdf(file_obj) -> List:
                 """)
             return "".join(details)
 
-        penalty_details = f"""
+        penalty_details_html = f"""
         {penalty_warning}
         <div class='penalty-box'>
             <div class='section-title'>💰 Penalty Clause Analysis</div>
@@ -787,7 +814,7 @@ def analyze_pdf(file_obj) -> List:
         </div>
         """
 
-        obligation_details = f"""
+        obligation_details_html = f"""
         {obligation_warning}
         <div class='obligation-box'>
             <div class='section-title'>📝 Obligation Clause Analysis</div>
@@ -795,7 +822,7 @@ def analyze_pdf(file_obj) -> List:
         </div>
         """
 
-        delay_details = f"""
+        delay_details_html = f"""
         {delay_warning}
         <div class='delay-box'>
             <div class='section-title'>⏱️ Delay Clause Analysis</div>
@@ -803,14 +830,14 @@ def analyze_pdf(file_obj) -> List:
         </div>
         """
 
-        penalty_amounts = "\n".join([f"<div class='count-item'><span class='count-label'>💰 Amount</span><span class='count-value'>${amt:,.2f}</span></div>" for amt in penalty_values[:5]]) if penalty_values else "<div class='success-box'>✅ No specific penalty amounts found - This is good news!</div>"
+        penalty_amounts_html = "\n".join([f"<div class='count-item'><span class='count-label'>💰 Amount</span><span class='count-value'>${amt:,.2f}</span></div>" for amt in penalty_values[:5]]) if penalty_values else "<div class='success-box'>✅ No specific penalty amounts found - This is good news!</div>"
        
         penalty_sentences = []
         for sentence in re.split(r'(?<=[.!?])\s+', text):
             if any(kw.lower() in sentence.lower() for kw in penalty_keywords.keys()):
                 penalty_sentences.append(sentence.strip())
        
-        penalty_examples = "\n".join([format_clause_example(sent, i+1) for i, sent in enumerate(penalty_sentences[:3])]) if penalty_sentences else "<div class='success-box'>✅ No penalty clauses found - Excellent contract terms!</div>"
+        penalty_examples_html = "\n".join([format_clause_example(sent, i+1) for i, sent in enumerate(penalty_sentences[:3])]) if penalty_sentences else "<div class='success-box'>✅ No penalty clauses found - Excellent contract terms!</div>"
        
         record_id = str(uuid.uuid4())
         sf_data = {
@@ -818,7 +845,7 @@ def analyze_pdf(file_obj) -> List:
             'risk_score': risk_score,
             'risk_level': risk_level,
             'record_id': record_id,
-            'penalty_examples': penalty_examples,
+            'penalty_examples': penalty_examples_html,
             'penalty_details': "\n".join([f"{kw}: {count}" for kw, count in penalty_counts.items()]),
             'penalty_amounts': "\n".join([f"${amt:,.2f}" for amt in penalty_values[:5]]) if penalty_values else "No specific penalty amounts found",
             'obligation_details': "\n".join([f"{kw}: {count}" for kw, count in obligation_counts.items()]),
@@ -863,21 +890,6 @@ def analyze_pdf(file_obj) -> List:
             "High": "This contract is high risk! Immediate legal review required."
         }
 
-        risk_summary_html = f"""
-        <div class='result-box'>
-            {sf_link_html}
-            <div class='section-title'>{risk_icon} Contract Risk Summary</div>
-            <div class='risk-row'>
-                <span class='risk-label'>Overall Risk Score</span>
-                <span class='risk-score risk-{risk_level.lower()}'>{risk_score:.1f}/100</span>
-            </div>
-            {risk_meter}
-            <div style='margin-top: 15px; font-size: 16px;'>
-                <strong>Assessment:</strong> {risk_advice[risk_level]}
-            </div>
-        </div>
-        """
-
         # Generate keyword matches section with highlighted context
         def format_keyword_matches(contexts_dict, title):
             sections = []
@@ -910,36 +922,44 @@ def analyze_pdf(file_obj) -> List:
         </div>
         """
 
-        # Prepare data for PDF report
-        analysis_results = {
-            'risk_score': risk_score,
-            'risk_level': risk_level,
-            'penalty_count': total_penalties,
-            'obligation_count': total_obligations,
-            'delay_count': total_delays,
-            'penalty_values': penalty_values,
-            'penalty_details': "Penalty clauses found in the document",
-            'obligation_details': "Obligation clauses found in the document",
-            'delay_details': "Delay clauses found in the document",
-            'sentiment_score': sentiment_score
-        }
+        risk_summary_html = f"""
+        <div class='result-box'>
+            {sf_link_html}
+            <div class='section-title'>{risk_icon} Contract Risk Summary</div>
+            <div class='risk-row'>
+                <span class='risk-label'>Overall Risk Score</span>
+                <span class='risk-score risk-{risk_level.lower()}'>{risk_score:.1f}/100</span>
+            </div>
+            {risk_meter}
+            <div style='margin-top: 15px; font-size: 16px;'>
+                <strong>Assessment:</strong> {risk_advice[risk_level]}
+            </div>
+        </div>
+        """
 
-        # Generate PDF report
-        pdf_buffer = generate_pdf_report(analysis_results, heatmap_buffer)
+        # Generate the PDF report
+        pdf_bytes = generate_full_pdf_report(
+            risk_summary=risk_summary_html,
+            penalty_details=penalty_details_html,
+            penalty_amounts=penalty_amounts_html,
+            obligation_details=obligation_details_html,
+            delay_details=delay_details_html,
+            penalty_examples=penalty_examples_html,
+            keyword_matches=keyword_matches_html,
+            file_name=os.path.basename(file_obj.name)
         
-        # Return all outputs including the PDF data
         return [
             risk_summary_html,
-            "",  # risk_visualization
-            penalty_details,
-            f"<div class='penalty-box'><div class='section-title'>💰 Penalty Amounts Found</div>{penalty_amounts}</div>",
-            obligation_details,
-            delay_details,
-            f"<div class='result-box'><div class='section-title'>📜 Text Extracted from PDF</div>{penalty_examples}</div>",
-            "",  # sentiment_analysis
-            keyword_matches_html,
-            pdf_buffer  # PDF data as the last output
-        ]
+            "",  # risk_visualization (hidden)
+            penalty_details_html,
+            penalty_amounts_html,
+            obligation_details_html,
+            delay_details_html,
+            penalty_examples_html,
+            "",  # sentiment_analysis (hidden)
+            keyword_matches_html
+        ], pdf_bytes
+        
     except Exception as e:
         logger.error(f"Analysis failed: {str(e)}")
         error_message = f"""
@@ -954,7 +974,8 @@ def analyze_pdf(file_obj) -> List:
             </div>
         </div>
         """
-        return [error_message] * 10  # Update to match total outputs including PDF
+        error_pdf = generate_error_pdf(str(e))
+        return [error_message] * 9, error_pdf
 
 # Create Gradio interface with blue theme and hidden elements
 with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Default(primary_hue="blue")) as demo:
@@ -983,7 +1004,6 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
        
         with gr.Column(scale=3):
             risk_summary = gr.HTML(label="Contract Risk Summary")
-            # Hidden risk visualization (kept in code but not displayed)
             risk_visualization = gr.HTML(label="Risk Visualization", visible=False)
    
     with gr.Row():
@@ -1000,15 +1020,13 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
     with gr.Row():
         penalty_examples = gr.HTML(label="Text Extracted from PDF")
    
-    # Hidden sentiment analysis (kept in code but not displayed)
     sentiment_analysis = gr.HTML(label="Contract Sentiment Analysis", visible=False)
-    
-    # Keyword matches output
     keyword_matches = gr.HTML(label="Keyword Matches in Document")
     
-    # PDF download button
-    pdf_download = gr.File(label="Download PDF Report", visible=False)
-   
+    # Add download button
+    download_btn = gr.Download(label="📥 Download Full Report as PDF")
+    
+    # Define the click event
     submit_btn.click(
         fn=analyze_pdf,
         inputs=file_input,
@@ -1017,18 +1035,12 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
             penalty_count, penalty_amounts,
             obligation_count, delay_count,
             penalty_examples, sentiment_analysis,
-            keyword_matches, pdf_download
+            keyword_matches
         ]
-    )
-    
-    # Show download button when PDF is ready
-    def toggle_download_visibility(pdf_data):
-        return gr.File.update(value=pdf_data, visible=pdf_data is not None)
-    
-    keyword_matches.change(
-        fn=toggle_download_visibility,
-        inputs=pdf_download,
-        outputs=pdf_download
+    ).then(
+        fn=lambda pdf_bytes: gr.File(value=pdf_bytes, visible=True),
+        inputs=submit_btn.outputs[1],  # The PDF bytes are the second return value
+        outputs=download_btn
     )
 
 if __name__ == "__main__":