Repository navigation
fix(proxy): authorize key model aliases the same way as team aliases #43049
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 8 commits
881b1cd
7a912f5
6633c31
4cd9616
387d60a
b6ee048
cf63470
d2d6c65
58eb2d4
86d5c5c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -89,6 +89,7 @@ | |
| _safe_get_request_headers, | ||
| _safe_get_request_query_params, | ||
| ) | ||
| from litellm.proxy.common_utils.model_listing_utils import alias_map | ||
| from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time | ||
| from litellm.proxy.common_utils.user_api_key_cache import ( | ||
| END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, | ||
|
|
@@ -722,6 +723,7 @@ async def _run_project_checks( | |
| model=_model, | ||
| project_object=project_object, | ||
| llm_router=llm_router, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
|
|
||
| if not skip_budget_checks: | ||
|
|
@@ -1018,6 +1020,7 @@ async def common_checks( | |
| team_object=team_object, | ||
| llm_router=llm_router, | ||
| team_model_aliases=(valid_token.team_model_aliases if valid_token else None), | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
| except ProxyException as team_denial: | ||
| if team_denial.type != ProxyErrorTypes.team_model_access_denied: | ||
|
|
@@ -1027,6 +1030,7 @@ async def common_checks( | |
| valid_token=valid_token, | ||
| team_object=team_object, | ||
| llm_router=llm_router, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ): | ||
| raise | ||
|
|
||
|
|
@@ -1043,6 +1047,7 @@ async def common_checks( | |
| proxy_logging_obj=proxy_logging_obj, | ||
| team_membership=loaded_team_membership, | ||
| team_membership_loaded=team_membership_loaded, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
|
|
||
| # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent | ||
|
|
@@ -1081,6 +1086,7 @@ async def common_checks( | |
| model=_model, | ||
| llm_router=llm_router, | ||
| user_object=user_object, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
|
|
||
| # 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget) | ||
|
|
@@ -4349,6 +4355,7 @@ def _can_object_call_model( | |
| models: list[str], | ||
| team_model_aliases: dict[str, str] | None = None, | ||
| team_id: str | None = None, | ||
| key_model_aliases: Mapping[str, str] | None = None, | ||
| object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user", | ||
| fallback_depth: int = 0, | ||
| ) -> Literal[True]: | ||
|
|
@@ -4378,6 +4385,7 @@ def _can_object_call_model( | |
| models=models, | ||
| team_model_aliases=team_model_aliases, | ||
| team_id=team_id, | ||
| key_model_aliases=key_model_aliases, | ||
| object_type=object_type, | ||
| fallback_depth=fallback_depth + 1, | ||
| ) | ||
|
|
@@ -4386,13 +4394,32 @@ def _can_object_call_model( | |
| from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent | ||
|
|
||
| compaction_parent: Final = native_compaction_parent(model) | ||
| potential_models: Final = [model, compaction_parent] if compaction_parent is not None else [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: Final = llm_router._get_model_from_alias(model) | ||
| if _model: | ||
| potential_models.append(_model) | ||
| global_or_router_alias_target: Final = ( | ||
| litellm.model_alias_map[model] | ||
| if model in litellm.model_alias_map | ||
| else ( | ||
| llm_router._get_model_from_alias(model) | ||
| if llm_router is not None and model in llm_router.model_group_alias | ||
| else None | ||
| ) | ||
| ) | ||
| after_team_alias: Final = team_model_aliases.get(model, model) if team_model_aliases else model | ||
|
greptile-apps[bot] marked this conversation as resolved.
Resolved
|
||
| after_key_alias: Final = ( | ||
| key_model_aliases.get(after_team_alias, after_team_alias) if key_model_aliases else after_team_alias | ||
| ) | ||
| after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias) | ||
| dispatched_model: Final = ( | ||
| key_model_aliases.get(after_global_alias, after_global_alias) if key_model_aliases else after_global_alias | ||
| ) | ||
| key_alias_applied: Final = after_key_alias != after_team_alias or dispatched_model != after_global_alias | ||
| potential_models: Final = ( | ||
| (dispatched_model,) | ||
| if key_alias_applied | ||
|
greptile-apps[bot] marked this conversation as resolved.
Resolved
|
||
| else ( | ||
| *((model, compaction_parent) if compaction_parent is not None else (model,)), | ||
| *((global_or_router_alias_target,) if global_or_router_alias_target else ()), | ||
| ) | ||
|
greptile-apps[bot] marked this conversation as resolved.
Resolved
|
||
| ) | ||
|
|
||
| ## check model access for alias + underlying model - allow if either is in allowed models | ||
| for m in potential_models: | ||
|
|
@@ -4444,6 +4471,8 @@ async def _check_agent_access_group_model_access( | |
| models=sorted(ceiling.models), | ||
| team_id=valid_token.team_id, | ||
| object_type="agent", | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
|
greptile-apps[bot] marked this conversation as resolved.
Resolved
cursor[bot] marked this conversation as resolved.
Resolved
|
||
| team_model_aliases=valid_token.team_model_aliases, | ||
|
greptile-apps[bot] marked this conversation as resolved.
Resolved
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Team aliases bypass agent model ceilingHigh Severity Passing Additional Locations (2)Reviewed by Cursor Bugbot for commit d2d6c65. Configure here.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed, fixed in 58eb2d4: the team alias is resolved to its target before the ceiling check, so no name admission remains |
||
| ) | ||
|
|
||
|
|
||
|
|
@@ -4471,12 +4500,14 @@ async def _check_agent_caller_model_access( | |
| if caller_auth is None: | ||
| return | ||
| caller_team: Final = await load_team(valid_token) | ||
| caller_key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token) | ||
| if caller_team is not None: | ||
| await can_team_access_model( | ||
| model=model, | ||
| team_object=caller_team, | ||
| llm_router=llm_router, | ||
| prisma_client=prisma_client, | ||
| key_model_aliases=caller_key_model_aliases, | ||
| ) | ||
| await _check_team_member_model_access( | ||
| model=model, | ||
|
|
@@ -4486,12 +4517,18 @@ async def _check_agent_caller_model_access( | |
| prisma_client=prisma_client, | ||
| user_api_key_cache=user_api_key_cache, | ||
| proxy_logging_obj=proxy_logging_obj, | ||
| key_model_aliases=caller_key_model_aliases, | ||
| ) | ||
| return | ||
| caller_user: Final = await load_user(valid_token) | ||
| if caller_user is None: | ||
| return | ||
| await can_user_call_model(model=model, llm_router=llm_router, user_object=caller_user) | ||
|
cursor[bot] marked this conversation as resolved.
Resolved
|
||
| await can_user_call_model( | ||
| model=model, | ||
| llm_router=llm_router, | ||
| user_object=caller_user, | ||
| key_model_aliases=caller_key_model_aliases, | ||
| ) | ||
|
|
||
|
|
||
| def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: | ||
|
|
@@ -4512,6 +4549,10 @@ def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None | |
| return False | ||
|
|
||
|
|
||
| def key_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth | None) -> Mapping[str, str] | None: | ||
| return alias_map(valid_token.aliases) if valid_token is not None and valid_token.aliases else None | ||
|
|
||
|
|
||
| def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]: | ||
| """ | ||
| Expand key model sentinels before auth checks. | ||
|
|
@@ -4831,6 +4872,7 @@ async def can_key_call_model( | |
| models=key_models, | ||
| team_model_aliases=valid_token.team_model_aliases, | ||
| team_id=valid_token.team_id, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| object_type="key", | ||
| ) | ||
| except ProxyException: | ||
|
|
@@ -4848,6 +4890,7 @@ async def can_key_call_model( | |
| models=models_from_groups, | ||
| team_model_aliases=valid_token.team_model_aliases, | ||
| team_id=valid_token.team_id, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| object_type="key", | ||
| ) | ||
| raise | ||
|
|
@@ -4906,6 +4949,7 @@ async def can_key_call_resolved_model( | |
| team_object=team_object, | ||
| llm_router=llm_router, | ||
| team_model_aliases=valid_token.team_model_aliases, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
| except ProxyException as team_denial: | ||
| if team_denial.type != ProxyErrorTypes.team_model_access_denied: | ||
|
|
@@ -4915,6 +4959,7 @@ async def can_key_call_resolved_model( | |
| valid_token=valid_token, | ||
| team_object=team_object, | ||
| llm_router=llm_router, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ): | ||
| raise | ||
|
|
||
|
|
@@ -4927,6 +4972,7 @@ async def can_key_call_resolved_model( | |
| prisma_client=prisma_client, | ||
| user_api_key_cache=user_api_key_cache, | ||
| proxy_logging_obj=proxy_logging_obj, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
|
|
||
| if valid_token.project_id is not None: | ||
|
|
@@ -4941,6 +4987,7 @@ async def can_key_call_resolved_model( | |
| model=model, | ||
| project_object=project_object, | ||
| llm_router=llm_router, | ||
| key_model_aliases=key_model_aliases_for_auth_check(valid_token), | ||
| ) | ||
|
|
||
|
|
||
|
|
@@ -4968,6 +5015,7 @@ async def can_team_access_model( | |
| team_object: LiteLLM_TeamTable | None, | ||
| llm_router: Router | None, | ||
| team_model_aliases: dict[str, str] | None = None, | ||
| key_model_aliases: Mapping[str, str] | None = None, | ||
| prisma_client: DatabaseClient | None = None, | ||
| ) -> Literal[True]: | ||
| """ | ||
|
|
@@ -4983,6 +5031,7 @@ async def can_team_access_model( | |
| models=team_object.models if team_object else [], | ||
| team_model_aliases=team_model_aliases, | ||
| team_id=team_object.team_id if team_object else None, | ||
| key_model_aliases=key_model_aliases, | ||
| object_type="team", | ||
| ) | ||
| except ProxyException: | ||
|
|
@@ -5000,6 +5049,7 @@ async def can_team_access_model( | |
| models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])), | ||
| team_model_aliases=team_model_aliases, | ||
| team_id=team_object.team_id if team_object else None, | ||
| key_model_aliases=key_model_aliases, | ||
| object_type="team", | ||
| ) | ||
| raise | ||
|
|
@@ -5058,6 +5108,7 @@ async def _key_access_group_grants_model( | |
| valid_token: UserAPIKeyAuth | None, | ||
| team_object: LiteLLM_TeamTable | None, | ||
| llm_router: Router | None, | ||
| key_model_aliases: Mapping[str, str] | None = None, | ||
| ) -> bool: | ||
| """ | ||
| Returns True if the key's `access_group_ids` expand to models that grant | ||
|
|
@@ -5078,6 +5129,7 @@ async def _key_access_group_grants_model( | |
| models=authorized_models, | ||
| team_model_aliases=valid_token.team_model_aliases if valid_token else None, | ||
| team_id=valid_token.team_id if valid_token else None, | ||
| key_model_aliases=key_model_aliases, | ||
| object_type="key", | ||
| ) | ||
| return True | ||
|
|
@@ -5089,6 +5141,7 @@ def can_project_access_model( | |
| model: str | list[str], | ||
| project_object: LiteLLM_ProjectTable, | ||
| llm_router: Router | None, | ||
| key_model_aliases: Mapping[str, str] | None = None, | ||
| ) -> Literal[True]: | ||
| """ | ||
| Returns True if the project can access a specific model. | ||
|
|
@@ -5099,6 +5152,7 @@ def can_project_access_model( | |
| model=model, | ||
| llm_router=llm_router, | ||
| models=project_object.models if project_object else [], | ||
| key_model_aliases=key_model_aliases, | ||
| object_type="project", | ||
| ) | ||
|
|
||
|
|
@@ -5107,6 +5161,7 @@ async def can_user_call_model( | |
| model: str | list[str], | ||
| llm_router: Router | None, | ||
| user_object: LiteLLM_UserTable | None, | ||
| key_model_aliases: Mapping[str, str] | None = None, | ||
| ) -> Literal[True]: | ||
| if user_object is None: | ||
| return True | ||
|
|
@@ -5128,6 +5183,7 @@ async def can_user_call_model( | |
| model=model, | ||
| llm_router=llm_router, | ||
| models=user_object.models, | ||
| key_model_aliases=key_model_aliases, | ||
| object_type="user", | ||
| ) | ||
|
|
||
|
|
@@ -5682,6 +5738,7 @@ async def _check_team_member_model_access( | |
| proxy_logging_obj: ProxyLogging, | ||
| team_membership: LiteLLM_TeamMembership | None = None, | ||
| team_membership_loaded: bool = False, | ||
| key_model_aliases: Mapping[str, str] | None = None, | ||
| ) -> None: | ||
| """ | ||
| Check if a team member's per-member model scope allows access to the requested model. | ||
|
|
@@ -5717,6 +5774,7 @@ async def _check_team_member_model_access( | |
| models=member_allowed_models, | ||
| object_type="team", | ||
| team_id=team_object.team_id, | ||
| key_model_aliases=key_model_aliases, | ||
| ) | ||
| except ProxyException: | ||
| internal_message: Final = ( | ||
|
|
||


Uh oh!
There was an error while loading. Please reload this page.