Update app.py

app.py

@@ -15,10 +15,11 @@ import textwrap
 from reportlab.lib.pagesizes import letter
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
 from reportlab.lib.styles import getSampleStyleSheet
-from datetime import datetime
+from reportlab.lib.units import inch
+from tenacity import retry, stop_after_attempt, wait_fixed, retry_if_exception_type
 
-# Set up logging with timestamp
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Set up logging
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Custom CSS for styling with dark mode compatibility
@@ -357,6 +358,7 @@ SF_USERNAME = "Kushalpavansekharm503@agentforce.com"
 SF_PASSWORD = "Kushal@123"
 SF_TOKEN = "WwUIFWBVUjeKn9VPKyWJmawY0"
 
+@retry(stop_max_attempt_number=3, wait_fixed=2000, retry=retry_if_exception_type(Exception))
 def authenticate_salesforce() -> Salesforce:
     """Authenticate with Salesforce and return a Salesforce client"""
     try:
@@ -369,7 +371,7 @@ def authenticate_salesforce() -> Salesforce:
         return sf
     except Exception as e:
         logger.error(f"Failed to authenticate with Salesforce: {str(e)}")
-        raise Exception(f"Salesforce authentication failed: {str(e)}")
+        raise
 
 def get_hugging_face_sentiment(text: str) -> float:
     """Get sentiment score using Hugging Face model"""
@@ -383,35 +385,110 @@ def get_hugging_face_sentiment(text: str) -> float:
         logger.error(f"Hugging Face sentiment analysis failed: {str(e)}. Using fallback score.")
         return 0.5
 
-def generate_sentiment_pdf(sentiment_score: float) -> BytesIO:
-    """Generate a PDF with sentiment analysis results"""
+def generate_contract_analysis_pdf(data: Dict, temp_dir: str) -> Tuple[str, str, BytesIO]:
+    """Generate a PDF with contract analysis results"""
     try:
-        buffer = BytesIO()
-        doc = SimpleDocTemplate(buffer, pagesize=letter)
-        styles = getSampleStyleSheet()
-        story = []
-        
-        story.append(Paragraph("Sentiment Analysis Report", styles['Title']))
-        story.append(Spacer(1, 12))
-        sentiment_text = f"Sentiment Score: {sentiment_score:.2f} (Scale: 0.0 Negative to 1.0 Positive)"
-        story.append(Paragraph(sentiment_text, styles['Normal']))
-        story.append(Spacer(1, 12))
-        interpretation = "Positive" if sentiment_score > 0.6 else "Negative" if sentiment_score < 0.4 else "Neutral"
-        story.append(Paragraph(f"Interpretation: {interpretation}", styles['Normal']))
-        current_time = datetime.now().strftime("%I:%M %p IST, %B %d, %Y")
-        story.append(Spacer(1, 12))
-        story.append(Paragraph(f"Generated on: {current_time}", styles['Normal']))
-        
-        doc.build(story)
-        buffer.seek(0)
-        return buffer
+        pdf_filename = f"contract_analysis_{data['record_id']}_{int(time.time())}.pdf"
+        pdf_path = os.path.join(temp_dir, pdf_filename)
+        pdf_file = BytesIO()
+        c = canvas.Canvas(pdf_file, pagesize=letter)
+        
+        c.setFont("Helvetica-Bold", 16)
+        c.drawString(1 * inch, 10 * inch, "Contract Analysis Report")
+        
+        c.setFont("Helvetica", 12)
+        c.drawString(1 * inch, 9.5 * inch, f"Date: {time.strftime('%Y-%m-%d')}")
+        c.drawString(1 * inch, 9.2 * inch, f"Time: {time.strftime('%H:%M:%S')} IST")
+        
+        c.setFont("Helvetica-Bold", 14)
+        c.drawString(1 * inch, 8.7 * inch, f"Risk Score: {data['risk_score']:.1f}/100 ({data['risk_level']})")
+        c.drawString(1 * inch, 8.4 * inch, f"Sentiment Score: {data['sentiment_score']:.2f} (0.0 Negative to 1.0 Positive)")
+        
+        y_position = 8.0 * inch
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(1 * inch, y_position, "Summary:")
+        y_position -= 0.3 * inch
+        
+        c.setFont("Helvetica", 10)
+        summary_data = {
+            "Total Penalties": sum(data['penalty_counts'].values()),
+            "Total Obligations": sum(data['obligation_counts'].values()),
+            "Total Delays": sum(data['delay_counts'].values()),
+            "Analysis Timestamp": time.strftime("%Y-%m-%d %H:%M:%S IST")
+        }
+        for key, value in summary_data.items():
+            c.drawString(1 * inch, y_position, f"{key}: {value}")
+            y_position -= 0.25 * inch
+        
+        y_position -= 0.5 * inch
+        c.setFont("Helvetica-Bold", 12)
+        c.drawString(1 * inch, y_position, "Detailed Analysis:")
+        y_position -= 0.3 * inch
+        
+        c.setFont("Helvetica", 10)
+        sections = [
+            ("Penalty Clauses", data['penalty_details']),
+            ("Obligation Clauses", data['obligation_details']),
+            ("Delay Clauses", data['delay_details']),
+            ("Penalty Amounts", data['penalty_amounts']),
+            ("Extracted Data", data['penalty_examples'])
+        ]
+        for title, content in sections:
+            c.drawString(1 * inch, y_position, f"{title}:")
+            y_position -= 0.2 * inch
+            lines = content.split("\n")
+            for line in lines:
+                if y_position < 1 * inch:
+                    c.showPage()
+                    c.setFont("Helvetica", 10)
+                    y_position = 10 * inch
+                c.drawString(1.2 * inch, y_position, line)
+                y_position -= 0.2 * inch
+        
+        c.save()
+        pdf_file.seek(0)
+        
+        with open(pdf_path, "wb") as f:
+            f.write(pdf_file.getvalue())
+        
+        logger.info(f"PDF generated: {pdf_path}")
+        return pdf_path, f"file://{pdf_path}", pdf_file
+    except Exception as e:
+        logger.error(f"Error generating PDF: {e}")
+        return "", "", None
+
+def upload_pdf_to_salesforce(sf: Salesforce, pdf_file: BytesIO, report_id: str) -> str:
+    """Upload PDF to Salesforce and return the download URL"""
+    try:
+        if not pdf_file:
+            logger.error("No PDF file provided for upload")
+            return ""
+        
+        encoded_pdf = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+        content_version_data = {
+            "Title": f"Contract_Analysis_{int(time.time())}",
+            "PathOnClient": f"contract_analysis_{int(time.time())}.pdf",
+            "VersionData": encoded_pdf,
+            "FirstPublishLocationId": report_id
+        }
+        content_version = sf.ContentVersion.create(content_version_data)
+        result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version['id']}'")
+        
+        if not result['records']:
+            logger.error("Failed to retrieve ContentVersion")
+            return ""
+        
+        file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version['id']}"
+        logger.info(f"PDF uploaded to Salesforce: {file_url}")
+        return file_url
     except Exception as e:
-        logger.error(f"PDF generation failed: {str(e)}")
-        raise Exception(f"PDF generation failed: {str(e)}")
+        logger.error(f"Error uploading PDF to Salesforce: {e}")
+        return ""
 
-def save_to_salesforce(sf: Salesforce, data: Dict) -> str:
-    """Save analysis results to Salesforce Custom_Risk_Analysis__c object"""
+def save_to_salesforce(sf: Salesforce, data: Dict, pdf_file: BytesIO) -> Dict:
+    """Save analysis results and PDF to Salesforce, return record and file URLs"""
     try:
+        # Save Custom_Risk_Analysis__c record
         record = {
             'Sentiment_Score__c': data['sentiment_score'],
             'Risk_Score__c': data['risk_score'],
@@ -425,10 +502,17 @@ def save_to_salesforce(sf: Salesforce, data: Dict) -> str:
         }
         result = sf.Custom_Risk_Analysis__c.create(record)
         logger.info(f"Successfully created Salesforce record: {result['id']}")
-        return result['id']
+        
+        # Upload PDF if available
+        download_url = upload_pdf_to_salesforce(sf, pdf_file, result['id']) if pdf_file else ""
+        
+        return {
+            'record_id': result['id'],
+            'download_url': download_url
+        }
     except Exception as e:
         logger.error(f"Failed to save to Salesforce: {str(e)}")
-        raise Exception(f"Salesforce record creation failed: {str(e)}")
+        raise Exception(f"Salesforce record or file creation failed: {str(e)}")
 
 def extract_text_from_pdf(pdf_path: str) -> str:
     """Extract text from PDF using pdfplumber"""
@@ -579,9 +663,12 @@ def format_clause_example(example: str, index: int) -> str:
     </div>
     """
 
-def analyze_pdf(file_obj) -> tuple:
+def analyze_pdf(file_obj) -> List:
     """Main analysis function for Gradio interface"""
     try:
+        import tempfile
+        temp_dir = tempfile.mkdtemp()
+        
         if not file_obj:
             raise Exception("No PDF file uploaded. Please upload a valid PDF file.")
 
@@ -676,18 +763,24 @@ def analyze_pdf(file_obj) -> tuple:
             'risk_score': risk_score,
             'risk_level': risk_level,
             'record_id': record_id,
-            'penalty_examples': extracted_data,
+            'penalty_counts': penalty_counts,
             '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_counts': obligation_counts,
             'obligation_details': "\n".join([f"{kw}: {count}" for kw, count in obligation_counts.items()]),
-            'delay_details': "\n".join([f"{kw}: {count}" for kw, count in delay_counts.items()])
+            'delay_counts': delay_counts,
+            'delay_details': "\n".join([f"{kw}: {count}" for kw, count in delay_counts.items()]),
+            'penalty_examples': extracted_data
         }
         
         try:
-            salesforce_record_id = save_to_salesforce(sf, sf_data)
-            logger.info(f"Saved to Salesforce with ID: {salesforce_record_id}")
+            pdf_path, pdf_url, pdf_file = generate_contract_analysis_pdf(sf_data, temp_dir)
+            salesforce_result = save_to_salesforce(sf, sf_data, pdf_file)
+            download_url = salesforce_result['download_url'] or pdf_url
+            logger.info(f"Saved to Salesforce with Record ID: {salesforce_result['record_id']}, PDF URL: {download_url}")
         except Exception as e:
             logger.error(f"Salesforce record creation failed: {str(e)}")
+            download_url = pdf_url if pdf_url else ""
 
         box_class = "success-box" if risk_level == "Low" else "warning-box" if risk_level == "Medium" else "danger-box"
         risk_icon = "✅" if risk_level == "Low" else "⚠" if risk_level == "Medium" else "🚨"
@@ -697,24 +790,7 @@ def analyze_pdf(file_obj) -> tuple:
             "High": "This contract is high risk! Immediate legal review required."
         }
         
-        # Generate PDF for download
-        pdf_buffer = generate_sentiment_pdf(sentiment_score)
-        pdf_file = (f"Sentiment_Analysis_{record_id}.pdf", pdf_buffer.getvalue(), "application/pdf")
-
-        risk_summary_output = f"""
-        <div class='result-box'>
-            <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>
-        """
-
+        # Add sentiment analysis PDF download link
         sentiment_analysis_output = f"""
         <div class='result-box'>
             <div class='section-title'>📊 Sentiment Analysis</div>
@@ -723,20 +799,35 @@ def analyze_pdf(file_obj) -> tuple:
                 <span class='risk-score'>{sentiment_score:.2f}</span>
             </div>
             {sentiment_meter}
+            <div style='margin-top: 15px;'>
+                <strong>Download Full Report:</strong>
+                {'<a href="' + download_url + '" download="contract_analysis.pdf" style="color: var(--primary-color); text-decoration: underline;">Download PDF</a>' if download_url else 'Failed to generate PDF download link'}
+            </div>
         </div>
         """
 
-        return (
-            risk_summary_output,
+        return [
+            f"""
+            <div class='result-box'>
+                <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>
+            """,
             "",  # Empty string for hidden 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'>📜 Extracted Data</div>{extracted_data}</div>",
-            sentiment_analysis_output,
-            pdf_file
-        )
+            sentiment_analysis_output
+        ]
     except Exception as e:
         logger.error(f"Analysis failed: {str(e)}")
         error_message = f"""
@@ -751,7 +842,11 @@ def analyze_pdf(file_obj) -> tuple:
             </div>
         </div>
         """
-        return [error_message] * 8 + (None,)
+        return [error_message] * 8
+    finally:
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            logger.info(f"Cleaned up temporary directory: {temp_dir}")
 
 # Create Gradio interface with dark mode compatibility
 with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Default(primary_hue="blue")) as demo:
@@ -799,9 +894,6 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
     with gr.Row():
         sentiment_analysis = gr.HTML(label="Sentiment Analysis")
     
-    with gr.Row():
-        pdf_download = gr.File(label="Download Sentiment Report", visible=True)
-    
     submit_btn.click(
         fn=analyze_pdf,
         inputs=[file_input],
@@ -809,8 +901,7 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
             risk_summary, risk_visualization,
             penalty_count, penalty_amounts,
             obligation_count, delay_count,
-            extracted_data, sentiment_analysis,
-            pdf_download
+            extracted_data, sentiment_analysis
         ]
     )