From f0d001658a2291894d9dc223cc0201a8691ebbb4 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 15 Jan 2024 10:43:26 -0800
Subject: [PATCH] (test) proxy exception mapping - exactly like OpenAI

---
 litellm/tests/test_proxy_exception_mapping.py | 73 ++++++++++++++++++-
 1 file changed, 71 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_proxy_exception_mapping.py b/litellm/tests/test_proxy_exception_mapping.py
index fcc0ad98c96..cab26f319ed 100644
--- a/litellm/tests/test_proxy_exception_mapping.py
+++ b/litellm/tests/test_proxy_exception_mapping.py
@@ -25,8 +25,8 @@ def client():
     filepath = os.path.dirname(os.path.abspath(__file__))
     config_fp = f"{filepath}/test_configs/test_bad_config.yaml"
     asyncio.run(initialize(config=config_fp))
-    app = FastAPI()
-    app.include_router(router)  # Include your router in the test app
+    from litellm.proxy.proxy_server import app
+
     return TestClient(app)
 
 
@@ -44,6 +44,10 @@ def test_chat_completion_exception(client):
 
         response = client.post("/chat/completions", json=test_data)
 
+        json_response = response.json()
+        print("keys in json response", json_response.keys())
+        assert json_response.keys() == {"error"}
+
         # make an openai client to call _make_status_error_from_response
         openai_client = openai.OpenAI(api_key="anything")
         openai_exception = openai_client._make_status_error_from_response(
@@ -69,6 +73,10 @@ def test_chat_completion_exception_azure(client):
 
         response = client.post("/chat/completions", json=test_data)
 
+        json_response = response.json()
+        print("keys in json response", json_response.keys())
+        assert json_response.keys() == {"error"}
+
         # make an openai client to call _make_status_error_from_response
         openai_client = openai.OpenAI(api_key="anything")
         openai_exception = openai_client._make_status_error_from_response(
@@ -90,6 +98,10 @@ def test_embedding_auth_exception_azure(client):
         response = client.post("/embeddings", json=test_data)
         print("Response from proxy=", response)
 
+        json_response = response.json()
+        print("keys in json response", json_response.keys())
+        assert json_response.keys() == {"error"}
+
         # make an openai client to call _make_status_error_from_response
         openai_client = openai.OpenAI(api_key="anything")
         openai_exception = openai_client._make_status_error_from_response(
@@ -117,6 +129,10 @@ def test_exception_openai_bad_model(client):
 
         response = client.post("/chat/completions", json=test_data)
 
+        json_response = response.json()
+        print("keys in json response", json_response.keys())
+        assert json_response.keys() == {"error"}
+
         # make an openai client to call _make_status_error_from_response
         openai_client = openai.OpenAI(api_key="anything")
         openai_exception = openai_client._make_status_error_from_response(
@@ -143,6 +159,10 @@ def test_chat_completion_exception_any_model(client):
 
         response = client.post("/chat/completions", json=test_data)
 
+        json_response = response.json()
+        print("keys in json response", json_response.keys())
+        assert json_response.keys() == {"error"}
+
         # make an openai client to call _make_status_error_from_response
         openai_client = openai.OpenAI(api_key="anything")
         openai_exception = openai_client._make_status_error_from_response(
@@ -163,6 +183,11 @@ def test_embedding_exception_any_model(client):
 
         response = client.post("/embeddings", json=test_data)
         print("Response from proxy=", response)
+        print(response.json())
+
+        json_response = response.json()
+        print("keys in json response", json_response.keys())
+        assert json_response.keys() == {"error"}
 
         # make an openai client to call _make_status_error_from_response
         openai_client = openai.OpenAI(api_key="anything")
@@ -174,3 +199,47 @@ def test_embedding_exception_any_model(client):
 
     except Exception as e:
         pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")
+
+
+# raise openai.BadRequestError
+def test_chat_completion_exception_azure_context_window(client):
+    try:
+        # Your test data
+        test_data = {
+            "model": "working-azure-gpt-3.5-turbo",
+            "messages": [
+                {"role": "user", "content": "hi" * 10000},
+            ],
+            "max_tokens": 10,
+        }
+        response = None
+
+        response = client.post("/chat/completions", json=test_data)
+        print("got response from server", response)
+
+        json_response = response.json()
+
+        print("keys in json response", json_response.keys())
+
+        assert json_response.keys() == {"error"}
+
+        assert json_response == {
+            "error": {
+                "message": "AzureException - Error code: 400 - {'error': {'message': \"This model's maximum context length is 4096 tokens. However, your messages resulted in 10007 tokens. Please reduce the length of the messages.\", 'type': 'invalid_request_error', 'param': 'messages', 'code': 'context_length_exceeded'}}",
+                "type": None,
+                "param": None,
+                "code": 400,
+            }
+        }
+
+        # make an openai client to call _make_status_error_from_response
+        openai_client = openai.OpenAI(api_key="anything")
+        openai_exception = openai_client._make_status_error_from_response(
+            response=response
+        )
+        print("exception from proxy", openai_exception)
+        assert isinstance(openai_exception, openai.BadRequestError)
+        print("passed exception is of type BadRequestError")
+
+    except Exception as e:
+        pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")