scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 128

fix(proxy): apply key and team router_settings to /v1/images/generations (#49286)

Browse source 
* fix(proxy): apply key and team router_settings to /v1/images/generations

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* style(proxy): format router settings helper

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): resolve key model_group_alias once on image generations

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
moe-berri 2026-10-30 12:38:21 -07:00 • committed by GitHub
parent 7b259690e7
commit 67954b5c3a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 389 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

93
litellm/proxy/common_request_processing.py
View file

@ -151,7 +151,7 @@ from litellm.router_utils.common_utils import resolve_model_group_alias, resolve
from litellm.types.extraction import ExtractionResponse
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.router_weights import validate_router_weights
from litellm.types.router_weights import RouterWeights, validate_router_weights
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
_LlmCallT = TypeVar("_LlmCallT")
@ -894,6 +894,61 @@ async def _resolve_per_request_model_group_alias(
return target
async def hierarchical_router_settings_updates(
requested_model: object,
user_api_key_dict: UserAPIKeyAuth,
llm_router: Router,
proxy_config: ProxyConfig,
proxy_logging_obj: ProxyLogging,
apply_model_group_alias: bool,
) -> Mapping[str, object]:
# Apply hierarchical router_settings (Key > Team)
# Global router_settings are already on the Router object itself.
# This sits with the other alias rewrites, and ahead of the guardrail
# 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.
from litellm.proxy.proxy_server import prisma_client
router_settings: Final[Mapping[str, object] | None] = _ROUTER_SETTINGS_ADAPTER.validate_python(
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 is None:
return MappingProxyType({})
# 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
def validated_router_weights() -> RouterWeights | None:
try:
return validate_router_weights(router_settings.get("weights"))
except ValidationError:
verbose_proxy_logger.warning("Ignoring invalid saved router weights; update team/key router_settings")
return None
router_weights: Final = validated_router_weights()
alias_target: Final = (
await _resolve_per_request_model_group_alias(
requested_model=requested_model,
router_settings=router_settings,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
if apply_model_group_alias
else None
)
settings_updates: Final[Mapping[str, object]] = {
"router_settings_override": router_settings,
"_router_weights": router_weights,
**({"model": alias_target} if alias_target is not None else {}),
}
return MappingProxyType(settings_updates)
_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)
@ -2495,41 +2550,17 @@ class ProxyBaseLLMRequestProcessing:
):
self.data["model"] = user_api_key_dict.aliases[self.data["model"]]
# Apply hierarchical router_settings (Key > Team)
# Global router_settings are already on the Router object itself.
# This sits with the other alias rewrites, and ahead of the guardrail
# 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(
self.data.update(
await hierarchical_router_settings_updates(
requested_model=self.data.get("model"),
router_settings=router_settings,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
proxy_config=proxy_config,
proxy_logging_obj=proxy_logging_obj,
apply_model_group_alias=True,
)
if alias_target is not None:
self.data["model"] = alias_target
)
aliased_model: Final = self.data.get("model")
if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str):

18
litellm/proxy/image_endpoints/endpoints.py
View file

@ -8,8 +8,10 @@ from typing import Final, get_type_hints
import orjson
from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status
from fastapi.responses import ORJSONResponse
from pydantic import TypeAdapter
import litellm
from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
@ -17,6 +19,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,
hierarchical_router_settings_updates,
log_llm_api_exception,
resolve_litellm_call_id,
)
@ -36,6 +39,7 @@ from litellm.types.images.main import ImageEditRequestParams
from litellm.types.llms.openai import ChatCompletionUserMessage
router: Final = APIRouter()
_IMAGE_GENERATION_MODEL_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object)
IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams))
IMAGE_EDIT_OPTIONAL_FIELD_DEFAULTS: Final = MappingProxyType({"prompt": None, "image": None})
@ -132,6 +136,19 @@ async def image_generation(
if data["model"] in litellm.model_alias_map:
data["model"] = litellm.model_alias_map[data["model"]]
if llm_router is not None:
requested_model: Final = _IMAGE_GENERATION_MODEL_ADAPTER.validate_python(data.get("model"))
data.update(
await hierarchical_router_settings_updates(
requested_model=requested_model,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
proxy_config=proxy_config,
proxy_logging_obj=proxy_logging_obj,
apply_model_group_alias=request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True,
)
)
### 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 +173,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

162
tests/integration/routing/test_key_router_settings_image_generation.py Normal file
View file

@ -0,0 +1,162 @@
from __future__ import annotations
import json
import uuid
from types import MappingProxyType
from typing import Final
from integration._support.client import Gateway, Scenario, gateway_from_environment
from integration._support.wire import Reply, Request, Wire, wire_server
FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"})
IMAGE_RESPONSE: Final = {
"created": 1,
"background": None,
"data": [
{
"b64_json": None,
"revised_prompt": None,
"url": "https://images.example/generated.png",
"provider_specific_fields": None,
}
],
"output_format": None,
"quality": None,
"size": None,
"usage": {
"total_tokens": 0,
"input_tokens": 0,
"input_tokens_details": {"image_tokens": 0, "text_tokens": 0},
"output_tokens": 0,
},
}
def _primary_unavailable(_: Request) -> Reply:
return Reply(
status=503,
body=json.dumps({"error": {"message": "primary unavailable", "type": "server_error"}}).encode(),
)
def _image_reply(_: Request) -> Reply:
return Reply(
body=json.dumps(
{
"created": 1,
"data": [{"url": "https://images.example/generated.png"}],
}
).encode()
)
def _image_model(scenario: Scenario, group: str, wire: Wire) -> str:
return scenario.model(
name=group,
model="openai/dall-e-3",
api_base=f"{wire.url}/v1",
api_key="integration-provider-key",
num_retries=0,
)
def _assert_image_fallback(
gateway: Gateway,
key: str,
primary: Wire,
fallback: Wire,
primary_model: str,
) -> None:
response: Final = gateway.request(
"POST",
"/v1/images/generations",
{"model": primary_model, "prompt": "draw a tree"},
key=key,
headers=FRESH_CONNECTION,
)
assert response.status_code == 200, response.text
assert response.json() == IMAGE_RESPONSE, response.text
assert tuple((request.method, request.target) for request in primary.drain()) == (
("POST", "/v1/images/generations"),
), response.text
assert tuple((request.method, request.target) for request in fallback.drain()) == (
("POST", "/v1/images/generations"),
), response.text
def test_key_router_settings_fallback_routes_image_generation() -> None:
with (
wire_server(_primary_unavailable) as primary_wire,
wire_server(_image_reply) as fallback_wire,
gateway_from_environment() as gateway,
gateway.scenario() as scenario,
):
primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire)
fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire)
key: Final = scenario.key(router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0})
_assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model)
def test_team_router_settings_fallback_routes_image_generation() -> None:
with (
wire_server(_primary_unavailable) as primary_wire,
wire_server(_image_reply) as fallback_wire,
gateway_from_environment() as gateway,
gateway.scenario() as scenario,
):
primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire)
fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire)
team: Final = scenario.team(
router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0}
)
key: Final = scenario.key(team_id=team)
_assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model)
def test_key_model_group_alias_resolves_once_for_image_generation() -> None:
with (
wire_server(_image_reply) as primary_wire,
wire_server(_image_reply) as target_wire,
wire_server(_image_reply) as chained_target_wire,
gateway_from_environment() as gateway,
gateway.scenario() as scenario,
):
primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire)
target_model: Final = _image_model(scenario, f"image-target-{uuid.uuid4().hex}", target_wire)
chained_target_model: Final = _image_model(
scenario, f"image-chained-target-{uuid.uuid4().hex}", chained_target_wire
)
alias_name: Final = f"image-alias-{uuid.uuid4().hex}"
key: Final = scenario.key(
router_settings={
"model_group_alias": {
alias_name: target_model,
target_model: chained_target_model,
}
}
)
response: Final = gateway.request(
"POST",
"/v1/images/generations",
{"model": alias_name, "prompt": "draw a tree"},
key=key,
headers=FRESH_CONNECTION,
)
primary_requests: Final = tuple((request.method, request.target) for request in primary_wire.drain())
target_requests: Final = tuple((request.method, request.target) for request in target_wire.drain())
chained_target_requests: Final = tuple(
(request.method, request.target) for request in chained_target_wire.drain()
)
assert alias_name not in {primary_model, target_model, chained_target_model}, response.text
assert response.status_code == 200, response.text
assert response.json() == IMAGE_RESPONSE, response.text
assert primary_requests == (), response.text
assert target_requests == (("POST", "/v1/images/generations"),), (
f"{response.text}; chained target requests: {chained_target_requests}"
)
assert chained_target_requests == (), response.text

148
tests/unit/proxy/image_endpoints/test_endpoints.py
View file

@ -3,7 +3,7 @@ import copy
import logging
from collections.abc import Iterator, Mapping
from types import SimpleNamespace
from typing import Any, Dict
from typing import Any, Dict, Final, cast
import orjson
import pytest
@ -12,10 +12,17 @@ from fastapi.testclient import TestClient
from starlette.requests import Request
from starlette.responses import Response
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.image_endpoints import endpoints
from litellm.proxy.proxy_server import PrismaClient, ProxyConfig
from litellm.proxy.utils import ProxyLogging
from litellm.types.utils import ImageObject, ImageResponse
from litellm.utils import custom_llm_setup
@pytest.mark.asyncio
@ -410,3 +417,142 @@ async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_i
assert raised.value.headers["x-litellm-call-id"] == call_id
assert [data["litellm_call_id"] for data in hook_request_data] == [call_id]
class _ImageRouterHandler(litellm.CustomLLM):
def __init__(self) -> None:
self.calls: list[str] = []
async def aimage_generation(
self,
model: str,
prompt: str,
model_response: ImageResponse,
api_key: str | None,
api_base: str | None,
optional_params: Mapping[str, object],
logging_obj: object,
timeout: object | None = None,
client: object | None = None,
) -> ImageResponse:
self.calls.append(model)
if model.endswith("failing"):
raise litellm.RateLimitError(message="primary failed", llm_provider="image-regression", model=model)
return ImageResponse(created=1, data=[ImageObject(url="https://images.example/generated.png")])
def _image_generation_client(
monkeypatch: pytest.MonkeyPatch,
user_api_key_dict: UserAPIKeyAuth,
team_settings: Mapping[str, object] | None = None,
) -> tuple[TestClient, _ImageRouterHandler]:
from litellm.proxy import proxy_server
handler: Final = _ImageRouterHandler()
cache: Final = UserApiKeyCache()
monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list])
monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers])
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-regression", "custom_handler": handler}])
custom_llm_setup()
if team_settings is not None:
team_id: Final = user_api_key_dict.team_id
assert team_id is not None
cache.set_cache(
key=f"team_id:{team_id}",
value=LiteLLM_TeamTableCachedObj(team_id=team_id, router_settings=dict(team_settings)),
local_only=True,
model_type=LiteLLM_TeamTableCachedObj,
)
monkeypatch.setattr(proxy_server, "prisma_client", cast(PrismaClient, object()))
router: Final = litellm.Router(
model_list=[
{"model_name": "A", "litellm_params": {"model": "image-regression/failing"}},
{"model_name": "B", "litellm_params": {"model": "image-regression/working"}},
],
num_retries=0,
)
monkeypatch.setattr(litellm, "model_alias_map", {})
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_config", ProxyConfig())
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache))
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "user_model", None)
monkeypatch.setattr(proxy_server, "version", "test-version")
app: Final = FastAPI()
app.include_router(endpoints.router)
app.dependency_overrides[user_api_key_auth] = lambda: user_api_key_dict
return TestClient(app, raise_server_exceptions=False), handler
def test_image_generation_uses_key_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None:
client, handler = _image_generation_client(
monkeypatch,
UserAPIKeyAuth(router_settings={"fallbacks": [{"A": ["B"]}]}),
)
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
assert response.status_code == 200, response.text
assert response.json() == ImageResponse(
created=1,
data=[ImageObject(url="https://images.example/generated.png")],
).model_dump(), response.text
assert handler.calls == ["failing", "working"], response.text
def test_image_generation_uses_team_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None:
team_id: Final = "image-regression-team"
client, handler = _image_generation_client(
monkeypatch,
UserAPIKeyAuth(team_id=team_id),
{"fallbacks": [{"A": ["B"]}]},
)
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
assert response.status_code == 200, response.text
assert response.json() == ImageResponse(
created=1,
data=[ImageObject(url="https://images.example/generated.png")],
).model_dump(), response.text
assert handler.calls == ["failing", "working"], response.text
def test_image_generation_without_router_settings_does_not_use_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth())
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
assert response.status_code == 500, response.text
assert handler.calls == ["failing"], response.text
def test_image_generation_uses_key_model_group_alias(monkeypatch: pytest.MonkeyPatch) -> None:
client, handler = _image_generation_client(
monkeypatch,
UserAPIKeyAuth(router_settings={"model_group_alias": {"alias": "B"}}),
)
response: Final = client.post("/v1/images/generations", json={"model": "alias", "prompt": "draw a tree"})
assert response.status_code == 200, response.text
assert response.json() == ImageResponse(
created=1,
data=[ImageObject(url="https://images.example/generated.png")],
).model_dump(), response.text
assert handler.calls == ["working"], response.text
def test_image_generation_uses_key_num_retries(monkeypatch: pytest.MonkeyPatch) -> None:
client, handler = _image_generation_client(
monkeypatch,
UserAPIKeyAuth(router_settings={"num_retries": 2}),
)
response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"})
assert response.status_code == 500, response.text
assert handler.calls == ["failing", "failing", "failing"], response.text