tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-21 08:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(auth_checks.py): ensure if key has access to aliased model, it still works

Browse Source
This commit is contained in:
Krrish Dholakia
2025-08-20 15:36:32 -07:00
parent b0c98efd14
commit 52cbedf044
3 changed files with 246 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_new_secret_config.yaml
+3 -11
View File
@@ -3,7 +3,7 @@ model_list:
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://webhook.site/4feb0d46-4b23-468c-bf55-7008b5deb36d
api_base: https://exampleopenaiendpoint-production.up.railway.app/
- model_name: gpt-5-mini
litellm_params:
model: azure/gpt-5-mini
@@ -14,13 +14,5 @@ model_list:
model_info:
mode: chat
litellm_settings:
cache: true
cache_params:
type: redis
ttl: 600
supported_call_types: ["acompletion", "completion"]
model_group_settings:
forward_client_headers_to_llm_api:
- fake-openai-endpoint
router_settings:
model_group_alias: {"my-fake-gpt-4": "fake-openai-endpoint"}
litellm/proxy/auth/auth_checks.py
+87 -58
View File
@@ -513,19 +513,19 @@ async def get_team_membership(
) -> Optional["LiteLLM_TeamMembership"]:
"""
Returns team membership object if user is member of team.
Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership).
"""
from litellm.proxy._types import LiteLLM_TeamMembership
if prisma_client is None:
raise Exception("No db connected")
if user_id is None or team_id is None:
return None
_key = "team_membership:{}:{}".format(user_id, team_id)
# check if in cache
cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key)
if cached_membership_obj is not None:
@@ -533,7 +533,7 @@ async def get_team_membership(
return LiteLLM_TeamMembership(**cached_membership_obj)
elif isinstance(cached_membership_obj, LiteLLM_TeamMembership):
return cached_membership_obj
# else, check db
try:
response = await prisma_client.db.litellm_teammembership.find_unique(
@@ -545,15 +545,17 @@ async def get_team_membership(
return None
# save the team membership object to cache
await user_api_key_cache.async_set_cache(
key=_key, value=response
)
await user_api_key_cache.async_set_cache(key=_key, value=response)
_response = LiteLLM_TeamMembership(**response.dict())
return _response
except Exception:
verbose_proxy_logger.exception("Error getting team membership for user_id: %s, team_id: %s", user_id, team_id)
verbose_proxy_logger.exception(
"Error getting team membership for user_id: %s, team_id: %s",
user_id,
team_id,
)
return None
@@ -1203,57 +1205,14 @@ async def get_org_object(
)
def _can_object_call_model(
model: Union[str, List[str]],
def _check_model_access_helper(
model: str,
llm_router: Optional[Router],
models: List[str],
team_model_aliases: Optional[Dict[str, str]] = None,
team_id: Optional[str] = None,
object_type: Literal["user", "team", "key", "org"] = "user",
fallback_depth: int = 0,
) -> Literal[True]:
"""
Checks if token can call a given model
Args:
- model: str
- llm_router: Optional[Router]
- models: List[str]
- team_model_aliases: Optional[Dict[str, str]]
- object_type: Literal["user", "team", "key", "org"]. We use the object type to raise the correct exception type
Returns:
- True: if token allowed to call model
Raises:
- Exception: If token not allowed to call model
"""
if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH:
raise Exception(
"Unable to parse model, max fallback depth exceeded - received model: {}".format(
model
)
)
if isinstance(model, list):
for m in model:
_can_object_call_model(
model=m,
llm_router=llm_router,
models=models,
team_model_aliases=team_model_aliases,
team_id=team_id,
object_type=object_type,
fallback_depth=fallback_depth + 1,
)
return True
if model in litellm.model_alias_map:
model = litellm.model_alias_map[model]
elif llm_router and model in llm_router.model_group_alias:
_model = llm_router._get_model_from_alias(model)
if _model:
model = _model
## check if model in allowed model names
from collections import defaultdict
@@ -1301,13 +1260,83 @@ def _can_object_call_model(
param="model",
code=status.HTTP_401_UNAUTHORIZED,
)
verbose_proxy_logger.debug(
f"filtered allowed_models: {filtered_models}; models: {models}"
)
return True
def _can_object_call_model(
model: Union[str, List[str]],
llm_router: Optional[Router],
models: List[str],
team_model_aliases: Optional[Dict[str, str]] = None,
team_id: Optional[str] = None,
object_type: Literal["user", "team", "key", "org"] = "user",
fallback_depth: int = 0,
) -> Literal[True]:
"""
Checks if token can call a given model
Args:
- model: str
- llm_router: Optional[Router]
- models: List[str]
- team_model_aliases: Optional[Dict[str, str]]
- object_type: Literal["user", "team", "key", "org"]. We use the object type to raise the correct exception type
Returns:
- True: if token allowed to call model
Raises:
- Exception: If token not allowed to call model
"""
if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH:
raise Exception(
"Unable to parse model, max fallback depth exceeded - received model: {}".format(
model
)
)
if isinstance(model, list):
for m in model:
_can_object_call_model(
model=m,
llm_router=llm_router,
models=models,
team_model_aliases=team_model_aliases,
team_id=team_id,
object_type=object_type,
fallback_depth=fallback_depth + 1,
)
return True
potential_models = [model]
if model in litellm.model_alias_map:
potential_models.append(litellm.model_alias_map[model])
elif llm_router and model in llm_router.model_group_alias:
_model = llm_router._get_model_from_alias(model)
if _model:
potential_models.append(_model)
## check model access for alias + underlying model - allow if either is in allowed models
for m in potential_models:
if _check_model_access_helper(
model=m,
llm_router=llm_router,
models=models,
team_model_aliases=team_model_aliases,
team_id=team_id,
object_type=object_type,
):
return True
raise ProxyException(
message=f"{object_type} not allowed to access model. This {object_type} can only access models={models}. Tried to access {model}",
type=ProxyErrorTypes.get_model_access_error_type_for_object(
object_type=object_type
),
param="model",
code=status.HTTP_401_UNAUTHORIZED,
)
def _model_in_team_aliases(
model: str, team_model_aliases: Optional[Dict[str, str]] = None
) -> bool:
litellm/proxy/proxy_server.py
+156 -128
View File
@@ -2186,6 +2186,7 @@ class ProxyConfig:
## ROUTER SETTINGS (e.g. routing_strategy, ...)
router_settings = config.get("router_settings", None)
if router_settings and isinstance(router_settings, dict):
arg_spec = inspect.getfullargspec(litellm.Router)
# model list already set
@@ -2664,6 +2665,11 @@ class ProxyConfig:
) -> None:
"""
Adds router settings from DB config to litellm proxy
1. Get router settings from DB
2. Get router settings from config
3. Combine both
4. Update router settings
"""
if llm_router is not None and prisma_client is not None:
db_router_settings = await prisma_client.db.litellm_config.find_first(
@@ -2958,81 +2964,96 @@ class ProxyConfig:
config_record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": "model_cost_map_reload_config"}
)
if config_record is None or config_record.param_value is None:
return # No configuration found, skip reload
config = config_record.param_value
interval_hours = config.get("interval_hours")
force_reload = config.get("force_reload", False)
if interval_hours is None and force_reload is False:
if interval_hours is None and force_reload is False:
return # No interval configured, skip reload
current_time = datetime.utcnow()
# Check if we need to reload based on interval or force reload
should_reload = False
if force_reload:
should_reload = True
verbose_proxy_logger.info("Model cost map reload triggered by force reload flag")
verbose_proxy_logger.info(
"Model cost map reload triggered by force reload flag"
)
elif interval_hours is not None:
# Use pod's in-memory last reload time
global last_model_cost_map_reload
if last_model_cost_map_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_model_cost_map_reload)
last_reload_time = datetime.fromisoformat(
last_model_cost_map_reload
)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
hours_since_last_reload = (
time_since_last_reload.total_seconds() / 3600
)
if hours_since_last_reload >= interval_hours:
should_reload = True
verbose_proxy_logger.info(f"Model cost map reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}")
verbose_proxy_logger.info(
f"Model cost map reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}"
)
except Exception as e:
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
verbose_proxy_logger.warning(
f"Error parsing last reload time: {e}"
)
# If we can't parse the last reload time, reload anyway
should_reload = True
else:
# No last reload time recorded, reload now
should_reload = True
verbose_proxy_logger.info("Model cost map reload triggered - no previous reload time recorded")
verbose_proxy_logger.info(
"Model cost map reload triggered - no previous reload time recorded"
)
if should_reload:
# Perform the reload
from litellm.litellm_core_utils.get_model_cost_map import (
get_model_cost_map,
)
model_cost_map_url = litellm.model_cost_map_url
new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
litellm.model_cost = new_model_cost_map
# Update pod's in-memory last reload time
last_model_cost_map_reload = current_time.isoformat()
# Clear force reload flag in database
await prisma_client.db.litellm_config.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": safe_dumps({
"interval_hours": interval_hours,
"force_reload": False
})
"param_value": safe_dumps(
{
"interval_hours": interval_hours,
"force_reload": False,
}
),
},
"update": {
"param_value": safe_dumps({
"force_reload": False
})
}
}
"update": {"param_value": safe_dumps({"force_reload": False})},
},
)
verbose_proxy_logger.info(f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}")
verbose_proxy_logger.info(
f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}"
)
except Exception as e:
verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {str(e)}")
verbose_proxy_logger.exception(
f"Error in _check_and_reload_model_cost_map: {str(e)}"
)
async def _init_prompts_in_db(self, prisma_client: PrismaClient):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
@@ -3631,7 +3652,6 @@ class ProxyStartupEvent:
args=[prisma_client, db_writer_client, proxy_logging_obj],
)
### ADD NEW MODELS ###
store_model_in_db = (
get_secret_bool("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db
@@ -3972,18 +3992,15 @@ async def model_info(
# Get provider information from the router deployment
if llm_router is None:
raise HTTPException(
status_code=500,
detail="Router not initialized"
)
raise HTTPException(status_code=500, detail="Router not initialized")
deployment = llm_router.get_deployment_by_model_group_name(model_id)
if deployment is None:
raise HTTPException(
status_code=404,
detail=f"Model '{model_id}' not found in router configuration"
detail=f"Model '{model_id}' not found in router configuration",
)
# Use the actual litellm model from the deployment to get provider info
_, provider, _, _ = litellm.get_llm_provider(model=deployment.litellm_params.model)
@@ -5750,7 +5767,9 @@ async def run_thread(
from litellm.llms.base_llm.base_utils import BaseTokenCounter
def _get_provider_token_counter(deployment: dict, model_to_use: str) -> Tuple[Optional[BaseTokenCounter], Optional[str], Optional[str]]:
def _get_provider_token_counter(
deployment: dict, model_to_use: str
) -> Tuple[Optional[BaseTokenCounter], Optional[str], Optional[str]]:
"""
Auto-route to the correct provider's token counter based on model/deployment.
Uses the existing get_provider_model_info infrastructure with switch-case pattern.
@@ -5788,7 +5807,11 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str) -> Tuple[Op
model=full_model, provider=llm_provider_enum
)
if provider_model_info is not None:
return provider_model_info.get_token_counter(), model, custom_llm_provider
return (
provider_model_info.get_token_counter(),
model,
custom_llm_provider,
)
except Exception:
# If provider detection fails, fall back to manual checks
@@ -5807,10 +5830,7 @@ def _get_provider_token_counter(deployment: dict, model_to_use: str) -> Tuple[Op
dependencies=[Depends(user_api_key_auth)],
response_model=TokenCountResponse,
)
async def token_counter(
request: TokenCountRequest,
call_endpoint: bool = False
):
async def token_counter(request: TokenCountRequest, call_endpoint: bool = False):
"""
Args:
request: TokenCountRequest
@@ -5865,12 +5885,19 @@ async def token_counter(
custom_llm_provider: Optional[str] = None
if call_endpoint is True and deployment is not None:
# Auto-route to the correct provider based on model
provider_counter, _model, custom_llm_provider = _get_provider_token_counter(deployment, model_to_use)
provider_counter, _model, custom_llm_provider = _get_provider_token_counter(
deployment, model_to_use
)
if _model is not None:
model_to_use = _model
if provider_counter is not None:
if provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider) is True:
if (
provider_counter.should_use_token_counting_api(
custom_llm_provider=custom_llm_provider
)
is True
):
result = await provider_counter.count_tokens(
model_to_use=model_to_use or "",
messages=messages, # type: ignore
@@ -9019,7 +9046,7 @@ async def get_litellm_model_cost_map(
status_code=403,
detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}",
)
try:
_model_cost_map = litellm.model_cost
return _model_cost_map
@@ -9041,7 +9068,7 @@ async def reload_model_cost_map(
):
"""
ADMIN ONLY / MASTER KEY Only Endpoint
Manually reload the model cost map from the remote source.
This will fetch fresh pricing data from the model_prices_and_context_window.json file.
"""
@@ -9051,59 +9078,55 @@ async def reload_model_cost_map(
status_code=403,
detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}",
)
try:
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail="Database connection not available"
status_code=500, detail="Database connection not available"
)
# Immediately reload the model cost map in the current pod
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
model_cost_map_url = litellm.model_cost_map_url
new_model_cost_map = get_model_cost_map(url=model_cost_map_url)
litellm.model_cost = new_model_cost_map
# Update pod's in-memory last reload time
global last_model_cost_map_reload
current_time = datetime.utcnow()
last_model_cost_map_reload = current_time.isoformat()
# Set force reload flag in database for other pods
await prisma_client.db.litellm_config.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": safe_dumps({
"interval_hours": None,
"force_reload": True
})
"param_value": safe_dumps(
{"interval_hours": None, "force_reload": True}
),
},
"update": {
"param_value": safe_dumps({
"force_reload": True
})
}
}
"update": {"param_value": safe_dumps({"force_reload": True})},
},
)
models_count = len(new_model_cost_map) if new_model_cost_map else 0
verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}")
verbose_proxy_logger.info(
f"Model cost map reloaded successfully in current pod. Models count: {models_count}"
)
return {
"message": f"Price data reloaded successfully! {models_count} models updated.",
"status": "success",
"models_count": models_count,
"timestamp": current_time.isoformat()
"timestamp": current_time.isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to reload model cost map: {str(e)}")
raise HTTPException(
status_code=500,
detail=f"Failed to reload model cost map: {str(e)}"
status_code=500, detail=f"Failed to reload model cost map: {str(e)}"
)
@@ -9119,7 +9142,7 @@ async def schedule_model_cost_map_reload(
):
"""
ADMIN ONLY / MASTER KEY Only Endpoint
Schedule periodic reload of the model cost map.
This will create a background job that reloads the model cost map every specified hours.
"""
@@ -9129,54 +9152,52 @@ async def schedule_model_cost_map_reload(
status_code=403,
detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}",
)
if hours <= 0:
raise HTTPException(
status_code=400,
detail="Hours must be greater than 0"
)
raise HTTPException(status_code=400, detail="Hours must be greater than 0")
try:
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail="Database connection not available"
status_code=500, detail="Database connection not available"
)
# Update database with new reload configuration
await prisma_client.db.litellm_config.upsert(
where={"param_name": "model_cost_map_reload_config"},
data={
"create": {
"param_name": "model_cost_map_reload_config",
"param_value": safe_dumps({
"interval_hours": hours,
"force_reload": False
})
"param_value": safe_dumps(
{"interval_hours": hours, "force_reload": False}
),
},
"update": {
"param_value": safe_dumps({
"interval_hours": hours,
"force_reload": False
})
}
}
"param_value": safe_dumps(
{"interval_hours": hours, "force_reload": False}
)
},
},
)
verbose_proxy_logger.info(f"Model cost map reload scheduled for every {hours} hours")
verbose_proxy_logger.info(
f"Model cost map reload scheduled for every {hours} hours"
)
return {
"message": f"Model cost map reload scheduled for every {hours} hours",
"status": "success",
"interval_hours": hours,
"timestamp": datetime.utcnow().isoformat()
"timestamp": datetime.utcnow().isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to schedule model cost map reload: {str(e)}")
verbose_proxy_logger.exception(
f"Failed to schedule model cost map reload: {str(e)}"
)
raise HTTPException(
status_code=500,
detail=f"Failed to schedule model cost map reload: {str(e)}"
detail=f"Failed to schedule model cost map reload: {str(e)}",
)
@@ -9191,7 +9212,7 @@ async def cancel_model_cost_map_reload(
):
"""
ADMIN ONLY / MASTER KEY Only Endpoint
Cancel the scheduled periodic reload of the model cost map.
"""
# Check if user is admin
@@ -9200,32 +9221,32 @@ async def cancel_model_cost_map_reload(
status_code=403,
detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}",
)
try:
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail="Database connection not available"
status_code=500, detail="Database connection not available"
)
# Remove reload configuration from database
await prisma_client.db.litellm_config.delete(
where={"param_name": "model_cost_map_reload_config"}
)
verbose_proxy_logger.info("Model cost map reload schedule cancelled")
return {
"message": "Model cost map reload schedule cancelled",
"status": "success",
"timestamp": datetime.utcnow().isoformat()
"timestamp": datetime.utcnow().isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to cancel model cost map reload: {str(e)}")
verbose_proxy_logger.exception(
f"Failed to cancel model cost map reload: {str(e)}"
)
raise HTTPException(
status_code=500,
detail=f"Failed to cancel model cost map reload: {str(e)}"
status_code=500, detail=f"Failed to cancel model cost map reload: {str(e)}"
)
@@ -9240,7 +9261,7 @@ async def get_model_cost_map_reload_status(
):
"""
ADMIN ONLY / MASTER KEY Only Endpoint
Get the status of the scheduled model cost map reload job.
"""
# Check if user is admin
@@ -9249,75 +9270,82 @@ async def get_model_cost_map_reload_status(
status_code=403,
detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}",
)
try:
global prisma_client, last_model_cost_map_reload
verbose_proxy_logger.info(f"Checking model cost map reload status. Last reload: {last_model_cost_map_reload}")
verbose_proxy_logger.info(
f"Checking model cost map reload status. Last reload: {last_model_cost_map_reload}"
)
if prisma_client is None:
verbose_proxy_logger.info("No database connection, returning not scheduled")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None
"next_run": None,
}
# Get reload configuration from database
config_record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": "model_cost_map_reload_config"}
)
if config_record is None or config_record.param_value is None:
verbose_proxy_logger.info("No model cost map reload configuration found")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None
"next_run": None,
}
config = config_record.param_value
interval_hours = config.get("interval_hours")
if interval_hours is None:
verbose_proxy_logger.info("No interval configured, returning not scheduled")
return {
"scheduled": False,
"interval_hours": None,
"last_run": None,
"next_run": None
"next_run": None,
}
current_time = datetime.utcnow()
next_run = None
# Use pod's in-memory last reload time
if last_model_cost_map_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_model_cost_map_reload)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600
if hours_since_last_reload < interval_hours:
next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat()
next_run = (
last_reload_time + timedelta(hours=interval_hours)
).isoformat()
except Exception as e:
verbose_proxy_logger.warning(f"Error parsing last reload time: {e}")
return {
"scheduled": True,
"interval_hours": interval_hours,
"last_run": last_model_cost_map_reload,
"next_run": next_run
"next_run": next_run,
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to get model cost map reload status: {str(e)}")
verbose_proxy_logger.exception(
f"Failed to get model cost map reload status: {str(e)}"
)
raise HTTPException(
status_code=500,
detail=f"Failed to get model cost map reload status: {str(e)}"
detail=f"Failed to get model cost map reload status: {str(e)}",
)
@router.get("/", dependencies=[Depends(user_api_key_auth)])
async def home(request: Request):
return "LiteLLM: RUNNING"