Update app.py

app.py

@@ -164,8 +164,15 @@ def save_to_salesforce(input_data, prediction, pdf_buffer):
         # Get the ContentVersion ID
         content_version_id = cv_result["id"]
 
-        # Construct the correct Salesforce URL for the ContentVersion
-        pdf_url = f"{os.environ.get('SF_INSTANCE_URL')}/sfc/servlet.shepherd/version/download/{content_version_id}"
+        # Query the ContentDocumentId from the ContentVersion
+        query = f"SELECT ContentDocumentId FROM ContentVersion WHERE Id = '{content_version_id}'"
+        query_result = sf.query(query)
+        if query_result["totalSize"] == 0:
+            return "Failed to retrieve ContentDocumentId for the ContentVersion"
+        content_document_id = query_result["records"][0]["ContentDocumentId"]
+
+        # Construct the correct Salesforce URL for the ContentDocument
+        pdf_url = f"{os.environ.get('SF_INSTANCE_URL')}/sfc/servlet.shepherd/document/download/{content_document_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})