Skip to content

Commit

Permalink
Adds error handling for openai's try again error in llms/openai modul…
Browse files Browse the repository at this point in the history
…e and its test

- Adds test for openai's try again error handling in chat_completion in llms/openai module
- Adds error handling for openai's try again error in chat_completion in llms/openai module
  • Loading branch information
aleric-cusher committed Nov 7, 2023
1 parent 636cfb7 commit 41e6b6a
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 2 deletions.
8 changes: 6 additions & 2 deletions superagi/llms/openai.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import openai
from openai import APIError, InvalidRequestError
from openai.error import RateLimitError, AuthenticationError, Timeout
from openai.error import RateLimitError, AuthenticationError, Timeout, TryAgain
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential

from superagi.config.config import get_config
Expand Down Expand Up @@ -58,7 +58,8 @@ def get_model(self):
@retry(
retry=(
retry_if_exception_type(RateLimitError) |
retry_if_exception_type(Timeout)
retry_if_exception_type(Timeout) |
retry_if_exception_type(TryAgain)
),
stop=stop_after_attempt(MAX_RETRY_ATTEMPTS), # Maximum number of retry attempts
wait=wait_random_exponential(min=MIN_WAIT, max=MAX_WAIT),
Expand Down Expand Up @@ -95,6 +96,9 @@ def chat_completion(self, messages, max_tokens=get_config("MAX_MODEL_TOKEN_LIMIT
except Timeout as timeout_error:
logger.info("OpenAi Timeout:", timeout_error)
raise Timeout(str(timeout_error))
except TryAgain as try_again_error:
logger.info("OpenAi TryAgain:", try_again_error)
raise TryAgain(str(try_again_error))
except AuthenticationError as auth_error:
logger.info("OpenAi AuthenticationError:", auth_error)
return {"error": "ERROR_AUTHENTICATION", "message": "Authentication error please check the api keys: "+str(auth_error)}
Expand Down
24 changes: 24 additions & 0 deletions tests/unit_tests/llms/test_open_ai.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,30 @@ def test_chat_completion_retry_timeout_error(mock_openai, mock_wait_random_expon
assert mock_openai.ChatCompletion.create.call_count == MAX_RETRY_ATTEMPTS


@patch('superagi.llms.openai.wait_random_exponential.__call__')
@patch('superagi.llms.openai.openai')
def test_chat_completion_retry_try_again_error(mock_openai, mock_wait_random_exponential):
# Arrange
model = 'gpt-4'
api_key = 'test_key'
openai_instance = OpenAi(api_key, model=model)

messages = [{"role": "system", "content": "You are a helpful assistant."}]
max_tokens = 100

mock_openai.ChatCompletion.create.side_effect = openai.error.TryAgain("Try Again")

# Mock sleep time
mock_wait_random_exponential.return_value = 0.1

# Act
with pytest.raises(tenacity.RetryError):
result = openai_instance.chat_completion(messages, max_tokens)

# Assert
assert mock_openai.ChatCompletion.create.call_count == MAX_RETRY_ATTEMPTS


def test_verify_access_key():
model = 'gpt-4'
api_key = 'test_key'
Expand Down

0 comments on commit 41e6b6a

Please sign in to comment.