You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
import typing
|
|
from functools import wraps
|
|
|
|
import openai
|
|
from dotenv import load_dotenv
|
|
from rich.progress import track
|
|
|
|
load_dotenv()
|
|
|
|
openai.api_key = os.getenv("OPENAI_API_KEY")
|
|
openai.debug = False
|
|
|
|
|
|
def setup_logging(verbosity: int) -> None:
|
|
logging_level = logging.WARNING
|
|
if verbosity == 1:
|
|
logging_level = logging.INFO
|
|
elif verbosity >= 2:
|
|
openai.debug = True
|
|
logging_level = logging.DEBUG
|
|
|
|
logging.basicConfig(
|
|
handlers=[
|
|
logging.StreamHandler(),
|
|
],
|
|
format="%(asctime)s - %(filename)s:%(lineno)d - %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
level=logging_level,
|
|
)
|
|
logging.captureWarnings(capture=True)
|
|
|
|
|
|
def retry(
|
|
exceptions: typing.Type[openai.error.RateLimitError], tries: int = 4, delay: int = 3, back_off: int = 2
|
|
) -> typing.Callable:
|
|
def deco_retry(f: typing.Any) -> typing.Callable:
|
|
@wraps(f)
|
|
def f_retry(*args, **kwargs): # type: ignore
|
|
m_retries, m_delay = tries, delay
|
|
while m_retries > 1:
|
|
try:
|
|
return f(*args, **kwargs)
|
|
except exceptions as e:
|
|
msg = f"🚨 {e} {os.linesep} "
|
|
logging.warning(msg)
|
|
for _ in track(range(m_delay), description="⌛ Retrying in ..."):
|
|
time.sleep(1)
|
|
m_retries -= 1
|
|
m_delay *= back_off
|
|
return f(*args, **kwargs)
|
|
|
|
return f_retry # true decorator
|
|
|
|
return deco_retry
|