Skip to content

Conversation

@codeflash-ai
Copy link

@codeflash-ai codeflash-ai bot commented Oct 30, 2025

📄 156% (1.56x) speedup for CustomGuardrail._guardrail_is_in_requested_guardrails in litellm/integrations/custom_guardrail.py

⏱️ Runtime : 821 microseconds 320 microseconds (best of 7 runs)

📝 Explanation and details

The optimization introduces two key performance improvements to the _guardrail_is_in_requested_guardrails method:

1. Fast-path optimization for string lists: The optimized code detects when requested_guardrails is a list of strings and uses Python's built-in in operator instead of iterating through each element. This leverages Python's highly optimized string membership testing, which is significantly faster than manual iteration with isinstance() checks.

2. Early None check: The code now checks if guardrail_name is None upfront and immediately returns False, avoiding any unnecessary iteration when the guardrail name is not set.

Why this leads to speedup:

  • The in operator for lists uses optimized C code that's much faster than Python loops with type checks
  • Early termination when guardrail_name is None eliminates wasted cycles
  • The optimization preserves the original fallback logic for mixed lists (strings + dicts), so correctness is maintained

Test case performance patterns:

  • Massive gains for string-only lists: Large-scale string tests show 1200-1500% speedups due to the fast-path optimization
  • Excellent performance when guardrail_name is None: 128-176% speedup from early termination
  • Slight slowdown for dict-heavy workloads: 20-35% slower for pure dictionary lists since they still require the original iteration logic, but mixed lists still benefit when strings are encountered first
  • Best suited for: Applications with predominantly string-based guardrail configurations and scenarios where guardrail names might be unset

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 10 Passed
🌀 Generated Regression Tests 126 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 6 Passed
📊 Tests Coverage 100.0%
⚙️ Existing Unit Tests and Runtime
Test File::Test Function Original ⏱️ Optimized ⏱️ Speedup
guardrails_tests/test_custom_guardrail.py::test_guardrail_is_in_requested_guardrails 3.20μs 3.47μs -7.77%⚠️
🌀 Generated Regression Tests and Runtime
from typing import Dict, List, Optional, Union

# imports
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail


class GuardrailEventHooks(str):
    pass

class DynamicGuardrailParams:
    pass

class Mode:
    def __init__(self, tags: Dict[str, GuardrailEventHooks], default: Optional[GuardrailEventHooks] = None):
        self.tags = tags
        self.default = default

class CustomLogger:
    def __init__(self, turn_off_message_logging: bool = False, message_logging: bool = True, **kwargs):
        self.message_logging = message_logging
        self.turn_off_message_logging = turn_off_message_logging
from litellm.integrations.custom_guardrail import CustomGuardrail

# ----------------------------
# Unit Tests for _guardrail_is_in_requested_guardrails
# ----------------------------

# --------- BASIC TEST CASES ---------
def test_basic_guardrail_name_in_str_list():
    # Test: guardrail_name matches one of the strings in the list
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = ["gr1", "gr2", "gr3"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.01μs -> 819ns (23.4% faster)

def test_basic_guardrail_name_not_in_str_list():
    # Test: guardrail_name does not match any string in the list
    cg = CustomGuardrail(guardrail_name="gr4")
    requested_guardrails = ["gr1", "gr2", "gr3"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.27μs -> 920ns (37.8% faster)

def test_basic_guardrail_name_in_dict_list():
    # Test: guardrail_name matches one of the keys in dicts in the list
    cg = CustomGuardrail(guardrail_name="gr2")
    requested_guardrails = [
        {"gr1": DynamicGuardrailParams()},
        {"gr2": DynamicGuardrailParams()},
        {"gr3": DynamicGuardrailParams()},
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.05μs -> 1.47μs (28.8% slower)

def test_basic_guardrail_name_not_in_dict_list():
    # Test: guardrail_name does not match any key in dicts in the list
    cg = CustomGuardrail(guardrail_name="gr4")
    requested_guardrails = [
        {"gr1": DynamicGuardrailParams()},
        {"gr2": DynamicGuardrailParams()},
        {"gr3": DynamicGuardrailParams()},
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.03μs -> 1.43μs (28.1% slower)

def test_basic_guardrail_name_in_mixed_list():
    # Test: guardrail_name matches in a mixed list of strings and dicts
    cg = CustomGuardrail(guardrail_name="gr2")
    requested_guardrails = [
        "gr1",
        {"gr2": DynamicGuardrailParams()},
        "gr3",
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.28μs -> 1.09μs (17.7% faster)

def test_basic_guardrail_name_not_in_mixed_list():
    # Test: guardrail_name does not match in a mixed list of strings and dicts
    cg = CustomGuardrail(guardrail_name="gr4")
    requested_guardrails = [
        "gr1",
        {"gr2": DynamicGuardrailParams()},
        "gr3",
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.39μs -> 1.11μs (25.7% faster)

# --------- EDGE TEST CASES ---------
def test_edge_empty_requested_guardrails():
    # Test: requested_guardrails is empty
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = []
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 467ns -> 660ns (29.2% slower)

def test_edge_guardrail_name_is_none():
    # Test: guardrail_name is None
    cg = CustomGuardrail(guardrail_name=None)
    requested_guardrails = ["gr1", "gr2"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.18μs -> 427ns (176% faster)

def test_edge_requested_guardrails_contains_none():
    # Test: requested_guardrails contains None values
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = [None, "gr1", None]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.12μs -> 1.44μs (22.2% slower)

def test_edge_requested_guardrails_contains_non_str_non_dict():
    # Test: requested_guardrails contains non-str, non-dict types (should ignore)
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = [123, 45.6, ["gr1"], {"gr1": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.52μs -> 1.79μs (15.4% slower)

def test_edge_requested_guardrails_dict_with_non_str_keys():
    # Test: dicts with non-str keys (should not match)
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = [{1: DynamicGuardrailParams()}, {2: DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.05μs -> 1.33μs (20.8% slower)

def test_edge_requested_guardrails_dict_with_none_key():
    # Test: dicts with None key
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = [{None: DynamicGuardrailParams()}, {"gr1": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 950ns -> 1.41μs (32.8% slower)

def test_edge_requested_guardrails_dict_with_empty_dict():
    # Test: dicts that are empty
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = [{}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 873ns -> 1.12μs (21.9% slower)

def test_edge_requested_guardrails_str_is_empty():
    # Test: empty string in requested_guardrails
    cg = CustomGuardrail(guardrail_name="")
    requested_guardrails = ["", "gr1"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 933ns -> 800ns (16.6% faster)

def test_edge_requested_guardrails_guardrail_name_empty_string():
    # Test: guardrail_name is empty string, requested_guardrails does not contain empty string
    cg = CustomGuardrail(guardrail_name="")
    requested_guardrails = ["gr1", "gr2"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.21μs -> 868ns (39.3% faster)

def test_edge_requested_guardrails_dict_with_multiple_keys():
    # Test: dicts with multiple keys, only one matches
    cg = CustomGuardrail(guardrail_name="gr2")
    requested_guardrails = [{"gr1": DynamicGuardrailParams(), "gr2": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 817ns -> 1.23μs (33.4% slower)

def test_edge_requested_guardrails_dict_with_multiple_keys_none_match():
    # Test: dicts with multiple keys, none match
    cg = CustomGuardrail(guardrail_name="gr3")
    requested_guardrails = [{"gr1": DynamicGuardrailParams(), "gr2": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 770ns -> 1.15μs (33.2% slower)

def test_edge_requested_guardrails_str_and_dict_with_same_name():
    # Test: both string and dict with same guardrail_name
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = ["gr1", {"gr1": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 980ns -> 868ns (12.9% faster)

def test_edge_requested_guardrails_str_and_dict_with_different_names():
    # Test: both string and dict, but only dict matches
    cg = CustomGuardrail(guardrail_name="gr2")
    requested_guardrails = ["gr1", {"gr2": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.15μs -> 1.06μs (8.66% faster)

def test_edge_requested_guardrails_str_and_dict_with_no_match():
    # Test: both string and dict, but neither matches
    cg = CustomGuardrail(guardrail_name="gr3")
    requested_guardrails = ["gr1", {"gr2": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.20μs -> 999ns (20.1% faster)

# --------- LARGE SCALE TEST CASES ---------
def test_large_scale_many_strings_match():
    # Test: large list of strings, guardrail_name is present
    cg = CustomGuardrail(guardrail_name="target_guardrail")
    requested_guardrails = ["gr" + str(i) for i in range(500)] + ["target_guardrail"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 42.5μs -> 3.18μs (1237% faster)

def test_large_scale_many_strings_no_match():
    # Test: large list of strings, guardrail_name is not present
    cg = CustomGuardrail(guardrail_name="not_present_guardrail")
    requested_guardrails = ["gr" + str(i) for i in range(1000)]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 83.3μs -> 5.15μs (1519% faster)

def test_large_scale_many_dicts_match():
    # Test: large list of dicts, guardrail_name is present as a key
    cg = CustomGuardrail(guardrail_name="target_guardrail")
    requested_guardrails = [{"gr" + str(i): DynamicGuardrailParams()} for i in range(500)]
    requested_guardrails.append({"target_guardrail": DynamicGuardrailParams()})
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 26.2μs -> 27.4μs (4.39% slower)

def test_large_scale_many_dicts_no_match():
    # Test: large list of dicts, guardrail_name is not present
    cg = CustomGuardrail(guardrail_name="not_present_guardrail")
    requested_guardrails = [{"gr" + str(i): DynamicGuardrailParams()} for i in range(1000)]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 52.6μs -> 50.4μs (4.26% faster)

def test_large_scale_mixed_list_match():
    # Test: large mixed list, guardrail_name matches in both string and dict
    cg = CustomGuardrail(guardrail_name="target_guardrail")
    requested_guardrails = (
        ["gr" + str(i) for i in range(500)] +
        [{"gr" + str(i): DynamicGuardrailParams()} for i in range(400)] +
        ["target_guardrail"]
    )
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 62.4μs -> 6.36μs (881% faster)

def test_large_scale_mixed_list_no_match():
    # Test: large mixed list, guardrail_name does not match
    cg = CustomGuardrail(guardrail_name="not_present_guardrail")
    requested_guardrails = (
        ["gr" + str(i) for i in range(500)] +
        [{"gr" + str(i): DynamicGuardrailParams()} for i in range(400)]
    )
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 65.8μs -> 6.30μs (943% faster)

def test_large_scale_dicts_with_multiple_keys_match():
    # Test: large list of dicts with multiple keys, guardrail_name matches one key
    cg = CustomGuardrail(guardrail_name="target_guardrail")
    requested_guardrails = [
        {"gr" + str(i): DynamicGuardrailParams(), "target_guardrail": DynamicGuardrailParams()}
        for i in range(500)
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 958ns -> 1.42μs (32.5% slower)

def test_large_scale_dicts_with_multiple_keys_no_match():
    # Test: large list of dicts with multiple keys, guardrail_name does not match
    cg = CustomGuardrail(guardrail_name="not_present_guardrail")
    requested_guardrails = [
        {"gr" + str(i): DynamicGuardrailParams(), "gr" + str(i+1): DynamicGuardrailParams()}
        for i in range(500)
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 26.5μs -> 27.1μs (2.11% slower)

# --------- MUTATION TESTING CASES ---------
def test_mutation_guardrail_name_partial_match():
    # Test: guardrail_name is a substring of a string in requested_guardrails, should not match
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = ["gr12", "gr13", "gr14"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.14μs -> 947ns (20.4% faster)

def test_mutation_guardrail_name_case_sensitive():
    # Test: guardrail_name matches case-insensitively, should not match (function is case-sensitive)
    cg = CustomGuardrail(guardrail_name="GR1")
    requested_guardrails = ["gr1", "gr2", "gr3"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.25μs -> 1.00μs (24.9% faster)

def test_mutation_guardrail_name_whitespace():
    # Test: guardrail_name matches with extra whitespace, should not match
    cg = CustomGuardrail(guardrail_name="gr1 ")
    requested_guardrails = ["gr1", "gr2"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 1.03μs -> 912ns (12.9% faster)

def test_mutation_requested_guardrails_dict_key_partial_match():
    # Test: guardrail_name is a substring of a dict key, should not match
    cg = CustomGuardrail(guardrail_name="gr1")
    requested_guardrails = [{"gr12": DynamicGuardrailParams()}, {"gr13": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested_guardrails) # 984ns -> 1.37μs (28.2% slower)
# 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 Dict, List, Optional, Union

# imports
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail


# Minimal stubs for required types/enums/classes
class GuardrailEventHooks(str):
    pass

class Mode:
    def __init__(self, tags: Dict[str, GuardrailEventHooks], default: Optional[GuardrailEventHooks] = None):
        self.tags = tags
        self.default = default

class DynamicGuardrailParams:
    pass

# CustomLogger stub from litellm/integrations/custom_logger.py
class CustomLogger:
    def __init__(
        self, 
        turn_off_message_logging: bool = False,
        message_logging: bool = True,
        **kwargs
    ) -> None:
        self.message_logging = message_logging
        self.turn_off_message_logging = turn_off_message_logging
from litellm.integrations.custom_guardrail import CustomGuardrail

# ------------------------
# Unit Tests
# ------------------------

# 1. Basic Test Cases

def test_single_string_match():
    # Guardrail name is present as a string in the list
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = ["guardrailA"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 882ns -> 775ns (13.8% faster)

def test_single_string_no_match():
    # Guardrail name is not present as a string in the list
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = ["guardrailB"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 985ns -> 884ns (11.4% faster)

def test_multiple_strings_with_match():
    # Guardrail name is present among several strings
    cg = CustomGuardrail(guardrail_name="guardrailC")
    requested = ["guardrailA", "guardrailB", "guardrailC"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.35μs -> 941ns (43.3% faster)

def test_multiple_strings_no_match():
    # Guardrail name is not present among several strings
    cg = CustomGuardrail(guardrail_name="guardrailD")
    requested = ["guardrailA", "guardrailB", "guardrailC"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.32μs -> 963ns (37.4% faster)

def test_single_dict_match():
    # Guardrail name is present as a key in a dict
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [{"guardrailA": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 797ns -> 1.18μs (32.2% slower)

def test_single_dict_no_match():
    # Guardrail name is not present as a key in a dict
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [{"guardrailB": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 785ns -> 1.18μs (33.7% slower)

def test_multiple_dicts_with_match():
    # Guardrail name is present among several dicts
    cg = CustomGuardrail(guardrail_name="guardrailC")
    requested = [
        {"guardrailA": DynamicGuardrailParams()},
        {"guardrailB": DynamicGuardrailParams()},
        {"guardrailC": DynamicGuardrailParams()},
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.15μs -> 1.47μs (21.8% slower)

def test_multiple_dicts_no_match():
    # Guardrail name is not present among several dicts
    cg = CustomGuardrail(guardrail_name="guardrailD")
    requested = [
        {"guardrailA": DynamicGuardrailParams()},
        {"guardrailB": DynamicGuardrailParams()},
        {"guardrailC": DynamicGuardrailParams()},
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.01μs -> 1.49μs (32.1% slower)

def test_mixed_strings_and_dicts_with_match():
    # Guardrail name is present as both string and dict key
    cg = CustomGuardrail(guardrail_name="guardrailX")
    requested = [
        "guardrailA",
        {"guardrailB": DynamicGuardrailParams()},
        "guardrailX",
        {"guardrailZ": DynamicGuardrailParams()},
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.28μs -> 1.00μs (27.8% faster)

def test_mixed_strings_and_dicts_no_match():
    # Guardrail name is not present in mixed list
    cg = CustomGuardrail(guardrail_name="guardrailY")
    requested = [
        "guardrailA",
        {"guardrailB": DynamicGuardrailParams()},
        "guardrailX",
        {"guardrailZ": DynamicGuardrailParams()},
    ]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.38μs -> 1.22μs (13.0% faster)

# 2. Edge Test Cases

def test_empty_requested_guardrails_list():
    # Empty list should always return False
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = []
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 471ns -> 702ns (32.9% slower)

def test_guardrail_name_is_none():
    # Guardrail name is None, should never match
    cg = CustomGuardrail(guardrail_name=None)
    requested = ["guardrailA", {"guardrailA": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.29μs -> 565ns (128% faster)

def test_requested_guardrails_contains_none():
    # None in the requested_guardrails list should be ignored
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [None, "guardrailA", None]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.19μs -> 1.53μs (22.4% slower)

def test_requested_guardrails_contains_empty_dict():
    # Empty dict in the requested_guardrails list should be ignored
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [{}, "guardrailA"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.17μs -> 1.60μs (26.5% slower)

def test_requested_guardrails_contains_empty_string():
    # Empty string in the requested_guardrails list should be ignored
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = ["", "guardrailA"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.23μs -> 932ns (31.7% faster)

def test_requested_guardrails_with_non_string_dict_key():
    # Dict with non-string key should not match
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [{123: DynamicGuardrailParams()}, "guardrailB"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.20μs -> 1.67μs (27.9% slower)

def test_requested_guardrails_with_duplicate_entries():
    # Multiple matching entries (should still return True)
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = ["guardrailA", "guardrailA", {"guardrailA": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 967ns -> 833ns (16.1% faster)

def test_requested_guardrails_with_case_sensitivity():
    # Should be case-sensitive
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = ["GUARDRAILA", {"GUARDRAILA": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.12μs -> 1.15μs (2.01% slower)

def test_requested_guardrails_with_whitespace():
    # Whitespace differences should not match
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [" guardrailA", "guardrailA ", {" guardrailA": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.20μs -> 1.03μs (16.9% faster)

def test_requested_guardrails_with_substring():
    # Substring should not match
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = ["guardrail", "A", {"guard": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.29μs -> 1.10μs (17.0% faster)

# 3. Large Scale Test Cases

def test_large_list_with_match_at_start():
    # Large list, match at the start
    cg = CustomGuardrail(guardrail_name="target")
    requested = ["target"] + [f"guardrail{i}" for i in range(999)]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.07μs -> 896ns (19.0% faster)

def test_large_list_with_match_at_end():
    # Large list, match at the end
    cg = CustomGuardrail(guardrail_name="target")
    requested = [f"guardrail{i}" for i in range(999)] + ["target"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 83.5μs -> 5.02μs (1564% faster)

def test_large_list_with_no_match():
    # Large list, no match
    cg = CustomGuardrail(guardrail_name="target")
    requested = [f"guardrail{i}" for i in range(1000)]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 83.3μs -> 5.35μs (1457% faster)

def test_large_list_of_dicts_with_match():
    # Large list of dicts, match present
    cg = CustomGuardrail(guardrail_name="target")
    requested = [{f"guardrail{i}": DynamicGuardrailParams()} for i in range(999)]
    requested.append({"target": DynamicGuardrailParams()})
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 50.0μs -> 53.7μs (6.96% slower)

def test_large_list_of_dicts_with_no_match():
    # Large list of dicts, no match
    cg = CustomGuardrail(guardrail_name="target")
    requested = [{f"guardrail{i}": DynamicGuardrailParams()} for i in range(1000)]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 49.6μs -> 53.2μs (6.67% slower)

def test_large_mixed_list_with_match():
    # Large mixed list, match present as string and dict
    cg = CustomGuardrail(guardrail_name="target")
    requested = (
        [f"guardrail{i}" for i in range(500)] +
        [{"guardrail"+str(i): DynamicGuardrailParams()} for i in range(499)] +
        ["target"] +
        [{"target": DynamicGuardrailParams()}]
    )
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 67.2μs -> 7.00μs (860% faster)

def test_large_mixed_list_with_no_match():
    # Large mixed list, no match
    cg = CustomGuardrail(guardrail_name="target")
    requested = (
        [f"guardrail{i}" for i in range(500)] +
        [{"guardrail"+str(i): DynamicGuardrailParams()} for i in range(500)]
    )
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 66.9μs -> 6.77μs (888% faster)

# 4. Additional Edge/Negative Cases

def test_requested_guardrails_with_non_str_non_dict():
    # List contains an int, should be ignored
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [1, 2.5, "guardrailA"]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.48μs -> 1.75μs (15.3% slower)

def test_requested_guardrails_with_all_non_str_non_dict():
    # List contains only non-str, non-dict types
    cg = CustomGuardrail(guardrail_name="guardrailA")
    requested = [1, 2.5, None, [], ()]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 1.79μs -> 2.01μs (11.0% slower)

def test_requested_guardrails_with_dict_with_multiple_keys():
    # Dict with multiple keys, one matches
    cg = CustomGuardrail(guardrail_name="guardrailB")
    requested = [{"guardrailA": DynamicGuardrailParams(), "guardrailB": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 858ns -> 1.22μs (29.7% slower)

def test_requested_guardrails_with_dict_with_multiple_keys_none_match():
    # Dict with multiple keys, none match
    cg = CustomGuardrail(guardrail_name="guardrailC")
    requested = [{"guardrailA": DynamicGuardrailParams(), "guardrailB": DynamicGuardrailParams()}]
    codeflash_output = cg._guardrail_is_in_requested_guardrails(requested) # 801ns -> 1.20μs (33.4% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.guardrails import DynamicGuardrailParams

def test_CustomGuardrail__guardrail_is_in_requested_guardrails():
    CustomGuardrail._guardrail_is_in_requested_guardrails(CustomGuardrail(guardrail_name='\x00', supported_event_hooks=None, event_hook=None, default_on=False, mask_request_content=False, mask_response_content=False), ['', '\x00'])

def test_CustomGuardrail__guardrail_is_in_requested_guardrails_2():
    CustomGuardrail._guardrail_is_in_requested_guardrails(CustomGuardrail(guardrail_name='', supported_event_hooks=[], event_hook=[], default_on=True, mask_request_content=True, mask_response_content=True), [{}, {'': DynamicGuardrailParams()}])

def test_CustomGuardrail__guardrail_is_in_requested_guardrails_3():
    CustomGuardrail._guardrail_is_in_requested_guardrails(CustomGuardrail(guardrail_name='', supported_event_hooks=[], event_hook=[], default_on=True, mask_request_content=False, mask_response_content=True), [])
🔎 Concolic Coverage Tests and Runtime
Test File::Test Function Original ⏱️ Optimized ⏱️ Speedup
codeflash_concolic_zbim32de/tmpir8ytiis/test_concolic_coverage.py::test_CustomGuardrail__guardrail_is_in_requested_guardrails 1.32μs 915ns 44.7%✅
codeflash_concolic_zbim32de/tmpir8ytiis/test_concolic_coverage.py::test_CustomGuardrail__guardrail_is_in_requested_guardrails_2 1.11μs 1.45μs -23.5%⚠️
codeflash_concolic_zbim32de/tmpir8ytiis/test_concolic_coverage.py::test_CustomGuardrail__guardrail_is_in_requested_guardrails_3 474ns 619ns -23.4%⚠️

To edit these changes git checkout codeflash/optimize-CustomGuardrail._guardrail_is_in_requested_guardrails-mhdw73k6 and push.

Codeflash Static Badge

The optimization introduces two key performance improvements to the `_guardrail_is_in_requested_guardrails` method:

**1. Fast-path optimization for string lists:** The optimized code detects when `requested_guardrails` is a list of strings and uses Python's built-in `in` operator instead of iterating through each element. This leverages Python's highly optimized string membership testing, which is significantly faster than manual iteration with `isinstance()` checks.

**2. Early None check:** The code now checks if `guardrail_name` is None upfront and immediately returns False, avoiding any unnecessary iteration when the guardrail name is not set.

**Why this leads to speedup:** 
- The `in` operator for lists uses optimized C code that's much faster than Python loops with type checks
- Early termination when `guardrail_name` is None eliminates wasted cycles
- The optimization preserves the original fallback logic for mixed lists (strings + dicts), so correctness is maintained

**Test case performance patterns:**
- **Massive gains for string-only lists:** Large-scale string tests show 1200-1500% speedups due to the fast-path optimization
- **Excellent performance when guardrail_name is None:** 128-176% speedup from early termination
- **Slight slowdown for dict-heavy workloads:** 20-35% slower for pure dictionary lists since they still require the original iteration logic, but mixed lists still benefit when strings are encountered first
- **Best suited for:** Applications with predominantly string-based guardrail configurations and scenarios where guardrail names might be unset
@codeflash-ai codeflash-ai bot requested a review from mashraf-222 October 30, 2025 20:44
@codeflash-ai codeflash-ai bot added ⚡️ codeflash Optimization PR opened by Codeflash AI 🎯 Quality: Medium Optimization Quality according to Codeflash labels Oct 30, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

⚡️ codeflash Optimization PR opened by Codeflash AI 🎯 Quality: Medium Optimization Quality according to Codeflash

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant