Custom Auth - bubble up custom exceptions (#13093)

* fix(enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py): bubble up exception if type is ProxyException

* docs(custom_auth.md): doc on bubbling up custom exceptions

docs/my-website/docs/proxy/custom_auth.md

@@ -60,9 +60,114 @@ Supported from v1.72.2+
 [Get free 7-day trial key](https://www.litellm.ai/enterprise#trial)
 :::
 
+### Usage
+
+1. Setup custom auth file
+
+```python
+"""
+Example custom auth function.
+
+This will allow all keys starting with "my-custom-key" to pass through.
+"""
+from typing import Union
+
+from fastapi import Request
+
+from litellm.proxy._types import UserAPIKeyAuth
+
+
+async def user_api_key_auth(
+    request: Request, api_key: str
+) -> Union[UserAPIKeyAuth, str]:
+    try:
+        if api_key.startswith("my-custom-key"):
+            return "sk-P1zJMdsqCPNN54alZd_ETw"
+        else:
+            raise Exception("Invalid API key")
+    except Exception:
+        raise Exception("Invalid API key")
+
+```
+
+2. Setup config.yaml
+
+Key change set `mode: auto`. This will check both litellm api key auth + custom auth.
+
 ```yaml
+model_list: 
+  - model_name: "openai-model"
+    litellm_params: 
+      model: "gpt-3.5-turbo"
+      api_key: os.environ/OPENAI_API_KEY
+
 general_settings:
   custom_auth: custom_auth_auto.user_api_key_auth
   custom_auth_settings:
     mode: "auto" # can be 'on', 'off', 'auto' - 'auto' checks both litellm api key auth + custom auth
+```
+
+Flow:
+1. Checks custom auth first
+2. If custom auth fails, checks litellm api key auth
+3. If both fail, returns 401
+
+
+3. Test it! 
+
+```bash
+curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
+-H 'Content-Type: application/json' \
+-H 'Authorization: Bearer sk-P1zJMdsqCPNN54alZd_ETw' \
+-d '{
+    "model": "openai-model",
+    "messages": [
+          {
+            "role": "user",
+            "content": "Hey! My name is John"
+          }
+        ]
+}'
+```
+
+
+
+
+#### Bubble up custom exceptions
+
+If you want to bubble up custom exceptions, you can do so by raising a `ProxyException`.
+
+```python
+"""
+Example custom auth function.
+
+This will allow all keys starting with "my-custom-key" to pass through.
+"""
+
+from typing import Union
+
+from fastapi import Request
+
+from litellm.proxy._types import UserAPIKeyAuth, ProxyException
+
+
+async def user_api_key_auth(
+    request: Request, api_key: str
+) -> Union[UserAPIKeyAuth, str]:
+    try:
+        if api_key.startswith("my-custom-key"):
+            return "sk-P1zJMdsqCPNN54alZd_ETw"
+        if api_key == "invalid-api-key":
+            # raise a custom exception back to the client
+            raise ProxyException(
+                message="Invalid API key",
+                type="invalid_request_error",
+                param="api_key",
+                code=401,
+            )
+        else:
+            raise Exception("Invalid API key")
+    except Exception:
+        raise Exception("Invalid API key")
+
 ```

enterprise/litellm_enterprise/proxy/auth/user_api_key_auth.py

@@ -3,7 +3,7 @@ from typing import Any, Optional
 from fastapi import Request
 
 from litellm._logging import verbose_proxy_logger
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 
 
 async def enterprise_custom_auth(
@@ -24,6 +24,8 @@ async def enterprise_custom_auth(
     elif custom_auth_settings["mode"] == "auto":
         try:
             return await user_custom_auth(request, api_key)
+        except ProxyException as e:
+            raise e
         except Exception as e:
             verbose_proxy_logger.debug(
                 f"Error in custom auth, checking litellm auth: {e}"

litellm/proxy/custom_auth_auto.py

@@ -3,11 +3,12 @@ Example custom auth function.
 
 This will allow all keys starting with "my-custom-key" to pass through.
 """
+
 from typing import Union
 
 from fastapi import Request
 
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 
 
 async def user_api_key_auth(
@@ -16,6 +17,14 @@ async def user_api_key_auth(
     try:
         if api_key.startswith("my-custom-key"):
             return "sk-P1zJMdsqCPNN54alZd_ETw"
+        if api_key == "invalid-api-key":
+            # raise a custom exception back to the client
+            raise ProxyException(
+                message="Invalid API key",
+                type="invalid_request_error",
+                param="api_key",
+                code=401,
+            )
         else:
             raise Exception("Invalid API key")
     except Exception: