Update main.py

main.py

@@ -78,19 +78,34 @@ def generate_recommendations_for_session(session_id):
         # Convert session data to a DataFrame
         raw_df = pd.DataFrame(session_data)
         
-        # Aggregate data by id and action
-        aggregated_data = raw_df.groupby(['id', 'action']).agg(
-            presence=('action', 'size'),
-            total_duration=('duration', 'sum')
-        ).reset_index()
+        # Check if 'duration' column exists
+        if 'duration' in raw_df.columns:
+            # Aggregate data by id and action, including duration
+            aggregated_data = raw_df.groupby(['id', 'action']).agg(
+                presence=('action', 'size'),
+                total_duration=('duration', 'sum')
+            ).reset_index()
+        else:
+            # Aggregate data by id and action, without duration
+            aggregated_data = raw_df.groupby(['id', 'action']).agg(
+                presence=('action', 'size')
+            ).reset_index()
         
         # Create a pivot table from the aggregated data
-        pivot_df = aggregated_data.pivot_table(
-            index=['id'],
-            columns='action',
-            values=['presence', 'total_duration'],
-            fill_value=0
-        )
+        if 'duration' in raw_df.columns:
+            pivot_df = aggregated_data.pivot_table(
+                index=['id'],
+                columns='action',
+                values=['presence', 'total_duration'],
+                fill_value=0
+            )
+        else:
+            pivot_df = aggregated_data.pivot_table(
+                index=['id'],
+                columns='action',
+                values=['presence'],
+                fill_value=0
+            )
         
         # Flatten column names
         pivot_df.columns = ['_'.join(col).strip() for col in pivot_df.columns.values]
@@ -99,7 +114,7 @@ def generate_recommendations_for_session(session_id):
         for col in ALL_COLUMNS:
             if f'presence_{col}' not in pivot_df.columns and col != 'time_spent':
                 pivot_df[f'presence_{col}'] = 0
-            elif col == 'time_spent' and 'total_duration_time_spent' not in pivot_df.columns:
+            elif col == 'time_spent' and 'duration' in raw_df.columns and 'total_duration_time_spent' not in pivot_df.columns:
                 pivot_df['total_duration_time_spent'] = 0
         
         # Calculate interaction score for each row