|
| 1 | +from typing import Callable, List, Union |
| 2 | + |
| 3 | +import litellm |
| 4 | +from litellm._logging import verbose_logger |
| 5 | +from litellm.integrations.custom_logger import CustomLogger |
| 6 | + |
| 7 | + |
| 8 | +class LoggingCallbackManager: |
| 9 | + """ |
| 10 | + A centralized class that allows easy add / remove callbacks for litellm. |
| 11 | +
|
| 12 | + Goals of this class: |
| 13 | + - Prevent adding duplicate callbacks / success_callback / failure_callback |
| 14 | + - Keep a reasonable MAX_CALLBACKS limit (this ensures callbacks don't exponentially grow and consume CPU Resources) |
| 15 | + """ |
| 16 | + |
| 17 | + # healthy maximum number of callbacks - unlikely someone needs more than 20 |
| 18 | + MAX_CALLBACKS = 30 |
| 19 | + |
| 20 | + def add_litellm_input_callback(self, callback: Union[CustomLogger, str]): |
| 21 | + """ |
| 22 | + Add a input callback to litellm.input_callback |
| 23 | + """ |
| 24 | + self._safe_add_callback_to_list( |
| 25 | + callback=callback, parent_list=litellm.input_callback |
| 26 | + ) |
| 27 | + |
| 28 | + def add_litellm_service_callback( |
| 29 | + self, callback: Union[CustomLogger, str, Callable] |
| 30 | + ): |
| 31 | + """ |
| 32 | + Add a service callback to litellm.service_callback |
| 33 | + """ |
| 34 | + self._safe_add_callback_to_list( |
| 35 | + callback=callback, parent_list=litellm.service_callback |
| 36 | + ) |
| 37 | + |
| 38 | + def add_litellm_callback(self, callback: Union[CustomLogger, str, Callable]): |
| 39 | + """ |
| 40 | + Add a callback to litellm.callbacks |
| 41 | +
|
| 42 | + Ensures no duplicates are added. |
| 43 | + """ |
| 44 | + self._safe_add_callback_to_list( |
| 45 | + callback=callback, parent_list=litellm.callbacks # type: ignore |
| 46 | + ) |
| 47 | + |
| 48 | + def add_litellm_success_callback( |
| 49 | + self, callback: Union[CustomLogger, str, Callable] |
| 50 | + ): |
| 51 | + """ |
| 52 | + Add a success callback to `litellm.success_callback` |
| 53 | + """ |
| 54 | + self._safe_add_callback_to_list( |
| 55 | + callback=callback, parent_list=litellm.success_callback |
| 56 | + ) |
| 57 | + |
| 58 | + def add_litellm_failure_callback( |
| 59 | + self, callback: Union[CustomLogger, str, Callable] |
| 60 | + ): |
| 61 | + """ |
| 62 | + Add a failure callback to `litellm.failure_callback` |
| 63 | + """ |
| 64 | + self._safe_add_callback_to_list( |
| 65 | + callback=callback, parent_list=litellm.failure_callback |
| 66 | + ) |
| 67 | + |
| 68 | + def add_litellm_async_success_callback( |
| 69 | + self, callback: Union[CustomLogger, Callable, str] |
| 70 | + ): |
| 71 | + """ |
| 72 | + Add a success callback to litellm._async_success_callback |
| 73 | + """ |
| 74 | + self._safe_add_callback_to_list( |
| 75 | + callback=callback, parent_list=litellm._async_success_callback |
| 76 | + ) |
| 77 | + |
| 78 | + def add_litellm_async_failure_callback( |
| 79 | + self, callback: Union[CustomLogger, Callable, str] |
| 80 | + ): |
| 81 | + """ |
| 82 | + Add a failure callback to litellm._async_failure_callback |
| 83 | + """ |
| 84 | + self._safe_add_callback_to_list( |
| 85 | + callback=callback, parent_list=litellm._async_failure_callback |
| 86 | + ) |
| 87 | + |
| 88 | + def _add_string_callback_to_list( |
| 89 | + self, callback: str, parent_list: List[Union[CustomLogger, Callable, str]] |
| 90 | + ): |
| 91 | + """ |
| 92 | + Add a string callback to a list, if the callback is already in the list, do not add it again. |
| 93 | + """ |
| 94 | + if callback not in parent_list: |
| 95 | + parent_list.append(callback) |
| 96 | + else: |
| 97 | + verbose_logger.debug( |
| 98 | + f"Callback {callback} already exists in {parent_list}, not adding again.." |
| 99 | + ) |
| 100 | + |
| 101 | + def _check_callback_list_size( |
| 102 | + self, parent_list: List[Union[CustomLogger, Callable, str]] |
| 103 | + ) -> bool: |
| 104 | + """ |
| 105 | + Check if adding another callback would exceed MAX_CALLBACKS |
| 106 | + Returns True if safe to add, False if would exceed limit |
| 107 | + """ |
| 108 | + if len(parent_list) >= self.MAX_CALLBACKS: |
| 109 | + verbose_logger.warning( |
| 110 | + f"Cannot add callback - would exceed MAX_CALLBACKS limit of {self.MAX_CALLBACKS}. Current callbacks: {len(parent_list)}" |
| 111 | + ) |
| 112 | + return False |
| 113 | + return True |
| 114 | + |
| 115 | + def _safe_add_callback_to_list( |
| 116 | + self, |
| 117 | + callback: Union[CustomLogger, Callable, str], |
| 118 | + parent_list: List[Union[CustomLogger, Callable, str]], |
| 119 | + ): |
| 120 | + """ |
| 121 | + Safe add a callback to a list, if the callback is already in the list, do not add it again. |
| 122 | +
|
| 123 | + Ensures no duplicates are added for `str`, `Callable`, and `CustomLogger` callbacks. |
| 124 | + """ |
| 125 | + # Check max callbacks limit first |
| 126 | + if not self._check_callback_list_size(parent_list): |
| 127 | + return |
| 128 | + |
| 129 | + if isinstance(callback, str): |
| 130 | + self._add_string_callback_to_list( |
| 131 | + callback=callback, parent_list=parent_list |
| 132 | + ) |
| 133 | + elif isinstance(callback, CustomLogger): |
| 134 | + self._add_custom_logger_to_list( |
| 135 | + custom_logger=callback, |
| 136 | + parent_list=parent_list, |
| 137 | + ) |
| 138 | + elif callable(callback): |
| 139 | + self._add_callback_function_to_list( |
| 140 | + callback=callback, parent_list=parent_list |
| 141 | + ) |
| 142 | + |
| 143 | + def _add_callback_function_to_list( |
| 144 | + self, callback: Callable, parent_list: List[Union[CustomLogger, Callable, str]] |
| 145 | + ): |
| 146 | + """ |
| 147 | + Add a callback function to a list, if the callback is already in the list, do not add it again. |
| 148 | + """ |
| 149 | + # Check if the function already exists in the list by comparing function objects |
| 150 | + if callback not in parent_list: |
| 151 | + parent_list.append(callback) |
| 152 | + else: |
| 153 | + verbose_logger.debug( |
| 154 | + f"Callback function {callback.__name__} already exists in {parent_list}, not adding again.." |
| 155 | + ) |
| 156 | + |
| 157 | + def _add_custom_logger_to_list( |
| 158 | + self, |
| 159 | + custom_logger: CustomLogger, |
| 160 | + parent_list: List[Union[CustomLogger, Callable, str]], |
| 161 | + ): |
| 162 | + """ |
| 163 | + Add a custom logger to a list, if another instance of the same custom logger exists in the list, do not add it again. |
| 164 | + """ |
| 165 | + # Check if an instance of the same class already exists in the list |
| 166 | + custom_logger_key = self._get_custom_logger_key(custom_logger) |
| 167 | + custom_logger_type_name = type(custom_logger).__name__ |
| 168 | + for existing_logger in parent_list: |
| 169 | + if ( |
| 170 | + isinstance(existing_logger, CustomLogger) |
| 171 | + and self._get_custom_logger_key(existing_logger) == custom_logger_key |
| 172 | + ): |
| 173 | + verbose_logger.debug( |
| 174 | + f"Custom logger of type {custom_logger_type_name}, key: {custom_logger_key} already exists in {parent_list}, not adding again.." |
| 175 | + ) |
| 176 | + return |
| 177 | + parent_list.append(custom_logger) |
| 178 | + |
| 179 | + def _get_custom_logger_key(self, custom_logger: CustomLogger): |
| 180 | + """ |
| 181 | + Get a unique key for a custom logger that considers only fundamental instance variables |
| 182 | +
|
| 183 | + Returns: |
| 184 | + str: A unique key combining the class name and fundamental instance variables (str, bool, int) |
| 185 | + """ |
| 186 | + key_parts = [type(custom_logger).__name__] |
| 187 | + |
| 188 | + # Add only fundamental type instance variables to the key |
| 189 | + for attr_name, attr_value in vars(custom_logger).items(): |
| 190 | + if not attr_name.startswith("_"): # Skip private attributes |
| 191 | + if isinstance(attr_value, (str, bool, int)): |
| 192 | + key_parts.append(f"{attr_name}={attr_value}") |
| 193 | + |
| 194 | + return "-".join(key_parts) |
| 195 | + |
| 196 | + def _reset_all_callbacks(self): |
| 197 | + """ |
| 198 | + Reset all callbacks to an empty list |
| 199 | +
|
| 200 | + Note: this is an internal function and should be used sparingly. |
| 201 | + """ |
| 202 | + litellm.input_callback = [] |
| 203 | + litellm.success_callback = [] |
| 204 | + litellm.failure_callback = [] |
| 205 | + litellm._async_success_callback = [] |
| 206 | + litellm._async_failure_callback = [] |
| 207 | + litellm.callbacks = [] |
0 commit comments