diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5f3f18c25ca..49d0775a94a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -647,11 +647,11 @@ class Logging(LiteLLMLoggingBaseClass): litellm_call_id: str, function_id: str, litellm_trace_id: str | None = None, - dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None, - dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, + dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, applied_guardrails: list[str] | None = None, kwargs: dict | None = None, log_raw_request_response: bool = False, @@ -724,15 +724,17 @@ class Logging(LiteLLMLoggingBaseClass): self.raw_request_only = raw_request_only # Initialize dynamic callbacks - self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks + self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_success_callbacks ) self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_success_callbacks ) - self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = ( + self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( + dynamic_failure_callbacks + ) + self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_failure_callbacks ) @@ -854,10 +856,26 @@ class Logging(LiteLLMLoggingBaseClass): litellm.logging_callback_manager.get_custom_logger_key(callback) } + def add_dynamic_function_callbacks( + self, + *, + async_success: Callable[..., object], + failure: Callable[..., object], + async_failure: Callable[..., object], + ) -> None: + self.dynamic_async_success_callbacks = self._with_dynamic_callback( + self.dynamic_async_success_callbacks, async_success + ) + self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure) + self.dynamic_async_failure_callbacks = self._with_dynamic_callback( + self.dynamic_async_failure_callbacks, async_failure + ) + @staticmethod def _with_dynamic_callback( - callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger - ) -> list[str | Callable | CustomLogger]: + callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None, + callback: Callable[..., object] | CustomLogger, + ) -> list[str | Callable[..., object] | CustomLogger]: existing: Final = tuple(callbacks or ()) return [*existing, *(() if callback in existing else (callback,))] @@ -894,9 +912,9 @@ class Logging(LiteLLMLoggingBaseClass): def _process_dynamic_callback_list( self, - callback_list: list[str | Callable | CustomLogger] | None, + callback_list: list[str | Callable[..., object] | CustomLogger] | None, dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"], - ) -> list[str | Callable | CustomLogger] | None: + ) -> list[str | Callable[..., object] | CustomLogger] | None: """ Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -908,7 +926,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback_list is None: return None - processed_list: Final[list[str | Callable | CustomLogger]] = [] + processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = [] for callback in callback_list: if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: for callback_instance in self._resolve_dynamic_callback_string(callback): diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fe7ac30e434..e1e8864678e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -880,8 +880,8 @@ def _build_passthrough_failure_request_payload( @dataclass(frozen=True, slots=True) class _TeamCallbackWiring: - success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg - failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg + success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg + failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg diff --git a/litellm/router.py b/litellm/router.py index 43a7b54c64f..4a52eabda29 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1709,6 +1709,11 @@ class Router: self._release_global_callbacks() def bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None: + logging_obj.add_dynamic_function_callbacks( + async_success=self.deployment_callback_on_success, + failure=self.deployment_callback_on_failure, + async_failure=self.async_deployment_callback_on_failure, + ) for selector in self._request_scoped_selectors: logging_obj.replace_global_callback(selector) for policy in self._request_policies_not_set_globally(): diff --git a/litellm/utils.py b/litellm/utils.py index 14ee35e571e..15692443ef0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1197,10 +1197,10 @@ def function_setup( for index in reversed(removed_async_items): litellm.failure_callback.pop(index) ### DYNAMIC CALLBACKS ### - dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None + dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index bc902233c3f..4839d3d3452 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -11,7 +11,7 @@ import redis import respx -from collections.abc import Coroutine, Iterator, Mapping +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping from datetime import datetime from types import SimpleNamespace from uuid import uuid4 @@ -2764,6 +2764,165 @@ async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_ await request() +async def _user_config_request(model: str, user_config: Mapping[str, object], arrives_with_logging_obj: bool) -> object: + return await ( + await route_request( + { + "model": model, + "messages": list(_HELLO_MESSAGES), + "num_retries": 0, + **({"litellm_logging_obj": _proxy_logging_obj(model)} if arrives_with_logging_obj else {}), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + + +def _ordered_pair(model_name: str, first_base: str, second_base: str) -> list[dict[str, object]]: + return [ + { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "api_base": api_base, "order": order}, + } + for order, api_base in ((1, first_base), (2, second_base)) + ] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_a_rate_limited_user_config_deployment_cools_down_for_the_next_requests( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + rate_limited_base: Final = "https://rate-limited.example/v1" + healthy_base: Final = "https://healthy.example/v1" + user_config: Final = {"model_list": _ordered_pair("flaky", rate_limited_base, healthy_base)} + + with respx.mock(assert_all_called=False) as respx_mock: + rate_limited: Final = respx_mock.post(f"{rate_limited_base}/chat/completions").mock( + return_value=httpx.Response(429, json={"error": {"message": "slow down"}}) + ) + respx_mock.post(f"{healthy_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(3): + await _user_config_request("flaky", user_config, arrives_with_logging_obj) + await asyncio.sleep(0.3) + + assert rate_limited.call_count == 1 + + +async def _remaining_usage_headers_after_429_then_200s( + send: Callable[[], Awaitable[object]], api_base: str +) -> list[tuple[object, object] | str]: + async def remaining_usage() -> tuple[object, object] | str: + try: + response: Final = await send() + except litellm.RateLimitError: + return "429" + finally: + await asyncio.sleep(0.3) + assert isinstance(response, litellm.ModelResponse) + headers: Final = response._hidden_params["additional_headers"] + return headers.get("x-ratelimit-remaining-requests"), headers.get("x-ratelimit-remaining-tokens") + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{api_base}/chat/completions").mock( + side_effect=[ + httpx.Response(429, json={"error": {"message": "slow down"}}), + _chat_completion_response("gpt-5.5"), + _chat_completion_response("gpt-5.5"), + ] + ) + return [await remaining_usage() for _ in range(3)] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_remaining_usage_headers_match_a_proxy_router_with_the_same_deployment( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_base: Final = "https://proxy-limited.example/v1" + user_config_base: Final = "https://user-config-limited.example/v1" + + def deployment(model_name: str, api_base: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": api_base, + "rpm": 10, + "tpm": 1000, + }, + } + + proxy_router: Final = litellm.Router(model_list=[deployment("proxy-limited", proxy_base)]) + try: + proxy_router_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: proxy_router.acompletion(model="proxy-limited", messages=list(_HELLO_MESSAGES), num_retries=0), + proxy_base, + ) + finally: + proxy_router.discard() + + user_config: Final = {"model_list": [deployment("user-config-limited", user_config_base)]} + user_config_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: _user_config_request("user-config-limited", user_config, arrives_with_logging_obj), user_config_base + ) + + assert user_config_headers == proxy_router_headers + + +async def _first_deployment_calls_over_successes_then_failures( + send: Callable[[], Awaitable[object]], first_base: str, second_base: str +) -> int: + successes: Final = [_chat_completion_response("gpt-5.5") for _ in range(6)] + failures: Final = [httpx.Response(500, json={"error": {"message": "boom"}}) for _ in range(6)] + with respx.mock(assert_all_called=False) as respx_mock: + first: Final = respx_mock.post(f"{first_base}/chat/completions").mock(side_effect=[*successes, *failures]) + respx_mock.post(f"{second_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(len(successes) + len(failures)): + await send() + await asyncio.sleep(0.3) + return first.call_count + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_successes_keep_a_deployment_out_of_cooldown_like_a_proxy_router( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_router: Final = litellm.Router( + model_list=_ordered_pair("proxy-pair", "https://proxy-first.example/v1", "https://proxy-second.example/v1") + ) + try: + proxy_router_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: proxy_router.acompletion(model="proxy-pair", messages=list(_HELLO_MESSAGES), num_retries=0), + "https://proxy-first.example/v1", + "https://proxy-second.example/v1", + ) + finally: + proxy_router.discard() + + user_config: Final = { + "model_list": _ordered_pair("user-pair", "https://user-first.example/v1", "https://user-second.example/v1") + } + + user_config_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: _user_config_request("user-pair", user_config, arrives_with_logging_obj), + "https://user-first.example/v1", + "https://user-second.example/v1", + ) + + assert user_config_calls == proxy_router_calls + + @pytest.mark.asyncio async def test_user_config_adaptive_router_learns_from_proxy_requests() -> None: user_config: Final = {