Update app.py

app.py

@@ -10,6 +10,9 @@ from simple_salesforce import Salesforce
 from dotenv import load_dotenv
 import requests
 from datetime import datetime
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from reportlab.lib.utils import ImageReader
 
 # Set up logging
 logging.basicConfig(level=logging.DEBUG)
@@ -45,26 +48,70 @@ def generate_plot(planned_cost, actual_spend, forecast_cost):
     fig, ax = plt.subplots(figsize=(8, 6))
     categories = ['Planned Cost', 'Actual Spend', 'Forecasted Cost']
     values = [planned_cost, actual_spend, forecast_cost]
-    ax.bar(categories, values, color=['blue', 'orange', 'green'])
-    ax.set_title("Budget Overview")
-    ax.set_ylabel("Amount ($)")
-    plt.xticks(rotation=45)
+    ax.bar(categories, values, color=['#1f77b4', '#ff7f0e', '#2ca02c'])  # Improved colors
+    ax.set_title("Budget Overview", fontsize=14, pad=15)
+    ax.set_ylabel("Amount ($)", fontsize=12)
+    ax.tick_params(axis='x', rotation=45)
+    ax.grid(True, axis='y', linestyle='--', alpha=0.7)
     
     # Save plot to a bytes buffer for Salesforce
     buf = io.BytesIO()
-    plt.savefig(buf, format='png', bbox_inches='tight')
+    plt.savefig(buf, format='png', bbox_inches='tight', dpi=100)
     buf.seek(0)
     image_base64 = base64.b64encode(buf.read()).decode('utf-8')
     
-    # Save plot to another buffer for Gradio
+    # Save plot to another buffer for Gradio and PDF
     buf_gradio = io.BytesIO()
-    plt.savefig(buf_gradio, format='png', bbox_inches='tight')
+    plt.savefig(buf_gradio, format='png', bbox_inches='tight', dpi=100)
     buf_gradio.seek(0)
     # Convert BytesIO to PIL Image for Gradio
     gradio_image = Image.open(buf_gradio)
     
+    # Save plot to a temporary file for PDF generation
+    buf_pdf = io.BytesIO()
+    plt.savefig(buf_pdf, format='png', bbox_inches='tight', dpi=100)
+    buf_pdf.seek(0)
+    
     plt.close()
-    return image_base64, gradio_image
+    return image_base64, gradio_image, buf_pdf
+
+# Function to generate a PDF report locally
+def generate_pdf(planned_cost, actual_spend, forecast_cost, total_risk, insights, status, top_causes, forecast_cost_eur, chart_image):
+    pdf_path = f"budget_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
+    c = canvas.Canvas(pdf_path, pagesize=letter)
+    width, height = letter
+
+    # Title
+    c.setFont("Helvetica-Bold", 16)
+    c.drawString(50, height - 50, "Budget Overrun Risk Report")
+
+    # Prediction Metrics
+    c.setFont("Helvetica", 12)
+    y_position = height - 100
+    c.drawString(50, y_position, f"Planned Cost: ${planned_cost:,.2f}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Actual Spend: ${actual_spend:,.2f}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Forecasted Cost: ${forecast_cost:,.2f}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Forecasted Cost (EUR): {forecast_cost_eur if isinstance(forecast_cost_eur, str) else f'{forecast_cost_eur:,.2f}'}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Total Risk: {total_risk}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Status: {status}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Insights: {insights}")
+    y_position -= 20
+    c.drawString(50, y_position, f"Top Causes: {top_causes}")
+    y_position -= 40
+
+    # Add the chart
+    chart_reader = ImageReader(chart_image)
+    c.drawImage(chart_reader, 50, y_position - 300, width=500, height=300)
+
+    c.showPage()
+    c.save()
+    return pdf_path
 
 # Function to generate an Excel file with prediction results
 def generate_excel(planned_cost, actual_spend, forecast_cost, total_risk, insights, status, top_causes, forecast_cost_eur):
@@ -137,8 +184,8 @@ def predict_risk(planned_cost, actual_spend, category, cement_index, labor_index
         logger.debug(f"Salesforce create result: {result}")
         record_id = result['id']
 
-        # Generate the plot for both Salesforce and Gradio
-        plot_data, chart_image = generate_plot(planned_cost, actual_spend, forecast_cost)
+        # Generate the plot for Salesforce, Gradio, and PDF
+        plot_data, chart_image, chart_image_pdf = generate_plot(planned_cost, actual_spend, forecast_cost)
         attachment = {
             "ParentId": record_id,
             "Name": f"Budget_Overview_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png",
@@ -168,6 +215,9 @@ def predict_risk(planned_cost, actual_spend, category, cement_index, labor_index
         logger.error(f"Failed to fetch exchange rates: {str(e)}")
         forecast_cost_eur = "N/A"
 
+    # Generate local PDF report
+    pdf_file = generate_pdf(planned_cost, actual_spend, forecast_cost, total_risk, insights, status, top_causes, forecast_cost_eur, chart_image_pdf)
+
     # Generate Excel file
     excel_file = generate_excel(planned_cost, actual_spend, forecast_cost, total_risk, insights, status, top_causes, forecast_cost_eur)
 
@@ -187,10 +237,11 @@ def predict_risk(planned_cost, actual_spend, category, cement_index, labor_index
         f"Total Risk: {total_risk}\n"
         f"Forecasted Cost: ${forecast_cost:,.2f} (EUR: {forecast_cost_eur if isinstance(forecast_cost_eur, str) else f'{forecast_cost_eur:,.2f}'})\n"
         f"Top Causes: {top_causes}\n"
-        f"PDF Report: {pdf_url}\n"
+        f"Salesforce PDF URL (requires login): {pdf_url}\n"
+        f"Local PDF Report: [Download link below]\n"
         f"Excel Report: [Download link below]"
     )
-    return output_text, chart_image, pdf_url, excel_file
+    return output_text, chart_image, pdf_file, excel_file
 
 # Gradio interface
 interface = gr.Interface(
@@ -206,7 +257,7 @@ interface = gr.Interface(
     outputs=[
         gr.Textbox(label="Prediction Results"),
         gr.Image(label="Forecast Chart"),
-        gr.File(label="Download PDF Report"),
+        gr.File(label="Download Local PDF Report"),
         gr.File(label="Download Excel Report")
     ],
     title="Budget Overrun Risk Estimator",