mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(proxy): register CustomLogger callbacks into all six callback lists
`initialize_callbacks_on_proxy` only appended user callbacks to `litellm.callbacks`. Several code paths read different lists: - pass-through endpoint logging reads `litellm._async_success_callback` (see `litellm_logging.py` ~L2640) - `log_pre_api_call` reads `litellm.input_callback` As a result, a CustomLogger registered via `litellm_settings.callbacks` in `config.yaml` silently never fires for pass-through requests and `log_pre_api_call` hooks, despite being correctly imported and added to `litellm.callbacks`. This change pushes every CustomLogger instance into all six dedicated lists (`input_callback`, `success_callback`, `failure_callback` and their `_async_*` counterparts) immediately after they land in `litellm.callbacks`. The push is idempotent so existing entries are not duplicated. Fixes #17310 Tests added in `tests/test_litellm/proxy/common_utils/test_callback_utils.py`: - list-branch registration covers all six lists - scalar-branch registration covers all six lists - idempotent on repeated calls - non-CustomLogger entries are ignored - helper-level idempotency check
This commit is contained in:
parent
e59e34bed3
commit
03477c1b8b
2 changed files with 225 additions and 0 deletions
|
|
@ -321,6 +321,16 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915
|
|||
else:
|
||||
litellm.callbacks = imported_list # type: ignore
|
||||
|
||||
# Also register CustomLogger instances into the dedicated
|
||||
# input/success/failure (sync + async) callback lists.
|
||||
# The pass-through endpoint logging chain reads
|
||||
# ``litellm._async_success_callback`` (see ``litellm_logging.py``
|
||||
# ~line 2640) rather than ``litellm.callbacks``, so a CustomLogger
|
||||
# registered only via ``litellm_settings.callbacks`` would silently
|
||||
# not fire for pass-through requests. Same goes for ``log_pre_api_call``
|
||||
# (uses ``litellm.input_callback``). See issue #17310.
|
||||
_register_custom_loggers_into_all_callback_lists(imported_list)
|
||||
|
||||
if "prometheus" in value:
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
|
|
@ -332,11 +342,42 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915
|
|||
config_file_path=config_file_path,
|
||||
)
|
||||
]
|
||||
_register_custom_loggers_into_all_callback_lists(litellm.callbacks)
|
||||
verbose_proxy_logger.debug(
|
||||
f"{blue_color_code} Initialized Callbacks - {litellm.callbacks} {reset_color_code}"
|
||||
)
|
||||
|
||||
|
||||
def _register_custom_loggers_into_all_callback_lists(
|
||||
callbacks: Iterable[Any],
|
||||
) -> None:
|
||||
"""Ensure every CustomLogger instance is in all six callback lists.
|
||||
|
||||
LiteLLM maintains separate lists for input / success / failure callbacks,
|
||||
each with sync and async variants. Different code paths read different
|
||||
lists (e.g. pass-through endpoints read ``_async_success_callback``,
|
||||
``log_pre_api_call`` reads ``input_callback``). Registering only into
|
||||
``litellm.callbacks`` is not enough for those paths to fire the callback.
|
||||
|
||||
Idempotent: callbacks already present in a list are not added again.
|
||||
"""
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
for callback in callbacks:
|
||||
if not isinstance(callback, CustomLogger):
|
||||
continue
|
||||
for parent_list in (
|
||||
litellm.input_callback,
|
||||
litellm.success_callback,
|
||||
litellm.failure_callback,
|
||||
litellm._async_input_callback,
|
||||
litellm._async_success_callback,
|
||||
litellm._async_failure_callback,
|
||||
):
|
||||
if callback not in parent_list:
|
||||
parent_list.append(callback)
|
||||
|
||||
|
||||
def get_model_group_from_litellm_kwargs(kwargs: dict) -> Optional[str]:
|
||||
_litellm_params = kwargs.get("litellm_params", None) or {}
|
||||
_metadata = (
|
||||
|
|
|
|||
|
|
@ -119,3 +119,187 @@ def test_initialize_callbacks_on_proxy_instantiates_compression_interception(
|
|||
assert "compression_interception" not in litellm.callbacks
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression tests for issue #17310: CustomLogger instances registered via
|
||||
# ``litellm_settings.callbacks`` must also land in input / success / failure
|
||||
# callback lists (sync + async). Otherwise pass-through endpoint logging and
|
||||
# ``log_pre_api_call`` silently skip the user callback.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
class _PRTestCustomLogger(CustomLogger):
|
||||
"""Minimal CustomLogger subclass used in the tests below."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
_ALL_CALLBACK_LIST_NAMES = (
|
||||
"input_callback",
|
||||
"success_callback",
|
||||
"failure_callback",
|
||||
"_async_input_callback",
|
||||
"_async_success_callback",
|
||||
"_async_failure_callback",
|
||||
)
|
||||
|
||||
|
||||
def _snapshot_callback_lists():
|
||||
return {name: list(getattr(litellm, name)) for name in _ALL_CALLBACK_LIST_NAMES} | {
|
||||
"callbacks": list(litellm.callbacks)
|
||||
if isinstance(litellm.callbacks, list)
|
||||
else []
|
||||
}
|
||||
|
||||
|
||||
def _restore_callback_lists(snap):
|
||||
for name in _ALL_CALLBACK_LIST_NAMES:
|
||||
getattr(litellm, name).clear()
|
||||
getattr(litellm, name).extend(snap[name])
|
||||
litellm.callbacks = snap["callbacks"]
|
||||
|
||||
|
||||
def test_initialize_callbacks_on_proxy_registers_custom_logger_into_all_lists(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A CustomLogger instance passed via the list form must land in
|
||||
every one of the six dedicated callback lists in addition to
|
||||
``litellm.callbacks``. Regression for issue #17310."""
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"litellm.proxy.proxy_server",
|
||||
SimpleNamespace(prisma_client=None),
|
||||
)
|
||||
snap = _snapshot_callback_lists()
|
||||
logger = _PRTestCustomLogger()
|
||||
try:
|
||||
initialize_callbacks_on_proxy(
|
||||
value=[logger],
|
||||
premium_user=False,
|
||||
config_file_path=".",
|
||||
litellm_settings={},
|
||||
callback_specific_params={},
|
||||
)
|
||||
assert logger in litellm.callbacks
|
||||
for name in _ALL_CALLBACK_LIST_NAMES:
|
||||
assert logger in getattr(
|
||||
litellm, name
|
||||
), f"CustomLogger missing from litellm.{name}"
|
||||
finally:
|
||||
_restore_callback_lists(snap)
|
||||
|
||||
|
||||
def test_initialize_callbacks_on_proxy_is_idempotent_for_custom_logger(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Re-running ``initialize_callbacks_on_proxy`` with the same
|
||||
CustomLogger must not duplicate it in any callback list."""
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"litellm.proxy.proxy_server",
|
||||
SimpleNamespace(prisma_client=None),
|
||||
)
|
||||
snap = _snapshot_callback_lists()
|
||||
logger = _PRTestCustomLogger()
|
||||
try:
|
||||
initialize_callbacks_on_proxy(
|
||||
value=[logger],
|
||||
premium_user=False,
|
||||
config_file_path=".",
|
||||
litellm_settings={},
|
||||
callback_specific_params={},
|
||||
)
|
||||
initialize_callbacks_on_proxy(
|
||||
value=[logger],
|
||||
premium_user=False,
|
||||
config_file_path=".",
|
||||
litellm_settings={},
|
||||
callback_specific_params={},
|
||||
)
|
||||
for name in _ALL_CALLBACK_LIST_NAMES:
|
||||
assert (
|
||||
getattr(litellm, name).count(logger) == 1
|
||||
), f"CustomLogger duplicated in litellm.{name}"
|
||||
finally:
|
||||
_restore_callback_lists(snap)
|
||||
|
||||
|
||||
def test_initialize_callbacks_on_proxy_scalar_value_registers_into_all_lists(
|
||||
monkeypatch,
|
||||
):
|
||||
"""The scalar (non-list) branch of ``initialize_callbacks_on_proxy``
|
||||
must also push CustomLogger instances into all six dedicated lists.
|
||||
Regression for issue #17310 (covers the ``else`` branch in addition
|
||||
to the list branch)."""
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"litellm.proxy.proxy_server",
|
||||
SimpleNamespace(prisma_client=None),
|
||||
)
|
||||
logger = _PRTestCustomLogger()
|
||||
# Make get_instance_fn return our logger regardless of input path.
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_utils.callback_utils.get_instance_fn",
|
||||
lambda value, config_file_path=None: logger,
|
||||
)
|
||||
snap = _snapshot_callback_lists()
|
||||
try:
|
||||
initialize_callbacks_on_proxy(
|
||||
value="my_module.callback_instance",
|
||||
premium_user=False,
|
||||
config_file_path=".",
|
||||
litellm_settings={},
|
||||
callback_specific_params={},
|
||||
)
|
||||
assert logger in litellm.callbacks
|
||||
for name in _ALL_CALLBACK_LIST_NAMES:
|
||||
assert logger in getattr(
|
||||
litellm, name
|
||||
), f"CustomLogger missing from litellm.{name} (scalar branch)"
|
||||
finally:
|
||||
_restore_callback_lists(snap)
|
||||
|
||||
|
||||
def test_register_custom_loggers_into_all_callback_lists_ignores_non_custom_logger():
|
||||
"""The helper must skip entries that are not CustomLogger instances."""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
_register_custom_loggers_into_all_callback_lists,
|
||||
)
|
||||
|
||||
snap = _snapshot_callback_lists()
|
||||
plain_obj = object()
|
||||
plain_str = "lago"
|
||||
try:
|
||||
_register_custom_loggers_into_all_callback_lists([plain_obj, plain_str])
|
||||
for name in _ALL_CALLBACK_LIST_NAMES:
|
||||
assert plain_obj not in getattr(
|
||||
litellm, name
|
||||
), f"object() should not be in litellm.{name}"
|
||||
assert plain_str not in getattr(
|
||||
litellm, name
|
||||
), f"plain string should not be in litellm.{name}"
|
||||
finally:
|
||||
_restore_callback_lists(snap)
|
||||
|
||||
|
||||
def test_register_custom_loggers_into_all_callback_lists_is_idempotent():
|
||||
"""Direct test of the helper: calling twice does not duplicate entries."""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
_register_custom_loggers_into_all_callback_lists,
|
||||
)
|
||||
|
||||
snap = _snapshot_callback_lists()
|
||||
logger = _PRTestCustomLogger()
|
||||
try:
|
||||
_register_custom_loggers_into_all_callback_lists([logger])
|
||||
_register_custom_loggers_into_all_callback_lists([logger])
|
||||
for name in _ALL_CALLBACK_LIST_NAMES:
|
||||
assert (
|
||||
getattr(litellm, name).count(logger) == 1
|
||||
), f"helper not idempotent for litellm.{name}"
|
||||
finally:
|
||||
_restore_callback_lists(snap)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue