From edfb0130d78732b00396f50b11a34033909c10a7 Mon Sep 17 00:00:00 2001 From: Jaeyeon Kim Date: Sat, 28 Feb 2026 02:37:48 +0900 Subject: [PATCH] feat: add setup_proxy(app) hook for CustomLogger callbacks Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/custom_logger.py | 12 ++ litellm/proxy/proxy_server.py | 24 ++++ .../proxy/test_callback_setup_proxy.py | 117 ++++++++++++++++++ 3 files changed, 153 insertions(+) create mode 100644 tests/test_litellm/proxy/test_callback_setup_proxy.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index c244363e389..d87a74ade41 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -82,6 +82,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self.turn_off_message_logging = turn_off_message_logging pass + def setup_proxy(self, app: Any) -> None: + """ + Called once during proxy startup after all callbacks are initialized. + + Override this method to interact with the FastAPI app during startup, + e.g., to add middleware, mount additional routes, or register shutdown hooks. + + Args: + app: The FastAPI application instance. + """ + pass + @staticmethod def get_callback_env_vars(callback_name: Optional[str] = None) -> List[str]: """ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bd5b5309e0f..e6a61b41d05 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -749,6 +749,27 @@ async def _initialize_shared_aiohttp_session(): return None +def _invoke_callback_setup_proxy(app: FastAPI) -> None: + """ + Call setup_proxy(app) on all registered CustomLogger callbacks. + + This gives callbacks an opportunity to interact with the FastAPI app + during startup, e.g., to add middleware or mount additional routes. + """ + from litellm.integrations.custom_logger import CustomLogger + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger) and hasattr(callback, "setup_proxy"): + try: + callback.setup_proxy(app) + except Exception as e: + verbose_proxy_logger.warning( + "Error calling setup_proxy on callback %s: %s", + type(callback).__name__, + str(e), + ) + + @asynccontextmanager async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session @@ -813,6 +834,9 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 if isinstance(worker_config, dict): await initialize(**worker_config) + # Call setup_proxy(app) on all registered CustomLogger callbacks + _invoke_callback_setup_proxy(app) + # check if DATABASE_URL in environment - load from there if prisma_client is None: _db_url: Optional[str] = get_secret("DATABASE_URL", None) # type: ignore diff --git a/tests/test_litellm/proxy/test_callback_setup_proxy.py b/tests/test_litellm/proxy/test_callback_setup_proxy.py new file mode 100644 index 00000000000..739682db9ef --- /dev/null +++ b/tests/test_litellm/proxy/test_callback_setup_proxy.py @@ -0,0 +1,117 @@ +""" +Tests for the CustomLogger.setup_proxy(app) hook. + +Verifies that: +1. CustomLogger has a setup_proxy method (no-op default) +2. _invoke_callback_setup_proxy calls setup_proxy on all CustomLogger callbacks +3. Non-CustomLogger callbacks are skipped +4. Exceptions in setup_proxy are caught and logged (don't crash startup) +""" + +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy.proxy_server import _invoke_callback_setup_proxy + + +class TestCustomLoggerSetupProxy: + """Tests for the setup_proxy method on CustomLogger.""" + + def test_custom_logger_has_setup_proxy_method(self): + """CustomLogger base class should have a setup_proxy method.""" + logger = CustomLogger() + assert hasattr(logger, "setup_proxy") + assert callable(logger.setup_proxy) + + def test_setup_proxy_default_is_noop(self): + """Default setup_proxy should do nothing and not raise.""" + logger = CustomLogger() + mock_app = MagicMock() + logger.setup_proxy(mock_app) # Should not raise + + +class TestInvokeCallbackSetupProxy: + """Tests for _invoke_callback_setup_proxy function.""" + + def test_calls_setup_proxy_on_custom_logger_callbacks(self): + """setup_proxy should be called on each CustomLogger in litellm.callbacks.""" + + class MyCallback(CustomLogger): + def __init__(self): + super().__init__() + self.setup_called_with = None + + def setup_proxy(self, app): + self.setup_called_with = app + + callback = MyCallback() + mock_app = MagicMock() + + original_callbacks = litellm.callbacks + litellm.callbacks = [callback] + try: + _invoke_callback_setup_proxy(mock_app) + assert callback.setup_called_with is mock_app + finally: + litellm.callbacks = original_callbacks + + def test_skips_non_custom_logger_callbacks(self): + """String callbacks and non-CustomLogger objects should be skipped.""" + mock_app = MagicMock() + + original_callbacks = litellm.callbacks + litellm.callbacks = ["langfuse", "sentry", 42] + try: + # Should not raise + _invoke_callback_setup_proxy(mock_app) + finally: + litellm.callbacks = original_callbacks + + def test_exception_in_setup_proxy_is_caught(self): + """If setup_proxy raises, it should be caught and logged, not crash.""" + + class BadCallback(CustomLogger): + def setup_proxy(self, app): + raise RuntimeError("setup failed") + + bad_callback = BadCallback() + good_callback = CustomLogger() + mock_app = MagicMock() + + original_callbacks = litellm.callbacks + litellm.callbacks = [bad_callback, good_callback] + try: + # Should not raise even though bad_callback.setup_proxy raises + _invoke_callback_setup_proxy(mock_app) + finally: + litellm.callbacks = original_callbacks + + def test_multiple_callbacks_all_called(self): + """All CustomLogger callbacks should have setup_proxy called.""" + call_order = [] + + class Callback1(CustomLogger): + def setup_proxy(self, app): + call_order.append("cb1") + + class Callback2(CustomLogger): + def setup_proxy(self, app): + call_order.append("cb2") + + mock_app = MagicMock() + original_callbacks = litellm.callbacks + litellm.callbacks = [Callback1(), "langfuse", Callback2()] + try: + _invoke_callback_setup_proxy(mock_app) + assert call_order == ["cb1", "cb2"] + finally: + litellm.callbacks = original_callbacks