Update app.py

app.py

@@ -8,24 +8,17 @@ from PIL import Image
 import io
 import os
 from dotenv import load_dotenv
+from simple_salesforce import Salesforce
 
 # Load environment variables
 load_dotenv()
 
 # Configuration
 AI_MODEL_ENDPOINT = os.getenv("AI_MODEL_ENDPOINT", "https://your-huggingface-or-lambda-endpoint/predict-milestone")
-SALESFORCE_API_URL = os.getenv("SALESFORCE_API_URL", "https://your-salesforce-instance/services/data/v61.0/sobjects/Construction_Project__c")
 S3_BUCKET = os.getenv("S3_BUCKET", "your-s3-bucket")
 AWS_REGION = os.getenv("AWS_REGION", "us-east-1")
 MAX_FILE_SIZE = 20 * 1024 * 1024  # 20MB in bytes
 MIN_IMAGE_DIMENSIONS = (800, 600)  # Minimum width and height for clarity
-SALESFORCE_AUTH = {
-    "access_token": os.getenv("SALESFORCE_ACCESS_TOKEN"),
-    "headers": {
-        "Authorization": f"Bearer {os.getenv('SALESFORCE_ACCESS_TOKEN')}",
-        "Content-Type": "application/json"
-    }
-}
 
 # Initialize AWS S3 client
 try:
@@ -39,33 +32,58 @@ except Exception as e:
     st.error(f"Failed to initialize S3 client: {str(e)}")
     st.stop()
 
+# Initialize Salesforce client
+try:
+    sf = Salesforce(
+        username=os.getenv("SF_USERNAME"),
+        password=os.getenv("SF_PASSWORD"),
+        security_token=os.getenv("SF_SECURITY_TOKEN"),
+        domain="login"  # Use "test" for sandbox
+    )
+except Exception as e:
+    st.error(f"Failed to connect to Salesforce: {str(e)}")
+    st.stop()
+
 # Streamlit app
 st.title("Construction Progress Verifier")
-st.write("Upload a construction site photo to track project milestones.")
-
-# Input for Salesforce Project ID
-project_id = st.text_input("Enter Salesforce Project ID (Construction_Project__c Record ID)", "")
-
-# File uploader
-uploaded_file = st.file_uploader("Choose an image (JPG/PNG, <20MB)", type=["jpg", "png"])
-
-if uploaded_file is not None and project_id:
-    # Validate file size
-    if uploaded_file.size > MAX_FILE_SIZE:
-        st.error("File size exceeds 20MB. Please upload a smaller image.")
-    else:
-        # Display progress bar and status
-        progress_bar = st.progress(0)
-        status_text = st.empty()
+st.write("Upload one or more construction site photos to track project milestones.")
+
+# Input for Salesforce Construction__c Record ID
+project_id = st.text_input("Enter Salesforce Construction ID (Construction__c Record ID)", "")
+
+# File uploader for multiple images
+uploaded_files = st.file_uploader("Choose images (JPG/PNG, <20MB each)", type=["jpg", "png"], accept_multiple_files=True)
+
+if uploaded_files and project_id:
+    # Initialize results list for output
+    results = []
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    total_images = len(uploaded_files)
+    progress_increment = 100 // (total_images * 3)  # Divide progress into upload, AI, and Salesforce steps
+
+    for idx, uploaded_file in enumerate(uploaded_files, 1):
+        # Validate file size
+        if uploaded_file.size > MAX_FILE_SIZE:
+            results.append({
+                "filename": uploaded_file.name,
+                "status": "Failed",
+                "message": "File size exceeds 20MB. Please upload a smaller image."
+            })
+            continue
 
         try:
             # Read and validate image
-            status_text.text("Validating image...")
-            progress_bar.progress(10)
+            status_text.text(f"Processing image {idx}/{total_images}: Validating {uploaded_file.name}...")
+            progress_bar.progress(min((idx - 1) * progress_increment * 3 + progress_increment, 100))
             image = Image.open(uploaded_file)
             if image.size < MIN_IMAGE_DIMENSIONS:
-                st.error("Image resolution too low. Please upload a higher resolution image (min 800x600).")
-                st.stop()
+                results.append({
+                    "filename": uploaded_file.name,
+                    "status": "Failed",
+                    "message": "Image resolution too low. Please upload a higher resolution image (min 800x600)."
+                })
+                continue
 
             # Convert image to base64
             buffered = io.BytesIO()
@@ -73,14 +91,14 @@ if uploaded_file is not None and project_id:
             img_str = base64.b64encode(buffered.getvalue()).decode()
 
             # Upload image to S3
-            status_text.text("Uploading image to storage...")
-            s3_key = f"construction_images/{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
+            status_text.text(f"Processing image {idx}/{total_images}: Uploading {uploaded_file.name} to storage...")
+            s3_key = f"construction_images/{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uploaded_file.name}"
             s3_client.put_object(Bucket=S3_BUCKET, Key=s3_key, Body=buffered.getvalue())
             s3_url = f"https://{S3_BUCKET}.s3.{AWS_REGION}.amazonaws.com/{s3_key}"
-            progress_bar.progress(30)
+            progress_bar.progress(min((idx - 1) * progress_increment * 3 + 2 * progress_increment, 100))
 
             # Call AI model endpoint
-            status_text.text("Analyzing image with AI model...")
+            status_text.text(f"Processing image {idx}/{total_images}: Analyzing {uploaded_file.name} with AI model...")
             payload = {"image": img_str}
             ai_response = requests.post(
                 AI_MODEL_ENDPOINT,
@@ -92,45 +110,55 @@ if uploaded_file is not None and project_id:
             if ai_response.status_code == 200:
                 ai_result = ai_response.json()
                 milestone = ai_result.get("milestone", "Unknown")
-                completion_percent = ai_result.get("percent_complete", 0)
-                progress_bar.progress(60)
 
-                # Update Salesforce record
-                status_text.text("Updating Salesforce record...")
-                sf_payload = {
+                # Update Salesforce record using simple-salesforce
+                status_text.text(f"Processing image {idx}/{total_images}: Updating Salesforce for {uploaded_file.name}...")
+                sf.Construction__c.update(project_id, {
                     "Current_Milestone__c": milestone,
-                    "Completion_Percentage__c": float(completion_percent),
                     "Last_Updated_Image__c": s3_url,
-                    "Last_Updated_On__c": datetime.utcnow().isoformat(),
-                    "Upload_Status__c": "Success",
-                    "Comments__c": "AI prediction completed successfully"
-                }
-
-                sf_url = f"{SALESFORCE_API_URL}/{project_id}"
-                sf_response = requests.patch(sf_url, headers=SALESFORCE_AUTH["headers"], json=sf_payload)
-
-                if sf_response.status_code in [200, 204]:
-                    status_text.text("Success! Salesforce record updated.")
-                    progress_bar.progress(100)
-                    st.success(f"Milestone: {milestone}, Completion: {completion_percent}%")
-                else:
-                    status_text.text("Failed to update Salesforce.")
-                    progress_bar.progress(0)
-                    st.error(f"Salesforce API error: {sf_response.text}")
-                    st.error("Please try again or contact support.")
+                    "Upload_Status__c": "Success"
+                })
+
+                results.append({
+                    "filename": uploaded_file.name,
+                    "status": "Success",
+                    "milestone": milestone,
+                    "s3_url": s3_url,
+                    "message": "Salesforce record updated successfully."
+                })
+                progress_bar.progress(min(idx * progress_increment * 3, 100))
             else:
-                status_text.text("AI model failed to process image.")
-                progress_bar.progress(0)
-                st.error(f"AI model error: {ai_response.text}")
-                st.error("Please retake the photo or try again later.")
+                results.append({
+                    "filename": uploaded_file.name,
+                    "status": "Failed",
+                    "message": f"AI model error: {ai_response.text}. Please retake the photo or try again."
+                })
+                progress_bar.progress(min(idx * progress_increment * 3, 100))
         except Exception as e:
-            status_text.text("An error occurred.")
-            progress_bar.progress(0)
-            st.error(f"Error: {str(e)}")
-            st.error("Please retake the photo or contact support.")
+            results.append({
+                "filename": uploaded_file.name,
+                "status": "Failed",
+                "message": f"Error processing image: {str(e)}. Please retake the photo or contact support."
+            })
+            progress_bar.progress(min(idx * progress_increment * 3, 100))
+
+    # Display results
+    st.markdown("### Processing Results")
+    for result in results:
+        if result["status"] == "Success":
+            st.success(
+                f"**{result['filename']}**: "
+                f"Milestone: {result['milestone']}, "
+                f"S3 URL: {result['s3_url']}, "
+                f"Message: {result['message']}"
+            )
+        else:
+            st.error(f"**{result['filename']}**: {result['message']}")
 
-elif uploaded_file is not None and not project_id:
-    st.warning("Please enter a valid Salesforce Project ID.")
+elif uploaded_files and not project_id:
+    st.warning("Please enter a valid Salesforce Construction ID.")
+elif project_id and not uploaded_files:
+    st.warning("Please upload at least one image.")
 
 # Footer
 st.markdown("---")