From afaca3f819db72b1bb78a2e7b2b85e6bb98e0042 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 25 Nov 2023 15:33:44 -0800
Subject: [PATCH] (fix) acompletion: Raise same error as completion

---
 litellm/llms/azure.py            | 10 ++++-----
 litellm/tests/test_exceptions.py | 38 ++++++++++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py
index 218762d753..2cf87ea240 100644
--- a/litellm/llms/azure.py
+++ b/litellm/llms/azure.py
@@ -196,13 +196,11 @@ class AzureChatCompletion(BaseLLM):
             response = await azure_client.chat.completions.create(**data) 
             response.model = "azure/" + str(response.model)
             return convert_to_model_response_object(response_object=json.loads(response.model_dump_json()), model_response_object=model_response)
+       except AzureOpenAIError as e: 
+            exception_mapping_worked = True
+            raise e
        except Exception as e: 
-           if isinstance(e,httpx.TimeoutException):
-                raise AzureOpenAIError(status_code=500, message="Request Timeout Error")
-           elif response is not None and hasattr(response, "text"):
-                raise AzureOpenAIError(status_code=500, message=f"{str(e)}\n\nOriginal Response: {response.text}")
-           else: 
-                raise AzureOpenAIError(status_code=500, message=f"{str(e)}")
+            raise e
 
     def streaming(self,
                   logging_obj,
diff --git a/litellm/tests/test_exceptions.py b/litellm/tests/test_exceptions.py
index c2a7a57ad9..ae7910e3bc 100644
--- a/litellm/tests/test_exceptions.py
+++ b/litellm/tests/test_exceptions.py
@@ -179,7 +179,7 @@ def test_completion_azure_exception():
         litellm.set_verbose=False
         ## Test azure call
         old_azure_key = os.environ["AZURE_API_KEY"]
-        os.environ["AZURE_API_KEY"] = ""
+        os.environ["AZURE_API_KEY"] = "good morning"
         response = completion(
             model="azure/chatgpt-v-2",
             messages=[
@@ -191,13 +191,47 @@ def test_completion_azure_exception():
         )
         print(f"response: {response}")
         print(response)
-    except openai.APIConnectionError as e:
+    except openai.AuthenticationError as e:
         os.environ["AZURE_API_KEY"] = old_azure_key
         print("good job got the correct error for azure when key not set")
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 test_completion_azure_exception()
 
+async def test_completion_azure_exception():
+    try:
+        import openai
+        import litellm
+        print("azure gpt-3.5 test\n\n")
+        litellm.set_verbose=False
+        ## Test azure call
+        old_azure_key = os.environ["AZURE_API_KEY"]
+        os.environ["AZURE_API_KEY"] = "good morning"
+        response = await litellm.acompletion(
+            model="azure/chatgpt-v-2",
+            messages=[
+                {
+                    "role": "user",
+                    "content": "hello"
+                }
+            ],
+        )
+        print(f"response: {response}")
+        print(response)
+    except openai.AuthenticationError as e:
+        os.environ["AZURE_API_KEY"] = old_azure_key
+        print("good job got the correct error for azure when key not set")
+        print(e)
+    except Exception as e:
+        print("Got wrong exception")
+        print("exception", e)
+        pytest.fail(f"Error occurred: {e}")
+
+import asyncio
+asyncio.run(
+    test_completion_azure_exception()
+)
+
 # test_invalid_request_error(model="command-nightly")
 # Test 3: Rate Limit Errors
 # def test_model_call(model):