feat(handle_jwt.py): support new custom validate function for jwt auth

litellm/proxy/_new_secret_config.yaml

@@ -1,5 +1,15 @@
 model_list:
   - model_name: azure-gpt-35-turbo
     litellm_params:
-      model: topaz/chatgpt-v-2
-      api_key: os.environ/AZURE_API_KEY
+      model: azure/chatgpt-v-2
+      api_key: os.environ/AZURE_API_KEY
+      api_base: os.environ/AZURE_API_BASE
+
+general_settings:
+  master_key: sk-1234
+  enable_jwt_auth: True
+  litellm_jwtauth:
+    user_id_jwt_field: "sub"
+    team_id_jwt_field: "client_id"
+    user_id_upsert: True
+    custom_validate: custom_validate.my_custom_validate

litellm/proxy/_types.py

@@ -2,7 +2,17 @@ import enum
 import json
 import uuid
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Union,
+    get_type_hints,
+)
 
 import httpx
 from pydantic import (
@@ -30,6 +40,8 @@ from litellm.types.utils import (
     TextCompletionResponse,
 )
 
+from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
+
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
 
@@ -2395,6 +2407,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     - public_key_ttl: Default - 600s. TTL for caching public JWT keys.
     - public_allowed_routes: list of allowed routes for authenticated but unknown litellm role jwt tokens.
     - enforce_rbac: If true, enforce RBAC for all routes.
+    - custom_validate: A custom function to validates the JWT token.
 
     See `auth_checks.py` for the specific routes
     """
@@ -2439,6 +2452,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     scope_mappings: Optional[List[ScopeMapping]] = None
     enforce_scope_based_access: bool = False
     enforce_team_based_model_access: bool = False
+    custom_validate: Optional[Callable[..., Literal[True]]] = None
 
     def __init__(self, **kwargs: Any) -> None:
         # get the attribute names for this Pydantic model
@@ -2451,6 +2465,12 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
         role_mappings = kwargs.get("role_mappings")
         scope_mappings = kwargs.get("scope_mappings")
         enforce_scope_based_access = kwargs.get("enforce_scope_based_access")
+        custom_validate = kwargs.get("custom_validate")
+
+        if custom_validate is not None:
+            fn = get_instance_fn(custom_validate)
+            validate_custom_validate_return_type(fn)
+            kwargs["custom_validate"] = fn
 
         if invalid_keys:
             raise ValueError(

litellm/proxy/auth/handle_jwt.py

@@ -862,6 +862,14 @@ class JWTAuthManager:
         """Main authentication and authorization builder"""
         jwt_valid_token: dict = await jwt_handler.auth_jwt(token=api_key)
 
+        # Check custom validate
+        if jwt_handler.litellm_jwtauth.custom_validate:
+            if not jwt_handler.litellm_jwtauth.custom_validate(jwt_valid_token):
+                raise HTTPException(
+                    status_code=403,
+                    detail="Invalid JWT token",
+                )
+
         # Check RBAC
         rbac_role = jwt_handler.get_rbac_role(token=jwt_valid_token)
         await JWTAuthManager.check_rbac_role(

litellm/proxy/common_utils/callback_utils.py

@@ -4,7 +4,7 @@ import litellm
 from litellm import get_secret
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
-from litellm.proxy.utils import get_instance_fn
+from litellm.proxy.types_utils.utils import get_instance_fn
 
 blue_color_code = "\033[94m"
 reset_color_code = "\033[0m"

litellm/proxy/custom_validate.py

@@ -0,0 +1,2 @@
+def my_custom_validate(token: str) -> bool:
+    return False

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -609,7 +609,7 @@ def create_pass_through_route(
     # check if target is an adapter.py or a url
     import uuid
 
-    from litellm.proxy.utils import get_instance_fn
+    from litellm.proxy.types_utils.utils import get_instance_fn
 
     try:
         if isinstance(target, CustomLogger):

litellm/proxy/proxy_server.py

@@ -232,6 +232,7 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import (
     router as spend_management_router,
 )
 from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
+from litellm.proxy.types_utils.utils import get_instance_fn
 from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
     router as ui_crud_endpoints_router,
 )
@@ -245,7 +246,6 @@ from litellm.proxy.utils import (
     _is_projected_spend_over_limit,
     _is_valid_team_configs,
     get_error_message_str,
-    get_instance_fn,
     hash_token,
     update_spend,
 )

litellm/proxy/types_utils/README.md

@@ -0,0 +1 @@
+Utility functions for proxy types.py

litellm/proxy/types_utils/utils.py

@@ -0,0 +1,64 @@
+import importlib
+import os
+from typing import Any, Callable, Literal, Optional, get_type_hints
+
+
+def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any:
+    module_name = value
+    instance_name = None
+    try:
+        # Split the path by dots to separate module from instance
+        parts = value.split(".")
+
+        # The module path is all but the last part, and the instance_name is the last part
+        module_name = ".".join(parts[:-1])
+        instance_name = parts[-1]
+
+        # If config_file_path is provided, use it to determine the module spec and load the module
+        if config_file_path is not None:
+            directory = os.path.dirname(config_file_path)
+            module_file_path = os.path.join(directory, *module_name.split("."))
+            module_file_path += ".py"
+
+            spec = importlib.util.spec_from_file_location(module_name, module_file_path)  # type: ignore
+            if spec is None:
+                raise ImportError(
+                    f"Could not find a module specification for {module_file_path}"
+                )
+            module = importlib.util.module_from_spec(spec)  # type: ignore
+            spec.loader.exec_module(module)  # type: ignore
+        else:
+            # Dynamically import the module
+            module = importlib.import_module(module_name)
+
+        # Get the instance from the module
+        instance = getattr(module, instance_name)
+
+        return instance
+    except ImportError as e:
+        # Re-raise the exception with a user-friendly message
+        if instance_name and module_name:
+            raise ImportError(
+                f"Could not import {instance_name} from {module_name}"
+            ) from e
+        else:
+            raise e
+    except Exception as e:
+        raise e
+
+
+def validate_custom_validate_return_type(
+    fn: Optional[Callable[..., Any]]
+) -> Optional[Callable[..., Literal[True]]]:
+    if fn is None:
+        return None
+
+    hints = get_type_hints(fn)
+    return_type = hints.get("return")
+
+    if return_type != Literal[True]:
+        raise TypeError(
+            f"Custom validator must be annotated to return Literal[True], got {return_type}"
+        )
+
+    return fn

litellm/proxy/utils.py

@@ -1,7 +1,6 @@
 import asyncio
 import copy
 import hashlib
-import importlib
 import json
 import os
 import smtplib
@@ -2222,51 +2221,6 @@ class PrismaClient:
         )
 
 
-### CUSTOM FILE ###
-def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any:
-    module_name = value
-    instance_name = None
-    try:
-        # Split the path by dots to separate module from instance
-        parts = value.split(".")
-
-        # The module path is all but the last part, and the instance_name is the last part
-        module_name = ".".join(parts[:-1])
-        instance_name = parts[-1]
-
-        # If config_file_path is provided, use it to determine the module spec and load the module
-        if config_file_path is not None:
-            directory = os.path.dirname(config_file_path)
-            module_file_path = os.path.join(directory, *module_name.split("."))
-            module_file_path += ".py"
-
-            spec = importlib.util.spec_from_file_location(module_name, module_file_path)  # type: ignore
-            if spec is None:
-                raise ImportError(
-                    f"Could not find a module specification for {module_file_path}"
-                )
-            module = importlib.util.module_from_spec(spec)  # type: ignore
-            spec.loader.exec_module(module)  # type: ignore
-        else:
-            # Dynamically import the module
-            module = importlib.import_module(module_name)
-
-        # Get the instance from the module
-        instance = getattr(module, instance_name)
-
-        return instance
-    except ImportError as e:
-        # Re-raise the exception with a user-friendly message
-        if instance_name and module_name:
-            raise ImportError(
-                f"Could not import {instance_name} from {module_name}"
-            ) from e
-        else:
-            raise e
-    except Exception as e:
-        raise e
-
-
 ### HELPER FUNCTIONS ###
 async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient):
     """

tests/proxy_unit_tests/test_proxy_custom_logger.py

@@ -51,7 +51,7 @@ print("Testing proxy custom logger")
 def test_embedding(client):
     try:
         litellm.set_verbose = False
-        from litellm.proxy.utils import get_instance_fn
+        from litellm.proxy.types_utils.utils import get_instance_fn
 
         my_custom_logger = get_instance_fn(
             value="custom_callbacks.my_custom_logger", config_file_path=python_file_path
@@ -122,7 +122,7 @@ def test_chat_completion(client):
     try:
         # Your test data
         litellm.set_verbose = False
-        from litellm.proxy.utils import get_instance_fn
+        from litellm.proxy.types_utils.utils import get_instance_fn
 
         my_custom_logger = get_instance_fn(
             value="custom_callbacks.my_custom_logger", config_file_path=python_file_path
@@ -217,7 +217,7 @@ def test_chat_completion_stream(client):
     try:
         # Your test data
         litellm.set_verbose = False
-        from litellm.proxy.utils import get_instance_fn
+        from litellm.proxy.types_utils.utils import get_instance_fn
 
         my_custom_logger = get_instance_fn(
             value="custom_callbacks.my_custom_logger", config_file_path=python_file_path