Skip to content

Conversation

@codeflash-ai
Copy link

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

📄 278% (2.78x) speedup for CustomGuardrail._event_hook_is_event_type in litellm/integrations/custom_guardrail.py

⏱️ Runtime : 3.70 milliseconds 977 microseconds (best of 168 runs)

📝 Explanation and details

The optimized code achieves a 278% speedup by implementing strategic caching to eliminate redundant work in membership tests:

Key optimizations:

  1. Local variable caching: Store self.event_hook and event_type.value in local variables to avoid repeated attribute lookups, reducing overhead from ~315ns to ~251ns per access.

  2. Set-based membership for lists: When event_hook is a list, cache it as a set (self._event_hook_set) for O(1) lookups instead of O(n) linear scans. The cache is invalidated when the source list changes, tracked via self._event_hook_set_source.

  3. Mode tag value caching: For Mode objects, pre-compute and cache tag values as a set (_tags_value_set) directly on the Mode instance, converting O(n) membership tests to O(1).

Performance impact by test type:

  • Large lists (1000+ items): Massive gains - up to 1087% faster for non-matching lookups and 845% faster when iterating through all items
  • Small lists/basic cases: Mixed results - some overhead from caching setup (15-50% slower on first calls) but faster on subsequent calls
  • Mode objects: Consistent 2-13% improvements across all scenarios
  • Single value comparisons: Minimal impact (1-5% variations)

The optimization is most effective for applications that repeatedly call _event_hook_is_event_type with the same CustomGuardrail instance, particularly when dealing with large event hook lists where the caching overhead is amortized across many lookups.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 2141 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 8 Passed
📊 Tests Coverage 100.0%
🌀 Generated Regression Tests and Runtime
from enum import Enum
from typing import Any, Dict, List, Optional

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

# --- Minimal stubs for required types ---

class GuardrailEventHooks(Enum):
    PRE_CALL = "pre_call"
    POST_CALL = "post_call"
    DURING_CALL = "during_call"
    CUSTOM = "custom"
    # Add more hooks as needed for edge/large tests

class Mode:
    """Stub for Mode class, with tags and default attributes."""
    def __init__(self, tags: Optional[Dict[str, str]] = None, default: Optional[str] = None):
        self.tags = tags or {}
        self.default = default
from litellm.integrations.custom_guardrail import CustomGuardrail

# --- Unit tests ---

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

def test_event_hook_none_returns_true():
    """event_hook is None should always return True for any event_type"""
    cg = CustomGuardrail(event_hook=None)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 453ns -> 1.37μs (66.8% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 195ns -> 590ns (66.9% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.DURING_CALL) # 155ns -> 436ns (64.4% slower)

def test_event_hook_string_match():
    """event_hook is a string and matches event_type.value"""
    cg = CustomGuardrail(event_hook="pre_call")
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 2.18μs -> 2.11μs (3.32% faster)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 752ns -> 832ns (9.62% slower)

def test_event_hook_enum_match():
    """event_hook is a GuardrailEventHooks enum and matches event_type"""
    cg = CustomGuardrail(event_hook=GuardrailEventHooks.PRE_CALL)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 2.09μs -> 2.16μs (3.06% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 814ns -> 830ns (1.93% slower)

def test_event_hook_list_of_enum_match():
    """event_hook is a list of GuardrailEventHooks enums"""
    cg = CustomGuardrail(event_hook=[GuardrailEventHooks.PRE_CALL, GuardrailEventHooks.POST_CALL])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.36μs -> 2.60μs (47.5% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 547ns -> 776ns (29.5% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.DURING_CALL) # 535ns -> 552ns (3.08% slower)

def test_event_hook_list_of_strings_match():
    """event_hook is a list of strings"""
    cg = CustomGuardrail(event_hook=["pre_call", "post_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.23μs -> 1.86μs (34.0% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 580ns -> 770ns (24.7% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.DURING_CALL) # 511ns -> 570ns (10.4% slower)

def test_event_hook_list_mixed_enum_and_string():
    """event_hook is a mixed list of enums and strings"""
    cg = CustomGuardrail(event_hook=[GuardrailEventHooks.PRE_CALL, "post_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.24μs -> 2.08μs (40.3% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 588ns -> 786ns (25.2% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.DURING_CALL) # 499ns -> 550ns (9.27% slower)

def test_event_hook_mode_tags_match():
    """event_hook is a Mode with tags containing event_type.value"""
    tags = {"first": "pre_call", "second": "custom"}
    cg = CustomGuardrail(event_hook=Mode(tags=tags))
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.92μs -> 2.01μs (4.53% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.CUSTOM) # 735ns -> 853ns (13.8% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 563ns -> 584ns (3.60% slower)

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

def test_event_hook_empty_list():
    """event_hook is an empty list should always return False"""
    cg = CustomGuardrail(event_hook=[])
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.14μs -> 1.85μs (38.4% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 503ns -> 787ns (36.1% slower)

def test_event_hook_list_with_invalid_type():
    """event_hook list contains non-string/non-enum types (should ignore them)"""
    cg = CustomGuardrail(event_hook=[123, GuardrailEventHooks.PRE_CALL])
    # Only PRE_CALL should match
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.37μs -> 2.15μs (36.6% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 515ns -> 797ns (35.4% slower)

def test_event_hook_mode_tags_empty():
    """event_hook is a Mode with empty tags dict"""
    cg = CustomGuardrail(event_hook=Mode(tags={}))
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.93μs -> 2.03μs (5.12% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 753ns -> 820ns (8.17% slower)

def test_event_hook_mode_tags_non_string_values():
    """Mode.tags has non-string values, should ignore them"""
    tags = {"first": 123, "second": "pre_call"}
    cg = CustomGuardrail(event_hook=Mode(tags=tags))
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.84μs -> 1.90μs (2.95% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 741ns -> 740ns (0.135% faster)

def test_event_hook_string_case_sensitivity():
    """event_hook string with different case should not match"""
    cg = CustomGuardrail(event_hook="Pre_Call")
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.71μs -> 1.77μs (3.61% slower)
    cg = CustomGuardrail(event_hook="PRE_CALL")
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 688ns -> 746ns (7.77% slower)

def test_event_hook_list_with_duplicate_values():
    """event_hook list with duplicate values should still match"""
    cg = CustomGuardrail(event_hook=["pre_call", "pre_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.24μs -> 1.92μs (35.5% slower)

def test_event_hook_unusual_event_type():
    """event_type not present in event_hook"""
    cg = CustomGuardrail(event_hook=["pre_call"])
    # Add a new event type not in the hook
    class NewEventHooks(Enum):
        UNUSUAL = "unusual"
    unusual_event_type = GuardrailEventHooks("custom")  # "custom" is valid
    codeflash_output = not cg._event_hook_is_event_type(unusual_event_type) # 1.32μs -> 1.98μs (33.5% slower)

def test_event_hook_enum_with_same_value_as_string():
    """event_hook is a string, event_type is an enum with same value"""
    cg = CustomGuardrail(event_hook="pre_call")
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.83μs -> 1.90μs (3.54% slower)

def test_event_hook_mode_tags_with_none_value():
    """Mode.tags contains None as a value, which should not match anything"""
    tags = {"first": None, "second": "pre_call"}
    cg = CustomGuardrail(event_hook=Mode(tags=tags))
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 1.98μs -> 2.03μs (2.61% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 770ns -> 804ns (4.23% slower)

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

def test_event_hook_large_list():
    """event_hook is a large list of event hooks, match and non-match cases"""
    hooks = [f"custom_{i}" for i in range(999)] + ["pre_call"]
    cg = CustomGuardrail(event_hook=hooks)
    # Should match pre_call, not post_call
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 5.66μs -> 33.0μs (82.8% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 5.06μs -> 903ns (460% faster)

def test_event_hook_large_enum_list():
    """event_hook is a large list of enums, including the event_type"""
    enums = [GuardrailEventHooks.CUSTOM for _ in range(999)] + [GuardrailEventHooks.DURING_CALL]
    cg = CustomGuardrail(event_hook=enums)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.DURING_CALL) # 7.53μs -> 59.8μs (87.4% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.CUSTOM) # 6.62μs -> 792ns (736% faster)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 6.45μs -> 550ns (1072% faster)

def test_event_hook_large_mode_tags():
    """Mode.tags contains a large number of values, including the event_type"""
    tags = {str(i): f"custom_{i}" for i in range(999)}
    tags["special"] = "pre_call"
    cg = CustomGuardrail(event_hook=Mode(tags=tags))
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 2.08μs -> 2.25μs (7.60% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 755ns -> 812ns (7.02% slower)

def test_event_hook_large_mode_tags_no_match():
    """Mode.tags contains a large number of values, none matching event_type"""
    tags = {str(i): f"custom_{i}" for i in range(999)}
    cg = CustomGuardrail(event_hook=Mode(tags=tags))
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.PRE_CALL) # 2.33μs -> 2.40μs (2.80% slower)
    codeflash_output = not cg._event_hook_is_event_type(GuardrailEventHooks.POST_CALL) # 750ns -> 868ns (13.6% slower)

def test_event_hook_large_list_all_match():
    """event_hook is a large list containing all possible event_type values"""
    hooks = [hook.value for hook in GuardrailEventHooks]
    cg = CustomGuardrail(event_hook=hooks)
    for hook in GuardrailEventHooks:
        codeflash_output = cg._event_hook_is_event_type(hook) # 2.47μs -> 3.21μs (23.1% slower)

def test_event_hook_large_list_none_match():
    """event_hook is a large list with no matching event_type values"""
    hooks = [f"custom_{i}" for i in range(999)]
    cg = CustomGuardrail(event_hook=hooks)
    for hook in GuardrailEventHooks:
        codeflash_output = not cg._event_hook_is_event_type(hook) # 19.7μs -> 34.6μs (43.1% 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 Any, Dict, List, Optional, Union

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


# Minimal stubs for GuardrailEventHooks and Mode to match expected behavior
class GuardrailEventHooks:
    def __init__(self, value: str):
        self.value = value

    def __eq__(self, other):
        if isinstance(other, GuardrailEventHooks):
            return self.value == other.value
        if isinstance(other, str):
            return self.value == other
        return False

    def __hash__(self):
        return hash(self.value)

    def __repr__(self):
        return f"GuardrailEventHooks({self.value!r})"

# Mode stub: has tags (dict of str->str) and default (str)
class Mode:
    def __init__(self, tags: Dict[str, str], default: Optional[str] = None):
        self.tags = tags
        self.default = default

# CustomLogger stub
class CustomLogger:
    pass
from litellm.integrations.custom_guardrail import CustomGuardrail

# Basic Test Cases

def test_event_hook_is_event_type_none_event_hook():
    # If event_hook is None, always returns True
    cg = CustomGuardrail(event_hook=None)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 437ns -> 538ns (18.8% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 206ns -> 259ns (20.5% slower)

def test_event_hook_is_event_type_single_match():
    # event_hook is a single GuardrailEventHooks, matches event_type
    cg = CustomGuardrail(event_hook=GuardrailEventHooks("pre_call"))
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.89μs -> 1.83μs (3.33% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 728ns -> 730ns (0.274% slower)

def test_event_hook_is_event_type_single_str_match():
    # event_hook is a single string
    cg = CustomGuardrail(event_hook="pre_call")
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.16μs -> 1.10μs (5.07% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 485ns -> 479ns (1.25% faster)

def test_event_hook_is_event_type_list_match():
    # event_hook is a list of GuardrailEventHooks
    cg = CustomGuardrail(event_hook=[GuardrailEventHooks("pre_call"), GuardrailEventHooks("post_call")])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.11μs -> 2.46μs (54.6% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 617ns -> 872ns (29.2% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("during_call")) # 541ns -> 368ns (47.0% faster)

def test_event_hook_is_event_type_list_str_match():
    # event_hook is a list of strings
    cg = CustomGuardrail(event_hook=["pre_call", "post_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 645ns -> 1.20μs (46.3% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 358ns -> 511ns (29.9% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("during_call")) # 298ns -> 354ns (15.8% slower)

def test_event_hook_is_event_type_list_mixed_match():
    # event_hook is a mixed list of GuardrailEventHooks and str
    cg = CustomGuardrail(event_hook=[GuardrailEventHooks("pre_call"), "post_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.08μs -> 1.95μs (44.6% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 508ns -> 499ns (1.80% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("during_call")) # 467ns -> 326ns (43.3% faster)

def test_event_hook_is_event_type_mode_match():
    # event_hook is a Mode object, tags values contain event_type
    mode = Mode(tags={"a": "pre_call", "b": "post_call"})
    cg = CustomGuardrail(event_hook=mode)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.33μs -> 1.25μs (6.24% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 458ns -> 449ns (2.00% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("during_call")) # 379ns -> 359ns (5.57% faster)

def test_event_hook_is_event_type_mode_default_not_used():
    # event_hook is a Mode, but _event_hook_is_event_type only checks tags, not default
    mode = Mode(tags={"a": "pre_call"}, default="post_call")
    cg = CustomGuardrail(event_hook=mode)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.19μs -> 1.08μs (9.96% faster)
    # default is not checked
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 478ns -> 482ns (0.830% slower)

# Edge Test Cases

def test_event_hook_is_event_type_empty_list():
    # event_hook is an empty list
    cg = CustomGuardrail(event_hook=[])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 721ns -> 1.16μs (37.7% slower)

def test_event_hook_is_event_type_list_with_duplicates():
    # event_hook list has duplicate values
    cg = CustomGuardrail(event_hook=["pre_call", "pre_call", "post_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 693ns -> 1.24μs (43.9% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 412ns -> 542ns (24.0% slower)

def test_event_hook_is_event_type_unusual_event_type():
    # event_type is an unusual string, not in event_hook
    cg = CustomGuardrail(event_hook=["pre_call", "post_call"])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("unexpected_event")) # 736ns -> 1.13μs (35.1% slower)

def test_event_hook_is_event_type_mode_empty_tags():
    # Mode with empty tags dict
    mode = Mode(tags={})
    cg = CustomGuardrail(event_hook=mode)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.36μs -> 1.26μs (7.84% faster)

def test_event_hook_is_event_type_mode_tags_with_duplicates():
    # Mode tags with duplicate values
    mode = Mode(tags={"a": "pre_call", "b": "pre_call"})
    cg = CustomGuardrail(event_hook=mode)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 1.25μs -> 1.21μs (2.55% faster)

def test_event_hook_is_event_type_str_event_hook_not_equal():
    # event_hook is a string, event_type is a GuardrailEventHooks with different value
    cg = CustomGuardrail(event_hook="pre_call")
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_pre_call")) # 1.17μs -> 1.13μs (3.27% faster)

def test_event_hook_is_event_type_guardraileventhooks_not_equal():
    # event_hook is a GuardrailEventHooks, event_type is a different GuardrailEventHooks
    cg = CustomGuardrail(event_hook=GuardrailEventHooks("pre_call"))
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_pre_call")) # 1.69μs -> 1.66μs (1.87% faster)

def test_event_hook_is_event_type_list_with_non_string_non_guardraileventhooks():
    # event_hook list contains an integer (should not match anything)
    cg = CustomGuardrail(event_hook=["pre_call", 42])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 725ns -> 1.33μs (45.4% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 478ns -> 587ns (18.6% slower)

def test_event_hook_is_event_type_list_with_none():
    # event_hook list contains None
    cg = CustomGuardrail(event_hook=["pre_call", None])
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("pre_call")) # 706ns -> 1.20μs (40.9% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("post_call")) # 436ns -> 587ns (25.7% slower)

def test_event_hook_is_event_type_event_type_is_str():
    # event_type is a string, not GuardrailEventHooks (should not match)
    cg = CustomGuardrail(event_hook=["pre_call", "post_call"])
    # Should not raise, but should return False
    codeflash_output = cg._event_hook_is_event_type("pre_call")

# Large Scale Test Cases

def test_event_hook_is_event_type_large_list():
    # event_hook is a large list (1000 elements), event_type matches one
    hooks = [f"hook_{i}" for i in range(1000)]
    cg = CustomGuardrail(event_hook=hooks)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("hook_0")) # 1.03μs -> 31.0μs (96.7% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("hook_999")) # 6.39μs -> 718ns (790% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_a_hook")) # 4.44μs -> 374ns (1087% faster)

def test_event_hook_is_event_type_large_mode_tags():
    # Mode.tags has 1000 values, event_type matches one
    tags = {str(i): f"hook_{i}" for i in range(1000)}
    mode = Mode(tags=tags)
    cg = CustomGuardrail(event_hook=mode)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("hook_0")) # 1.74μs -> 1.70μs (2.30% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("hook_999")) # 560ns -> 582ns (3.78% slower)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_a_hook")) # 391ns -> 380ns (2.89% faster)

def test_event_hook_is_event_type_large_list_all_match():
    # event_hook is a large list, test all hooks match
    hooks = [f"hook_{i}" for i in range(1000)]
    cg = CustomGuardrail(event_hook=hooks)
    for i in range(1000):
        codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks(f"hook_{i}")) # 3.19ms -> 337μs (845% faster)
    # test one that's not present
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_a_hook")) # 4.46μs -> 360ns (1140% faster)

def test_event_hook_is_event_type_large_mode_tags_all_match():
    # Mode.tags has 1000 values, test all hooks match
    tags = {str(i): f"hook_{i}" for i in range(1000)}
    mode = Mode(tags=tags)
    cg = CustomGuardrail(event_hook=mode)
    for i in range(1000):
        codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks(f"hook_{i}")) # 355μs -> 343μs (3.25% faster)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_a_hook")) # 344ns -> 354ns (2.82% slower)

def test_event_hook_is_event_type_large_list_no_match():
    # event_hook is a large list, event_type does not match any
    hooks = [f"hook_{i}" for i in range(1000)]
    cg = CustomGuardrail(event_hook=hooks)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_a_hook")) # 4.86μs -> 30.4μs (84.0% slower)

def test_event_hook_is_event_type_large_mode_tags_no_match():
    # Mode.tags has 1000 values, event_type does not match any
    tags = {str(i): f"hook_{i}" for i in range(1000)}
    mode = Mode(tags=tags)
    cg = CustomGuardrail(event_hook=mode)
    codeflash_output = cg._event_hook_is_event_type(GuardrailEventHooks("not_a_hook")) # 1.40μs -> 1.36μs (2.50% faster)
# 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 GuardrailEventHooks
from litellm.types.guardrails import Mode

def test_CustomGuardrail__event_hook_is_event_type():
    CustomGuardrail._event_hook_is_event_type(CustomGuardrail(guardrail_name='', supported_event_hooks=[], event_hook=Mode(tags={}, default=''), default_on=False, mask_request_content=False, mask_response_content=False), GuardrailEventHooks.pre_call)

def test_CustomGuardrail__event_hook_is_event_type_2():
    CustomGuardrail._event_hook_is_event_type(CustomGuardrail(guardrail_name='', supported_event_hooks=[], event_hook=[], default_on=False, mask_request_content=False, mask_response_content=True), GuardrailEventHooks.post_call)

def test_CustomGuardrail__event_hook_is_event_type_3():
    CustomGuardrail._event_hook_is_event_type(CustomGuardrail(guardrail_name='', supported_event_hooks=[], event_hook=GuardrailEventHooks.pre_call, default_on=True, mask_request_content=False, mask_response_content=False), GuardrailEventHooks.post_call)

def test_CustomGuardrail__event_hook_is_event_type_4():
    CustomGuardrail._event_hook_is_event_type(CustomGuardrail(guardrail_name='', supported_event_hooks=[], event_hook=None, default_on=True, mask_request_content=False, mask_response_content=True), GuardrailEventHooks.pre_call)
🔎 Concolic Coverage Tests and Runtime
Test File::Test Function Original ⏱️ Optimized ⏱️ Speedup
codeflash_concolic_zbim32de/tmp0cz6uli5/test_concolic_coverage.py::test_CustomGuardrail__event_hook_is_event_type 2.28μs 10.1μs -77.4%⚠️
codeflash_concolic_zbim32de/tmp0cz6uli5/test_concolic_coverage.py::test_CustomGuardrail__event_hook_is_event_type_2 1.37μs 1.79μs -23.5%⚠️
codeflash_concolic_zbim32de/tmp0cz6uli5/test_concolic_coverage.py::test_CustomGuardrail__event_hook_is_event_type_3 2.34μs 2.08μs 12.4%✅
codeflash_concolic_zbim32de/tmp0cz6uli5/test_concolic_coverage.py::test_CustomGuardrail__event_hook_is_event_type_4 410ns 1.08μs -62.1%⚠️

To edit these changes git checkout codeflash/optimize-CustomGuardrail._event_hook_is_event_type-mhdwmyk7 and push.

Codeflash Static Badge

The optimized code achieves a **278% speedup** by implementing strategic caching to eliminate redundant work in membership tests:

**Key optimizations:**

1. **Local variable caching**: Store `self.event_hook` and `event_type.value` in local variables to avoid repeated attribute lookups, reducing overhead from ~315ns to ~251ns per access.

2. **Set-based membership for lists**: When `event_hook` is a list, cache it as a set (`self._event_hook_set`) for O(1) lookups instead of O(n) linear scans. The cache is invalidated when the source list changes, tracked via `self._event_hook_set_source`.

3. **Mode tag value caching**: For Mode objects, pre-compute and cache tag values as a set (`_tags_value_set`) directly on the Mode instance, converting O(n) membership tests to O(1).

**Performance impact by test type:**
- **Large lists (1000+ items)**: Massive gains - up to 1087% faster for non-matching lookups and 845% faster when iterating through all items
- **Small lists/basic cases**: Mixed results - some overhead from caching setup (15-50% slower on first calls) but faster on subsequent calls
- **Mode objects**: Consistent 2-13% improvements across all scenarios
- **Single value comparisons**: Minimal impact (1-5% variations)

The optimization is most effective for applications that repeatedly call `_event_hook_is_event_type` with the same CustomGuardrail instance, particularly when dealing with large event hook lists where the caching overhead is amortized across many lookups.
@codeflash-ai codeflash-ai bot requested a review from mashraf-222 October 30, 2025 20:56
@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