Fix timeout issue on data inserting script (#411)

* Fix timeout

* Update script

* Add more checks

Author: lickem22

core_backend/add_new_data_to_db.py

@@ -1,6 +1,7 @@
 import argparse
 import json
 import random
+import time
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timedelta
@@ -11,10 +12,17 @@
 from app.config import (
     LITELLM_API_KEY,
     LITELLM_ENDPOINT,
-    LITELLM_MODEL_DEFAULT,
+    LITELLM_MODEL_GENERATION,
 )
+from app.contents.models import ContentDB
 from app.database import get_session
-from app.question_answer.models import ContentFeedbackDB, QueryDB, ResponseFeedbackDB
+from app.llm_call.utils import remove_json_markdown
+from app.question_answer.models import (
+    ContentFeedbackDB,
+    QueryDB,
+    QueryResponseContentDB,
+    ResponseFeedbackDB,
+)
 from app.urgency_detection.models import UrgencyQueryDB
 from app.users.models import UserDB
 from app.utils import get_key_hash
@@ -29,6 +37,7 @@
 
 try:
     import requests  # type: ignore
+
 except ImportError:
     print(
         "Please install requests library using `pip install requests` "
@@ -39,6 +48,7 @@
     (QueryDB, "query_datetime_utc"),
     (ResponseFeedbackDB, "feedback_datetime_utc"),
     (ContentFeedbackDB, "feedback_datetime_utc"),
+    (QueryResponseContentDB, "created_datetime_utc"),
     (UrgencyQueryDB, "message_datetime_utc"),
 ]
 
@@ -106,47 +116,54 @@ def generate_feedback(question_text: str, faq_text: str, sentiment: str) -> dict
     """
 
     response = completion(
-        model=LITELLM_MODEL_DEFAULT,
+        model=LITELLM_MODEL_GENERATION,
         api_base=LITELLM_ENDPOINT,
         api_key=LITELLM_API_KEY,
         messages=[{"role": "user", "content": prompt}],
         max_tokens=100,
         temperature=0.7,
     )
 
-    # Extract the output from the response
-    feedback_output = response["choices"][0]["message"]["content"].strip()
-    feedback_output = feedback_output.replace("json", "")
-    feedback_output = feedback_output.replace("\n", "").strip()
-
     try:
+        # Extract the output from the response
+        feedback_output = response["choices"][0]["message"]["content"].strip()
+        feedback_output = remove_json_markdown(feedback_output)
         feedback_dict = json.loads(feedback_output)
         if isinstance(feedback_dict, dict) and "output" in feedback_dict:
-
             return feedback_dict
         else:
             raise ValueError("Output is not in the correct format.")
-    except (SyntaxError, ValueError) as e:
+    except Exception as e:
         print(f"Output is not in the correct format.{e}")
         return None
 
 
-def save_single_row(endpoint: str, data: dict) -> dict:
+def save_single_row(endpoint: str, data: dict, retries: int = 2) -> dict | None:
     """
     Save a single row in the database.
     """
-
-    response = requests.post(
-        endpoint,
-        headers={
-            "accept": "application/json",
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {API_KEY}",
-        },
-        json=data,
-        verify=False,
-    )
-    return response.json()
+    try:
+        response = requests.post(
+            endpoint,
+            headers={
+                "accept": "application/json",
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {API_KEY}",
+            },
+            json=data,
+            verify=False,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    except Exception as e:
+        if retries > 0:
+            # Implement exponential wait before retrying
+            time.sleep(2 ** (2 - retries))
+            return save_single_row(endpoint, data, retries=retries - 1)
+        else:
+            print(f"Request failed after retries: {e}")
+            return None
 
 
 def process_search(_id: int, text: str) -> tuple | None:
@@ -161,7 +178,7 @@ def process_search(_id: int, text: str) -> tuple | None:
         "generate_tts": False,
     }
     response = save_single_row(endpoint, data)
-    if "search_results" in response:
+    if response and isinstance(response, dict) and "search_results" in response:
         return (
             _id,
             response["query_id"],
@@ -215,7 +232,13 @@ def process_content_feedback(
     if is_off_topic and feedback_sentiment == "positive":
         return None
     # randomly get a content from the search results to provide feedback on
-    content = search_results[str(random.randint(0, 3))]
+    content_num = str(random.randint(0, 3))
+    if not search_results or not isinstance(search_results, dict):
+        return None
+    if content_num not in search_results:
+        return None
+
+    content = search_results[content_num]
 
     # Get content text and use to generate feedback text using LLMs
     content_text = content["title"] + " " + content["text"]
@@ -253,19 +276,16 @@ def process_urgency_detection(_id: int, text: str) -> tuple | None:
     }
 
     response = save_single_row(endpoint, data)
-    if "is_urgent" in response:
+    if response and "is_urgent" in response:
         return (response["is_urgent"],)
     return None
 
 
-def create_random_datetime_from_string(date_string: str) -> datetime:
+def create_random_datetime_from_string(start_date: datetime) -> datetime:
     """
     Create a random datetime from a date in the format "%d-%m-%y
     to today
     """
-    date_format = "%d-%m-%y"
-
-    start_date = datetime.strptime(date_string, date_format)
 
     time_difference = datetime.now() - start_date
     random_number_of_days = random.randint(0, time_difference.days)
@@ -296,6 +316,7 @@ def update_date_of_records(models: list, random_dates: list, api_key: str) -> No
     # Create a dictionary to map the query_id to the random date
     date_map_dic = {queries[i].query_id: random_dates[i] for i in range(len(queries))}
     for model in models:
+        print(f"Updating the date of the records for {model[0].__name__}...")
         session = next(get_session())
 
         rows = [c for c in session.query(model[0]).all() if c.user_id == user.user_id]
@@ -312,12 +333,31 @@ def update_date_of_records(models: list, random_dates: list, api_key: str) -> No
         session.commit()
 
 
+def update_date_of_contents(date: datetime) -> None:
+    """
+    Update the date of the content records in the database for consistency
+    """
+    session = next(get_session())
+    contents = session.query(ContentDB).all()
+    for content in contents:
+        content.created_datetime_utc = date
+        content.updated_datetime_utc = date
+        session.merge(content)
+    session.commit()
+
+
 if __name__ == "__main__":
     HOST = args.host
     NB_WORKERS = int(args.nb_workers) if args.nb_workers else 8
     API_KEY = args.api_key if args.api_key else ADMIN_API_KEY
 
-    start_date = args.start_date if args.start_date else "01-08-23"
+    date_string = args.start_date if args.start_date else "01-08-23"
+    date_format = "%d-%m-%y"
+    start_date = datetime.strptime(date_string, date_format)
+    assert (
+        start_date and start_date < datetime.now()
+    ), "Invalid start date. Please provide a valid start date."
+
     path = args.csv
     df = pd.read_csv(path)
     saved_queries = defaultdict(list)
@@ -409,5 +449,8 @@ def update_date_of_records(models: list, random_dates: list, api_key: str) -> No
     ]
     print("Updating the date of the records...")
     update_date_of_records(MODELS, random_dates, API_KEY)
+
+    print("Updating the date of the content records...")
+    update_date_of_contents(start_date)
     print("All records dates updated successfully.")
     print("All records added successfully.")