fix: Fix missing error in openai_request retry strategy (#4802)

* Fix missing error in openai_request retry strategy

* Correctly handle OpenAIUnauthorizedError

Co-authored-by: bogdankostic <bogdankostic@web.de>

---------

Co-authored-by: bogdankostic <bogdankostic@web.de>
This commit is contained in:
Silvano Cerza 2023-05-10 10:31:07 +02:00 committed by GitHub
parent c734c58b4b
commit f12e5a0127
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 60 additions and 1 deletions

View File

@ -128,7 +128,9 @@ def _openai_text_completion_tokenization_details(model_name: str):
@tenacity.retry(
retry=tenacity.retry_if_exception_type(OpenAIRateLimitError),
reraise=True,
retry=tenacity.retry_if_exception_type(OpenAIError)
and tenacity.retry_if_not_exception_type(OpenAIUnauthorizedError),
wait=tenacity.wait_exponential(multiplier=OPENAI_BACKOFF),
stop=tenacity.stop_after_attempt(OPENAI_MAX_RETRIES),
)

View File

@ -0,0 +1,57 @@
from unittest.mock import patch
import pytest
from tenacity import wait_none
from haystack.errors import OpenAIError, OpenAIRateLimitError, OpenAIUnauthorizedError
from haystack.utils.openai_utils import openai_request
@pytest.mark.unit
@patch("haystack.utils.openai_utils.requests")
def test_openai_request_retries_generic_error(mock_requests):
mock_requests.request.return_value.status_code = 418
with pytest.raises(OpenAIError):
# We need to use a custom wait amount otherwise the test would take forever to run
# as the original wait time is exponential
openai_request.retry_with(wait=wait_none())(url="some_url", headers={}, payload={}, read_response=False)
assert mock_requests.request.call_count == 5
@pytest.mark.unit
@patch("haystack.utils.openai_utils.requests")
def test_openai_request_retries_on_rate_limit_error(mock_requests):
mock_requests.request.return_value.status_code = 429
with pytest.raises(OpenAIRateLimitError):
# We need to use a custom wait amount otherwise the test would take forever to run
# as the original wait time is exponential
openai_request.retry_with(wait=wait_none())(url="some_url", headers={}, payload={}, read_response=False)
assert mock_requests.request.call_count == 5
@pytest.mark.unit
@patch("haystack.utils.openai_utils.requests")
def test_openai_request_does_not_retry_on_unauthorized_error(mock_requests):
mock_requests.request.return_value.status_code = 401
with pytest.raises(OpenAIUnauthorizedError):
# We need to use a custom wait amount otherwise the test would take forever to run
# as the original wait time is exponential
openai_request.retry_with(wait=wait_none())(url="some_url", headers={}, payload={}, read_response=False)
assert mock_requests.request.call_count == 1
@pytest.mark.unit
@patch("haystack.utils.openai_utils.requests")
def test_openai_request_does_not_retry_on_success(mock_requests):
mock_requests.request.return_value.status_code = 200
# We need to use a custom wait amount otherwise the test would take forever to run
# as the original wait time is exponential
openai_request.retry_with(wait=wait_none())(url="some_url", headers={}, payload={}, read_response=False)
assert mock_requests.request.call_count == 1