phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

use safe get_internal_user_object

Browse source
This commit is contained in:
Ishaan Jaff 2024-09-26 17:23:33 -07:00
parent 8900255bc7
commit e34d9c51c9
1 changed files with 25 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -747,19 +747,33 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
user_id: str, user_id: str,
user_api_key_dict: UserAPIKeyAuth, user_api_key_dict: UserAPIKeyAuth,
) -> Optional[dict]: ) -> Optional[dict]:
"""
Helper to get the 'Internal User Object'
It uses the `get_user_object` function from `litellm.proxy.auth.auth_checks`
We need this because the UserApiKeyAuth object does not contain the rpm/tpm limits for a User AND there could be a perf impact by additionally reading the UserTable.
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.auth.auth_checks import get_user_object
from litellm.proxy.proxy_server import prisma_client from litellm.proxy.proxy_server import prisma_client
_user_id_rate_limits = await get_user_object( try:
user_id=user_id, _user_id_rate_limits = await get_user_object(
prisma_client=prisma_client, user_id=user_id,
user_api_key_cache=self.internal_usage_cache.dual_cache, prisma_client=prisma_client,
user_id_upsert=False, user_api_key_cache=self.internal_usage_cache.dual_cache,
parent_otel_span=user_api_key_dict.parent_otel_span, user_id_upsert=False,
proxy_logging_obj=None, parent_otel_span=user_api_key_dict.parent_otel_span,
) proxy_logging_obj=None,
)
if _user_id_rate_limits is None: if _user_id_rate_limits is None:
return None
return _user_id_rate_limits.model_dump()
except Exception as e:
verbose_proxy_logger.exception(
"Parallel Request Limiter: Error getting user object", str(e)
)
return None return None
return _user_id_rate_limits.model_dump()