⚡️ Speed up function get_role_based_models by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for get_role_based_models in litellm/proxy/auth/auth_checks.py

⏱️ Runtime : 263 microseconds → 202 microseconds (best of 61 runs)

📝 Explanation and details

The optimization removes an unnecessary cast() operation and simplifies the default parameter in the dictionary lookup. Specifically:

Key Changes:

1. Removed cast() operation: The original code used cast(Optional[List[RoleBasedPermissions]], ...) which adds runtime overhead for type checking that provides no functional benefit
2. Simplified dictionary get: Changed from general_settings.get("role_permissions", []) to general_settings.get("role_permissions"), removing the default empty list parameter

Performance Impact:
The line profiler shows the cast operation consumed 12.5% of total runtime (306,989ns out of 2,453,950ns total). By eliminating this unnecessary type conversion, the function achieves a 29% speedup (263μs → 202μs).

Why This Optimization Works:

• cast() in Python is not a no-op - it has measurable overhead even though it doesn't change the runtime value
• The simplified .get() call is more efficient since we immediately check for None anyway, making the default empty list redundant
• These micro-optimizations compound when the function is called frequently in authentication flows

Real-World Benefit:
Based on the function reference, this optimization directly benefits JWT authentication flows where get_role_based_models() is called via can_rbac_role_call_model() for every model access check. The test results show consistent 2-3x speedups across all scenarios, with particularly strong gains (100-170% faster) for common cases like missing roles or empty permissions, making authentication checks more responsive.

Test Case Performance:
The optimization shows excellent results across all test scenarios, with the most significant improvements in edge cases (165-171% faster for missing/None role permissions) and solid gains even in large-scale scenarios (4-116% faster with 1000 roles).

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from dataclasses import dataclass
from typing import Any, List, Literal, Optional

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

# --- Function and supporting types to test ---


# Simulate RBAC_ROLES and RoleBasedPermissions as per the original code
RBAC_ROLES = Literal["admin", "user", "viewer", "guest", "custom"]

@dataclass
class RoleBasedPermissions:
    role: RBAC_ROLES
    models: Optional[List[str]] = None
    routes: Optional[List[str]] = None
from litellm.proxy.auth.auth_checks import get_role_based_models

# --- Unit tests ---

# 1. Basic Test Cases

def test_basic_single_role_models():
    # Basic: Only one role present, should return its models
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4", "gpt-3.5"])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.88μs -> 1.47μs (95.6% faster)

def test_basic_multiple_roles():
    # Basic: Multiple roles, should return correct models for each
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4", "gpt-3.5"]),
            RoleBasedPermissions(role="user", models=["gpt-3.5"]),
            RoleBasedPermissions(role="viewer", models=["gpt-2"])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 2.48μs -> 1.28μs (93.8% faster)
    codeflash_output = get_role_based_models("user", general_settings) # 1.21μs -> 597ns (103% faster)
    codeflash_output = get_role_based_models("viewer", general_settings) # 1.07μs -> 493ns (117% faster)

def test_basic_role_not_found():
    # Basic: Role not present, should return None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4"])
        ]
    }
    codeflash_output = get_role_based_models("user", general_settings) # 2.20μs -> 1.05μs (110% faster)

def test_basic_models_none():
    # Basic: Role present, but models is None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=None)
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 2.33μs -> 1.17μs (98.3% faster)

def test_basic_empty_models_list():
    # Basic: Role present, models is empty list
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=[])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 2.44μs -> 1.13μs (116% faster)

def test_basic_empty_role_permissions():
    # Basic: role_permissions is empty list
    general_settings = {
        "role_permissions": []
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 2.13μs -> 911ns (133% faster)

def test_basic_missing_role_permissions_key():
    # Basic: general_settings has no 'role_permissions' key
    general_settings = {}
    codeflash_output = get_role_based_models("admin", general_settings) # 2.15μs -> 791ns (171% faster)

# 2. Edge Test Cases

def test_edge_role_permissions_none():
    # Edge: role_permissions is None
    general_settings = {
        "role_permissions": None
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 1.99μs -> 750ns (165% faster)



def test_edge_role_with_empty_string():
    # Edge: role is empty string
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="", models=["gpt-4"])
        ]
    }
    codeflash_output = get_role_based_models("", general_settings) # 3.15μs -> 1.74μs (81.2% faster)
    codeflash_output = get_role_based_models("admin", general_settings) # 1.23μs -> 590ns (109% faster)

def test_edge_models_with_none_and_empty_string():
    # Edge: models contains None or empty string
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=[None, ""])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 2.43μs -> 1.25μs (95.4% faster)

def test_edge_role_permissions_contains_duplicates():
    # Edge: role_permissions contains duplicate roles
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4"]),
            RoleBasedPermissions(role="admin", models=["gpt-3.5"])
        ]
    }
    # Should return the first match
    codeflash_output = get_role_based_models("admin", general_settings) # 2.45μs -> 1.23μs (98.9% faster)

def test_edge_role_case_sensitivity():
    # Edge: role case sensitivity
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="Admin", models=["gpt-4"]),
            RoleBasedPermissions(role="admin", models=["gpt-3.5"])
        ]
    }
    # Should match exact string
    codeflash_output = get_role_based_models("admin", general_settings) # 2.60μs -> 1.31μs (98.6% faster)
    codeflash_output = get_role_based_models("Admin", general_settings) # 1.07μs -> 474ns (126% faster)

def test_edge_models_is_not_list():
    # Edge: models is not a list
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models="gpt-4")
        ]
    }
    # Should return the string as-is
    codeflash_output = get_role_based_models("admin", general_settings) # 2.31μs -> 1.17μs (97.2% faster)

def test_edge_roles_with_special_characters():
    # Edge: role names with special characters
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin$", models=["gpt-4"]),
            RoleBasedPermissions(role="user@", models=["gpt-3.5"])
        ]
    }
    codeflash_output = get_role_based_models("admin$", general_settings) # 2.41μs -> 1.10μs (119% faster)
    codeflash_output = get_role_based_models("user@", general_settings) # 1.35μs -> 633ns (114% faster)
    codeflash_output = get_role_based_models("admin", general_settings) # 1.08μs -> 502ns (115% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_roles():
    # Large Scale: 1000 roles, each with one model
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=f"role{i}", models=[f"model{i}"])
            for i in range(1000)
        ]
    }
    # Test a few random roles
    codeflash_output = get_role_based_models("role0", general_settings) # 3.34μs -> 1.55μs (116% faster)
    codeflash_output = get_role_based_models("role999", general_settings) # 35.4μs -> 33.9μs (4.20% faster)
    codeflash_output = get_role_based_models("role500", general_settings) # 17.9μs -> 17.1μs (4.55% faster)
    # Test a role not present
    codeflash_output = get_role_based_models("role1001", general_settings) # 31.3μs -> 30.4μs (2.91% faster)

def test_large_scale_many_models_per_role():
    # Large Scale: One role, 1000 models
    models = [f"model{i}" for i in range(1000)]
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=models)
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings) # 2.29μs -> 1.12μs (104% faster)
    # Test a role not present
    codeflash_output = get_role_based_models("user", general_settings) # 1.15μs -> 482ns (139% faster)

def test_large_scale_roles_with_mixed_models_none_and_empty():
    # Large Scale: 500 roles, some with None, some with empty, some with lists
    perms = []
    for i in range(500):
        if i % 3 == 0:
            perms.append(RoleBasedPermissions(role=f"role{i}", models=None))
        elif i % 3 == 1:
            perms.append(RoleBasedPermissions(role=f"role{i}", models=[]))
        else:
            perms.append(RoleBasedPermissions(role=f"role{i}", models=[f"model{i}"]))
    general_settings = {"role_permissions": perms}
    # Check a few cases
    codeflash_output = get_role_based_models("role0", general_settings) # 2.72μs -> 1.25μs (117% faster)
    codeflash_output = get_role_based_models("role1", general_settings) # 1.40μs -> 603ns (132% faster)
    codeflash_output = get_role_based_models("role2", general_settings) # 1.13μs -> 512ns (121% faster)
    codeflash_output = get_role_based_models("role499", general_settings) # 17.7μs -> 17.0μs (4.41% faster)

def test_large_scale_roles_with_duplicate_names():
    # Large Scale: Many roles with duplicate names, should get first match
    perms = []
    for i in range(500):
        perms.append(RoleBasedPermissions(role="dup", models=[f"model{i}"]))
    general_settings = {"role_permissions": perms}
    # Should return the first model
    codeflash_output = get_role_based_models("dup", general_settings) # 2.47μs -> 1.15μs (115% faster)

def test_large_scale_roles_with_non_list_models():
    # Large Scale: 100 roles, models is a string
    perms = []
    for i in range(100):
        perms.append(RoleBasedPermissions(role=f"role{i}", models=f"model{i}"))
    general_settings = {"role_permissions": perms}
    # Should return the string as-is
    codeflash_output = get_role_based_models("role50", general_settings) # 4.39μs -> 3.13μs (40.2% faster)
    codeflash_output = get_role_based_models("role99", general_settings) # 4.82μs -> 4.04μs (19.3% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from dataclasses import dataclass
from typing import List, Literal, Optional

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

# --- Function and dependencies to test ---

# Simulate RBAC_ROLES and RoleBasedPermissions as per the original code
RBAC_ROLES = Literal["admin", "user", "viewer", "guest", "custom"]

@dataclass
class RoleBasedPermissions:
    role: RBAC_ROLES
    models: Optional[List[str]] = None
    routes: Optional[List[str]] = None
from litellm.proxy.auth.auth_checks import get_role_based_models

# --- Unit tests ---

# 1. Basic Test Cases

def test_basic_admin_role_returns_models():
    # Admin role with models
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4", "gpt-3.5"])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.34μs -> 1.05μs (123% faster)

def test_basic_user_role_returns_models():
    # User role with models
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="user", models=["gpt-3.5"])
        ]
    }
    codeflash_output = get_role_based_models("user", general_settings); result = codeflash_output # 2.33μs -> 1.13μs (106% faster)

def test_basic_role_not_present_returns_none():
    # Role not present in permissions
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4"])
        ]
    }
    codeflash_output = get_role_based_models("user", general_settings); result = codeflash_output # 2.23μs -> 843ns (164% faster)

def test_basic_multiple_roles_returns_correct_models():
    # Multiple roles, select correct one
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4"]),
            RoleBasedPermissions(role="user", models=["gpt-3.5"]),
            RoleBasedPermissions(role="viewer", models=["gpt-3"])
        ]
    }
    codeflash_output = get_role_based_models("viewer", general_settings); result = codeflash_output # 2.45μs -> 1.26μs (94.2% faster)

def test_basic_models_none_returns_none():
    # Role present, but models is None
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=None)
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.39μs -> 1.26μs (90.0% faster)

def test_basic_empty_models_returns_empty_list():
    # Role present, models is empty list
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=[])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.43μs -> 1.14μs (114% faster)

# 2. Edge Test Cases

def test_edge_no_role_permissions_key_returns_none():
    # general_settings missing "role_permissions" key
    general_settings = {}
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.07μs -> 786ns (163% faster)

def test_edge_role_permissions_is_none_returns_none():
    # "role_permissions" explicitly set to None
    general_settings = {
        "role_permissions": None
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 1.93μs -> 727ns (165% faster)

def test_edge_role_permissions_is_empty_list_returns_none():
    # "role_permissions" is empty list
    general_settings = {
        "role_permissions": []
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.05μs -> 858ns (139% faster)


def test_edge_role_case_sensitivity():
    # Role names are case sensitive
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="Admin", models=["gpt-4"]),
            RoleBasedPermissions(role="admin", models=["gpt-3.5"])
        ]
    }
    codeflash_output = get_role_based_models("Admin", general_settings); result = codeflash_output # 3.06μs -> 1.65μs (86.1% faster)
    codeflash_output = get_role_based_models("admin", general_settings); result2 = codeflash_output # 1.45μs -> 668ns (117% faster)

def test_edge_role_with_extra_fields_ignored():
    # RoleBasedPermissions with extra fields should not affect result
    class RoleBasedPermissionsExtra(RoleBasedPermissions):
        def __init__(self, role, models, routes, extra):
            super().__init__(role, models, routes)
            self.extra = extra
    general_settings = {
        "role_permissions": [
            RoleBasedPermissionsExtra(role="user", models=["gpt-3.5"], routes=None, extra="foo")
        ]
    }
    codeflash_output = get_role_based_models("user", general_settings); result = codeflash_output # 2.60μs -> 1.26μs (106% faster)

def test_edge_models_is_not_list_returns_value():
    # models is not a list (e.g., string), should return as is
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models="gpt-4")
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.45μs -> 1.26μs (94.2% faster)



def test_large_many_roles_and_models():
    # Test with 1000 roles and models
    n = 1000
    roles = [f"custom{i}" for i in range(n)]
    permissions = [
        RoleBasedPermissions(role=role, models=[f"model_{i}"]) for i, role in enumerate(roles)
    ]
    general_settings = {
        "role_permissions": permissions
    }
    # Pick a random role to test
    idx = 567
    codeflash_output = get_role_based_models(f"custom{idx}", general_settings); result = codeflash_output # 23.4μs -> 21.2μs (10.2% faster)

def test_large_many_roles_some_missing_models():
    # 1000 roles, some with missing models
    n = 1000
    permissions = []
    for i in range(n):
        if i % 10 == 0:
            permissions.append(RoleBasedPermissions(role=f"custom{i}", models=None))
        else:
            permissions.append(RoleBasedPermissions(role=f"custom{i}", models=[f"model_{i}"]))
    general_settings = {
        "role_permissions": permissions
    }
    # Test a role with missing models
    codeflash_output = get_role_based_models("custom20", general_settings); result = codeflash_output # 4.25μs -> 2.47μs (71.9% faster)
    # Test a role with present models
    codeflash_output = get_role_based_models("custom21", general_settings); result2 = codeflash_output # 2.23μs -> 1.32μs (69.6% faster)

def test_large_role_permissions_is_large_empty_list():
    # role_permissions is a large empty list
    general_settings = {
        "role_permissions": []
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.13μs -> 947ns (125% faster)

def test_large_role_permissions_with_duplicate_roles():
    # Duplicate roles, should return the first match
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role="admin", models=["gpt-4"]),
            RoleBasedPermissions(role="admin", models=["gpt-3.5"]),
            RoleBasedPermissions(role="admin", models=["gpt-2"])
        ]
    }
    codeflash_output = get_role_based_models("admin", general_settings); result = codeflash_output # 2.52μs -> 1.32μs (90.9% faster)

def test_large_role_permissions_with_all_roles_missing():
    # None of the roles match
    general_settings = {
        "role_permissions": [
            RoleBasedPermissions(role=f"custom{i}", models=[f"model_{i}"]) for i in range(1000)
        ]
    }
    codeflash_output = get_role_based_models("guest", general_settings); result = codeflash_output # 31.5μs -> 29.4μs (7.42% 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_models-mhwwix6c and push.