close #402

microsoft · Oct 24, 2023 · de23d55 · de23d55
1 parent c8f8cbc
commit de23d55
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git 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
@@ -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}])