add openai error handler with retry and logger (#148)

This commit is contained in:
LiviaSun
2024-02-06 12:38:41 +08:00
committed by GitHub
parent 322421fae3
commit e2bf732bc3

View File

@@ -1,3 +1,5 @@
import logging
import time
from typing import Callable, List, Optional, Union from typing import Callable, List, Optional, Union
import numpy as np import numpy as np
@@ -13,6 +15,9 @@ except ImportError as e:
openai = tiktoken = e openai = tiktoken = e
logger = logging.getLogger("openai")
def create_logit_bias_int(tokenizer): def create_logit_bias_int(tokenizer):
"""Get logit bias for integer numbers.""" """Get logit bias for integer numbers."""
int_token_ids = [] int_token_ids = []
@@ -199,10 +204,11 @@ class OpenAI(BaseBackend):
return decision, scores, scores return decision, scores, scores
def openai_completion(client, is_chat=None, prompt=None, **kwargs): def openai_completion(client, retries=3, is_chat=None, prompt=None, **kwargs):
for attempt in range(retries):
try: try:
if is_chat: if is_chat:
if kwargs["stop"] is None: if "stop" in kwargs and kwargs["stop"] is None:
kwargs.pop("stop") kwargs.pop("stop")
ret = client.chat.completions.create(messages=prompt, **kwargs) ret = client.chat.completions.create(messages=prompt, **kwargs)
comp = ret.choices[0].message.content comp = ret.choices[0].message.content
@@ -212,17 +218,24 @@ def openai_completion(client, is_chat=None, prompt=None, **kwargs):
comp = [c.text for c in ret.choices] comp = [c.text for c in ret.choices]
else: else:
comp = ret.choices[0].text comp = ret.choices[0].text
except openai.OpenAIError as e: break
print(f"OpenAI Error: {e}") except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e:
logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...")
time.sleep(5)
if attempt == retries - 1:
raise e
except Exception as e:
logger.error(f"RuntimeError {e}.")
raise e raise e
return comp return comp
def openai_completion_stream(client, is_chat=None, prompt=None, **kwargs): def openai_completion_stream(client, retries=3, is_chat=None, prompt=None, **kwargs):
for attempt in range(retries):
try: try:
if is_chat: if is_chat:
if kwargs["stop"] is None: if "stop" in kwargs and kwargs["stop"] is None:
kwargs.pop("stop") kwargs.pop("stop")
generator = client.chat.completions.create( generator = client.chat.completions.create(
messages=prompt, stream=True, **kwargs messages=prompt, stream=True, **kwargs
@@ -231,10 +244,18 @@ def openai_completion_stream(client, is_chat=None, prompt=None, **kwargs):
content = ret.choices[0].delta.content content = ret.choices[0].delta.content
yield content or "", {} yield content or "", {}
else: else:
generator = client.completions.create(prompt=prompt, stream=True, **kwargs) generator = client.completions.create(
prompt=prompt, stream=True, **kwargs
)
for ret in generator: for ret in generator:
content = ret.choices[0].text content = ret.choices[0].text
yield content or "", {} yield content or "", {}
except openai.OpenAIError as e: break
print(f"OpenAI Error: {e}") except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e:
logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...")
time.sleep(5)
if attempt == retries - 1:
raise e
except Exception as e:
logger.error(f"RuntimeError {e}.")
raise e raise e