Update app.py

app.py

@@ -10,16 +10,37 @@ from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.lib.units import inch
 from io import BytesIO
 from simple_salesforce import Salesforce
-import os
 
-# Salesforce credentials (from environment variables or directly from app.py)
-SF_USERNAME = os.getenv('SF_USERNAME', 'your_salesforce_username')
-SF_PASSWORD = os.getenv('SF_PASSWORD', 'your_salesforce_password')
-SF_SECURITY_TOKEN = os.getenv('SF_SECURITY_TOKEN', 'your_salesforce_security_token')
-SF_INSTANCE_URL = os.getenv('SF_INSTANCE_URL', 'https://your_instance.salesforce.com')
+# Streamlit app configuration
+st.set_page_config(page_title="Delay 🚀", layout="wide")
+
+# Salesforce connection (using Streamlit secrets)
+try:
+    sf = Salesforce(
+        username=st.secrets["SF_USERNAME"],
+        password=st.secrets["SF_PASSWORD"],
+        security_token=st.secrets["SF_SECURITY_TOKEN"],
+        instance_url=st.secrets["SF_INSTANCE_URL"]
+    )
+except Exception as e:
+    st.error(f"Failed to connect to Salesforce: {str(e)}")
+    sf = None
+
+# Title
+st.title("Project Delay Predictor 🚀")
+
+# Task options per phase
+task_options = {
+    "Planning": ["Define Scope", "Resource Allocation", "Permit Acquisition"],
+    "Design": ["Architectural Drafting", "Engineering Analysis", "Design Review"],
+    "Construction": ["Foundation Work", "Structural Build", "Utility Installation"]
+}
 
-# Initialize Salesforce connection
-sf = Salesforce(username=SF_USERNAME, password=SF_PASSWORD, security_token=SF_SECURITY_TOKEN, instance_url=SF_INSTANCE_URL)
+# Initialize session state for phase and task
+if 'phase' not in st.session_state:
+    st.session_state.phase = ""
+if 'task' not in st.session_state:
+    st.session_state.task = ""
 
 # Function to format high_risk_phases with flag and alert
 def format_high_risk_phases(high_risk_phases):
@@ -63,7 +84,6 @@ def generate_pdf(input_data, prediction, heatmap_fig):
 
     # Prediction Results
     story.append(Paragraph("Prediction Results", styles['Heading2']))
-    # Format high_risk_phases for PDF
     high_risk_text = "<br/>".join(format_high_risk_phases(prediction['high_risk_phases']))
     prediction_fields = [
         f"Delay Probability: {prediction['delay_probability']:.2f}%",
@@ -77,65 +97,45 @@ def generate_pdf(input_data, prediction, heatmap_fig):
 
     # Heatmap
     story.append(Paragraph("Delay Risk Heatmap", styles['Heading2']))
-    # Save heatmap figure to BytesIO as PNG
     img_buffer = BytesIO()
     heatmap_fig.savefig(img_buffer, format='png', bbox_inches='tight')
     img_buffer.seek(0)
-    # Embed image in PDF (scale to fit page width)
     story.append(Image(img_buffer, width=6*inch, height=2*inch))
-    plt.close(heatmap_fig)  # Close figure to free memory
+    plt.close(heatmap_fig)
 
     doc.build(story)
     buffer.seek(0)
     return buffer
 
-# Function to save data in Salesforce custom object (Delay_Predictor__c)
+# Function to save data to Salesforce
 def save_to_salesforce(input_data, prediction):
+    if sf is None:
+        return "Salesforce connection not established."
     try:
-        # Create record in Delay_Predictor__c object
-        record = {
-            'Project_Name__c': input_data['project_name'],
-            'Phase__c': input_data['phase'],
-            'Task__c': input_data['task'],
-            'Current_Progress__c': input_data['current_progress'],
-            'Task_Expected_Duration__c': input_data['task_expected_duration'],
-            'Task_Actual_Duration__c': input_data['task_actual_duration'],
-            'Workforce_Gap__c': input_data['workforce_gap'],
-            'Workforce_Skill_Level__c': input_data['workforce_skill_level'],
-            'Workforce_Shift_Hours__c': input_data['workforce_shift_hours'],
-            'Weather_Impact_Score__c': input_data['weather_impact_score'],
-            'Weather_Condition__c': input_data['weather_condition'],
-            'Weather_Forecast_Date__c': input_data['weather_forecast_date'],
-            'Delay_Probability__c': prediction['delay_probability'],
-            'AI_Insights__c': prediction['ai_insights'],
-            'High_Risk_Phases__c': str(prediction['high_risk_phases']),  # Storing as a string (can be a complex field if needed)
-            'Weather_Condition_Forecast__c': prediction['weather_condition']
+        # Prepare data for Delay_Predictor__c object
+        sf_data = {
+            "Project_Name__c": input_data["project_name"],
+            "Phase__c": input_data["phase"],
+            "Task__c": input_data["task"],
+            "Current_Progress__c": input_data["current_progress"],
+            "Task_Expected_Duration__c": input_data["task_expected_duration"],
+            "Task_Actual_Duration__c": input_data["task_actual_duration"],
+            "Workforce_Gap__c": input_data["workforce_gap"],
+            "Workforce_Skill_Level__c": input_data["workforce_skill_level"],
+            "Workforce_Shift_Hours__c": input_data["workforce_shift_hours"],
+            "Weather_Impact_Score__c": input_data["weather_impact_score"],
+            "Weather_Condition__c": input_data["weather_condition"],
+            "Weather_Forecast_Date__c": input_data["weather_forecast_date"],
+            "Delay_Probability__c": prediction["delay_probability"],
+            "AI_Insights__c": prediction["ai_insights"],
+            # Store high_risk_phases as a formatted string
+            "High_Risk_Phases__c": "; ".join(format_high_risk_phases(prediction["high_risk_phases"]))
         }
-
-        # Insert record into Salesforce
-        sf.Delay_Predictor__c.create(record)
-        print("Record saved to Salesforce successfully!")
+        # 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']}"
     except Exception as e:
-        print(f"Error saving to Salesforce: {str(e)}")
-
-# Streamlit app configuration
-st.set_page_config(page_title="Delay 🚀", layout="wide")
-
-# Title
-st.title("Project Delay Predictor 🚀")
-
-# Task options per phase
-task_options = {
-    "Planning": ["Define Scope", "Resource Allocation", "Permit Acquisition"],
-    "Design": ["Architectural Drafting", "Engineering Analysis", "Design Review"],
-    "Construction": ["Foundation Work", "Structural Build", "Utility Installation"]
-}
-
-# Initialize session state for phase and task
-if 'phase' not in st.session_state:
-    st.session_state.phase = ""
-if 'task' not in st.session_state:
-    st.session_state.task = ""
+        return f"Error saving to Salesforce: {str(e)}"
 
 # Input form
 with st.form("project_form"):
@@ -145,10 +145,9 @@ with st.form("project_form"):
         project_name = st.text_input("Project Name")
         phase = st.selectbox("Phase", [""] + ["Planning", "Design", "Construction"], index=0, key="phase_select")
         
-        # Update task options based on phase
         if phase != st.session_state.phase:
             st.session_state.phase = phase
-            st.session_state.task = ""  # Reset task when phase changes
+            st.session_state.task = ""
         task_options_list = [""] + task_options.get(phase, []) if phase else [""]
         task = st.selectbox("Task", task_options_list, index=0, key="task_select")
         current_progress = st.number_input("Current Progress (%)", min_value=0.0, max_value=100.0, step=1.0)
@@ -159,10 +158,8 @@ with st.form("project_form"):
         workforce_gap = st.number_input("Workforce Gap (%)", min_value=0.0, max_value=100.0, step=1.0)
         workforce_skill_level = st.selectbox("Workforce Skill Level", ["Low", "Medium", "High"], index=1)
         workforce_shift_hours = st.number_input("Workforce Shift Hours", min_value=0, step=1, value=8)
-        # Debug output for shift hours
         st.write(f"**Selected Shift Hours**: {workforce_shift_hours}")
         weather_impact_score = st.number_input("Weather Impact Score (0-100)", min_value=0, max_value=100, step=1)
-        # Display computed weather condition
         weather_condition = get_weather_condition(weather_impact_score)
         st.write(f"**Weather Condition**: {weather_condition}")
         weather_forecast_date = st.date_input("Weather Forecast Date", min_value=datetime(2025, 1, 1))
@@ -171,7 +168,6 @@ with st.form("project_form"):
 
 # Process form submission
 if submit_button:
-    # Prepare input data
     input_data = {
         "project_name": project_name,
         "phase": phase,
@@ -187,19 +183,16 @@ if submit_button:
         "weather_forecast_date": weather_forecast_date.strftime("%Y-%m-%d")
     }
     
-    # Validate inputs
     error = validate_inputs(input_data)
     if error:
         st.error(error)
     else:
-        # Get prediction
         with st.spinner("Predicting..."):
             prediction = predict_delay(input_data)
         
         if "error" in prediction:
             st.error(prediction["error"])
         else:
-            # Display results
             st.subheader("Prediction Results")
             st.write(f"**Delay Probability**: {prediction['delay_probability']:.2f}%")
             st.write("**High Risk Phases**:")
@@ -208,14 +201,9 @@ if submit_button:
             st.write(f"**AI Insights**: {prediction['ai_insights']}")
             st.write(f"**Weather Condition**: {prediction['weather_condition']}")
             
-            # Generate and display heatmap
             fig = generate_heatmap(prediction['delay_probability'], f"{phase}: {task}")
             st.pyplot(fig)
             
-            # Save to Salesforce after prediction
-            save_to_salesforce(input_data, prediction)
-            
-            # Generate PDF with heatmap and provide download link
             pdf_buffer = generate_pdf(input_data, prediction, fig)
             st.download_button(
                 label="Download Prediction Report (PDF)",
@@ -224,6 +212,12 @@ if submit_button:
                 mime="application/pdf"
             )
             
-            # Store prediction in session state
+            # Save to Salesforce
+            sf_error = save_to_salesforce(input_data, prediction)
+            if sf_error:
+                st.error(sf_error)
+            else:
+                st.success("Prediction data successfully saved to Salesforce!")
+            
             st.session_state.prediction = prediction
-            st.session_state.input_data = input_data
+            st.session_state.input_data = input_data