mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(proxy): filter unresolved key/team model strings from /v1/models
This commit is contained in:
omarsherif0
parent
4e12d3c562
commit
832210de09
2 changed files with 193 additions and 17 deletions
|
|
@ -4,10 +4,12 @@ from typing import Dict, List, Optional, Set
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
from litellm.types.utils import LlmProvidersSet
|
||||
from litellm.utils import get_valid_models
|
||||
|
||||
|
||||
|
|
@ -64,6 +66,82 @@ def _get_models_from_access_groups(
|
|||
return all_models
|
||||
|
||||
|
||||
def _is_known_provider_qualified_model(model: str) -> bool:
|
||||
"""
|
||||
Return True for provider-qualified model identifiers such as:
|
||||
- openai/gpt-4o-mini
|
||||
- bedrock/us.amazon.nova-micro-v1:0
|
||||
- openrouter/auto
|
||||
|
||||
This intentionally allows provider-prefixed models even when the exact
|
||||
model name is not yet in LiteLLM's static model list, since those routes
|
||||
may still be valid for pass-through / BYOK flows.
|
||||
"""
|
||||
if "/" not in model:
|
||||
return False
|
||||
|
||||
provider_prefix = model.split("/", 1)[0]
|
||||
return provider_prefix in LlmProvidersSet or JSONProviderRegistry.exists(
|
||||
provider_prefix
|
||||
)
|
||||
|
||||
|
||||
def _should_include_model_in_complete_list(
|
||||
model: str,
|
||||
proxy_model_list: List[str],
|
||||
model_access_groups: Dict[str, List[str]],
|
||||
) -> bool:
|
||||
"""
|
||||
Filter out unresolved strings from the final model list returned by
|
||||
/v1/models and related endpoints.
|
||||
|
||||
Keep:
|
||||
- configured proxy model groups
|
||||
- configured access group names
|
||||
- known base model IDs (e.g. gpt-4o-mini)
|
||||
- known provider-qualified routes (e.g. openai/gpt-4o-mini, bedrock/*)
|
||||
|
||||
Drop:
|
||||
- arbitrary strings that don't resolve to a proxy model, access group, or
|
||||
a recognized LiteLLM/provider model route.
|
||||
"""
|
||||
if model in (
|
||||
SpecialModelNames.all_proxy_models.value,
|
||||
SpecialModelNames.all_team_models.value,
|
||||
):
|
||||
return False
|
||||
|
||||
if model in proxy_model_list or model in model_access_groups:
|
||||
return True
|
||||
|
||||
if model in litellm.model_list_set:
|
||||
return True
|
||||
|
||||
if model == "*":
|
||||
return True
|
||||
|
||||
if _is_known_provider_qualified_model(model):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _filter_complete_model_list(
|
||||
models: List[str],
|
||||
proxy_model_list: List[str],
|
||||
model_access_groups: Dict[str, List[str]],
|
||||
) -> List[str]:
|
||||
return [
|
||||
model
|
||||
for model in models
|
||||
if _should_include_model_in_complete_list(
|
||||
model=model,
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
async def get_mcp_server_ids(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> List[str]:
|
||||
|
|
@ -211,6 +289,12 @@ def get_complete_model_list(
|
|||
valid_models = get_valid_models()
|
||||
append_unique(valid_models)
|
||||
|
||||
unique_models = _filter_complete_model_list(
|
||||
models=unique_models,
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
)
|
||||
|
||||
if only_model_access_groups:
|
||||
model_access_groups_to_return: List[str] = []
|
||||
for model in unique_models:
|
||||
|
|
|
|||
|
|
@ -37,7 +37,10 @@ def test_get_team_models_all_proxy_models_includes_access_groups():
|
|||
}
|
||||
|
||||
result = get_team_models(
|
||||
team_models, proxy_model_list, model_access_groups, include_model_access_groups=True
|
||||
team_models,
|
||||
proxy_model_list,
|
||||
model_access_groups,
|
||||
include_model_access_groups=True,
|
||||
)
|
||||
assert "group-a" in result
|
||||
assert "group-b" in result
|
||||
|
|
@ -61,7 +64,10 @@ def test_get_team_models_all_proxy_models_without_include_flag():
|
|||
}
|
||||
|
||||
result = get_team_models(
|
||||
team_models, proxy_model_list, model_access_groups, include_model_access_groups=False
|
||||
team_models,
|
||||
proxy_model_list,
|
||||
model_access_groups,
|
||||
include_model_access_groups=False,
|
||||
)
|
||||
assert "group-a" not in result
|
||||
assert "group-b" not in result
|
||||
|
|
@ -159,43 +165,66 @@ def test_get_key_models_does_not_mutate_input():
|
|||
"key_models,team_models,proxy_model_list,model_list,expected",
|
||||
[
|
||||
(
|
||||
["anthropic/claude-3-haiku-20240307", "anthropic/claude-3-5-haiku-20241022"],
|
||||
[
|
||||
"anthropic/claude-3-haiku-20240307",
|
||||
"anthropic/claude-3-5-haiku-20241022",
|
||||
],
|
||||
[],
|
||||
[],
|
||||
[{"model_name": "anthropic/*", "litellm_params": {"model": "anthropic/*"}}],
|
||||
["anthropic/claude-3-haiku-20240307", "anthropic/claude-3-5-haiku-20241022"]
|
||||
[
|
||||
"anthropic/claude-3-haiku-20240307",
|
||||
"anthropic/claude-3-5-haiku-20241022",
|
||||
],
|
||||
),
|
||||
(
|
||||
[],
|
||||
["anthropic/claude-3-haiku-20240307", "anthropic/claude-3-5-haiku-20241022"],
|
||||
[
|
||||
"anthropic/claude-3-haiku-20240307",
|
||||
"anthropic/claude-3-5-haiku-20241022",
|
||||
],
|
||||
[],
|
||||
[{"model_name": "anthropic/*", "litellm_params": {"model": "anthropic/*"}}],
|
||||
["anthropic/claude-3-haiku-20240307", "anthropic/claude-3-5-haiku-20241022"]
|
||||
[
|
||||
"anthropic/claude-3-haiku-20240307",
|
||||
"anthropic/claude-3-5-haiku-20241022",
|
||||
],
|
||||
),
|
||||
(
|
||||
[],
|
||||
[],
|
||||
["anthropic/claude-3-haiku-20240307", "anthropic/claude-3-5-haiku-20241022"],
|
||||
[
|
||||
"anthropic/claude-3-haiku-20240307",
|
||||
"anthropic/claude-3-5-haiku-20241022",
|
||||
],
|
||||
[{"model_name": "anthropic/*", "litellm_params": {"model": "anthropic/*"}}],
|
||||
["anthropic/claude-3-haiku-20240307", "anthropic/claude-3-5-haiku-20241022"]
|
||||
[
|
||||
"anthropic/claude-3-haiku-20240307",
|
||||
"anthropic/claude-3-5-haiku-20241022",
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_complete_model_list_order(key_models, team_models, proxy_model_list, model_list, expected):
|
||||
def test_get_complete_model_list_order(
|
||||
key_models, team_models, proxy_model_list, model_list, expected
|
||||
):
|
||||
"""
|
||||
Test that get_complete_model_list preserves order
|
||||
"""
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list
|
||||
from litellm import Router
|
||||
|
||||
assert get_complete_model_list(
|
||||
proxy_model_list=proxy_model_list,
|
||||
key_models=key_models,
|
||||
team_models=team_models,
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
llm_router=Router(model_list=model_list),
|
||||
) == expected
|
||||
assert (
|
||||
get_complete_model_list(
|
||||
proxy_model_list=proxy_model_list,
|
||||
key_models=key_models,
|
||||
team_models=team_models,
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
llm_router=Router(model_list=model_list),
|
||||
)
|
||||
== expected
|
||||
)
|
||||
|
||||
|
||||
def test_get_complete_model_list_byok_wildcard_expansion():
|
||||
|
|
@ -220,3 +249,66 @@ def test_get_complete_model_list_byok_wildcard_expansion():
|
|||
assert len(result) > 0
|
||||
assert all(m.startswith("openai/") for m in result)
|
||||
assert "openai/*" not in result
|
||||
|
||||
|
||||
def test_get_complete_model_list_filters_unknown_non_model_strings():
|
||||
"""
|
||||
If a key contains an arbitrary string that is neither:
|
||||
- a configured proxy model
|
||||
- a configured access group
|
||||
- a known LiteLLM model id
|
||||
- nor a recognized provider-qualified route
|
||||
|
||||
it should not leak into the final /v1/models response.
|
||||
"""
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list
|
||||
|
||||
result = get_complete_model_list(
|
||||
key_models=["team-sales-api"],
|
||||
team_models=[],
|
||||
proxy_model_list=["gpt-4o-mini"],
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
model_access_groups={},
|
||||
)
|
||||
|
||||
assert "team-sales-api" not in result
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_get_complete_model_list_keeps_known_base_model_ids():
|
||||
"""
|
||||
Exact model IDs can be valid even when they are not configured as proxy
|
||||
model groups, so known LiteLLM model ids should remain in the final list.
|
||||
"""
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list
|
||||
|
||||
result = get_complete_model_list(
|
||||
key_models=["gpt-4o-mini"],
|
||||
team_models=[],
|
||||
proxy_model_list=[],
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
model_access_groups={},
|
||||
)
|
||||
|
||||
assert result == ["gpt-4o-mini"]
|
||||
|
||||
|
||||
def test_get_complete_model_list_keeps_provider_qualified_models():
|
||||
"""
|
||||
Provider-qualified model identifiers should survive filtering even if the
|
||||
exact model name is newer than LiteLLM's baked-in model list.
|
||||
"""
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list
|
||||
|
||||
result = get_complete_model_list(
|
||||
key_models=["bedrock/very_new_model"],
|
||||
team_models=[],
|
||||
proxy_model_list=[],
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
model_access_groups={},
|
||||
)
|
||||
|
||||
assert result == ["bedrock/very_new_model"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue