diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cec176ac559..2b22400a567 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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") diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..02d84c86761 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/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 diff --git a/tests/unit/proxy/image_endpoints/test_image_generation_router_settings.py b/tests/unit/proxy/image_endpoints/test_image_generation_router_settings.py new file mode 100644 index 00000000000..150e5c54414 --- /dev/null +++ b/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