Update main.py

main.py

@@ -78,34 +78,19 @@ def generate_recommendations_for_session(session_id):
         # Convert session data to a DataFrame
         raw_df = pd.DataFrame(session_data)
         
-        # 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()
+        # Aggregate data by id and action
+        aggregated_data = raw_df.groupby(['id', 'action']).agg(
+            presence=('action', 'size'),
+            total_duration=('duration', 'sum')
+        ).reset_index()
         
         # Create a pivot table from the aggregated data
-        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
-            )
+        pivot_df = aggregated_data.pivot_table(
+            index=['id'],
+            columns='action',
+            values=['presence', 'total_duration'],
+            fill_value=0
+        )
         
         # Flatten column names
         pivot_df.columns = ['_'.join(col).strip() for col in pivot_df.columns.values]
@@ -114,7 +99,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 'duration' in raw_df.columns and 'total_duration_time_spent' not in pivot_df.columns:
+            elif col == 'time_spent' and 'total_duration_time_spent' not in pivot_df.columns:
                 pivot_df['total_duration_time_spent'] = 0
         
         # Calculate interaction score for each row
@@ -144,6 +129,7 @@ def generate_recommendations_for_session(session_id):
         logger.error(f"Error in generate_recommendations_for_session: {e}")
         return None
 
+
 def calculate_interaction_score(row):
     try:
         # Calculate the score based on the presence of different actions