scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

test fix test_aaauser_personal_budgets

Browse source
This commit is contained in:
Ishaan Jaff 2025-06-06 20:55:27 -07:00
parent bc835c6044
commit 18081cf250
1 changed files with 22 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -51,7 +51,10 @@ from litellm.proxy.auth.auth_utils import (
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.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
@ -587,7 +590,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
_end_user_object = None
end_user_params = {}
end_user_id = get_end_user_id_from_request_body(request_data, dict(request.headers))
end_user_id = get_end_user_id_from_request_body(
request_data, _safe_get_request_headers(request)
)
if end_user_id:
try:
end_user_params["end_user_id"] = end_user_id
@ -601,23 +606,23 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
proxy_logging_obj=proxy_logging_obj,
)
if _end_user_object is not None:
end_user_params[
"allowed_model_region"
] = _end_user_object.allowed_model_region
end_user_params["allowed_model_region"] = (
_end_user_object.allowed_model_region
)
if _end_user_object.litellm_budget_table is not None:
budget_info = _end_user_object.litellm_budget_table
if budget_info.tpm_limit is not None:
end_user_params[
"end_user_tpm_limit"
] = budget_info.tpm_limit
end_user_params["end_user_tpm_limit"] = (
budget_info.tpm_limit
)
if budget_info.rpm_limit is not None:
end_user_params[
"end_user_rpm_limit"
] = budget_info.rpm_limit
end_user_params["end_user_rpm_limit"] = (
budget_info.rpm_limit
)
if budget_info.max_budget is not None:
end_user_params[
"end_user_max_budget"
] = budget_info.max_budget
end_user_params["end_user_max_budget"] = (
budget_info.max_budget
)
except Exception as e:
if isinstance(e, litellm.BudgetExceededError):
raise e
@ -1140,7 +1145,9 @@ async def user_api_key_auth(
custom_litellm_key_header=custom_litellm_key_header,
)
end_user_id = get_end_user_id_from_request_body(request_data, dict(request.headers))
end_user_id = get_end_user_id_from_request_body(
request_data, _safe_get_request_headers(request)
)
if end_user_id is not None:
user_api_key_auth_obj.end_user_id = end_user_id