Merge pull request #1449 from BerriAI/litellm_proxy_improve_exception_mapping

[Feat] Litellm Proxy improve exception mapping

litellm/proxy/proxy_server.py

@@ -94,7 +94,12 @@ from fastapi import (
 from fastapi.routing import APIRouter
 from fastapi.security import OAuth2PasswordBearer
 from fastapi.encoders import jsonable_encoder
-from fastapi.responses import StreamingResponse, FileResponse, ORJSONResponse
+from fastapi.responses import (
+    StreamingResponse,
+    FileResponse,
+    ORJSONResponse,
+    JSONResponse,
+)
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.security.api_key import APIKeyHeader
 import json
@@ -106,6 +111,42 @@ app = FastAPI(
     title="LiteLLM API",
     description="Proxy Server to call 100+ LLMs in the OpenAI format\n\nAdmin Panel on [https://dashboard.litellm.ai/admin](https://dashboard.litellm.ai/admin)",
 )
+
+
+class ProxyException(Exception):
+    # NOTE: DO NOT MODIFY THIS
+    # This is used to map exactly to OPENAI Exceptions
+    def __init__(
+        self,
+        message: str,
+        type: str,
+        param: Optional[str],
+        code: Optional[int],
+    ):
+        self.message = message
+        self.type = type
+        self.param = param
+        self.code = code
+
+
+@app.exception_handler(ProxyException)
+async def openai_exception_handler(request: Request, exc: ProxyException):
+    # NOTE: DO NOT MODIFY THIS, its crucial to map to Openai exceptions
+    return JSONResponse(
+        status_code=int(exc.code)
+        if exc.code
+        else status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={
+            "error": {
+                "message": exc.message,
+                "type": exc.type,
+                "param": exc.param,
+                "code": exc.code,
+            }
+        },
+    )
+
+
 router = APIRouter()
 origins = ["*"]
 
@@ -1423,11 +1464,12 @@ async def completion(
         traceback.print_exc()
         error_traceback = traceback.format_exc()
         error_msg = f"{str(e)}\n\n{error_traceback}"
-        try:
-            status = e.status_code  # type: ignore
-        except:
-            status = 500
-        raise HTTPException(status_code=status, detail=error_msg)
+        raise ProxyException(
+            message=getattr(e, "message", error_msg),
+            type=getattr(e, "type", "None"),
+            param=getattr(e, "param", "None"),
+            code=getattr(e, "status_code", 500),
+        )
 
 
 @router.post(
@@ -1611,11 +1653,13 @@ async def chat_completion(
         else:
             error_traceback = traceback.format_exc()
             error_msg = f"{str(e)}\n\n{error_traceback}"
-            try:
-                status = e.status_code  # type: ignore
-            except:
-                status = 500
-            raise HTTPException(status_code=status, detail=error_msg)
+
+        raise ProxyException(
+            message=getattr(e, "message", error_msg),
+            type=getattr(e, "type", "None"),
+            param=getattr(e, "param", "None"),
+            code=getattr(e, "status_code", 500),
+        )
 
 
 @router.post(
@@ -1751,11 +1795,12 @@ async def embeddings(
         else:
             error_traceback = traceback.format_exc()
             error_msg = f"{str(e)}\n\n{error_traceback}"
-            try:
-                status = e.status_code  # type: ignore
-            except:
-                status = 500
-            raise HTTPException(status_code=status, detail=error_msg)
+            raise ProxyException(
+                message=getattr(e, "message", error_msg),
+                type=getattr(e, "type", "None"),
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", 500),
+            )
 
 
 @router.post(
@@ -1865,11 +1910,12 @@ async def image_generation(
         else:
             error_traceback = traceback.format_exc()
             error_msg = f"{str(e)}\n\n{error_traceback}"
-            try:
-                status = e.status_code  # type: ignore
-            except:
-                status = 500
-            raise HTTPException(status_code=status, detail=error_msg)
+            raise ProxyException(
+                message=getattr(e, "message", error_msg),
+                type=getattr(e, "type", "None"),
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", 500),
+            )
 
 
 #### KEY MANAGEMENT ####

litellm/proxy/tests/test_openai_exception_request.py

@@ -0,0 +1,51 @@
+import openai, httpx, os
+from dotenv import load_dotenv
+
+load_dotenv()
+client = openai.OpenAI(
+    api_key="anything",
+    base_url="http://0.0.0.0:8000",
+    http_client=httpx.Client(verify=False),
+)
+
+try:
+    # request sent to model set on litellm proxy, `litellm --model`
+    response = client.chat.completions.create(
+        model="azure-gpt-3.5",
+        messages=[
+            {
+                "role": "user",
+                "content": "this is a test request, write a short poem" * 2000,
+            }
+        ],
+    )
+
+    print(response)
+except Exception as e:
+    print(e)
+    variables_proxy_exception = vars(e)
+    print("proxy exception variables", variables_proxy_exception.keys())
+    print(variables_proxy_exception["body"])
+
+
+api_key = os.getenv("AZURE_API_KEY")
+azure_endpoint = os.getenv("AZURE_API_BASE")
+print(api_key, azure_endpoint)
+client = openai.AzureOpenAI(
+    api_key=os.getenv("AZURE_API_KEY"),
+    azure_endpoint=os.getenv("AZURE_API_BASE", "default"),
+)
+try:
+    response = client.chat.completions.create(
+        model="chatgpt-v-2",
+        messages=[
+            {
+                "role": "user",
+                "content": "this is a test request, write a short poem" * 2000,
+            }
+        ],
+    )
+except Exception as e:
+    print(e)
+    variables_exception = vars(e)
+    print("openai client exception variables", variables_exception.keys())

litellm/proxy/tests/test_openai_js.js

@@ -0,0 +1,31 @@
+const openai = require('openai');
+
+// set DEBUG=true in env
+process.env.DEBUG=false;
+async function runOpenAI() {
+  const client = new openai.OpenAI({
+    apiKey: 'your_api_key_here',
+    baseURL: 'http://0.0.0.0:8000'
+  });
+
+  try {
+    const response = await client.chat.completions.create({
+      model: 'azure-gpt-3.5',
+      messages: [
+        {
+          role: 'user',
+          content: 'this is a test request, write a short poem'.repeat(2000),
+        },
+      ],
+    });
+
+    console.log(response);
+  } catch (error) {
+    console.log("got this exception from server");
+    console.error(error);
+    console.log("done with exception from proxy");
+  }
+}
+
+// Call the asynchronous function
+runOpenAI();

litellm/tests/test_configs/test_bad_config.yaml

@@ -3,6 +3,11 @@ model_list:
     litellm_params:
       api_key: bad-key
       model: gpt-3.5-turbo
+  - model_name: working-azure-gpt-3.5-turbo
+    litellm_params:
+      model: azure/chatgpt-v-2
+      api_base: os.environ/AZURE_API_BASE
+      api_key: os.environ/AZURE_API_KEY
   - model_name: azure-gpt-3.5-turbo
     litellm_params:
       model: azure/chatgpt-v-2

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)}")