Update app.py

app.py

@@ -11,6 +11,7 @@ from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.lib.units import inch
 from io import BytesIO
 from simple_salesforce import Salesforce
+import base64
 
 # Streamlit app configuration
 st.set_page_config(page_title="Delay 🚀", layout="wide")
@@ -108,8 +109,8 @@ def generate_pdf(input_data, prediction, heatmap_fig):
     buffer.seek(0)
     return buffer
 
-# Function to save data to Salesforce
-def save_to_salesforce(input_data, prediction):
+# Function to save data to Salesforce, including PDF
+def save_to_salesforce(input_data, prediction, pdf_buffer):
     if sf is None:
         return "Salesforce connection not established."
     try:
@@ -133,7 +134,37 @@ def save_to_salesforce(input_data, prediction):
         }
         # Create a new record in Delay_Predictor__c
         result = sf.Delay_Predictor__c.create(sf_data)
-        return None if result["success"] else f"Salesforce save failed: {result['errors']}"
+        if not result["success"]:
+            return f"Salesforce save failed: {result['errors']}"
+
+        # Get the record ID
+        record_id = result["id"]
+
+        # Upload PDF as ContentVersion
+        pdf_data = pdf_buffer.getvalue()
+        pdf_base64 = base64.b64encode(pdf_data).decode('utf-8')
+        content_version = {
+            "Title": f"Delay_Prediction_Report_{input_data['project_name']}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
+            "PathOnClient": "project_delay_report.pdf",
+            "VersionData": pdf_base64,
+            "FirstPublishLocationId": record_id  # Link to the Delay_Predictor__c record
+        }
+        cv_result = sf.ContentVersion.create(content_version)
+        if not cv_result["success"]:
+            return f"Failed to upload PDF to Salesforce: {cv_result['errors']}"
+
+        # Get the ContentVersion ID
+        content_version_id = cv_result["id"]
+
+        # Construct the internal Salesforce URL for the ContentVersion
+        pdf_url = f"{os.environ.get('SF_INSTANCE_URL')}/sfc/servlet.shepherd/version/download/{content_version_id}"
+
+        # Update the Delay_Predictor__c record with the PDF URL
+        update_result = sf.Delay_Predictor__c.update(record_id, {"PDF_Report__c": pdf_url})
+        if not update_result == 204:  # 204 indicates success for updates
+            return f"Failed to update PDF_Report__c field: {update_result}"
+
+        return None
     except Exception as e:
         return f"Error saving to Salesforce: {str(e)}"
 
@@ -212,12 +243,12 @@ if submit_button:
                 mime="application/pdf"
             )
             
-            # Save to Salesforce
-            sf_error = save_to_salesforce(input_data, prediction)
+            # Save to Salesforce, including PDF
+            sf_error = save_to_salesforce(input_data, prediction, pdf_buffer)
             if sf_error:
                 st.error(sf_error)
             else:
-                st.success("Prediction data successfully saved to Salesforce!")
+                st.success("Prediction data and PDF successfully saved to Salesforce!")
             
             st.session_state.prediction = prediction
             st.session_state.input_data = input_data