⚡️ Speed up function _get_role_based_permissions by 32%

[codeflash-ai]

📄 32% (0.32x) speedup for _get_role_based_permissions in litellm/proxy/auth/auth_checks.py

⏱️ Runtime : 230 microseconds → 174 microseconds (best of 250 runs)

📝 Explanation and details

The optimization achieves a 31% speedup by eliminating unnecessary type casting overhead in Python's function call stack. Here are the key changes:

What was optimized:

1. Removed expensive cast() operation - The original code used cast(Optional[List[RoleBasedPermissions]], general_settings.get("role_permissions", [])) which adds function call overhead without runtime benefit
2. Simplified None check - Changed from is None to if not role_based_permissions to handle both None and empty list cases uniformly
3. Removed default empty list - Using .get("role_permissions") instead of .get("role_permissions", []) since the truthiness check handles missing keys correctly

Why this is faster:

• The cast() function, while a no-op at runtime, still incurs Python function call overhead (~280ns per the profiler)
• Removing the default [] parameter eliminates unnecessary list creation when the key is missing
• The not check is slightly faster than is None comparison and handles more cases

Performance characteristics:

• Best gains on edge cases (empty/missing permissions): 243-257% faster
• Consistent improvements across all test scenarios: 100-150% faster for typical cases
• Minimal impact on large-scale searches: 7-20% faster when iterating through many roles

Impact on workloads:
Based on the function references, this optimization directly benefits JWT authentication flows where get_role_based_models() and get_role_based_routes() are called frequently during permission checks. Since auth validation is typically in the hot path for API requests, even microsecond improvements compound significantly under load.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 44 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from typing import List, Literal, Optional

# imports
import pytest  # used for our unit tests
from litellm.proxy.auth.auth_checks import _get_role_based_permissions


# function to test
class RoleBasedPermissions:
    """
    Mimics the structure used in the actual implementation.
    """
    def __init__(self, role: str, models: Optional[List[str]] = None, routes: Optional[List[str]] = None):
        self.role = role
        self.models = models
        self.routes = routes

RBAC_ROLES = Literal["admin", "user", "guest", "superuser", "readonly"]
from litellm.proxy.auth.auth_checks import _get_role_based_permissions

# unit tests

# --------------------------- BASIC TEST CASES ---------------------------

def test_basic_single_role_models():
    # Test with one role, models key
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["modelA", "modelB"], routes=["routeA"])
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 2.33μs -> 1.02μs (129% faster)

def test_basic_single_role_routes():
    # Test with one role, routes key
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["modelA"], routes=["routeA", "routeB"])
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "routes"); result = codeflash_output # 2.05μs -> 975ns (110% faster)

def test_basic_multiple_roles():
    # Test with multiple roles, pick correct one
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="user", models=["modelX"], routes=["routeX"]),
            RoleBasedPermissions(role="admin", models=["modelA"], routes=["routeA"]),
            RoleBasedPermissions(role="guest", models=["modelG"], routes=["routeG"]),
        ]
    }
    codeflash_output = _get_role_based_permissions("guest", general_settings, "models"); result = codeflash_output # 2.34μs -> 1.15μs (103% faster)

def test_basic_no_matching_role():
    # No role matches
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="user", models=["modelX"], routes=["routeX"]),
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 1.98μs -> 781ns (153% faster)

# --------------------------- EDGE TEST CASES ---------------------------

def test_edge_empty_role_permissions_list():
    # Empty list of role_permissions
    general_settings = {
        "role_permissions": []
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 1.90μs -> 547ns (248% faster)

def test_edge_role_permissions_none():
    # role_permissions is None
    general_settings = {
        "role_permissions": None
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 1.75μs -> 501ns (250% faster)

def test_edge_missing_role_permissions_key():
    # general_settings missing role_permissions key
    general_settings = {}
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 1.83μs -> 512ns (257% faster)

def test_edge_role_with_empty_models_and_routes():
    # Role exists but models and routes are empty lists
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=[], routes=[])
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result_models = codeflash_output # 2.09μs -> 988ns (112% faster)
    codeflash_output = _get_role_based_permissions("admin", general_settings, "routes"); result_routes = codeflash_output # 1.07μs -> 409ns (163% faster)

def test_edge_role_with_none_models_and_routes():
    # Role exists but models and routes are None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=None, routes=None)
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result_models = codeflash_output # 2.17μs -> 867ns (150% faster)
    codeflash_output = _get_role_based_permissions("admin", general_settings, "routes"); result_routes = codeflash_output # 1.05μs -> 422ns (148% faster)

def test_edge_multiple_roles_with_same_name():
    # Multiple entries for the same role; should return the first match
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["modelA"], routes=["routeA"]),
            RoleBasedPermissions(role="admin", models=["modelB"], routes=["routeB"]),
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 2.08μs -> 838ns (148% faster)

def test_edge_role_name_case_sensitivity():
    # Role names are case sensitive
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="Admin", models=["modelA"], routes=["routeA"]),
            RoleBasedPermissions(role="admin", models=["modelB"], routes=["routeB"]),
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 2.27μs -> 1.06μs (113% faster)

def test_edge_role_with_extra_attributes():
    # RoleBasedPermissions with extra attributes should not affect outcome
    class ExtendedRoleBasedPermissions(RoleBasedPermissions):
        def __init__(self, role, models, routes, extra):
            super().__init__(role, models, routes)
            self.extra = extra

    general_settings = {
        "role_permissions": [
            ExtendedRoleBasedPermissions(role="admin", models=["modelA"], routes=["routeA"], extra="extra_info")
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 2.16μs -> 894ns (142% faster)

def test_edge_role_permissions_not_a_list():
    # role_permissions is not a list (e.g., a dict, string, or int)
    general_settings = {
        "role_permissions": "not_a_list"
    }
    # Should not crash, but will iterate over string (which has no .role)
    with pytest.raises(AttributeError):
        _get_role_based_permissions("admin", general_settings, "models") # 2.99μs -> 1.71μs (74.8% faster)

def test_edge_role_permissions_list_with_non_role_objects():
    # role_permissions contains non-RoleBasedPermissions objects
    general_settings = {
        "role_permissions": [
            {"role": "admin", "models": ["modelA"], "routes": ["routeA"]},
            123,
            "random_string"
        ]
    }
    # Should raise AttributeError when trying to access .role
    with pytest.raises(AttributeError):
        _get_role_based_permissions("admin", general_settings, "models") # 2.84μs -> 1.64μs (73.6% faster)

# --------------------------- LARGE SCALE TEST CASES ---------------------------

def test_large_scale_many_roles():
    # Test with a large number of roles
    num_roles = 500
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=f"user{i}", models=[f"model{i}"], routes=[f"route{i}"]) for i in range(num_roles)
        ]
    }
    # Pick a role in the middle
    codeflash_output = _get_role_based_permissions(f"user250", general_settings, "models"); result = codeflash_output # 10.9μs -> 9.09μs (20.3% faster)

def test_large_scale_many_models_and_routes():
    # Role with many models and routes
    models = [f"model{i}" for i in range(1000)]
    routes = [f"route{i}" for i in range(1000)]
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=models, routes=routes)
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result_models = codeflash_output # 2.31μs -> 1.01μs (129% faster)
    codeflash_output = _get_role_based_permissions("admin", general_settings, "routes"); result_routes = codeflash_output # 951ns -> 323ns (194% faster)
    # Check performance by ensuring it doesn't time out (pytest will handle this)

def test_large_scale_first_and_last_role():
    # Ensure first and last roles are accessible
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="first", models=["firstModel"], routes=["firstRoute"]),
        ] + [
            RoleBasedPermissions(role=f"user{i}", models=[f"model{i}"], routes=[f"route{i}"]) for i in range(1, 998)
        ] + [
            RoleBasedPermissions(role="last", models=["lastModel"], routes=["lastRoute"]),
        ]
    }
    codeflash_output = _get_role_based_permissions("first", general_settings, "models"); result_first = codeflash_output # 2.65μs -> 1.03μs (156% faster)
    codeflash_output = _get_role_based_permissions("last", general_settings, "routes"); result_last = codeflash_output # 29.8μs -> 28.9μs (3.06% faster)

def test_large_scale_no_matching_role():
    # Large list, no matching role
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=f"user{i}", models=[f"model{i}"], routes=[f"route{i}"]) for i in range(1000)
        ]
    }
    codeflash_output = _get_role_based_permissions("admin", general_settings, "models"); result = codeflash_output # 31.4μs -> 29.2μs (7.72% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import List, Literal, Optional

# imports
import pytest
from litellm.proxy.auth.auth_checks import _get_role_based_permissions

# Mock types to match the function's signature
RBAC_ROLES = Literal["admin", "user", "viewer", "editor", "custom"]
# Simulate the RoleBasedPermissions dataclass
class RoleBasedPermissions:
    def __init__(self, role: str, models: Optional[List[str]] = None, routes: Optional[List[str]] = None):
        self.role = role
        self.models = models
        self.routes = routes
from litellm.proxy.auth.auth_checks import _get_role_based_permissions

# --------------------- UNIT TESTS ---------------------

# 1. Basic Test Cases

def test_basic_single_role_models():
    # Test with one role, key="models"
    perms = [RoleBasedPermissions(role="admin", models=["gpt-3", "gpt-4"], routes=["/a", "/b"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.17μs -> 875ns (148% faster)

def test_basic_single_role_routes():
    # Test with one role, key="routes"
    perms = [RoleBasedPermissions(role="admin", models=["gpt-3"], routes=["/admin", "/dashboard"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "routes"); result = codeflash_output # 2.07μs -> 887ns (134% faster)

def test_basic_multiple_roles():
    # Test with multiple roles, pick the correct one
    perms = [
        RoleBasedPermissions(role="user", models=["gpt-3"], routes=["/user"]),
        RoleBasedPermissions(role="admin", models=["gpt-4"], routes=["/admin"]),
        RoleBasedPermissions(role="viewer", models=[], routes=[]),
    ]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.18μs -> 1.00μs (117% faster)

def test_basic_no_match():
    # Role not present in permissions
    perms = [
        RoleBasedPermissions(role="user", models=["gpt-3"], routes=["/user"]),
        RoleBasedPermissions(role="viewer", models=[], routes=[]),
    ]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.04μs -> 818ns (150% faster)

def test_basic_empty_permissions_list():
    # Empty role_permissions list
    settings = {"role_permissions": []}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 1.85μs -> 541ns (243% faster)

# 2. Edge Test Cases

def test_edge_role_permissions_none():
    # role_permissions is None
    settings = {"role_permissions": None}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 1.75μs -> 523ns (234% faster)

def test_edge_no_role_permissions_key():
    # role_permissions key missing
    settings = {}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 1.81μs -> 509ns (255% faster)

def test_edge_models_none():
    # models is None for the matched role
    perms = [RoleBasedPermissions(role="admin", models=None, routes=["/admin"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.13μs -> 969ns (120% faster)

def test_edge_routes_none():
    # routes is None for the matched role
    perms = [RoleBasedPermissions(role="admin", models=["gpt-3"], routes=None)]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "routes"); result = codeflash_output # 2.12μs -> 975ns (117% faster)

def test_edge_empty_models_and_routes():
    # models and routes are empty lists
    perms = [RoleBasedPermissions(role="admin", models=[], routes=[])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result_models = codeflash_output # 2.12μs -> 887ns (139% faster)
    codeflash_output = _get_role_based_permissions("admin", settings, "routes"); result_routes = codeflash_output # 1.06μs -> 410ns (160% faster)

def test_edge_multiple_same_roles():
    # Multiple entries for the same role, should return first match
    perms = [
        RoleBasedPermissions(role="admin", models=["gpt-3"], routes=["/admin1"]),
        RoleBasedPermissions(role="admin", models=["gpt-4"], routes=["/admin2"]),
    ]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.05μs -> 849ns (142% faster)

def test_edge_role_case_sensitivity():
    # Role match should be case-sensitive
    perms = [RoleBasedPermissions(role="Admin", models=["gpt-3"], routes=["/admin"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 1.95μs -> 785ns (149% faster)

def test_edge_non_list_models_routes():
    # models/routes are not lists (should be returned as is)
    perms = [RoleBasedPermissions(role="admin", models="gpt-3", routes="/admin")]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result_models = codeflash_output # 2.07μs -> 840ns (146% faster)
    codeflash_output = _get_role_based_permissions("admin", settings, "routes"); result_routes = codeflash_output # 1.06μs -> 424ns (151% faster)

def test_edge_custom_role():
    # Role not in default set but present in permissions
    perms = [RoleBasedPermissions(role="custom", models=["special"], routes=["/special"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("custom", settings, "models"); result = codeflash_output # 1.99μs -> 856ns (133% faster)

def test_edge_permission_object_missing_attribute():
    # RoleBasedPermissions object missing the attribute (should raise AttributeError)
    class IncompletePerm:
        def __init__(self, role):
            self.role = role
    perms = [IncompletePerm("admin")]
    settings = {"role_permissions": perms}
    with pytest.raises(AttributeError):
        _get_role_based_permissions("admin", settings, "models") # 3.23μs -> 1.98μs (63.2% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_roles():
    # Test with 1000 roles, pick the last one
    perms = [RoleBasedPermissions(role=f"user{i}", models=[f"model{i}"], routes=[f"/route{i}"]) for i in range(999)]
    perms.append(RoleBasedPermissions(role="admin", models=["gpt-4"], routes=["/admin"]))
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 31.7μs -> 29.5μs (7.30% faster)

def test_large_scale_first_match():
    # 1000 entries, multiple for same role, should return first
    perms = [RoleBasedPermissions(role="admin", models=[f"model{i}"], routes=[f"/route{i}"]) for i in range(500)]
    perms += [RoleBasedPermissions(role="admin", models=["final"], routes=["/final"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.44μs -> 977ns (150% faster)

def test_large_scale_no_match():
    # 1000 entries, none match
    perms = [RoleBasedPermissions(role=f"user{i}", models=[f"model{i}"], routes=[f"/route{i}"]) for i in range(1000)]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 31.5μs -> 29.3μs (7.26% faster)

def test_large_scale_all_empty_models():
    # 1000 entries, all with empty models, match on 500th
    perms = [RoleBasedPermissions(role=f"user{i}", models=[], routes=[]) for i in range(500)]
    perms.append(RoleBasedPermissions(role="admin", models=[], routes=[]))
    perms += [RoleBasedPermissions(role=f"user{i}", models=[], routes=[]) for i in range(499)]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 17.2μs -> 15.4μs (12.0% faster)

def test_large_scale_large_models_list():
    # One role with a large models list
    big_models = [f"model{i}" for i in range(999)]
    perms = [RoleBasedPermissions(role="admin", models=big_models, routes=["/admin"])]
    settings = {"role_permissions": perms}
    codeflash_output = _get_role_based_permissions("admin", settings, "models"); result = codeflash_output # 2.20μs -> 956ns (130% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_get_role_based_permissions-mhww9hl2 and push.