From 52cbedf044f3094026ded2b16b9fe01d1b2c6244 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 20 Aug 2025 15:36:32 -0700 Subject: [PATCH] fix(auth_checks.py): ensure if key has access to aliased model, it still works --- litellm/proxy/_new_secret_config.yaml | 14 +- litellm/proxy/auth/auth_checks.py | 145 +++++++------ litellm/proxy/proxy_server.py | 284 ++++++++++++++------------ 3 files changed, 246 insertions(+), 197 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index d02002aeff..2675c2f687 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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"} \ No newline at end of file diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index cac24812ed..aa7a6a568a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 81ed9e68bf..591747237e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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"