Retry on rate limit error on OpenAI calls

This commit is contained in:
Ori Kotek
2023-07-20 15:01:12 +03:00
parent 6134c2ff61
commit 2f9546e144

View File

@ -1,12 +1,12 @@
import logging
import openai
from openai.error import APIError, Timeout, TryAgain
from openai.error import APIError, Timeout, TryAgain, RateLimitError
from retry import retry
from pr_agent.config_loader import settings
OPENAI_RETRIES=2
OPENAI_RETRIES=5
class AiHandler:
"""
@ -34,7 +34,7 @@ class AiHandler:
except AttributeError as e:
raise ValueError("OpenAI key is required") from e
@retry(exceptions=(APIError, Timeout, TryAgain, AttributeError),
@retry(exceptions=(APIError, Timeout, TryAgain, AttributeError, RateLimitError),
tries=OPENAI_RETRIES, delay=2, backoff=2, jitter=(1, 3))
async def chat_completion(self, model: str, temperature: float, system: str, user: str):
"""
@ -69,6 +69,9 @@ class AiHandler:
except (APIError, Timeout, TryAgain) as e:
logging.error("Error during OpenAI inference: ", e)
raise
except (RateLimitError) as e:
logging.error("Rate limit error during OpenAI inference: ", e)
raise
if response is None or len(response.choices) == 0:
raise TryAgain
resp = response.choices[0]['message']['content']