mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(proxy): apply key/team level router_settings.model_group_alias
Key and team router_settings reach the request as router_settings_override, but route_llm_request only forwarded an allowlist of per-request Router kwargs, which never included model_group_alias; the Router itself only ever reads its own global alias map. Resolve the per-request alias by rewriting the requested model group before routing. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
daf22ec871
commit
3695ac5b13
2 changed files with 118 additions and 1 deletions
|
|
@ -3,10 +3,12 @@ from typing import TYPE_CHECKING, Any, Literal, Mapping, Optional
|
|||
|
||||
import httpx
|
||||
from fastapi import HTTPException, status
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router_utils.common_utils import _is_proxy_admin_request
|
||||
from litellm.types.router import RouterModelGroupAliasItem
|
||||
|
||||
# Router-internal mock_testing_* flag names — kept in sync with
|
||||
# ``litellm.types.router.MockRouterTestingParams`` by the test
|
||||
|
|
@ -28,6 +30,41 @@ else:
|
|||
LitellmRouter = Any
|
||||
|
||||
|
||||
_MODEL_GROUP_ALIAS_ADAPTER = TypeAdapter(dict[str, str | RouterModelGroupAliasItem])
|
||||
|
||||
|
||||
def _resolve_per_request_model_group_alias(override_settings: Mapping[str, object], model: object) -> str | None:
|
||||
"""
|
||||
Resolve ``model`` through a ``model_group_alias`` map coming from key/team level
|
||||
``router_settings``.
|
||||
|
||||
Returns the aliased model group name, or None when there is nothing to rewrite.
|
||||
"""
|
||||
if not isinstance(model, str):
|
||||
return None
|
||||
|
||||
raw_alias_map = override_settings.get("model_group_alias")
|
||||
if not raw_alias_map:
|
||||
return None
|
||||
|
||||
try:
|
||||
alias_map = _MODEL_GROUP_ALIAS_ADAPTER.validate_python(raw_alias_map)
|
||||
except ValidationError:
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring malformed model_group_alias in key/team router_settings: %s",
|
||||
raw_alias_map,
|
||||
)
|
||||
return None
|
||||
|
||||
alias_item = alias_map.get(model)
|
||||
if alias_item is None:
|
||||
return None
|
||||
|
||||
return alias_item if isinstance(alias_item, str) else alias_item["model"]
|
||||
|
||||
|
||||
def _route_user_config_request(data: dict, route_type: str):
|
||||
"""Route a request using the user-provided router config."""
|
||||
router_config = data.pop("user_config")
|
||||
|
|
@ -446,6 +483,10 @@ async def route_request(
|
|||
if key in override_settings and key not in data:
|
||||
data[key] = override_settings[key]
|
||||
|
||||
aliased_model = _resolve_per_request_model_group_alias(override_settings, data.get("model"))
|
||||
if aliased_model is not None:
|
||||
data["model"] = aliased_model
|
||||
|
||||
# Use main router with overridden kwargs
|
||||
if llm_router is not None:
|
||||
return getattr(llm_router, f"{route_type}")(**data)
|
||||
|
|
|
|||
|
|
@ -384,7 +384,7 @@ async def test_route_request_with_router_settings_override():
|
|||
"timeout": 30,
|
||||
"model_group_retry_policy": {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}},
|
||||
"routing_strategy": "least-busy",
|
||||
# This setting should be ignored (not in per_request_settings list)
|
||||
# applied by rewriting the model, never forwarded as a router kwarg
|
||||
"model_group_alias": {"alias": "real_model"},
|
||||
},
|
||||
}
|
||||
|
|
@ -404,6 +404,7 @@ async def test_route_request_with_router_settings_override():
|
|||
assert call_kwargs["routing_strategy"] == "least-busy"
|
||||
# Verify unsupported settings were NOT merged
|
||||
assert "model_group_alias" not in call_kwargs
|
||||
assert call_kwargs["model"] == "gpt-3.5-turbo"
|
||||
# Verify router_settings_override was removed from data
|
||||
assert "router_settings_override" not in call_kwargs
|
||||
|
||||
|
|
@ -471,6 +472,81 @@ async def test_route_request_with_router_settings_override_preserves_existing():
|
|||
assert call_kwargs["timeout"] == 30
|
||||
|
||||
|
||||
def _router_with_two_model_groups():
|
||||
import litellm
|
||||
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "group-a",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"api_key": "fake",
|
||||
"mock_response": "from-group-a",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "group-b",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "fake",
|
||||
"mock_response": "from-group-b",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"alias_map",
|
||||
[
|
||||
{"group-a": "group-b"},
|
||||
{"group-a": {"model": "group-b", "hidden": True}},
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_route_request_applies_model_group_alias_from_router_settings_override(alias_map):
|
||||
"""
|
||||
Key/team level ``router_settings.model_group_alias`` must rewrite the model group,
|
||||
the same way the global ``router_settings.model_group_alias`` does. Regression test
|
||||
for the alias being silently dropped because it is not a per-request Router kwarg.
|
||||
"""
|
||||
router = _router_with_two_model_groups()
|
||||
data = {
|
||||
"model": "group-a",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"router_settings_override": {"model_group_alias": alias_map},
|
||||
}
|
||||
|
||||
response = await (await route_request(data, router, None, "acompletion"))
|
||||
|
||||
assert response.choices[0].message.content == "from-group-b"
|
||||
assert response.model == "gpt-4o-mini"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"alias_map",
|
||||
[
|
||||
{"some-other-group": "group-b"},
|
||||
{"group-a": ["group-b"]},
|
||||
"not-a-map",
|
||||
{},
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_route_request_leaves_model_untouched_for_irrelevant_or_malformed_alias(alias_map):
|
||||
router = _router_with_two_model_groups()
|
||||
data = {
|
||||
"model": "group-a",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"router_settings_override": {"model_group_alias": alias_map},
|
||||
}
|
||||
|
||||
response = await (await route_request(data, router, None, "acompletion"))
|
||||
|
||||
assert response.choices[0].message.content == "from-group-a"
|
||||
|
||||
|
||||
def test_mock_testing_kwarg_names_matches_dataclass():
|
||||
"""``_MOCK_TESTING_KWARG_NAMES`` is hardcoded to avoid a cyclic import
|
||||
against ``litellm.types.router``. This test guards against drift —
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue