mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 10:44:24 +00:00
Litellm dev 01 31 2025 p2 (#8164)
* docs(token_auth.md): clarify title * refactor(handle_jwt.py): add jwt auth manager + refactor to handle groups allows user to call model if user belongs to group with model access * refactor(handle_jwt.py): refactor to first check if service call then check user call * feat(handle_jwt.py): new `enforce_team_access` param only allows user to call model if a team they belong to has model access allows controlling user model access by team * fix(handle_jwt.py): fix error string, remove unecessary param * docs(token_auth.md): add controlling model access for jwt tokens via teams to docs * test: fix tests post refactor * fix: fix linting errors * fix: fix linting error * test: fix import error
This commit is contained in:
Krish Dholakia
GitHub
parent
795a71ce1e
commit
2147cad307
6 changed files with 447 additions and 242 deletions
|
|
@ -26,14 +26,10 @@ from litellm.proxy.auth.auth_checks import (
|
|||
_handle_failed_db_connection_for_get_key_object,
|
||||
_virtual_key_max_budget_check,
|
||||
_virtual_key_soft_budget_check,
|
||||
allowed_routes_check,
|
||||
can_key_call_model,
|
||||
common_checks,
|
||||
get_actual_routes,
|
||||
get_end_user_object,
|
||||
get_key_object,
|
||||
get_org_object,
|
||||
get_role_based_models,
|
||||
get_team_object,
|
||||
get_user_object,
|
||||
is_valid_fallback_model,
|
||||
|
|
@ -47,7 +43,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
route_in_additonal_public_routes,
|
||||
should_run_auth_on_pass_through_provider_route,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
|
||||
from litellm.proxy.auth.oauth2_check import check_oauth2_token
|
||||
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
|
|
@ -282,201 +278,6 @@ def get_rbac_role(jwt_handler: JWTHandler, scopes: List[str]) -> str:
|
|||
return LitellmUserRoles.TEAM
|
||||
|
||||
|
||||
def can_rbac_role_call_model(
|
||||
rbac_role: RBAC_ROLES,
|
||||
general_settings: dict,
|
||||
model: Optional[str],
|
||||
) -> Literal[True]:
|
||||
"""
|
||||
Checks if user is allowed to access the model, based on their role.
|
||||
"""
|
||||
role_based_models = get_role_based_models(
|
||||
rbac_role=rbac_role, general_settings=general_settings
|
||||
)
|
||||
if role_based_models is None or model is None:
|
||||
return True
|
||||
|
||||
if model not in role_based_models:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User role={rbac_role} not allowed to call model={model}. Allowed models={role_based_models}",
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
async def _jwt_auth_user_api_key_auth_builder(
|
||||
api_key: str,
|
||||
jwt_handler: JWTHandler,
|
||||
request_data: dict,
|
||||
general_settings: dict,
|
||||
route: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
parent_otel_span: Optional[Span],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> JWTAuthBuilderResult:
|
||||
|
||||
# check if valid token
|
||||
jwt_valid_token: dict = await jwt_handler.auth_jwt(token=api_key)
|
||||
|
||||
# check if unmatched token and enforce_rbac is true
|
||||
if jwt_handler.litellm_jwtauth.enforce_rbac is True:
|
||||
rbac_role = jwt_handler.get_rbac_role(token=jwt_valid_token)
|
||||
if rbac_role is None:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="Unmatched token passed in. enforce_rbac is set to True. Token must belong to a proxy admin, team, or user. See how to set roles in config here: https://docs.litellm.ai/docs/proxy/token_auth#advanced---spend-tracking-end-users--internal-users--team--org",
|
||||
)
|
||||
else:
|
||||
# run rbac validation checks
|
||||
can_rbac_role_call_model(
|
||||
rbac_role=rbac_role,
|
||||
general_settings=general_settings,
|
||||
model=request_data.get("model"),
|
||||
)
|
||||
# get scopes
|
||||
scopes = jwt_handler.get_scopes(token=jwt_valid_token)
|
||||
|
||||
# [OPTIONAL] allowed user email domains
|
||||
valid_user_email: Optional[bool] = None
|
||||
user_email: Optional[str] = None
|
||||
if jwt_handler.is_enforced_email_domain():
|
||||
"""
|
||||
if 'allowed_email_subdomains' is set,
|
||||
|
||||
- checks if token contains 'email' field
|
||||
- checks if 'email' is from an allowed domain
|
||||
"""
|
||||
user_email = jwt_handler.get_user_email(
|
||||
token=jwt_valid_token, default_value=None
|
||||
)
|
||||
if user_email is None:
|
||||
valid_user_email = False
|
||||
else:
|
||||
valid_user_email = jwt_handler.is_allowed_domain(user_email=user_email)
|
||||
|
||||
# [OPTIONAL] track spend against an internal employee - `LiteLLM_UserTable`
|
||||
user_object = None
|
||||
user_id = jwt_handler.get_user_id(token=jwt_valid_token, default_value=user_email)
|
||||
|
||||
# get org id
|
||||
org_id = jwt_handler.get_org_id(token=jwt_valid_token, default_value=None)
|
||||
# get team id
|
||||
team_id = jwt_handler.get_team_id(token=jwt_valid_token, default_value=None)
|
||||
# get end user id
|
||||
end_user_id = jwt_handler.get_end_user_id(token=jwt_valid_token, default_value=None)
|
||||
|
||||
# check if admin
|
||||
is_admin = jwt_handler.is_admin(scopes=scopes)
|
||||
# if admin return
|
||||
if is_admin:
|
||||
# check allowed admin routes
|
||||
is_allowed = allowed_routes_check(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_route=route,
|
||||
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
|
||||
)
|
||||
if is_allowed:
|
||||
return JWTAuthBuilderResult(
|
||||
is_proxy_admin=True,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
org_object=None,
|
||||
token=api_key,
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
end_user_id=end_user_id,
|
||||
org_id=org_id,
|
||||
)
|
||||
else:
|
||||
allowed_routes: List[Any] = jwt_handler.litellm_jwtauth.admin_allowed_routes
|
||||
actual_routes = get_actual_routes(allowed_routes=allowed_routes)
|
||||
raise Exception(
|
||||
f"Admin not allowed to access this route. Route={route}, Allowed Routes={actual_routes}"
|
||||
)
|
||||
|
||||
if team_id is None and jwt_handler.is_required_team_id() is True:
|
||||
raise Exception(
|
||||
f"No team id passed in. Field checked in jwt token - '{jwt_handler.litellm_jwtauth.team_id_jwt_field}'"
|
||||
)
|
||||
|
||||
team_object: Optional[LiteLLM_TeamTable] = None
|
||||
if team_id is not None:
|
||||
# check allowed team routes
|
||||
is_allowed = allowed_routes_check(
|
||||
user_role=LitellmUserRoles.TEAM,
|
||||
user_route=route,
|
||||
litellm_proxy_roles=jwt_handler.litellm_jwtauth,
|
||||
)
|
||||
if is_allowed is False:
|
||||
allowed_routes = jwt_handler.litellm_jwtauth.team_allowed_routes # type: ignore
|
||||
actual_routes = get_actual_routes(allowed_routes=allowed_routes)
|
||||
raise Exception(
|
||||
f"Team not allowed to access this route. Route={route}, Allowed Routes={actual_routes}"
|
||||
)
|
||||
|
||||
# check if team in db
|
||||
team_object = await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# [OPTIONAL] track spend for an org id - `LiteLLM_OrganizationTable`
|
||||
|
||||
org_object: Optional[LiteLLM_OrganizationTable] = None
|
||||
if org_id is not None:
|
||||
org_object = await get_org_object(
|
||||
org_id=org_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if user_id is not None:
|
||||
# get the user object
|
||||
user_object = await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=jwt_handler.is_upsert_user_id(
|
||||
valid_user_email=valid_user_email
|
||||
),
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
# [OPTIONAL] track spend against an external user - `LiteLLM_EndUserTable`
|
||||
end_user_object = None
|
||||
|
||||
if end_user_id is not None:
|
||||
# get the end-user object
|
||||
end_user_object = await get_end_user_object(
|
||||
end_user_id=end_user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
return JWTAuthBuilderResult(
|
||||
is_proxy_admin=False,
|
||||
team_id=team_id,
|
||||
team_object=team_object,
|
||||
user_id=user_id,
|
||||
user_object=user_object,
|
||||
org_id=org_id,
|
||||
org_object=org_object,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
token=api_key,
|
||||
)
|
||||
|
||||
|
||||
async def _user_api_key_auth_builder( # noqa: PLR0915
|
||||
request: Request,
|
||||
api_key: str,
|
||||
|
|
@ -612,7 +413,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
is_jwt = jwt_handler.is_jwt(token=api_key)
|
||||
verbose_proxy_logger.debug("is_jwt: %s", is_jwt)
|
||||
if is_jwt:
|
||||
result = await _jwt_auth_user_api_key_auth_builder(
|
||||
result = await JWTAuthManager.auth_builder(
|
||||
request_data=request_data,
|
||||
general_settings=general_settings,
|
||||
api_key=api_key,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue