mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix(proxy): constrain wildcard credential hydration
This commit is contained in:
parent
12eeece17a
commit
c08357d57e
2 changed files with 109 additions and 2 deletions
|
|
@ -8,10 +8,13 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
|||
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.router import CredentialLiteLLMParams, LiteLLM_Params
|
||||
from litellm.utils import get_valid_models
|
||||
|
||||
|
||||
_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)
|
||||
|
||||
|
||||
def _check_wildcard_routing(model: str) -> bool:
|
||||
"""
|
||||
Returns True if a model is a provider wildcard.
|
||||
|
|
@ -246,7 +249,11 @@ def _hydrate_litellm_credential_name(
|
|||
|
||||
litellm_params = litellm_params.model_copy()
|
||||
for key, value in credential_values.items():
|
||||
setattr(litellm_params, key, value)
|
||||
if (
|
||||
key in _CREDENTIAL_LITELLM_PARAM_FIELDS
|
||||
and getattr(litellm_params, key, None) is None
|
||||
):
|
||||
setattr(litellm_params, key, value)
|
||||
litellm_params.litellm_credential_name = None
|
||||
return litellm_params
|
||||
|
||||
|
|
|
|||
|
|
@ -326,6 +326,106 @@ def test_get_complete_model_list_expands_team_scoped_wildcard_with_stored_creden
|
|||
}
|
||||
|
||||
|
||||
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_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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue