add check for _model_is_within_list_of_allowed_models

This commit is contained in:
Ishaan Jaff 2024-11-25 13:58:08 -08:00
parent 0e36333051
commit e9cdbff75a

View file

@ -92,17 +92,15 @@ def common_checks( # noqa: PLR0915
): ):
# this means the team has access to all models on the proxy # this means the team has access to all models on the proxy
if ( if (
"all-proxy-models" in team_object.models _model_is_within_list_of_allowed_models(
or "*" in team_object.models model=_model, allowed_models=team_object.models
or "openai/*" in team_object.models )
is True
): ):
# this means the team has access to all models on the proxy
pass pass
# check if the team model is an access_group # check if the team model is an access_group
elif model_in_access_group(_model, team_object.models) is True: elif model_in_access_group(_model, team_object.models) is True:
pass pass
elif _model and "*" in _model:
pass
else: else:
raise Exception( raise Exception(
f"Team={team_object.team_id} not allowed to call model={_model}. Allowed team models = {team_object.models}" f"Team={team_object.team_id} not allowed to call model={_model}. Allowed team models = {team_object.models}"
@ -868,18 +866,16 @@ async def can_key_call_model(
filtered_models += models_in_current_access_groups filtered_models += models_in_current_access_groups
verbose_proxy_logger.debug(f"model: {model}; allowed_models: {filtered_models}") verbose_proxy_logger.debug(f"model: {model}; allowed_models: {filtered_models}")
# Check for universal access patterns if (
if len(filtered_models) == 0: _model_is_within_list_of_allowed_models(
return True model=model, allowed_models=filtered_models
if "*" in filtered_models:
return True
if model_matches_patterns(model=model, allowed_models=filtered_models) is True:
return True
if model is not None and model not in filtered_models:
raise ValueError(
f"API Key not allowed to access model. This token can only access models={valid_token.models}. Tried to access {model}"
) )
is False
):
raise ValueError(
f"API Key not allowed to access model. List of allowed models={filtered_models}. Tried to access {model}"
)
valid_token.models = filtered_models valid_token.models = filtered_models
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
f"filtered allowed_models: {filtered_models}; valid_token.models: {valid_token.models}" f"filtered allowed_models: {filtered_models}; valid_token.models: {valid_token.models}"
@ -887,6 +883,22 @@ async def can_key_call_model(
return True return True
def _model_is_within_list_of_allowed_models(
model: str, allowed_models: List[str]
) -> bool:
# Check for universal access patterns
if len(allowed_models) == 0:
return True
if "*" in allowed_models:
return True
if "all-proxy-models" in allowed_models:
return True
if model_matches_patterns(model=model, allowed_models=allowed_models) is True:
return True
return False
def model_matches_patterns(model: str, allowed_models: List[str]) -> bool: def model_matches_patterns(model: str, allowed_models: List[str]) -> bool:
""" """
Helper function to check if a model matches any of the allowed model patterns. Helper function to check if a model matches any of the allowed model patterns.