Update app.py

app.py

@@ -389,7 +389,7 @@ def generate_sentiment_pdf(sentiment_score: float) -> 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)"
@@ -397,7 +397,7 @@ def generate_sentiment_pdf(sentiment_score: float) -> BytesIO:
         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']))
-       
+        
         doc.build(story)
         buffer.seek(0)
         return buffer
@@ -405,10 +405,9 @@ def generate_sentiment_pdf(sentiment_score: float) -> BytesIO:
         logger.error(f"PDF generation failed: {str(e)}")
         raise Exception(f"PDF generation failed: {str(e)}")
 
-def save_to_salesforce(sf: Salesforce, data: Dict) -> Dict:
-    """Save analysis results and PDF to Salesforce, return record and file IDs"""
+def save_to_salesforce(sf: Salesforce, data: Dict) -> str:
+    """Save analysis results to Salesforce Custom_Risk_Analysis__c object"""
     try:
-        # Save Custom_Risk_Analysis__c record
         record = {
             'Sentiment_Score__c': data['sentiment_score'],
             'Risk_Score__c': data['risk_score'],
@@ -422,36 +421,10 @@ def save_to_salesforce(sf: Salesforce, data: Dict) -> Dict:
         }
         result = sf.Custom_Risk_Analysis__c.create(record)
         logger.info(f"Successfully created Salesforce record: {result['id']}")
-       
-        # Generate and save PDF as ContentVersion
-        pdf_buffer = generate_sentiment_pdf(data['sentiment_score'])
-        pdf_data = base64.b64encode(pdf_buffer.getvalue()).decode('utf-8')
-       
-        content_version = {
-            'Title': f"Sentiment_Analysis_{data['record_id']}.pdf",
-            'PathOnClient': f"Sentiment_Analysis_{data['record_id']}.pdf",
-            'VersionData': pdf_data,
-            'FirstPublishLocationId': result['id']
-        }
-        file_result = sf.ContentVersion.create(content_version)
-        logger.info(f"Successfully uploaded PDF to Salesforce: {file_result['id']}")
-       
-        # Get the ContentDocumentId
-        content_version_data = sf.ContentVersion.get(file_result['id'])
-        content_document_id = content_version_data['ContentDocumentId']
-       
-        # Construct the download URL
-        instance_url = sf.base_url.rstrip('/s')
-        download_url = f"{instance_url}/sfc/servlet.shepherd/version/download/{file_result['id']}"
-       
-        return {
-            'record_id': result['id'],
-            'file_id': file_result['id'],
-            'download_url': download_url
-        }
+        return result['id']
     except Exception as e:
         logger.error(f"Failed to save to Salesforce: {str(e)}")
-        raise Exception(f"Salesforce record or file creation failed: {str(e)}")
+        raise Exception(f"Salesforce record creation failed: {str(e)}")
 
 def extract_text_from_pdf(pdf_path: str) -> str:
     """Extract text from PDF using pdfplumber"""
@@ -482,7 +455,7 @@ def find_penalty_values(text: str) -> List[float]:
         r'\d+\s*(?:percent|%)',
         r'(?:\b[a-z]+\s*)+dollars',
     ]
-   
+    
     penalties = []
     for pattern in patterns:
         matches = re.finditer(pattern, text, flags=re.IGNORECASE)
@@ -502,7 +475,7 @@ def calculate_risk_score(penalty_count: int, penalty_values: List[float], obliga
     """Calculate risk score based on various factors"""
     score = 0
     score += min(penalty_count * 5, 30)
-   
+    
     if penalty_values:
         avg_penalty = sum(penalty_values) / len(penalty_values)
         if avg_penalty > 1000000:
@@ -513,11 +486,11 @@ def calculate_risk_score(penalty_count: int, penalty_values: List[float], obliga
             score += 15
         else:
             score += 5
-   
+    
     score += min(obligation_count * 2, 20)
     score += min(delay_count * 10, 30)
     score = min(score, 100)
-   
+    
     if score < 30:
         return score, "Low"
     elif score < 70:
@@ -557,7 +530,7 @@ def generate_heatmap(risk_level: str):
     """Generate a simple heatmap based on risk level"""
     try:
         fig, ax = plt.subplots(figsize=(8, 2))
-       
+        
         if risk_level == "Low":
             cmap = plt.cm.Blues
             color = '#4CAF50'
@@ -567,14 +540,14 @@ def generate_heatmap(risk_level: str):
         else:
             cmap = plt.cm.Reds
             color = '#F44336'
-       
+        
         gradient = np.linspace(0, 1, 256).reshape(1, -1)
         gradient = np.vstack((gradient, gradient))
-       
+        
         ax.imshow(gradient, aspect='auto', cmap=cmap)
         ax.text(128, 0.5, f"{risk_level} Risk", color='white' if risk_level in ["High", "Medium"] else 'black',
                 ha='center', va='center', fontsize=24, fontweight='bold')
-       
+        
         ax.set_axis_off()
         plt.tight_layout()
         return fig
@@ -602,7 +575,7 @@ def format_clause_example(example: str, index: int) -> str:
     </div>
     """
 
-def analyze_pdf(file_obj) -> List:
+def analyze_pdf(file_obj) -> tuple:
     """Main analysis function for Gradio interface"""
     try:
         if not file_obj:
@@ -629,17 +602,17 @@ def analyze_pdf(file_obj) -> List:
         penalty_keywords = ["penalty", "fine", "forfeit", "liquidated damages", "breach"]
         obligation_keywords = ["shall", "must", "required to", "obligated to", "duty"]
         delay_keywords = ["delay", "late", "overdue", "extension", "time is of the essence"]
-       
+        
         penalty_counts = count_keywords(text, penalty_keywords)
         obligation_counts = count_keywords(text, obligation_keywords)
         delay_counts = count_keywords(text, delay_keywords)
-       
+        
         penalty_values = find_penalty_values(text)
-       
+        
         total_penalties = sum(penalty_counts.values())
         total_obligations = sum(obligation_counts.values())
         total_delays = sum(delay_counts.values())
-       
+        
         # Generate warning messages with emojis
         penalty_warning = format_warning_message(total_penalties, "penalty", "💰")
         obligation_warning = format_warning_message(total_obligations, "obligation", "📝")
@@ -685,14 +658,14 @@ def analyze_pdf(file_obj) -> List:
         """
 
         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_sentences = []
         for sentence in re.split(r'(?<=[.!?])\s+', text):
             if any(kw.lower() in sentence.lower() for kw in penalty_keywords):
                 penalty_sentences.append(sentence.strip())
-       
+        
         extracted_data = "\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 = {
             'sentiment_score': sentiment_score,
@@ -705,15 +678,12 @@ def analyze_pdf(file_obj) -> List:
             '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()])
         }
-       
+        
         try:
-            salesforce_result = save_to_salesforce(sf, sf_data)
-            salesforce_record_id = salesforce_result['record_id']
-            download_url = salesforce_result['download_url']
-            logger.info(f"Saved to Salesforce with Record ID: {salesforce_record_id}, PDF URL: {download_url}")
+            salesforce_record_id = save_to_salesforce(sf, sf_data)
+            logger.info(f"Saved to Salesforce with ID: {salesforce_record_id}")
         except Exception as e:
             logger.error(f"Salesforce record creation failed: {str(e)}")
-            download_url = None
 
         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 "🚨"
@@ -722,8 +692,25 @@ def analyze_pdf(file_obj) -> List:
             "Medium": "This contract has moderate risk. Careful review advised.",
             "High": "This contract is high risk! Immediate legal review required."
         }
-       
-        # Add sentiment analysis PDF download link
+        
+        # Generate PDF for download
+        pdf_buffer = generate_sentiment_pdf(sentiment_score)
+        pdf_file = (f"Sentiment_Analysis_{record_id}.pdf", pdf_buffer.read(), "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>
+        """
+
         sentiment_analysis_output = f"""
         <div class='result-box'>
             <div class='section-title'>📊 Sentiment Analysis</div>
@@ -732,35 +719,20 @@ def analyze_pdf(file_obj) -> List:
                 <span class='risk-score'>{sentiment_score:.2f}</span>
             </div>
             {sentiment_meter}
-            <div style='margin-top: 15px;'>
-                <strong>Download Sentiment Report:</strong>
-                {'<a href="' + download_url + '" target="_blank" style="color: var(--primary-color); text-decoration: underline;">Download PDF</a>' if download_url else 'Failed to generate PDF download link'}
-            </div>
         </div>
         """
 
-        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>
-            """,
+        return (
+            risk_summary_output,
             "",  # 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
-        ]
+            sentiment_analysis_output,
+            pdf_file
+        )
     except Exception as e:
         logger.error(f"Analysis failed: {str(e)}")
         error_message = f"""
@@ -775,7 +747,7 @@ def analyze_pdf(file_obj) -> List:
             </div>
         </div>
         """
-        return [error_message] * 8
+        return [error_message] * 8 + (None,)
 
 # 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:
@@ -787,7 +759,7 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
             </p>
         </div>
     """)
-   
+    
     with gr.Row():
         with gr.Column(scale=1):
             file_input = gr.File(
@@ -801,28 +773,31 @@ with gr.Blocks(css=css, title="PDF Contract Risk Analyzer", theme=gr.themes.Defa
                 </div>
             """)
             submit_btn = gr.Button("Analyze Contract", variant="primary")
-       
+        
         with gr.Column(scale=3):
             risk_summary = gr.HTML(label="Contract Risk Summary")
             risk_visualization = gr.HTML(label="Risk Visualization", visible=False, elem_id="risk-visualization")
-   
+    
     with gr.Row():
         with gr.Column():
             penalty_count = gr.HTML(label="Penalty Clauses Analysis")
             penalty_amounts = gr.HTML(label="Penalty Amounts Found")
-       
+        
         with gr.Column():
             obligation_count = gr.HTML(label="Obligation Clauses Analysis")
-       
+        
         with gr.Column():
             delay_count = gr.HTML(label="Delay Clauses Analysis")
-   
+    
     with gr.Row():
         extracted_data = gr.HTML(label="Extracted Data")
-   
+    
     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],
@@ -830,7 +805,8 @@ 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
+            extracted_data, sentiment_analysis,
+            pdf_download
         ]
     )