feat(handle_jwt.py): initial commit adding custom RBAC support on jwt… (#8037)

* feat(handle_jwt.py): initial commit adding custom RBAC support on jwt auth

allows admin to define user role field and allowed roles which map to 'internal_user' on litellm

* fix(auth_checks.py): ensure user allowed to access model, when calling via personal keys

Fixes https://github.com/BerriAI/litellm/issues/8029

* feat(handle_jwt.py): support role based access with model permission control on proxy

Allows admin to just grant users roles on IDP (e.g. Azure AD/Keycloak) and user can immediately start calling models

* docs(rbac): add docs on rbac for model access control

make it clear how admin can use roles to control model access on proxy

* fix: fix linting errors

* test(test_user_api_key_auth.py): add unit testing to ensure rbac role is correctly enforced

* test(test_user_api_key_auth.py): add more testing

* test(test_users.py): add unit testing to ensure user model access is always checked for new keys

Resolves https://github.com/BerriAI/litellm/issues/8029

* test: fix unit test

* fix(dot_notation_indexing.py): fix typing to work with python 3.8

tests/proxy_unit_tests/test_auth_checks.py

@@ -508,3 +508,43 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
     assert (
         alert_triggered == expect_alert
     ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
+
+
+@pytest.mark.asyncio
+async def test_can_user_call_model():
+    from litellm.proxy.auth.auth_checks import can_user_call_model
+    from litellm.proxy._types import ProxyException
+    from litellm import Router
+
+    router = Router(
+        model_list=[
+            {
+                "model_name": "anthropic-claude",
+                "litellm_params": {"model": "anthropic/anthropic-claude"},
+            },
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "test-api-key"},
+            },
+        ]
+    )
+
+    args = {
+        "model": "anthropic-claude",
+        "llm_router": router,
+        "user_object": LiteLLM_UserTable(
+            user_id="testuser21@mycompany.com",
+            max_budget=None,
+            spend=0.0042295,
+            model_max_budget={},
+            model_spend={},
+            user_email="testuser@mycompany.com",
+            models=["gpt-3.5-turbo"],
+        ),
+    }
+
+    with pytest.raises(ProxyException) as e:
+        await can_user_call_model(**args)
+
+    args["model"] = "gpt-3.5-turbo"
+    await can_user_call_model(**args)

tests/proxy_unit_tests/test_user_api_key_auth.py

@@ -855,6 +855,8 @@ async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypa
         "user_api_key_cache": Mock(),
         "parent_otel_span": None,
         "proxy_logging_obj": Mock(),
+        "request_data": {},
+        "general_settings": {},
     }
 
     if enforce_rbac:
@@ -877,3 +879,55 @@ def test_user_api_key_auth_end_user_str():
 
     user_api_key_auth = UserAPIKeyAuth(**user_api_key_args)
     assert user_api_key_auth.end_user_id == "1"
+
+
+def test_can_rbac_role_call_model():
+    from litellm.proxy.auth.user_api_key_auth import can_rbac_role_call_model
+    from litellm.proxy._types import RoleBasedPermissions
+
+    roles_based_permissions = [
+        RoleBasedPermissions(
+            role=LitellmUserRoles.INTERNAL_USER,
+            models=["gpt-4"],
+        ),
+        RoleBasedPermissions(
+            role=LitellmUserRoles.PROXY_ADMIN,
+            models=["anthropic-claude"],
+        ),
+    ]
+
+    assert can_rbac_role_call_model(
+        rbac_role=LitellmUserRoles.INTERNAL_USER,
+        general_settings={"role_permissions": roles_based_permissions},
+        model="gpt-4",
+    )
+
+    with pytest.raises(HTTPException):
+        can_rbac_role_call_model(
+            rbac_role=LitellmUserRoles.INTERNAL_USER,
+            general_settings={"role_permissions": roles_based_permissions},
+            model="gpt-4o",
+        )
+
+    with pytest.raises(HTTPException):
+        can_rbac_role_call_model(
+            rbac_role=LitellmUserRoles.PROXY_ADMIN,
+            general_settings={"role_permissions": roles_based_permissions},
+            model="gpt-4o",
+        )
+
+
+def test_can_rbac_role_call_model_no_role_permissions():
+    from litellm.proxy.auth.user_api_key_auth import can_rbac_role_call_model
+
+    assert can_rbac_role_call_model(
+        rbac_role=LitellmUserRoles.INTERNAL_USER,
+        general_settings={},
+        model="gpt-4",
+    )
+
+    assert can_rbac_role_call_model(
+        rbac_role=LitellmUserRoles.PROXY_ADMIN,
+        general_settings={"role_permissions": []},
+        model="anthropic-claude",
+    )