File size: 2,261 Bytes
744eb4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import openai
import time
import random
import os
def retry_with_exponential_backoff(
func,
initial_delay: float = 1,
exponential_base: float = 2,
jitter: bool = True,
max_retries: int = 40,
max_delay: int = 30,
errors: tuple = (openai.error.RateLimitError, openai.error.ServiceUnavailableError, openai.error.Timeout),
):
"""Retry a function with exponential backoff."""
def wrapper(*args, **kwargs):
num_retries = 0
delay = initial_delay
while True:
try:
return func(*args, **kwargs)
except errors as e:
# * print the error info
num_retries += 1
if num_retries > max_retries:
print(f"[OPENAI] Encounter error: {e}.")
raise Exception(
f"[OPENAI] Maximum number of retries ({max_retries}) exceeded."
)
delay *= exponential_base * (1 + jitter * random.random())
time.sleep(min(delay, max_delay))
except Exception as e:
raise e
return wrapper
class OpenAIGPT():
def __init__(self, model="gpt-3.5-turbo-0613", temperature=1, top_p=1, max_tokens=2048, **kwargs) -> None:
setup_openai(model)
self.default_chat_parameters = {
"model": model,
"temperature": temperature,
"top_p": top_p,
"max_tokens": max_tokens,
**kwargs
}
@retry_with_exponential_backoff
def safe_chat_complete(self, messages, content_only=True, **kwargs):
chat_parameters = self.default_chat_parameters.copy()
if len(kwargs) > 0:
chat_parameters.update(**kwargs)
response = openai.ChatCompletion.create(
messages=messages,
**chat_parameters
)
if content_only:
response = response['choices'][0]["message"]['content']
return response
def setup_openai(model_name):
# Setup OpenAI API Key
print("[OPENAI] Setting OpenAI api_key...")
openai.api_key = os.getenv('OPENAI_API_KEY')
print(f"[OPENAI] OpenAI organization: {openai.organization}")
print(f"[OPENAI] Using MODEL: {model_name}") |