Update app.py

app.py

@@ -46,7 +46,11 @@ def connect_to_salesforce():
         logger.error(f"Salesforce connection failed: {e}")
         raise
 
-sf = connect_to_salesforce()
+try:
+    sf = connect_to_salesforce()
+except Exception as e:
+    logger.error(f"Failed to initialize Salesforce connection: {e}")
+    sf = None
 
 # Choices
 equipment_choices = [
@@ -77,66 +81,104 @@ def call_ai_model(usage, idle, freq, cost, last):
 
 # Core processing
 def process_equipment_utilization(equip, proj, use_h, idle_h, move_f, cost_h, last_maint, ai_sug):
-    if not ai_sug:
-        ai_sug, conf, score = call_ai_model(use_h, idle_h, move_f, cost_h, last_maint)
-    else:
-        conf, score = 0.9, 85.0
-    summary = {
-        "Equipment Name": equip,
-        "Project": proj,
-        "Usage Hours": use_h,
-        "Idle Hours": idle_h,
-        "Suggestion": ai_sug,
-        "Confidence": conf,
-        "Utilization Score": score,
-        "Cost per Hour": cost_h,
-        "Last Maintenance": last_maint or "N/A"
-    }
-    record_data = {
-        "Equipment_Name__c": equip,
-        "Project_Name__c": proj,
-        "Usage_Hours__c": use_h,
-        "Idle_Hours__c": idle_h,
-        "AI_Suggestion__c": ai_sug,
-        "Suggestion_Confidence__c": conf * 100,
-        "Utilization_Score__c": score,
-        "Cost_per_Hour__c": cost_h,
-        "Report_Link__c": "Pending",
-        "Last_Maintenance__c": last_maint if last_maint != "N/A" else None,
-        "Dashboard_Flag__c": False
-    }
-    resp = sf.Equipment_Utilization_Record__c.create(record_data)
-    rec_id = resp.get("id")
-    # Prepare paths
-    uid = uuid.uuid4().hex[:8]
-    pdf_path = Path(f"static/reports/report_{uid}.pdf")
-    csv_path = Path(f"static/reports/report_{uid}.csv")
-    pdf_path.parent.mkdir(parents=True, exist_ok=True)
-    # Generate PDF
-    c = canvas.Canvas(str(pdf_path), pagesize=letter)
-    c.setFont("Helvetica", 12)
-    c.drawString(100, 750, "Equipment Utilization Report")
-    c.drawString(100, 735, f"Record ID: {rec_id}")
-    y = 710
-    for k, v in summary.items(): c.drawString(100, y, f"{k}: {v}"); y -= 20
-    c.save()
-    # Generate CSV
-    with open(csv_path, "w", newline="") as f:
-        w = csv.writer(f)
-        w.writerow(["Field", "Value"])
-        w.writerow(["Record ID", rec_id])
-        [w.writerow([k, v]) for k, v in summary.items()]
-    # Upload PDF
-    encoded = base64.b64encode(pdf_path.read_bytes()).decode()
-    cv = sf.ContentVersion.create({
-        "Title": "UtilReport",
-        "PathOnClient": os.path.basename(str(pdf_path)),
-        "VersionData": encoded,
-        "FirstPublishLocationId": Report_File_Path
-    })
-    pdf_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{cv['id']}"
-    sf.Equipment_Utilization_Record__c.update(rec_id, {"Report_Link__c": pdf_url})
-    return {"Salesforce_Record_Id": rec_id, "Summary": summary, "Report_Link": pdf_url, "CSV_Report_Link": str(csv_path), "Report_File_Path": str(pdf_path)}
+    try:
+        if not sf:
+            raise ValueError("Salesforce connection is not initialized. Please check credentials and try again.")
+
+        if not ai_sug:
+            ai_sug, conf, score = call_ai_model(use_h, idle_h, move_f, cost_h, last_maint)
+        else:
+            conf, score = 0.9, 85.0
+        summary = {
+            "Equipment Name": equip,
+            "Project": proj,
+            "Usage Hours": use_h,
+            "Idle Hours": idle_h,
+            "Suggestion": ai_sug,
+            "Confidence": conf,
+            "Utilization Score": score,
+            "Cost per Hour": cost_h,
+            "Last Maintenance": last_maint or "N/A"
+        }
+        record_data = {
+            "Equipment_Name__c": equip,
+            "Project_Name__c": proj,
+            "Usage_Hours__c": use_h,
+            "Idle_Hours__c": idle_h,
+            "AI_Suggestion__c": ai_sug,
+            "Suggestion_Confidence__c": conf * 100,
+            "Utilization_Score__c": score,
+            "Cost_per_Hour__c": cost_h,
+            "Report_Link__c": "Pending",
+            "Last_Maintenance__c": last_maint if last_maint != "N/A" else None,
+            "Dashboard_Flag__c": False
+        }
+        resp = sf.Equipment_Utilization_Record__c.create(record_data)
+        if not resp.get("success"):
+            raise ValueError(f"Failed to create Salesforce record: {resp}")
+        rec_id = resp.get("id")
+        if not rec_id:
+            raise ValueError("Salesforce record creation succeeded, but no record ID was returned.")
+
+        # Prepare paths
+        uid = uuid.uuid4().hex[:8]
+        pdf_path = Path(f"static/reports/report_{uid}.pdf")
+        csv_path = Path(f"static/reports/report_{uid}.csv")
+        try:
+            pdf_path.parent.mkdir(parents=True, exist_ok=True)
+        except Exception as e:
+            raise ValueError(f"Failed to create directory {pdf_path.parent}: {e}")
+
+        # Generate PDF
+        try:
+            c = canvas.Canvas(str(pdf_path), pagesize=letter)
+            c.setFont("Helvetica", 12)
+            c.drawString(100, 750, "Equipment Utilization Report")
+            c.drawString(100, 735, f"Record ID: {rec_id}")
+            y = 710
+            for k, v in summary.items():
+                c.drawString(100, y, f"{k}: {v}")
+                y -= 20
+            c.save()
+        except Exception as e:
+            raise ValueError(f"Failed to generate PDF: {e}")
+
+        # Generate CSV
+        try:
+            with open(csv_path, "w", newline="") as f:
+                w = csv.writer(f)
+                w.writerow(["Field", "Value"])
+                w.writerow(["Record ID", rec_id])
+                [w.writerow([k, v]) for k, v in summary.items()]
+        except Exception as e:
+            raise ValueError(f"Failed to generate CSV: {e}")
+
+        # Upload PDF to Salesforce
+        try:
+            encoded = base64.b64encode(pdf_path.read_bytes()).decode()
+            cv = sf.ContentVersion.create({
+                "Title": "UtilReport",
+                "PathOnClient": os.path.basename(str(pdf_path)),
+                "VersionData": encoded,
+                "FirstPublishLocationId": rec_id
+            })
+            if not cv.get("success"):
+                raise ValueError(f"Failed to upload PDF to Salesforce: {cv}")
+            pdf_url = f"https://{sf.sf_instance}/sfc/servlet.shepherd/version/download/{cv['id']}"
+            sf.Equipment_Utilization_Record__c.update(rec_id, {"Report_Link__c": pdf_url})
+        except Exception as e:
+            raise ValueError(f"Failed to upload PDF to Salesforce: {e}")
+
+        return {
+            "Salesforce_Record_Id": rec_id,
+            "Summary": summary,
+            "Report_Link": pdf_url,
+            "CSV_Report_Link": str(csv_path),
+            "Report_File_Path": str(pdf_path)
+        }
+    except Exception as e:
+        logger.error(f"Error in process_equipment_utilization: {e}")
+        raise ValueError(f"Processing failed: {str(e)}")
 
 # Format output
 def format_output(result):
@@ -172,24 +214,52 @@ def format_output(result):
 
 # Gradio callbacks
 def manual_input(equipment, project, usage, idle, freq, cost, last, ai_suggestion):
-    last_val = last or "N/A"
-    res = process_equipment_utilization(equipment, project, usage, idle, freq, cost, last_val, ai_suggestion)
-    formatted = format_output(res)
-    return formatted, res.get("Report_File_Path")
+    try:
+        # Validate inputs
+        if not equipment or equipment not in equipment_choices:
+            raise ValueError("Please select a valid Equipment Name.")
+        if not project or project not in project_choices:
+            raise ValueError("Please select a valid Project Name.")
+        if usage is None or usage < 0:
+            raise ValueError("Usage Hours must be a non-negative number.")
+        if idle is None or idle < 0:
+            raise ValueError("Idle Hours must be a non-negative number.")
+        if freq is None or freq < 0:
+            raise ValueError("Movement Frequency must be a non-negative number.")
+        if cost is None or cost < 0:
+            raise ValueError("Cost per Hour must be a non-negative number.")
+        
+        last_val = last or "N/A"
+        res = process_equipment_utilization(equipment, project, usage, idle, freq, cost, last_val, ai_suggestion)
+        formatted = format_output(res)
+        return formatted, res.get("Report_File_Path")
+    except Exception as e:
+        logger.error(f"Error in manual_input: {e}")
+        return f"Error: {str(e)}", None
 
 def batch_upload(csv_file):
-    if not csv_file: return {}
-    df = pd.read_csv(csv_file.name)
-    ids = []
-    for _, row in df.iterrows():
-        rec = process_equipment_utilization(
-            row['equipment_name'], row['project_name'],
-            float(row['usage_hours']), float(row['idle_hours']),
-            float(row['movement_frequency']), float(row['cost_per_hour']),
-            row.get('last_maintenance','N/A'), row.get('ai_suggestion','')
-        )
-        ids.append(rec['Salesforce_Record_Id'])
-    return {"records": ids}
+    try:
+        if not csv_file:
+            raise ValueError("Please upload a CSV file.")
+        df = pd.read_csv(csv_file.name)
+        required_columns = ['equipment_name', 'project_name', 'usage_hours', 'idle_hours', 'movement_frequency', 'cost_per_hour']
+        missing_columns = [col for col in required_columns if col not in df.columns]
+        if missing_columns:
+            raise ValueError(f"CSV file is missing required columns: {', '.join(missing_columns)}")
+        
+        ids = []
+        for _, row in df.iterrows():
+            rec = process_equipment_utilization(
+                row['equipment_name'], row['project_name'],
+                float(row['usage_hours']), float(row['idle_hours']),
+                float(row['movement_frequency']), float(row['cost_per_hour']),
+                row.get('last_maintenance', 'N/A'), row.get('ai_suggestion', '')
+            )
+            ids.append(rec['Salesforce_Record_Id'])
+        return {"records": ids}
+    except Exception as e:
+        logger.error(f"Error in batch_upload: {e}")
+        return {"error": str(e)}
 
 # Interface
 with gr.Blocks() as app: