diff --git a/src/openai_pygenerator/openai_pygenerator.py b/src/openai_pygenerator/openai_pygenerator.py index f69e162..50dd993 100644 --- a/src/openai_pygenerator/openai_pygenerator.py +++ b/src/openai_pygenerator/openai_pygenerator.py @@ -96,7 +96,7 @@ def generate_completions( APIError, ServiceUnavailableError, ) as err: - if isinstance(err, APIError) and not (err.http_status in [524, 502]): + if isinstance(err, APIError) and not (err.http_status in [524, 502, 500]): raise logger.warning("Error returned from openai API: %s", err) logger.debug("retries = %d", retries) diff --git a/tests/test_gpt.py b/tests/test_gpt.py index 255665f..4349a12 100644 --- a/tests/test_gpt.py +++ b/tests/test_gpt.py @@ -82,6 +82,7 @@ def test_generate_completion(mock_openai, mock_sleep, error): [ RateLimitError("rate limited", http_status=429), APIError("Gateway Timeout", http_status=524), + APIError("Server shutdown", http_status=500), ServiceUnavailableError("Service unavailable"), urlex.ReadTimeoutError("test-pool", "http://test", "read timeout"), openai.error.Timeout,