Update app.py

app.py

@@ -223,15 +223,37 @@ h1 {
 
 def validate_csv(df):
     """
-    Validate that the CSV has the required columns.
+    Validate that the CSV has the required columns, handling both original and renamed columns.
     Returns True if valid, False otherwise with an error message.
     """
     # Strip whitespace from column names
     df.columns = df.columns.str.strip()
-    required_columns = ['device_id', 'usage_hours', 'amc_date', 'status']
-    missing_columns = [col for col in required_columns if col not in df.columns]
-    if missing_columns:
-        return False, f"Missing required columns: {', '.join(missing_columns)}. Found columns: {', '.join(df.columns)}"
+
+    # Define expected original and renamed columns
+    original_columns = ['device_id', 'usage_hours', 'amc_date', 'status']
+    renamed_columns = ['equipment', 'usage_count', 'amc_expiry', 'status']
+    
+    # Check for original columns
+    missing_original = [col for col in original_columns if col not in df.columns]
+    # Check for renamed columns
+    missing_renamed = [col for col in renamed_columns if col not in df.columns]
+
+    # If original columns are present, proceed as is
+    if not missing_original:
+        logging.info("Found original columns in CSV. Proceeding with validation.")
+    # If renamed columns are present, map them back to original for validation
+    elif not missing_renamed:
+        logging.info("Found renamed columns in CSV. Mapping back to original names for validation.")
+        df.rename(columns={
+            'equipment': 'device_id',
+            'usage_count': 'usage_hours',
+            'amc_expiry': 'amc_date'
+        }, inplace=True)
+    else:
+        # If neither set is fully present, report missing columns
+        found_columns = ', '.join(df.columns)
+        return False, f"Missing required columns. Expected either {', '.join(original_columns)} or {', '.join(renamed_columns)}. Found columns: {found_columns}"
+
     # Validate data types
     try:
         df['usage_hours'] = pd.to_numeric(df['usage_hours'], errors='raise')
@@ -242,6 +264,14 @@ def validate_csv(df):
             df['downtime'] = pd.to_numeric(df['downtime'], errors='raise')
     except Exception as e:
         return False, f"Invalid data types: {str(e)}"
+    
+    # Rename columns to internal names after validation
+    df.rename(columns={
+        'device_id': 'equipment',
+        'usage_hours': 'usage_count',
+        'amc_date': 'amc_expiry'
+    }, inplace=True)
+    
     return True, ""
 
 def generate_device_cards(df, anomaly_df):
@@ -377,13 +407,7 @@ def process_files(uploaded_files):
             df = pd.read_csv(file.name, delimiter=',', skipinitialspace=True)
             # Log the columns for debugging
             logging.info(f"Columns in {file.name}: {', '.join(df.columns)}")
-            # Rename columns to match expected names
-            df = df.rename(columns={
-                'device_id': 'equipment',
-                'usage_hours': 'usage_count',
-                'amc_date': 'amc_expiry'
-            })
-            # Validate CSV structure
+            # Validate CSV structure (renaming happens inside validate_csv now)
             is_valid, error_msg = validate_csv(df)
             if not is_valid:
                 logging.error(f"Failed to load {file.name}: {error_msg}")