Work-around openai.RateLimitError with exponential backoff

boteval-darma-task/endpoints/_utils.py

@@ -0,0 +1,59 @@
+# imports
+import random
+import time
+import openai
+from boteval import log 
+
+
+# define a retry decorator
+def retry_with_exponential_backoff(
+    func,
+    initial_delay: float = 1,
+    exponential_base: float = 2,
+    jitter: bool = True,
+    max_retries: int = -1,
+    errors: tuple = (openai.error.RateLimitError,),
+):
+    """Retry a function with exponential backoff.
+
+    Ex.
+        @retry_with_exponential_backoff
+        def completions_with_backoff(**kwargs):
+            return openai.Completion.create(**kwargs)
+
+        completions_with_backoff(model="text-davinci-002", prompt="Once upon a time,")
+    """
+
+    def wrapper(*args, **kwargs):
+        # Initialize variables
+        num_retries = 0
+        delay = initial_delay
+
+        # Loop until a successful response or max_retries is hit or an exception is raised
+        while True:
+            try:
+                return func(*args, **kwargs)
+
+            # Retry on specified errors
+            except errors as e:
+                # Increment retries
+                num_retries += 1
+
+                # Check if max retries has been reached
+                if num_retries > max_retries:
+                    raise Exception(
+                        f"Maximum number of retries ({max_retries}) exceeded."
+                    )
+
+                # Increment the delay
+                delay *= exponential_base * (1 + jitter * random.random())
+
+                # Sleep for the delay
+                log.warning(f'OpenAI RateLimitError - sleeping for {delay} then retrying..')
+                time.sleep(delay)
+
+            # Raise exceptions for any errors not specified
+            except Exception as e:
+                raise e
+
+    return wrapper

boteval-darma-task/endpoints/chatgpt.py

@@ -6,6 +6,8 @@
 from openai.error import RateLimitError
 from boteval import log 
 from typing import Union, List, Dict
+
+from ._utils import retry_with_exponential_backoff
 from . import Endpoint
 
 class ChatGPT(Endpoint):
@@ -121,6 +123,7 @@ def _messages_compose(self, instruction:str, turns: List[Dict], turn_idx:int, **
 
         return messages
 
+
     @staticmethod
     def query_completion_api(
             messages: List[Dict[str,str]], engine:str,
@@ -129,6 +132,11 @@ def query_completion_api(
             max_timeout_rounds = 10,
             **kwargs
         ):
+
+        @retry_with_exponential_backoff
+        def _complete_with_backoff(**kwargs):
+            return openai.ChatCompletion.create(**kwargs)
+
         log.debug(f"Using engine: {engine}")
 
         for i in range(max_timeout_rounds):        
@@ -137,23 +145,17 @@ def query_completion_api(
 
             log.debug(f"Input messages: {messages}")
 
-            try:
-                response = openai.ChatCompletion.create(
-                    model=engine,
-                    messages = messages, 
-                    temperature=temperature,
-                    max_tokens=1024,
-                    top_p=1,
-                    n=n,
-                    frequency_penalty=frequency_penalty,
-                    presence_penalty=presence_penalty,
-                    stop=["user A:", "user B:", "user C:", "user D:"]
-                )
-            except RateLimitError:
-                time_to_sleep = 0.15
-                log.critical(f'OpenAI RateLimitError - sleeping for {time_to_sleep} then retrying..')
-                sleep(time_to_sleep)
-                continue
+            response = _complete_with_backoff(
+                model=engine,
+                messages = messages, 
+                temperature=temperature,
+                max_tokens=1024,
+                top_p=1,
+                n=n,
+                frequency_penalty=frequency_penalty,
+                presence_penalty=presence_penalty,
+                stop=["user A:", "user B:", "user C:", "user D:"]
+            )
 
 
             response_text = response.choices[0]['message']['content'].strip() 

boteval-darma-task/endpoints/gpt3.py

@@ -4,6 +4,7 @@
 from openai.error import RateLimitError
 from boteval import log
 from typing import List, Dict
+from ._utils import retry_with_exponential_backoff
 from . import Endpoint
 
 class GPT3(Endpoint):
@@ -124,30 +125,29 @@ def query_completion_api(
             **kwargs
         ):
 
+        @retry_with_exponential_backoff
+        def _complete_with_backoff(**kwargs):
+            return  openai.Completion.create(**kwargs)
+
+
         for i in range(max_timeout_rounds):
             # GPT-3 Generation
             if i > 0:
                 log.critical(f'GPT timeout - retry #{i}')
 
             log.debug(f"Input prompt: {prompt}")
 
-            try:
-                response = openai.Completion.create(
-                    model=engine,
-                    prompt=prompt,
-                    temperature=temperature,
-                    max_tokens=1024,
-                    top_p=1,
-                    n=n,
-                    frequency_penalty=frequency_penalty,
-                    presence_penalty=presence_penalty,
-                    stop=["user A:", "user B:", "user C:", "user D:"]
-                )
-            except RateLimitError:
-                time_to_sleep = 0.15
-                log.critical(f'OpenAI RateLimitError - sleeping for {time_to_sleep} then retrying..')
-                sleep(time_to_sleep)
-                continue
+            response = _complete_with_backoff(
+                model=engine,
+                prompt=prompt,
+                temperature=temperature,
+                max_tokens=1024,
+                top_p=1,
+                n=n,
+                frequency_penalty=frequency_penalty,
+                presence_penalty=presence_penalty,
+                stop=["user A:", "user B:", "user C:", "user D:"]
+            )
 
             # Toxicity Classification
             # https://beta.openai.com/docs/models/content-filter