Added simple retry on Open AI chat if a Rate Limit or 502 Bad Gateway error received

scripts/llm_utils.py

@@ -1,26 +1,52 @@
+import time
 import openai
+from colorama import Fore
 from config import Config
+
 cfg = Config()
 
 openai.api_key = cfg.openai_api_key
 
+
 # Overly simple abstraction until we create something better
+# simple retry mechanism when getting a rate error or a bad gateway
 def create_chat_completion(messages, model=None, temperature=cfg.temperature, max_tokens=None)->str:
     """Create a chat completion using the OpenAI API"""
-    if cfg.use_azure:
+    response = None
-        response = openai.ChatCompletion.create(
+    num_retries = 5
-            deployment_id=cfg.get_azure_deployment_id_for_model(model),
+    for attempt in range(num_retries):
-            model=model,
+        try:
-            messages=messages,
+            if cfg.use_azure:
-            temperature=temperature,
+                response = openai.ChatCompletion.create(
-            max_tokens=max_tokens
+                    deployment_id=cfg.get_azure_deployment_id_for_model(model),
-        )
+                    model=model,
-    else:
+                    messages=messages,
-        response = openai.ChatCompletion.create(
+                    temperature=temperature,
-            model=model,
+                    max_tokens=max_tokens
-            messages=messages,
+                )
-            temperature=temperature,
+            else:
-            max_tokens=max_tokens
+                response = openai.ChatCompletion.create(
-        )
+                    model=model,
+                    messages=messages,
+                    temperature=temperature,
+                    max_tokens=max_tokens
+                )
+            break
+        except openai.error.RateLimitError:
+            if cfg.debug_mode:
+                print(Fore.RED + "Error: ", "API Rate Limit Reached. Waiting 20 seconds..." + Fore.RESET)
+            time.sleep(20)
+        except openai.error.APIError as e:
+            if e.http_status == 502:
+                if cfg.debug_mode:
+                    print(Fore.RED + "Error: ", "API Bad gateway. Waiting 20 seconds..." + Fore.RESET)
+                time.sleep(20)
+            else:
+                raise
+            if attempt == num_retries - 1:
+                raise
+
+    if response is None:
+        raise RuntimeError("Failed to get response after 5 retries")
 
     return response.choices[0].message["content"]