mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
d2ac51893b
Regenerated every touched file from origin/main applying only the B1 test deletions and the unused import and helper cleanup they leave behind, without running the formatter across untouched code. CI only checks ruff format under litellm/, so the earlier reflows of test files were pure diff noise for reviewers Also drops the tests/local_testing/test_prompt_caching.py entry from the caching-local shard in test-unit.yml since that file is deleted Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
881 lines
29 KiB
Python
881 lines
29 KiB
Python
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
|
|
def test_get_team_models_for_all_models_and_team_only_models():
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
|
|
team_models = ["all-proxy-models", "team-only-model", "team-only-model-2"]
|
|
proxy_model_list = ["model1", "model2", "model3"]
|
|
model_access_groups = {}
|
|
include_model_access_groups = False
|
|
|
|
result = get_team_models(
|
|
team_models, proxy_model_list, model_access_groups, include_model_access_groups
|
|
)
|
|
combined_models = team_models + proxy_model_list
|
|
assert set(result) == set(combined_models)
|
|
|
|
|
|
def test_get_team_models_all_proxy_models_includes_access_groups():
|
|
"""
|
|
When a team has 'all-proxy-models' and include_model_access_groups=True,
|
|
the result should include model access group names (e.g. 'claude-model-group')
|
|
in addition to individual model names.
|
|
"""
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
|
|
team_models = ["all-proxy-models"]
|
|
proxy_model_list = ["model1", "model2"]
|
|
model_access_groups = {
|
|
"group-a": ["model1"],
|
|
"group-b": ["model2"],
|
|
}
|
|
|
|
result = get_team_models(
|
|
team_models,
|
|
proxy_model_list,
|
|
model_access_groups,
|
|
include_model_access_groups=True,
|
|
)
|
|
assert "group-a" in result
|
|
assert "group-b" in result
|
|
assert "model1" in result
|
|
assert "model2" in result
|
|
assert len(result) == len(set(result)), "result should have no duplicates"
|
|
|
|
|
|
def test_get_team_models_all_proxy_models_without_include_flag():
|
|
"""
|
|
When include_model_access_groups=False, access group names should NOT
|
|
appear in the result even with 'all-proxy-models'.
|
|
"""
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
|
|
team_models = ["all-proxy-models"]
|
|
proxy_model_list = ["model1", "model2"]
|
|
model_access_groups = {
|
|
"group-a": ["model1"],
|
|
"group-b": ["model2"],
|
|
}
|
|
|
|
result = get_team_models(
|
|
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
|
|
assert "model1" in result
|
|
assert "model2" in result
|
|
|
|
|
|
def test_get_key_models_all_proxy_models_includes_access_groups():
|
|
"""
|
|
When a key has 'all-proxy-models' and include_model_access_groups=True,
|
|
the result should include model access group names.
|
|
"""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(
|
|
models=["all-proxy-models"],
|
|
api_key="test-key",
|
|
)
|
|
proxy_model_list = ["model1", "model2"]
|
|
model_access_groups = {
|
|
"group-a": ["model1"],
|
|
}
|
|
|
|
result = get_key_models(
|
|
user_api_key_dict=user_api_key_dict,
|
|
proxy_model_list=proxy_model_list,
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=True,
|
|
)
|
|
assert "group-a" in result
|
|
assert "model1" in result
|
|
assert "model2" in result
|
|
assert len(result) == len(set(result)), "result should have no duplicates"
|
|
|
|
|
|
def test_get_key_models_passes_include_model_access_groups():
|
|
"""
|
|
When a key explicitly has an access group name in its models list and
|
|
include_model_access_groups=True, the group name should be retained
|
|
(not stripped by _get_models_from_access_groups).
|
|
"""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(
|
|
models=["group-a"],
|
|
api_key="test-key",
|
|
)
|
|
proxy_model_list = ["model1", "model2"]
|
|
model_access_groups = {
|
|
"group-a": ["model1", "model2"],
|
|
}
|
|
|
|
result = get_key_models(
|
|
user_api_key_dict=user_api_key_dict,
|
|
proxy_model_list=proxy_model_list,
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=True,
|
|
)
|
|
assert "group-a" in result
|
|
assert "model1" in result
|
|
assert "model2" in result
|
|
|
|
|
|
def test_get_key_models_keeps_literal_model_colliding_with_group_name():
|
|
"""A name that is BOTH a deployed model and an access group grants both at
|
|
runtime (_check_model_access_helper unions them), so the listing must keep
|
|
the literal model alongside the group members instead of dropping it."""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
|
|
user_api_key_dict = UserAPIKeyAuth(models=["beta-models"], api_key="test-key")
|
|
|
|
result = get_key_models(
|
|
user_api_key_dict=user_api_key_dict,
|
|
proxy_model_list=["beta-models", "member-a", "unrelated"],
|
|
model_access_groups={"beta-models": ["member-a"]},
|
|
include_model_access_groups=False,
|
|
)
|
|
assert sorted(result) == ["beta-models", "member-a"]
|
|
|
|
|
|
def test_get_team_models_keeps_literal_model_colliding_with_group_name():
|
|
"""Team flavor of the collision case: literal deployment survives group expansion."""
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
|
|
result = get_team_models(
|
|
team_models=["beta-models"],
|
|
proxy_model_list=["beta-models", "member-a", "unrelated"],
|
|
model_access_groups={"beta-models": ["member-a"]},
|
|
include_model_access_groups=False,
|
|
)
|
|
assert sorted(result) == ["beta-models", "member-a"]
|
|
|
|
|
|
def test_get_team_models_drops_group_name_that_is_not_a_deployed_model():
|
|
"""No collision: a pure access-group name is still replaced by its members."""
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
|
|
result = get_team_models(
|
|
team_models=["beta-models"],
|
|
proxy_model_list=["member-a", "unrelated"],
|
|
model_access_groups={"beta-models": ["member-a"]},
|
|
include_model_access_groups=False,
|
|
)
|
|
assert result == ["member-a"]
|
|
|
|
|
|
def test_get_key_models_does_not_mutate_input():
|
|
"""
|
|
get_key_models must not mutate user_api_key_dict.models in-place.
|
|
_get_models_from_access_groups uses .pop()/.extend() which would corrupt
|
|
cached UserAPIKeyAuth objects if all_models were an alias instead of a copy.
|
|
"""
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
|
|
original_models = ["group-a", "extra-model"]
|
|
user_api_key_dict = UserAPIKeyAuth(
|
|
models=list(original_models), # give it a list
|
|
api_key="test-key",
|
|
)
|
|
model_access_groups = {
|
|
"group-a": ["model1", "model2"],
|
|
}
|
|
|
|
_ = get_key_models(
|
|
user_api_key_dict=user_api_key_dict,
|
|
proxy_model_list=["model1", "model2"],
|
|
model_access_groups=model_access_groups,
|
|
include_model_access_groups=False,
|
|
)
|
|
# The original models list on the auth object must be unchanged
|
|
assert user_api_key_dict.models == original_models
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"key_models,team_models,proxy_model_list,model_list,expected",
|
|
[
|
|
(
|
|
[
|
|
"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",
|
|
],
|
|
[],
|
|
[{"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",
|
|
],
|
|
[{"model_name": "anthropic/*", "litellm_params": {"model": "anthropic/*"}}],
|
|
[
|
|
"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
|
|
):
|
|
"""
|
|
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
|
|
)
|
|
|
|
|
|
def test_get_complete_model_list_byok_wildcard_expansion():
|
|
"""
|
|
Test that wildcard models (e.g., openai/*) are expanded when the router has
|
|
no deployment for them - BYOK case where team has openai/* but proxy has
|
|
no openai config.
|
|
"""
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
from litellm import Router
|
|
|
|
# Router with empty model_list - no openai/* deployment (BYOK scenario)
|
|
result = get_complete_model_list(
|
|
key_models=[],
|
|
team_models=["openai/*"],
|
|
proxy_model_list=[],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
llm_router=Router(model_list=[]),
|
|
)
|
|
# Should expand openai/* to actual OpenAI models
|
|
assert len(result) > 0
|
|
assert all(m.startswith("openai/") for m in result)
|
|
assert "openai/*" not in result
|
|
|
|
|
|
def test_get_complete_model_list_expands_team_scoped_wildcard_with_stored_credential(
|
|
monkeypatch,
|
|
):
|
|
"""
|
|
Team-scoped BYOK wildcard deployments are stored under an internal model_name,
|
|
with the public wildcard name in model_info.team_public_model_name.
|
|
"""
|
|
import litellm
|
|
from litellm import Router
|
|
from litellm.proxy.auth import model_checks
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
from litellm.types.utils import CredentialItem
|
|
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"credential_list",
|
|
[
|
|
CredentialItem(
|
|
credential_name="openai-credential",
|
|
credential_info={"provider": "openai"},
|
|
credential_values={
|
|
"api_key": "stored-openai-key",
|
|
"api_base": "https://example.openai.test/v1",
|
|
},
|
|
)
|
|
],
|
|
)
|
|
|
|
captured_params = {}
|
|
|
|
def fake_get_provider_models(provider, litellm_params=None):
|
|
captured_params["provider"] = provider
|
|
captured_params["api_key"] = litellm_params.api_key
|
|
captured_params["api_base"] = litellm_params.api_base
|
|
captured_params["credential_name"] = litellm_params.litellm_credential_name
|
|
return ["gpt-4o"]
|
|
|
|
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "model_name_team-1_generated",
|
|
"litellm_params": {
|
|
"model": "openai/*",
|
|
"custom_llm_provider": "openai",
|
|
"litellm_credential_name": "openai-credential",
|
|
},
|
|
"model_info": {
|
|
"team_id": "team-1",
|
|
"team_public_model_name": "openai/*",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
result = get_complete_model_list(
|
|
key_models=[],
|
|
team_models=["openai/*"],
|
|
proxy_model_list=[],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
llm_router=router,
|
|
team_id="team-1",
|
|
)
|
|
|
|
assert "openai/gpt-4o" in result
|
|
assert captured_params == {
|
|
"provider": "openai",
|
|
"api_key": "stored-openai-key",
|
|
"api_base": "https://example.openai.test/v1",
|
|
"credential_name": None,
|
|
}
|
|
|
|
|
|
def test_wildcard_credential_hydration_preserves_deployment_params(
|
|
monkeypatch,
|
|
):
|
|
import litellm
|
|
from litellm.proxy.auth import model_checks
|
|
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
|
from litellm.types.router import LiteLLM_Params
|
|
from litellm.types.utils import CredentialItem
|
|
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"credential_list",
|
|
[
|
|
CredentialItem(
|
|
credential_name="openai-credential",
|
|
credential_info={"provider": "openai"},
|
|
credential_values={
|
|
"api_key": "stored-openai-key",
|
|
"api_version": "credential-version",
|
|
"model": "openai/wrong-model",
|
|
"unexpected_field": "unexpected-value",
|
|
},
|
|
)
|
|
],
|
|
)
|
|
|
|
captured_params = {}
|
|
|
|
def fake_get_provider_models(provider, litellm_params=None):
|
|
captured_params["provider"] = provider
|
|
captured_params["model"] = litellm_params.model
|
|
captured_params["api_key"] = litellm_params.api_key
|
|
captured_params["api_version"] = litellm_params.api_version
|
|
captured_params["credential_name"] = litellm_params.litellm_credential_name
|
|
captured_params["has_unexpected_field"] = hasattr(
|
|
litellm_params, "unexpected_field"
|
|
)
|
|
return ["gpt-4o"]
|
|
|
|
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
|
|
|
|
result = get_known_models_from_wildcard(
|
|
wildcard_model="openai/*",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/*",
|
|
custom_llm_provider="openai",
|
|
api_version="deployment-version",
|
|
litellm_credential_name="openai-credential",
|
|
),
|
|
)
|
|
|
|
assert result == ["openai/gpt-4o"]
|
|
assert captured_params == {
|
|
"provider": "openai",
|
|
"model": "openai/*",
|
|
"api_key": "stored-openai-key",
|
|
"api_version": "deployment-version",
|
|
"credential_name": None,
|
|
"has_unexpected_field": False,
|
|
}
|
|
|
|
|
|
def test_wildcard_custom_prefix_does_not_stack_provider_prefix(monkeypatch):
|
|
"""Regression test for #30358.
|
|
|
|
A wildcard with a custom prefix (e.g. ``ollama_server1/*`` to distinguish multiple Ollama
|
|
instances) must not stack the provider's own prefix onto the expanded model ids. The expanded
|
|
ids should be ``ollama_server1/gemma3:1b`` rather than ``ollama_server1/ollama/gemma3:1b``.
|
|
"""
|
|
from litellm.proxy.auth import model_checks
|
|
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
|
from litellm.types.router import LiteLLM_Params
|
|
|
|
monkeypatch.setattr(
|
|
model_checks,
|
|
"get_provider_models",
|
|
lambda provider, litellm_params=None: ["ollama/gemma3:1b", "ollama/llama3:8b"],
|
|
)
|
|
|
|
result = get_known_models_from_wildcard(
|
|
wildcard_model="ollama_server1/*",
|
|
litellm_params=LiteLLM_Params(
|
|
model="ollama_chat/*", custom_llm_provider="ollama_chat"
|
|
),
|
|
)
|
|
|
|
assert result == ["ollama_server1/gemma3:1b", "ollama_server1/llama3:8b"]
|
|
|
|
|
|
def test_wildcard_custom_prefix_keeps_org_segment_for_non_provider_first_segment(
|
|
monkeypatch,
|
|
):
|
|
"""Only a known provider prefix should be stripped before re-prefixing.
|
|
|
|
If ``get_provider_models`` returns ids whose first segment is an org rather than a litellm
|
|
provider (e.g. ``meta-llama/Llama-3-8B``), stripping the first slash segment would drop the
|
|
org and produce an uncallable id. The org segment must be preserved.
|
|
"""
|
|
from litellm.proxy.auth import model_checks
|
|
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
|
from litellm.types.router import LiteLLM_Params
|
|
|
|
monkeypatch.setattr(
|
|
model_checks,
|
|
"get_provider_models",
|
|
lambda provider, litellm_params=None: ["meta-llama/Llama-3-8B"],
|
|
)
|
|
|
|
result = get_known_models_from_wildcard(
|
|
wildcard_model="my_hf/*",
|
|
litellm_params=LiteLLM_Params(
|
|
model="huggingface/*", custom_llm_provider="huggingface"
|
|
),
|
|
)
|
|
|
|
assert result == ["my_hf/meta-llama/Llama-3-8B"]
|
|
|
|
|
|
def test_wildcard_credential_hydration_preserves_missing_credential_name(
|
|
monkeypatch,
|
|
):
|
|
import litellm
|
|
from litellm.proxy.auth import model_checks
|
|
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
|
from litellm.types.router import LiteLLM_Params
|
|
|
|
monkeypatch.setattr(litellm, "credential_list", [])
|
|
|
|
captured_params = {}
|
|
|
|
def fake_get_provider_models(provider, litellm_params=None):
|
|
captured_params["provider"] = provider
|
|
captured_params["api_key"] = litellm_params.api_key
|
|
captured_params["credential_name"] = litellm_params.litellm_credential_name
|
|
return ["gpt-4o"]
|
|
|
|
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
|
|
|
|
result = get_known_models_from_wildcard(
|
|
wildcard_model="openai/*",
|
|
litellm_params=LiteLLM_Params(
|
|
model="openai/*",
|
|
custom_llm_provider="openai",
|
|
api_key=None,
|
|
litellm_credential_name="missing-credential",
|
|
),
|
|
)
|
|
|
|
assert result == ["openai/gpt-4o"]
|
|
assert captured_params == {
|
|
"provider": "openai",
|
|
"api_key": None,
|
|
"credential_name": "missing-credential",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_available_models_for_user_expands_query_team_wildcard(
|
|
monkeypatch,
|
|
):
|
|
import litellm
|
|
from litellm import Router
|
|
from litellm.proxy.auth import model_checks
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import get_available_models_for_user
|
|
from litellm.types.utils import CredentialItem
|
|
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"credential_list",
|
|
[
|
|
CredentialItem(
|
|
credential_name="openai-credential",
|
|
credential_info={"provider": "openai"},
|
|
credential_values={"api_key": "stored-openai-key"},
|
|
)
|
|
],
|
|
)
|
|
|
|
def fake_get_provider_models(provider, litellm_params=None):
|
|
assert litellm_params.api_key == "stored-openai-key"
|
|
assert litellm_params.litellm_credential_name is None
|
|
return ["gpt-4o-mini"]
|
|
|
|
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "model_name_team-1_generated",
|
|
"litellm_params": {
|
|
"model": "openai/*",
|
|
"custom_llm_provider": "openai",
|
|
"litellm_credential_name": "openai-credential",
|
|
},
|
|
"model_info": {
|
|
"team_id": "team-1",
|
|
"team_public_model_name": "openai/*",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
result = await get_available_models_for_user(
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-test",
|
|
models=[],
|
|
team_id="team-1",
|
|
team_models=["openai/*"],
|
|
),
|
|
llm_router=router,
|
|
general_settings={},
|
|
user_model=None,
|
|
team_id="team-1",
|
|
)
|
|
|
|
assert "openai/gpt-4o-mini" in result
|
|
|
|
|
|
def test_get_key_models_all_team_models_recursive_team():
|
|
"""GH#30619: when key and team both have all-team-models,
|
|
the sentinel should expand to proxy_model_list."""
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
from litellm.proxy._types import SpecialModelNames
|
|
|
|
user_api_key_dict = type(
|
|
"obj",
|
|
(object,),
|
|
{
|
|
"models": [SpecialModelNames.all_team_models.value],
|
|
"team_id": "team-1",
|
|
"team_models": [SpecialModelNames.all_team_models.value],
|
|
},
|
|
)()
|
|
proxy_model_list = ["model-a", "model-b"]
|
|
result = get_key_models(user_api_key_dict, proxy_model_list, {})
|
|
assert SpecialModelNames.all_team_models.value not in result
|
|
assert set(result) == {"model-a", "model-b"}
|
|
|
|
|
|
def test_get_key_models_all_team_models_keeps_mixed_team_entries():
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
from litellm.proxy._types import SpecialModelNames
|
|
|
|
user_api_key_dict = type(
|
|
"obj",
|
|
(object,),
|
|
{
|
|
"models": [SpecialModelNames.all_team_models.value],
|
|
"team_id": "team-1",
|
|
"team_models": [
|
|
SpecialModelNames.all_team_models.value,
|
|
"restricted-model",
|
|
],
|
|
},
|
|
)()
|
|
result = get_key_models(user_api_key_dict, ["model-a", "model-b"], {})
|
|
assert SpecialModelNames.all_team_models.value not in result
|
|
assert set(result) == {"model-a", "model-b", "restricted-model"}
|
|
|
|
|
|
def test_get_team_models_all_team_models_expands():
|
|
"""GH#30619: all-team-models in team_models should expand."""
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
from litellm.proxy._types import SpecialModelNames
|
|
|
|
result = get_team_models(
|
|
[SpecialModelNames.all_team_models.value],
|
|
["model-a", "model-b"],
|
|
{},
|
|
)
|
|
assert SpecialModelNames.all_team_models.value not in result
|
|
assert set(result) == {"model-a", "model-b"}
|
|
|
|
|
|
def test_get_team_models_all_team_models_expands_with_access_groups():
|
|
"""GH#30619: all-team-models with include_model_access_groups
|
|
should include access group keys."""
|
|
from litellm.proxy.auth.model_checks import get_team_models
|
|
from litellm.proxy._types import SpecialModelNames
|
|
|
|
result = get_team_models(
|
|
[SpecialModelNames.all_team_models.value],
|
|
["model-a", "model-b"],
|
|
{"group-1": ["g1-model"], "group-2": ["g2-model"]},
|
|
include_model_access_groups=True,
|
|
)
|
|
assert SpecialModelNames.all_team_models.value not in result
|
|
assert "model-a" in result
|
|
assert "model-b" in result
|
|
assert "group-1" in result
|
|
assert "group-2" in result
|
|
|
|
|
|
def test_get_key_models_teamless_all_team_models_returns_unrestricted():
|
|
"""Teamless key with all-team-models must resolve the same as leaving the
|
|
models field empty ([] = unrestricted). The sentinel must not leak into
|
|
the returned list. Fails if someone adds a team_id guard to the sentinel
|
|
expansion in get_key_models."""
|
|
from litellm.proxy._types import SpecialModelNames
|
|
from litellm.proxy.auth.model_checks import get_key_models
|
|
|
|
user_api_key_dict = type(
|
|
"obj",
|
|
(object,),
|
|
{
|
|
"models": [SpecialModelNames.all_team_models.value],
|
|
"team_id": None,
|
|
"team_models": [],
|
|
},
|
|
)()
|
|
proxy_model_list = ["gpt-4o", "claude-sonnet-4-20250514"]
|
|
result = get_key_models(user_api_key_dict, proxy_model_list, {})
|
|
assert SpecialModelNames.all_team_models.value not in result
|
|
assert result == [], "should return [] (unrestricted), same as an unscoped key"
|
|
|
|
|
|
def test_expand_wildcard_deployments_non_wildcard_passthrough():
|
|
"""Non-wildcard deployments must be returned unchanged."""
|
|
from litellm.proxy.auth.model_checks import (
|
|
expand_wildcard_deployments_for_model_info,
|
|
)
|
|
|
|
deployment = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}
|
|
result = expand_wildcard_deployments_for_model_info([deployment])
|
|
assert result == [deployment]
|
|
|
|
|
|
def test_expand_wildcard_deployments_openai_wildcard():
|
|
"""openai/* should expand into ≥1 known openai model entries."""
|
|
|
|
from litellm.proxy.auth.model_checks import (
|
|
expand_wildcard_deployments_for_model_info,
|
|
)
|
|
|
|
fake_models = ["openai/gpt-4o", "openai/gpt-4o-mini"]
|
|
deployment = {
|
|
"model_name": "openai/*",
|
|
"litellm_params": {"model": "openai/*"},
|
|
}
|
|
with patch(
|
|
"litellm.proxy.auth.model_checks.get_known_models_from_wildcard",
|
|
return_value=fake_models,
|
|
):
|
|
result = expand_wildcard_deployments_for_model_info([deployment])
|
|
|
|
assert len(result) == 2
|
|
assert all(r["model_name"] in fake_models for r in result)
|
|
assert all(r["litellm_params"]["model"] in fake_models for r in result)
|
|
|
|
|
|
def test_expand_wildcard_concrete_model_name_with_wildcard_litellm_params():
|
|
"""Concrete model_name must not be overwritten when only litellm_params.model is wildcard."""
|
|
from litellm.proxy.auth.model_checks import (
|
|
expand_wildcard_deployments_for_model_info,
|
|
)
|
|
|
|
deployment = {
|
|
"model_name": "my-custom-alias",
|
|
"litellm_params": {"model": "openai/*"},
|
|
}
|
|
result = expand_wildcard_deployments_for_model_info([deployment])
|
|
# model_name is not a wildcard, so the deployment passes through unchanged
|
|
assert result == [deployment]
|
|
|
|
|
|
def test_expand_wildcard_invalid_litellm_params_passthrough():
|
|
"""Deployments with invalid litellm_params must pass through unchanged (no 500)."""
|
|
from litellm.proxy.auth.model_checks import (
|
|
expand_wildcard_deployments_for_model_info,
|
|
)
|
|
|
|
deployment = {
|
|
"model_name": "openai/*",
|
|
"litellm_params": {
|
|
"model": "openai/*",
|
|
"max_retries": "not-an-int-field-that-breaks",
|
|
},
|
|
}
|
|
# Even if LiteLLM_Params construction fails the deployment should survive
|
|
result = expand_wildcard_deployments_for_model_info([deployment])
|
|
assert result == [deployment]
|
|
|
|
|
|
def test_get_complete_model_list_excludes_wildcard_routes_by_default():
|
|
"""Regression (LIT-4108): a wildcard with a matching router deployment leaked into /v1/models."""
|
|
from litellm import Router
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "bedrock/*",
|
|
"litellm_params": {"model": "bedrock/*"},
|
|
},
|
|
{
|
|
"model_name": "gpt-4",
|
|
"litellm_params": {"model": "openai/gpt-4"},
|
|
},
|
|
]
|
|
)
|
|
|
|
result = get_complete_model_list(
|
|
key_models=[],
|
|
team_models=[],
|
|
proxy_model_list=["bedrock/*", "gpt-4"],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
return_wildcard_routes=False,
|
|
llm_router=router,
|
|
)
|
|
|
|
assert "bedrock/*" not in result
|
|
assert "gpt-4" in result
|
|
assert any(m.startswith("bedrock/") for m in result)
|
|
|
|
|
|
def test_get_complete_model_list_excludes_wildcard_routes_without_router():
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
|
|
result = get_complete_model_list(
|
|
key_models=[],
|
|
team_models=[],
|
|
proxy_model_list=["bedrock/*", "gpt-4"],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
return_wildcard_routes=False,
|
|
llm_router=None,
|
|
)
|
|
|
|
assert "bedrock/*" not in result
|
|
assert "gpt-4" in result
|
|
assert any(m.startswith("bedrock/") for m in result)
|
|
|
|
|
|
def test_get_complete_model_list_includes_wildcard_routes_when_requested():
|
|
from litellm import Router
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "bedrock/*",
|
|
"litellm_params": {"model": "bedrock/*"},
|
|
},
|
|
]
|
|
)
|
|
|
|
result = get_complete_model_list(
|
|
key_models=[],
|
|
team_models=[],
|
|
proxy_model_list=["bedrock/*"],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
return_wildcard_routes=True,
|
|
llm_router=router,
|
|
)
|
|
|
|
assert result.count("bedrock/*") == 1
|
|
assert any(m.startswith("bedrock/") and m != "bedrock/*" for m in result)
|
|
|
|
|
|
def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion():
|
|
"""models_by_provider was a frozen import-time snapshot of set unions, so cost map
|
|
reloads (which call add_known_models) never reached wildcard expansion until a
|
|
process restart (LIT-4947)."""
|
|
import litellm
|
|
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
|
|
|
fake_model = "vertex_ai/gemini-lit4947-regression"
|
|
captured_reference = litellm.models_by_provider
|
|
assert fake_model not in litellm.models_by_provider["vertex_ai"]
|
|
try:
|
|
litellm.add_known_models(
|
|
model_cost_map={
|
|
fake_model: {"litellm_provider": "vertex_ai-language-models", "mode": "chat"}
|
|
}
|
|
)
|
|
assert fake_model in litellm.models_by_provider["vertex_ai"]
|
|
assert litellm.models_by_provider is captured_reference
|
|
assert fake_model in captured_reference["vertex_ai"]
|
|
assert fake_model in get_known_models_from_wildcard("vertex_ai/*")
|
|
finally:
|
|
litellm.vertex_language_models.discard(fake_model)
|
|
litellm.add_known_models(model_cost_map={})
|
|
assert fake_model not in litellm.models_by_provider["vertex_ai"]
|
|
|
|
|
|
def test_get_complete_model_list_drops_no_default_models_sentinel():
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
|
|
result = get_complete_model_list(
|
|
key_models=["no-default-models", "model-a"],
|
|
team_models=[],
|
|
proxy_model_list=["model-a", "model-b"],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
)
|
|
assert result == ["model-a"]
|
|
|
|
|
|
def test_get_complete_model_list_sentinel_only_grants_nothing():
|
|
from litellm.proxy.auth.model_checks import get_complete_model_list
|
|
|
|
result = get_complete_model_list(
|
|
key_models=["no-default-models"],
|
|
team_models=["no-default-models"],
|
|
proxy_model_list=["model-a", "model-b"],
|
|
user_model=None,
|
|
infer_model_from_keys=False,
|
|
)
|
|
assert result == []
|