diff --git a/litellm/router.py b/litellm/router.py index 4a52eabda29..6ae768c281b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10611,6 +10611,7 @@ class Router: changes, otherwise a released router keeps recording turns through its hook.""" from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, ) for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): @@ -10620,6 +10621,8 @@ class Router: litellm.logging_callback_manager.add_litellm_callback( AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) ) + if self.adaptive_routers: + litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ diff --git a/litellm/router_strategy/adaptive_router/__init__.py b/litellm/router_strategy/adaptive_router/__init__.py index d7f55ebced9..1f7e6b93169 100644 --- a/litellm/router_strategy/adaptive_router/__init__.py +++ b/litellm/router_strategy/adaptive_router/__init__.py @@ -1,6 +1,9 @@ """Adaptive router strategy. See README.md for design overview.""" from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, +) -__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"] +__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"] diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 92207e07c52..3427b81be68 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen return content, tool_calls -class AdaptiveRouterPostCallHook(CustomLogger): - """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" - - def __init__(self, adaptive_router: AdaptiveRouter) -> None: - self.adaptive_router = adaptive_router - +class AdaptiveRouterResponseHeaderHook(CustomLogger): async def async_post_call_response_headers_hook( self, data: Mapping[str, object], @@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger): return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} + +class AdaptiveRouterPostCallHook(CustomLogger): + """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" + + def __init__(self, adaptive_router: AdaptiveRouter) -> None: + self.adaptive_router = adaptive_router + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._record(kwargs, response_obj, response_status=200) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index ec09fc44cab..311f102c0ae 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -154,10 +154,9 @@ class RouterBudgetLimiting(CustomLogger): self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) async def _drain_increments_after_retire(self) -> None: - if not self.redis_increment_operation_queue: - return - if await self._push_in_memory_increments_to_redis(): - await self._drain_increments_after_retire() + while self.redis_increment_operation_queue: + if not await self._push_in_memory_increments_to_redis(): + await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL) @with_service_target("router_budgets") async def async_filter_deployments( diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 4839d3d3452..6044475652c 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -21,7 +21,9 @@ from unittest.mock import MagicMock from fakeredis import TcpFakeServer from fastapi import HTTPException +from litellm.caching.caching import DualCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.route_llm_request import ( JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, @@ -29,6 +31,11 @@ from litellm.proxy.route_llm_request import ( route_request, user_config_routers, ) +from litellm.proxy.utils import ProxyLogging +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ADAPTIVE_ROUTER_RESPONSE_HEADER, +) from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError @@ -2565,6 +2572,41 @@ async def _proxy_router_outcomes( router.discard() +@pytest.mark.asyncio +async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None: + data: Final = { + "model": "smart-cheap-router", + "messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}], + "metadata": {}, + "user_config": { + "model_list": [ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast", "smart"]}, + }, + }, + { + **_mock_deployment("fast"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}}, + }, + { + **_mock_deployment("smart"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + }, + ] + }, + } + + await (await route_request(data, None, None, "acompletion")) + headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=None + ) + + assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]} + + def _global_router_callbacks() -> tuple[list[object], list[object]]: return list(litellm.callbacks), list(litellm.input_callback) diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index ad61f43c5a0..0ed22658144 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import ( ) from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, _recent_tool_results, _resolve_session_key, ) @@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): """The header hook returns the `x-litellm-adaptive-router-model` header so proxy header construction picks it up (works for both streaming and non-streaming; `async_post_call_success_hook` is too late for streaming).""" - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"adaptive_router_chosen_model": "smart"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"litellm_session_id": "sess-A"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_no_metadata(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_not_dict(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": "not-a-dict"}, user_api_key_dict=MagicMock(), response=MagicMock(), diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 4efcdb51bc5..ff6de3b9b86 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -786,3 +786,28 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: await limiter._push_in_memory_increments_to_redis() assert redis_cache.values[_SPEND_KEY] == 10.0 assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None: + pipeline_started: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + pipeline_completed=pipeline_completed, + should_fail_pipeline=True, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600)) + + budget_limiter.retire() + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert redis_cache.values == {} + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == []