Update app.py

app.py

@@ -46,10 +46,11 @@ def update_news():
     status_code = 200
     try:
         db_read = DBRead()
-        db_write = DBWrite()
+        db_write = DBWrite(db_type="production")
+        prediction_db_write = DBWrite(db_type="prediction")
         old_news = db_read.read_news_from_db()
         new_news = get_news()
-        news_df, is_db_updation_required = predict_news_category_similar_news(old_news, new_news, interpreter, label_encoder, 
+        news_df, prediction_df, is_db_updation_required = predict_news_category_similar_news(old_news, new_news, interpreter, label_encoder, 
                                                                               tokenizer, collection, vectorizer, sent_model, ce_model)
         if news_df is None:
             raise Exception('Could not generate category predictions. Aborting the database insertion. No new articles are inserted into the collection.')
@@ -62,7 +63,11 @@ def update_news():
 
         if is_db_updation_required:
             news_json = [*json.loads(news_df.reset_index(drop=True).to_json(orient="index")).values()]
+            prediction_json = [*json.loads(prediction_df.reset_index(drop=True).to_json(orient="index")).values()]
+            
             db_write.insert_news_into_db(news_json)
+            prediction_db_write.insert_news_into_db(prediction_json)
+            
         else:
             logger.warning('DB is not updated as it is not required.')
     except Exception as e: