⚡️ Speed up function _custom_logger_class_exists_in_success_callbacks by 87%

[codeflash-ai]

📄 87% (0.87x) speedup for _custom_logger_class_exists_in_success_callbacks in litellm/utils.py

⏱️ Runtime : 295 microseconds → 158 microseconds (best of 33 runs)

📝 Explanation and details

The optimization eliminates the expensive list concatenation operation (litellm.success_callback + litellm._async_success_callback) that was creating a temporary combined list on every function call. This concatenation dominated the runtime, consuming 97.9% of execution time according to the line profiler.

Key optimizations:

1. Eliminated list concatenation: The original code created a new temporary list by concatenating two callback lists, which required memory allocation and copying all elements
2. Added early termination: Two separate loops now check each callback list individually, allowing the function to return immediately upon finding the first match
3. Cached type lookup: Added cb_type = type(callback_class) to avoid repeated type() calls during isinstance() checks

Performance benefits:

• 86% speedup overall (295μs → 158μs)
• Early termination is especially effective when matches are found early in the first list (up to 417% faster in some test cases)
• Reduced memory pressure by avoiding temporary list creation
• Better cache locality by processing one list at a time

The optimization is particularly effective for scenarios with:

• Large callback lists where matches occur early
• Cases with many callbacks in the first list (success_callback) since it's checked first
• Workloads that frequently call this function, as the memory allocation overhead is eliminated entirely

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
litellm_utils_tests/test_utils.py::test_custom_logger_exists_in_callbacks_individual_functions	6.18μs	3.07μs	101%✅

🌀 Generated Regression Tests and Runtime

import datetime
# Patch the global litellm reference in the function to our dummy module
import sys
import types

# imports
import pytest  # used for our unit tests
from litellm.utils import _custom_logger_class_exists_in_success_callbacks


# Simulate litellm module and CustomLogger for testing
class CustomLogger:
    pass

class AnotherLogger:
    pass

class YetAnotherLogger(CustomLogger):
    pass
from litellm.utils import _custom_logger_class_exists_in_success_callbacks

# 1. Basic Test Cases

def test_empty_callbacks_returns_false():
    # No callbacks present
    import litellm
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.00μs -> 836ns (139% faster)

def test_single_success_callback_match():
    # CustomLogger instance present in success_callback
    import litellm
    logger = CustomLogger()
    litellm.success_callback = [logger]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(logger) # 2.08μs -> 760ns (174% faster)

def test_single_async_success_callback_match():
    # CustomLogger instance present in _async_success_callback
    import litellm
    logger = CustomLogger()
    litellm._async_success_callback = [logger]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(logger) # 2.80μs -> 1.13μs (147% faster)

def test_single_success_callback_no_match():
    # AnotherLogger instance present, not CustomLogger
    import litellm
    litellm.success_callback = [AnotherLogger()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 1.95μs -> 962ns (103% faster)

def test_single_async_success_callback_no_match():
    # AnotherLogger instance present, not CustomLogger
    import litellm
    litellm._async_success_callback = [AnotherLogger()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.13μs -> 882ns (142% faster)

def test_multiple_callbacks_one_match():
    # Multiple callbacks, one is matching type
    import litellm
    litellm.success_callback = [AnotherLogger(), CustomLogger(), YetAnotherLogger()]
    litellm._async_success_callback = [AnotherLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 1.98μs -> 843ns (135% faster)

def test_multiple_callbacks_no_match():
    # Multiple callbacks, none are matching type
    import litellm
    litellm.success_callback = [AnotherLogger(), YetAnotherLogger()]
    litellm._async_success_callback = [AnotherLogger()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 1.97μs -> 886ns (122% faster)

def test_subclass_instance_does_not_match():
    # Subclass instance should not match base class instance
    import litellm
    litellm.success_callback = [YetAnotherLogger()]
    # Even though YetAnotherLogger is a subclass of CustomLogger, the function checks for exact type
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.16μs -> 757ns (186% faster)

def test_base_class_instance_does_not_match_subclass():
    # Base class instance should not match subclass instance
    import litellm
    litellm.success_callback = [CustomLogger()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(YetAnotherLogger()) # 2.10μs -> 1.13μs (86.1% faster)

def test_multiple_matches():
    # More than one matching instance present
    import litellm
    litellm.success_callback = [CustomLogger(), CustomLogger()]
    litellm._async_success_callback = [CustomLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 1.79μs -> 756ns (137% faster)

def test_mixed_types():
    # Mixed types, only one matches
    import litellm
    litellm.success_callback = [AnotherLogger(), CustomLogger(), YetAnotherLogger()]
    litellm._async_success_callback = [AnotherLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.07μs -> 922ns (124% faster)
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(AnotherLogger()) # 783ns -> 397ns (97.2% faster)
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(YetAnotherLogger()) # 776ns -> 402ns (93.0% faster)

# 2. Edge Test Cases

def test_callbacks_with_none_and_non_logger_objects():
    # Callbacks contain None and unrelated objects
    import litellm
    litellm.success_callback = [None, 42, "string", CustomLogger()]
    litellm._async_success_callback = [object()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.18μs -> 1.17μs (85.9% faster)

def test_callbacks_with_only_none_and_non_logger_objects():
    # Callbacks contain only None and unrelated objects
    import litellm
    litellm.success_callback = [None, 42, "string"]
    litellm._async_success_callback = [object()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.14μs -> 1.25μs (71.5% faster)

def test_callbacks_with_same_type_different_instance():
    # Instance passed is different from the one in callbacks, but same type
    import litellm
    logger1 = CustomLogger()
    logger2 = CustomLogger()
    litellm.success_callback = [logger1]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(logger2) # 2.11μs -> 739ns (186% faster)

def test_callbacks_with_different_subclass_instance():
    # Instance passed is subclass, but callback is base class
    import litellm
    litellm.success_callback = [CustomLogger()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(YetAnotherLogger()) # 1.77μs -> 1.05μs (67.9% faster)

def test_callbacks_with_multiple_types_and_subclasses():
    # Multiple types, including subclass and base class
    import litellm
    litellm.success_callback = [CustomLogger(), YetAnotherLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.15μs -> 751ns (187% faster)
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(YetAnotherLogger()) # 1.09μs -> 529ns (107% faster)

def test_callbacks_with_mutable_objects():
    # Callbacks contain mutable objects (lists, dicts)
    import litellm
    litellm.success_callback = [[], {}, CustomLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.20μs -> 989ns (122% faster)

def test_callbacks_with_callable_objects():
    # Callbacks contain callable objects (functions, lambdas)
    import litellm
    def dummy_func(): pass
    litellm.success_callback = [dummy_func, lambda x: x, CustomLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.35μs -> 992ns (137% faster)

def test_callbacks_with_large_number_of_non_matching_types():
    # Many non-matching types, only one matches
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(100)] + [CustomLogger()]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 7.22μs -> 4.00μs (80.7% faster)

def test_callbacks_with_large_number_of_non_matching_types_no_match():
    # Many non-matching types, none matches
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(100)]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 6.57μs -> 4.00μs (64.3% faster)

def test_callbacks_with_object_of_same_name_different_class():
    # Object with same class name, but different class object
    import litellm
    class CustomLogger2:
        pass
    litellm.success_callback = [CustomLogger2()]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 2.03μs -> 1.14μs (78.4% faster)

# 3. Large Scale Test Cases

def test_large_success_callback_with_one_match():
    # Large callback list, one matching instance at the end
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(999)] + [CustomLogger()]
    litellm._async_success_callback = []
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 48.5μs -> 28.9μs (68.2% faster)

def test_large_async_success_callback_with_one_match():
    # Large async callback list, one matching instance at the beginning
    import litellm
    litellm.success_callback = []
    litellm._async_success_callback = [CustomLogger()] + [AnotherLogger() for _ in range(999)]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 4.52μs -> 1.15μs (293% faster)

def test_large_success_and_async_callbacks_no_match():
    # Large lists, no matching instance
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(500)]
    litellm._async_success_callback = [YetAnotherLogger() for _ in range(500)]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 26.4μs -> 15.5μs (70.2% faster)

def test_large_success_and_async_callbacks_multiple_matches():
    # Large lists, multiple matching instances scattered
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(300)] + [CustomLogger()] + [AnotherLogger() for _ in range(300)]
    litellm._async_success_callback = [CustomLogger() for _ in range(10)]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 17.0μs -> 9.37μs (81.2% faster)

def test_large_success_and_async_callbacks_all_match():
    # Large lists, all matching instances
    import litellm
    litellm.success_callback = [CustomLogger() for _ in range(500)]
    litellm._async_success_callback = [CustomLogger() for _ in range(500)]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 4.80μs -> 927ns (417% faster)

def test_large_success_and_async_callbacks_all_non_match():
    # Large lists, all non-matching instances
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(500)]
    litellm._async_success_callback = [YetAnotherLogger() for _ in range(500)]
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 27.3μs -> 15.2μs (80.0% faster)

def test_large_success_and_async_callbacks_mixed_types():
    # Large lists, mixed types, only a few matches
    import litellm
    litellm.success_callback = [AnotherLogger() for _ in range(250)] + [CustomLogger() for _ in range(5)] + [YetAnotherLogger() for _ in range(245)]
    litellm._async_success_callback = [AnotherLogger() for _ in range(495)] + [CustomLogger() for _ in range(5)]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 15.6μs -> 7.89μs (98.1% faster)

def test_large_success_and_async_callbacks_with_none_and_other_objects():
    # Large lists, mixed with None and unrelated objects
    import litellm
    litellm.success_callback = [None for _ in range(100)] + [CustomLogger()] + [42 for _ in range(899)]
    litellm._async_success_callback = [object() for _ in range(1000)]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 10.8μs -> 3.74μs (188% faster)

def test_large_success_and_async_callbacks_with_callable_objects():
    # Large lists, mixed with callable objects
    import litellm
    litellm.success_callback = [lambda x: x for _ in range(500)] + [CustomLogger()]
    litellm._async_success_callback = [lambda x: x for _ in range(500)]
    codeflash_output = _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 27.2μs -> 15.0μs (81.2% faster)

def test_large_success_and_async_callbacks_with_varied_types_no_match():
    # Large lists, varied types, no match
    import litellm
    litellm.success_callback = [None, 42, "string", object()] * 250
    litellm._async_success_callback = [AnotherLogger(), YetAnotherLogger()] * 250
    codeflash_output = not _custom_logger_class_exists_in_success_callbacks(CustomLogger()) # 49.8μs -> 29.9μs (66.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import sys
# function to test
# (copied from the prompt for completeness)
import types

# imports
import pytest
from litellm.utils import _custom_logger_class_exists_in_success_callbacks


class CustomLogger:
    pass
from litellm.utils import \
    _custom_logger_class_exists_in_success_callbacks  # Now, all references to litellm in this module will point to our DummyLitellm instance


# Helper classes for testing
class LoggerA(CustomLogger):
    pass

class LoggerB(CustomLogger):
    pass

class NotALogger:
    pass

# ========== BASIC TEST CASES ==========

def setup_function():
    """Reset the litellm callback lists before each test."""
    litellm.success_callback.clear()
    litellm._async_success_callback.clear()

To edit these changes git checkout codeflash/optimize-_custom_logger_class_exists_in_success_callbacks-mhditqlz and push.