Update main.py

main.py

@@ -175,6 +175,9 @@ def process_product(product_name, data):
         # Concatenate the original DataFrame with the new fake transaction DataFrame
         data = pd.concat([data, fake_transaction_df], ignore_index=True)
         
+        # Extract the 'sub_sku' column
+        sub_sku = df['sub_sku']
+        
         # Summarize the sales count per month
         data['transaction_date'] = pd.to_datetime(data['transaction_date'])
         data.set_index('transaction_date', inplace=True)
@@ -186,6 +189,7 @@ def process_product(product_name, data):
         result_dict = rounded_value.to_dict(orient="records")[0]
 
         return {
+            "sub_sku": sub_sku,
             "Product Name": product_name,
             "next_month": str(result_dict["next_month"]),
             "predicted_count": result_dict["predicted_count"]
@@ -193,6 +197,7 @@ def process_product(product_name, data):
         }
     except Exception as e:
         return {
+            "sub_sku": sub_sku,
             "Product Name": product_name,
             "next_month": str(e),
             "predicted_count": "not predicted"
@@ -215,8 +220,10 @@ async def generate_product_count_prediction(b_id: int):
                 for future in as_completed(futures):
                     results.append(future.result())
 
-
-            delete_json(b_id)
+            try:
+               delete_json(b_id)
+            except:
+               print("no errors")
             insert_data(b_id,results)
             return {"status": "success",
                     "b_id":b_id,