scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge 5a7383ce0a into 0c98afa780

Browse source
This commit is contained in:
Nischalgouda 2026-09-12 18:30:46 +00:00 • committed by GitHub
commit 453b3777c7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 87 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -3003,6 +3003,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
team_alias: str | None = None
team_tpm_limit: int | None = None
team_rpm_limit: int | None = None
team_max_parallel_requests: int | None = None
team_max_budget: float | None = None
team_soft_budget: float | None = None
team_models: list = []
@ -3022,6 +3023,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
end_user_id: str | None = None
end_user_tpm_limit: int | None = None
end_user_rpm_limit: int | None = None
end_user_max_parallel_requests: int | None = None
end_user_max_budget: float | None = None
end_user_model_max_budget: dict | None = None
@ -3071,6 +3073,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
api_key: str | None = None
user_role: LitellmUserRoles | None = None
user_max_parallel_requests: int | None = None
allowed_model_region: AllowedModelRegion | None = None
parent_otel_span: Span | None = None
rpm_limit_per_model: dict[str, int] | None = None

17
litellm/proxy/auth/user_api_key_auth.py
View file

@ -526,6 +526,9 @@ def _apply_budget_limits_to_end_user_params(
if budget_info.rpm_limit is not None:
end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit
if budget_info.max_parallel_requests is not None:
end_user_params["end_user_max_parallel_requests"] = budget_info.max_parallel_requests
if budget_info.max_budget is not None:
end_user_params["end_user_max_budget"] = budget_info.max_budget
@ -610,6 +613,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"]
if end_user_params.get("end_user_rpm_limit") is not None:
valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"]
if end_user_params.get("end_user_max_parallel_requests") is not None:
valid_token.end_user_max_parallel_requests = end_user_params["end_user_max_parallel_requests"]
if end_user_params.get("allowed_model_region") is not None:
valid_token.allowed_model_region = end_user_params["allowed_model_region"]
if end_user_params.get("end_user_model_max_budget") is not None:
@ -1525,6 +1530,12 @@ async def _user_api_key_auth_builder(
user_id=user_id,
user_email=user_email,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_max_parallel_requests=(team_object.max_parallel_requests if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
team_metadata=(team_object.metadata if team_object is not None else None),
org_id=org_id,
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
@ -1535,11 +1546,17 @@ async def _user_api_key_auth_builder(
valid_token = UserAPIKeyAuth(
api_key=None,
team_id=team_id,
team_alias=(team_object.team_alias if team_object is not None else None),
team_tpm_limit=(team_object.tpm_limit if team_object is not None else None),
team_rpm_limit=(team_object.rpm_limit if team_object is not None else None),
team_max_parallel_requests=(team_object.max_parallel_requests if team_object is not None else None),
team_models=(team_object.models if team_object is not None else []),
user_role=(
LitellmUserRoles(user_object.user_role)
if user_object is not None and user_object.user_role is not None
else LitellmUserRoles.INTERNAL_USER
),
user_max_parallel_requests=(user_object.max_parallel_requests if user_object is not None else None),
user_id=user_id,
user_email=user_email,
org_id=org_id,

70
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -407,6 +407,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
if user_rpm_limit is None:
user_rpm_limit = sys.maxsize
user_max_parallel_requests = getattr(user_api_key_dict, "user_max_parallel_requests", sys.maxsize)
if user_max_parallel_requests is None:
user_max_parallel_requests = sys.maxsize
request_count_api_key = f"{user_id}::{precise_minute}::request_count"
# print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}")
await self.check_key_in_limits(
@ -414,7 +418,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a user
max_parallel_requests=user_max_parallel_requests, #TODO:Support max parallel requests for a user(has been implemented)
current=cache_objects["request_count_user_id"],
request_count_api_key=request_count_api_key,
tpm_limit=user_tpm_limit,
@ -435,6 +439,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
if team_rpm_limit is None:
team_rpm_limit = sys.maxsize
team_max_parallel_requests = getattr(user_api_key_dict, "team_max_parallel_requests", sys.maxsize)
if team_max_parallel_requests is None:
team_max_parallel_requests = sys.maxsize
request_count_api_key = f"{team_id}::{precise_minute}::request_count"
# print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}")
await self.check_key_in_limits(
@ -442,7 +450,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a team
max_parallel_requests=team_max_parallel_requests, # TODO: Support max parallel requests for a team(has been implemented)
current=cache_objects["request_count_team_id"],
request_count_api_key=request_count_api_key,
tpm_limit=team_tpm_limit,
@ -462,6 +470,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
if end_user_rpm_limit is None:
end_user_rpm_limit = sys.maxsize
end_user_max_parallel_requests = getattr(user_api_key_dict, "end_user_max_parallel_requests", sys.maxsize)
if end_user_max_parallel_requests is None:
end_user_max_parallel_requests = sys.maxsize
# now do the same tpm/rpm checks
request_count_api_key = f"{user_api_key_dict.end_user_id}::{precise_minute}::request_count"
@ -471,7 +483,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
cache=cache,
data=data,
call_type=call_type,
max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for an End-User
max_parallel_requests=end_user_max_parallel_requests, # TODO: Support max parallel requests for an End-User(has been implemented)
request_count_api_key=request_count_api_key,
current=cache_objects["request_count_end_user_id"],
tpm_limit=end_user_tpm_limit,
@ -762,6 +774,58 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
ttl=60,
litellm_parent_otel_span=litellm_parent_otel_span,
) # save in cache for up to 1 min.
user_api_key_user_id = _metadata.get("user_api_key_user_id", None)
if user_api_key_user_id is not None:
request_count_user = f"{user_api_key_user_id}::{precise_minute}::request_count"
current_user = await self.internal_usage_cache.async_get_cache(
key=request_count_user,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {"current_requests": 1, "current_tpm": 0, "current_rpm": 0}
new_val_user = {
"current_requests": max(current_user["current_requests"] - 1, 0),
"current_tpm": current_user["current_tpm"],
"current_rpm": current_user["current_rpm"],
}
await self.internal_usage_cache.async_set_cache(
request_count_user, new_val_user, ttl=60,
litellm_parent_otel_span=litellm_parent_otel_span,
)
user_api_key_team_id = _metadata.get("user_api_key_team_id", None)
if user_api_key_team_id is not None:
request_count_team = f"{user_api_key_team_id}::{precise_minute}::request_count"
current_team = await self.internal_usage_cache.async_get_cache(
key=request_count_team,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {"current_requests": 1, "current_tpm": 0, "current_rpm": 0}
new_val_team = {
"current_requests": max(current_team["current_requests"] - 1, 0),
"current_tpm": current_team["current_tpm"],
"current_rpm": current_team["current_rpm"],
}
await self.internal_usage_cache.async_set_cache(
request_count_team, new_val_team, ttl=60,
litellm_parent_otel_span=litellm_parent_otel_span,
)
user_api_key_end_user_id = _metadata.get("user_api_key_end_user_id", None)
if user_api_key_end_user_id is not None:
request_count_end_user = f"{user_api_key_end_user_id}::{precise_minute}::request_count"
current_end_user = await self.internal_usage_cache.async_get_cache(
key=request_count_end_user,
litellm_parent_otel_span=litellm_parent_otel_span,
) or {"current_requests": 1, "current_tpm": 0, "current_rpm": 0}
new_val_end_user = {
"current_requests": max(current_end_user["current_requests"] - 1, 0),
"current_tpm": current_end_user["current_tpm"],
"current_rpm": current_end_user["current_rpm"],
}
await self.internal_usage_cache.async_set_cache(
request_count_end_user, new_val_end_user, ttl=60,
litellm_parent_otel_span=litellm_parent_otel_span,
)
except Exception as e:
verbose_proxy_logger.exception("Inside Parallel Request Limiter: An exception occurred - %s", e)