Update app.py

app.py

@@ -1,60 +1,211 @@
-import json
+import gradio as gr
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from io import BytesIO
+import os
+import logging
+import base64
+import shutil
+import tempfile
+from simple_salesforce import Salesforce
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib import colors
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
+from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from transformers import BertTokenizer, BertForSequenceClassification
 import torch
-from flask import Flask, request, jsonify
 
-# Initialize Flask app
-app = Flask(__name__)
+# Configure logging to show detailed messages
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
-# Load the pre-trained BERT model and tokenizer (Replace with your own fine-tuned model path)
-model = BertForSequenceClassification.from_pretrained('path_to_your_finetuned_model')  # Replace with your model path
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  # Replace with your tokenizer path if custom
+# Salesforce credentials (use environment variables in production)
+SALESFORCE_USERNAME = os.getenv("SALESFORCE_USERNAME", "user@example.com")
+SALESFORCE_PASSWORD = os.getenv("SALESFORCE_PASSWORD", "password")
+SALESFORCE_SECURITY_TOKEN = os.getenv("SALESFORCE_SECURITY_TOKEN", "security_token")
+SALESFORCE_DOMAIN = os.getenv("SALESFORCE_DOMAIN", "login")
 
-# Define the possible risk levels (These are the output classes from your fine-tuned BERT model)
-risk_labels = ["low", "medium", "high"]
+# Load the BERT model and tokenizer for risk classification (fine-tuned for contract clauses)
+model = BertForSequenceClassification.from_pretrained('path_to_finetuned_model')  # Replace with your model path
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
-# Function to process contract text through the BERT model
+# Function to authenticate with Salesforce
+def get_salesforce_connection():
+    try:
+        sf = Salesforce(
+            username=SALESFORCE_USERNAME,
+            password=SALESFORCE_PASSWORD,
+            security_token=SALESFORCE_SECURITY_TOKEN,
+            domain=SALESFORCE_DOMAIN
+        )
+        return sf
+    except Exception as e:
+        logger.error(f"Failed to connect to Salesforce: {str(e)}", exc_info=True)
+        return None
+
+# Function to parse contract and predict risk score using BERT model
 def process_contract(contract_text):
-    # Tokenize input contract text
     inputs = tokenizer(contract_text, return_tensors="pt", truncation=True, padding=True, max_length=512)
     
-    # Perform inference with the BERT model to get the logits (raw prediction scores)
     with torch.no_grad():
         outputs = model(**inputs)
     
-    # Extract logits (raw prediction scores)
     logits = outputs.logits
-    
-    # Predict the risk level (index of max logit score)
     predicted_class = torch.argmax(logits, dim=1).item()
     
-    # Map prediction index to risk level
+    risk_labels = ["low", "medium", "high"]
     risk_tag = risk_labels[predicted_class]
     
-    # Return the predicted risk tag and the confidence score (highest logit)
-    return {
-        "risk_tag": risk_tag,
-        "score": logits.max().item(),  # Confidence score
-        "raw_scores": logits.squeeze().tolist()  # Optionally return the raw logits for more insight
+    return risk_tag, logits.max().item()
+
+# Function to generate a heatmap of the contract with section-wise risk levels
+def generate_heatmap(contract_text):
+    # Assuming the contract is split into sections; this is a simplified approach
+    sections = contract_text.split("\n\n")  # Split by paragraphs/sections
+    risks = []
+    for section in sections:
+        risk_tag, score = process_contract(section)
+        risks.append((section, risk_tag, score))
+    
+    # Create a heatmap visualization
+    fig, ax = plt.subplots(figsize=(10, len(sections) * 0.5))
+    ax.barh(range(len(sections)), [r[2] for r in risks], color='red', height=0.4)
+    
+    ax.set_yticks(range(len(sections)))
+    ax.set_yticklabels([r[0][:50] for r in risks])  # Display first 50 characters of each section as label
+    ax.set_xlabel('Risk Score')
+    ax.set_title('Risk Heatmap of Contract Sections')
+    
+    # Adjust layout and return the figure
+    plt.tight_layout()
+    return fig
+
+# Function to generate comprehensive PDF report
+def generate_pdf_report(project_title, risk_tags, ai_plan_score, estimated_duration, location, weather, gantt_chart_path=None):
+    pdf_file = BytesIO()
+    doc = SimpleDocTemplate(pdf_file, pagesize=letter)
+    styles = getSampleStyleSheet()
+    elements = []
+    
+    title_style = ParagraphStyle('Title', parent=styles['Heading1'], fontSize=18, alignment=1, spaceAfter=20)
+    elements.append(Paragraph(f"Project Report: {project_title}", title_style))
+    
+    details_style = styles['BodyText']
+    details = [
+        f"<b>Location:</b> {location}",
+        f"<b>Weather:</b> {weather.capitalize()}",
+        f"<b>Estimated Duration:</b> {estimated_duration} days",
+        f"<b>AI Plan Score:</b> {ai_plan_score:.1f}%",
+    ]
+    for detail in details:
+        elements.append(Paragraph(detail, details_style))
+    
+    elements.append(Spacer(1, 12))
+    elements.append(Paragraph("<b>Risk Assessment:</b>", styles['Heading2']))
+    
+    for risk in risk_tags.split("\n"):
+        elements.append(Paragraph(f"• {risk}", details_style))
+    
+    if gantt_chart_path:
+        elements.append(Spacer(1, 24))
+        elements.append(Paragraph("<b>Project Timeline:</b>", styles['Heading2']))
+        img = Image(gantt_chart_path, width=6 * inch, height=4 * inch)
+        elements.append(img)
+    
+    doc.build(elements)
+    pdf_file.seek(0)
+    return pdf_file
+
+# Function to upload the generated PDF to Salesforce
+def upload_pdf_to_salesforce(pdf_file, project_title):
+    sf = get_salesforce_connection()
+    if not sf:
+        logger.error("Salesforce connection failed. Cannot upload PDF.")
+        return None, None
+
+    encoded_pdf_data = base64.b64encode(pdf_file.getvalue()).decode('utf-8')
+    content_version_data = {
+        "Title": f"{project_title} - Comprehensive Report",
+        "PathOnClient": f"{project_title}_Report.pdf",
+        "VersionData": encoded_pdf_data,
     }
+    content_version = sf.ContentVersion.create(content_version_data)
+    content_version_id = content_version["id"]
+    result = sf.query(f"SELECT Id, ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'")
+    content_document_id = result['records'][0]['ContentDocumentId']
+    file_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{content_version_id}"
+    return content_version_id, file_url
+
+# Function to log project data to Salesforce
+def send_to_salesforce(project_title, gantt_chart_url, ai_plan_score, estimated_duration, risk_tags, status="Draft", record_id=None, location="", weather_type=""):
+    sf = get_salesforce_connection()
+    if not sf:
+        logger.error("Salesforce connection failed. Cannot proceed with record creation/update.")
+        return None
+    
+    sf_data = {
+        "Name": project_title[:80],
+        "Project_Title__c": project_title,
+        "Estimated_Duration__c": estimated_duration,
+        "AI_Plan_Score__c": ai_plan_score,
+        "Status__c": status,
+        "Location__c": location,
+        "Weather_Type__c": weather_type,
+        "Risk_Tags__c": risk_tags,
+    }
+    
+    if gantt_chart_url:
+        sf_data["Gantt_Chart_PDF__c"] = gantt_chart_url
+    
+    if record_id:
+        sf.AI_Project_Timeline__c.update(record_id, sf_data)
+        return record_id
+    else:
+        project_record = sf.AI_Project_Timeline__c.create(sf_data)
+        return project_record['id']
 
-# Flask route to process the contract text (POST request)
-@app.route('/process_contract', methods=['POST'])
-def analyze_contract():
+# Gradio interface function
+def gradio_interface(boq_file, weather, location, project_title):
     try:
-        # Get JSON data from the incoming POST request
-        data = request.json
-        contract_text = data['contract_text']
+        if not boq_file:
+            return None, "Error: No BOQ file uploaded", None, None
+        
+        fig = generate_heatmap(boq_file)
+        risk_tags = "Risk tags will be displayed here..."  # Generate risk tags logic based on contract analysis
         
-        # Process the contract text and extract risk details
-        result = process_contract(contract_text)
+        # Generating PDF report
+        pdf_report = generate_pdf_report(project_title, risk_tags, ai_plan_score=90, estimated_duration=30, location=location, weather=weather)
         
-        # Return the result as JSON
-        return jsonify(result), 200
+        # Upload to Salesforce
+        pdf_content_id, pdf_url = upload_pdf_to_salesforce(pdf_report, project_title)
+        
+        return fig, risk_tags, pdf_url, pdf_report
     except Exception as e:
-        # Return error message in case of failure
-        return jsonify({"error": str(e)}), 500
+        logger.error(f"Error in Gradio interface: {str(e)}")
+        return None, f"Error in Gradio interface: {str(e)}", None, None
+
+# Create Gradio interface
+demo = gr.Blocks()
+with demo:
+    gr.Markdown("## Contract Risk Analyzer")
+    gr.Markdown("Upload a contract, and the system will generate a heatmap and PDF report highlighting risk-prone clauses.")
+    
+    with gr.Row():
+        with gr.Column():
+            contract_file = gr.File(label="Upload Contract (PDF or Text)")
+            weather = gr.Dropdown(label="Weather", choices=["sunny", "rainy", "cloudy"], value="sunny")
+            location = gr.Textbox(label="Location", placeholder="Enter project location")
+            project_title = gr.Textbox(label="Project Title", placeholder="Enter project title")
+            submit_btn = gr.Button("Analyze Contract")
+        
+        with gr.Column():
+            plot_output = gr.Plot(label="Heatmap Visualization")
+            risk_tags_output = gr.Textbox(label="Risk Tags")
+            download_pdf = gr.File(label="Download Full Report (PDF)")
+
+    submit_btn.click(fn=gradio_interface, inputs=[contract_file, weather, location, project_title], outputs=[plot_output, risk_tags_output, download_pdf])
 
-# Run the Flask app
-if __name__ == '__main__':
-    app.run(debug=True, host="0.0.0.0", port=5000)
+if __name__ == "__main__":
+    demo.launch()