phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix _return_user_api_key_auth_obj (#7591)

Browse source
This commit is contained in:
Ishaan Jaff 2025-01-06 16:43:14 -08:00 • committed by GitHub
parent 7ac1c0c1b2
commit 1bea935889
2 changed files with 17 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -732,7 +732,7 @@ async def user_api_key_auth( # noqa: PLR0915
) )
if is_master_key_valid: if is_master_key_valid:
_user_api_key_obj = _return_user_api_key_auth_obj( _user_api_key_obj = await _return_user_api_key_auth_obj(
user_obj=None, user_obj=None,
user_role=LitellmUserRoles.PROXY_ADMIN, user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=master_key, api_key=master_key,
@ -1177,7 +1177,7 @@ async def user_api_key_auth( # noqa: PLR0915
# No token was found when looking up in the DB # No token was found when looking up in the DB
raise Exception("Invalid proxy server token passed") raise Exception("Invalid proxy server token passed")
if valid_token_dict is not None: if valid_token_dict is not None:
return _return_user_api_key_auth_obj( return await _return_user_api_key_auth_obj(
user_obj=user_obj, user_obj=user_obj,
api_key=api_key, api_key=api_key,
parent_otel_span=parent_otel_span, parent_otel_span=parent_otel_span,
@ -1240,7 +1240,7 @@ async def user_api_key_auth( # noqa: PLR0915
) )
def _return_user_api_key_auth_obj( async def _return_user_api_key_auth_obj(
user_obj: Optional[LiteLLM_UserTable], user_obj: Optional[LiteLLM_UserTable],
api_key: str, api_key: str,
parent_otel_span: Optional[Span], parent_otel_span: Optional[Span],
@ -1250,14 +1250,18 @@ def _return_user_api_key_auth_obj(
user_role: Optional[LitellmUserRoles] = None, user_role: Optional[LitellmUserRoles] = None,
) -> UserAPIKeyAuth: ) -> UserAPIKeyAuth:
end_time = datetime.now() end_time = datetime.now()
user_api_key_service_logger_obj.service_success_hook(
service=ServiceTypes.AUTH, asyncio.create_task(
call_type=route, user_api_key_service_logger_obj.async_service_success_hook(
start_time=start_time, service=ServiceTypes.AUTH,
end_time=end_time, call_type=route,
duration=end_time.timestamp() - start_time.timestamp(), start_time=start_time,
parent_otel_span=parent_otel_span, end_time=end_time,
duration=end_time.timestamp() - start_time.timestamp(),
parent_otel_span=parent_otel_span,
)
) )
retrieved_user_role = ( retrieved_user_role = (
user_role or _get_user_role(user_obj=user_obj) or LitellmUserRoles.INTERNAL_USER user_role or _get_user_role(user_obj=user_obj) or LitellmUserRoles.INTERNAL_USER
) )

5
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -141,12 +141,13 @@ async def test_check_blocked_team():
("proxy_admin_viewer", "proxy_admin_viewer"), ("proxy_admin_viewer", "proxy_admin_viewer"),
], ],
) )
def test_returned_user_api_key_auth(user_role, expected_role): @pytest.mark.asyncio
async def test_returned_user_api_key_auth(user_role, expected_role):
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles
from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj
from datetime import datetime from datetime import datetime
new_obj = _return_user_api_key_auth_obj( new_obj = await _return_user_api_key_auth_obj(
user_obj=LiteLLM_UserTable( user_obj=LiteLLM_UserTable(
user_role=user_role, user_id="", max_budget=None, user_email="" user_role=user_role, user_id="", max_budget=None, user_email=""
), ),