Merge d3c46c82cc into dddbaa803d

litellm/proxy/common_request_processing.py

@@ -839,6 +839,47 @@ async def _resolve_per_request_model_group_alias(
     return target
 
 
+async def apply_hierarchical_router_settings(
+    data: dict,
+    user_api_key_dict: UserAPIKeyAuth | None,
+    llm_router: Router | None,
+    proxy_config: ProxyConfig | None,
+    proxy_logging_obj: ProxyLogging | None = None,
+) -> None:
+    """
+    Apply hierarchical router_settings (Key > Team) to the request data dictionary.
+    Passes settings as per-request overrides (router_settings_override) and resolves
+    any per-request model group alias.
+    """
+    if llm_router is not None and proxy_config is not None and user_api_key_dict is not None:
+        from litellm.proxy.proxy_server import prisma_client
+
+        router_settings: Final = await proxy_config._get_hierarchical_router_settings(
+            user_api_key_dict=user_api_key_dict,
+            prisma_client=prisma_client,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+        # If router_settings found (from key or team), apply them
+        # Pass settings as per-request overrides instead of creating a new Router
+        # This avoids expensive Router instantiation on each request
+        if router_settings is not None:
+            data["router_settings_override"] = router_settings
+            try:
+                data["_router_weights"] = validate_router_weights(router_settings.get("weights"))
+            except ValidationError:
+                data["_router_weights"] = None
+                verbose_proxy_logger.warning("Ignoring invalid saved router weights; update team/key router_settings")
+            alias_target: Final = await _resolve_per_request_model_group_alias(
+                requested_model=data.get("model"),
+                router_settings=router_settings,
+                user_api_key_dict=user_api_key_dict,
+                llm_router=llm_router,
+            )
+            if alias_target is not None:
+                data["model"] = alias_target
+
+
 _REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
 _DECODED_JSON: Final = TypeAdapter(object)
 _ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@@ -2377,35 +2418,13 @@ class ProxyBaseLLMRequestProcessing:
         # merge and the pre-call hooks, so everything that keys off the model
         # group -- model-level guardrails, per-model budgets and rate limits,
         # the logging object -- sees the group that will actually serve.
-        if llm_router is not None and proxy_config is not None:
-            from litellm.proxy.proxy_server import prisma_client
-
-            router_settings: Final = await proxy_config._get_hierarchical_router_settings(
-                user_api_key_dict=user_api_key_dict,
-                prisma_client=prisma_client,
-                proxy_logging_obj=proxy_logging_obj,
-            )
-
-            # If router_settings found (from key or team), apply them
-            # Pass settings as per-request overrides instead of creating a new Router
-            # This avoids expensive Router instantiation on each request
-            if router_settings is not None:
-                self.data["router_settings_override"] = router_settings
-                try:
-                    self.data["_router_weights"] = validate_router_weights(router_settings.get("weights"))
-                except ValidationError:
-                    self.data["_router_weights"] = None
-                    verbose_proxy_logger.warning(
-                        "Ignoring invalid saved router weights; update team/key router_settings"
-                    )
-                alias_target: Final = await _resolve_per_request_model_group_alias(
-                    requested_model=self.data.get("model"),
-                    router_settings=router_settings,
-                    user_api_key_dict=user_api_key_dict,
-                    llm_router=llm_router,
-                )
-                if alias_target is not None:
-                    self.data["model"] = alias_target
+        await apply_hierarchical_router_settings(
+            data=self.data,
+            user_api_key_dict=user_api_key_dict,
+            llm_router=llm_router,
+            proxy_config=proxy_config,
+            proxy_logging_obj=proxy_logging_obj,
+        )
 
         aliased_model: Final = self.data.get("model")
 

litellm/proxy/image_endpoints/endpoints.py

@@ -17,6 +17,7 @@ from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
 from litellm.proxy.common_request_processing import (
     ProxyBaseLLMRequestProcessing,
+    apply_hierarchical_router_settings,
     log_llm_api_exception,
     resolve_litellm_call_id,
 )
@@ -132,6 +133,15 @@ async def image_generation(
         if data["model"] in litellm.model_alias_map:
             data["model"] = litellm.model_alias_map[data["model"]]
 
+        # Apply hierarchical router_settings (Key > Team)
+        await apply_hierarchical_router_settings(
+            data=data,
+            user_api_key_dict=user_api_key_dict,
+            llm_router=llm_router,
+            proxy_config=proxy_config,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
         ### CALL HOOKS ### - modify incoming data / reject request before calling the model
         prompt_value: Final = data.get("prompt")
         if prompt_value is not None:
@@ -156,6 +166,7 @@ async def image_generation(
             route_type="aimage_generation",
             llm_router=llm_router,
             user_model=user_model,
+            user_api_key_dict=user_api_key_dict,
         )
         response = await llm_call
 

tests/unit/proxy/image_endpoints/test_image_generation_router_settings.py

@@ -0,0 +1,122 @@
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+import litellm
+from litellm import CustomLLM, Router
+from litellm.proxy import proxy_server
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+from litellm.proxy.image_endpoints import endpoints
+from litellm.types.utils import ImageObject
+
+
+class Upstream(CustomLLM):
+    def __init__(self, fail: bool = False):
+        super().__init__()
+        self.calls = 0
+        self.fail = fail
+
+    def _maybe_fail(self):
+        self.calls += 1
+        if self.fail:
+            raise litellm.ServiceUnavailableError("503", model="x", llm_provider="x")
+
+    async def acompletion(self, *args, **kwargs):
+        self._maybe_fail()
+        return litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "ok"}}])
+
+    async def aimage_generation(self, model, prompt, model_response, **kwargs):
+        self._maybe_fail()
+        model_response.data = [ImageObject(url="https://example.com/out.png")]
+        return model_response
+
+
+@pytest.mark.asyncio
+async def test_image_generation_applies_key_router_settings_fallbacks(monkeypatch):
+    """Ensure key router_settings fallbacks are applied to /v1/images/generations."""
+    primary, fallback = Upstream(fail=True), Upstream(fail=False)
+    monkeypatch.setattr(
+        litellm,
+        "custom_provider_map",
+        [
+            {"provider": "failing", "custom_handler": primary},
+            {"provider": "ok", "custom_handler": fallback},
+        ],
+    )
+    litellm.utils.custom_llm_setup()
+
+    edge = [{"A": ["B"]}]
+    monkeypatch.setattr(
+        proxy_server,
+        "llm_router",
+        Router(
+            model_list=[
+                {"model_name": "A", "litellm_params": {"model": "failing/x", "num_retries": 0}},
+                {"model_name": "B", "litellm_params": {"model": "ok/x", "num_retries": 0}},
+            ],
+            num_retries=0,
+            fallbacks=[],
+        ),
+    )
+    monkeypatch.setattr(
+        proxy_server,
+        "proxy_logging_obj",
+        proxy_server.ProxyLogging(user_api_key_cache=proxy_server.user_api_key_cache),
+    )
+
+    key = UserAPIKeyAuth(api_key="sk-test", router_settings={"fallbacks": edge})
+    app = FastAPI()
+    app.include_router(endpoints.router)
+    app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
+    app.dependency_overrides[endpoints.user_api_key_auth] = lambda: key
+    client = TestClient(app)
+
+    r = client.post("/v1/images/generations", json={"model": "A", "prompt": "hi"})
+    assert r.status_code == 200, r.text
+    assert primary.calls == 1
+    assert fallback.calls == 1
+
+
+@pytest.mark.asyncio
+async def test_image_generation_applies_key_router_settings_model_alias(monkeypatch):
+    """Ensure key router_settings model_group_alias is applied to /v1/images/generations."""
+    upstream = Upstream()
+    monkeypatch.setattr(
+        litellm,
+        "custom_provider_map",
+        [
+            {"provider": "ok", "custom_handler": upstream},
+        ],
+    )
+    litellm.utils.custom_llm_setup()
+
+    monkeypatch.setattr(
+        proxy_server,
+        "llm_router",
+        Router(
+            model_list=[
+                {"model_name": "real-model", "litellm_params": {"model": "ok/x", "num_retries": 0}},
+            ],
+            num_retries=0,
+        ),
+    )
+    monkeypatch.setattr(
+        proxy_server,
+        "proxy_logging_obj",
+        proxy_server.ProxyLogging(user_api_key_cache=proxy_server.user_api_key_cache),
+    )
+
+    key = UserAPIKeyAuth(
+        api_key="sk-test",
+        models=["real-model"],
+        router_settings={"model_group_alias": {"alias-model": "real-model"}},
+    )
+    app = FastAPI()
+    app.include_router(endpoints.router)
+    app.add_exception_handler(ProxyException, proxy_server.openai_exception_handler)
+    app.dependency_overrides[endpoints.user_api_key_auth] = lambda: key
+    client = TestClient(app)
+
+    r = client.post("/v1/images/generations", json={"model": "alias-model", "prompt": "hi"})
+    assert r.status_code == 200, r.text
+    assert upstream.calls == 1