diff --git a/litellm/router.py b/litellm/router.py index 2d66bc3158d..4e535e575e6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4313,6 +4313,7 @@ class Router: passthrough_on_no_deployment = kwargs.pop("passthrough_on_no_deployment", False) function_name = "_ageneric_api_call_with_fallbacks" + deployment = None try: parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) try: @@ -4392,6 +4393,9 @@ class Router: ) if model is not None: self.fail_calls[model] += 1 + if deployment is not None: + self._set_deployment_num_retries_on_exception(e, deployment) + self._set_failed_deployment_id_on_exception(e, deployment) raise e async def _aresponses_with_streaming_fallbacks( diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 8faf6bcd9cf..cdc664693dc 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -769,3 +769,85 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo ) assert response._hidden_params["model_id"] == "fallback" + + +# --------------------------------------------------------------------------- +# Regression: generic API path (used by /v1/messages, image_edit, etc.) +# must also stamp failed_deployment_id for weighted failover to work +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_generic_api_path_sets_failed_deployment_id(): + """_ageneric_api_call_with_fallbacks_helper (used by /v1/messages and + other factory_function endpoints) did not call + _set_failed_deployment_id_on_exception, so weighted failover always + bailed out. Verify the attribute is now stamped.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("region down"), + "weight": 1, + }, + "model_info": {"id": "dep-A"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + enable_weighted_failover=True, + ) + + with pytest.raises(Exception) as exc_info: + await router._ageneric_api_call_with_fallbacks_helper( + model="test-model", + original_generic_function=router._acompletion, + messages=[{"role": "user", "content": "hi"}], + ) + + assert getattr(exc_info.value, "failed_deployment_id", None) == "dep-A" + + +@pytest.mark.asyncio +async def test_weighted_failover_works_on_generic_api_path(): + """End-to-end: weighted failover via the _ageneric_api_call_with_fallbacks + path (used by /v1/messages, image_edit, generate_content, etc.) should + re-pick a healthy sibling in the same model group before falling back + to a cross-group fallback. Uses anthropic_messages to exercise the + exact code path reported in the bug.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("region-A down"), + "weight": 1, + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from B", + "weight": 0, + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + enable_weighted_failover=True, + ) + + response = await router.anthropic_messages( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert response._hidden_params["model_id"] == "B"