diff --git a/autogen/oai/completion.py b/autogen/oai/completion.py index ce9d50f8690..29ed334bd82 100644 --- a/autogen/oai/completion.py +++ b/autogen/oai/completion.py @@ -13,8 +13,7 @@ try: import openai - from openai.error import ( - ServiceUnavailableError, + from openai import ( RateLimitError, APIError, BadRequestError, @@ -221,10 +220,7 @@ def _get_response(cls, config: Dict, raise_on_ratelimit_or_timeout=False, use_ca response = openai_completion.create(**config) else: response = openai_completion.create(request_timeout=request_timeout, **config) - except ( - ServiceUnavailableError, - APIConnectionError, - ): + except APIConnectionError: # transient error logger.info(f"retrying in {retry_wait_time} seconds...", exc_info=1) sleep(retry_wait_time) diff --git a/test/oai/test_completion.py b/test/oai/test_completion.py index f33533b7bf8..dd03864de8e 100644 --- a/test/oai/test_completion.py +++ b/test/oai/test_completion.py @@ -272,7 +272,7 @@ def test_humaneval(num_samples=1): ) response = autogen.ChatCompletion.create(context=test_data[0], config_list=config_list, **config) print(response) - from openai.error import RateLimitError + from openai import RateLimitError try: code, cost, selected = implement(tune_data[1], [{**config_list[-1], **config}])