diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 80113fff38c..d1e0caf13d0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -386,8 +386,6 @@ async def get_user_object( - if valid, return LiteLLM_UserTable object with defined limits - if not, then raise an error """ - if prisma_client is None: - raise Exception("No db connected") if user_id is None: return None @@ -400,6 +398,8 @@ async def get_user_object( elif isinstance(cached_user_obj, LiteLLM_UserTable): return cached_user_obj # else, check db + if prisma_client is None: + raise Exception("No db connected") try: response = await prisma_client.db.litellm_usertable.find_unique( diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 7b3c0960ff0..a35fea324f3 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -322,10 +322,9 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # check if REQUEST ALLOWED for user_id user_id = user_api_key_dict.user_id if user_id is not None: - request_count_api_key = f"{user_id}::{precise_minute}::request_count" - _user_id_rate_limits = await self.internal_usage_cache.async_get_cache( - key=request_count_api_key, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + _user_id_rate_limits = await self.get_internal_user_object( + user_id=user_id, + user_api_key_dict=user_api_key_dict, ) # get user tpm/rpm limits if _user_id_rate_limits is not None and isinstance( @@ -742,3 +741,25 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): str(e) ) ) + + async def get_internal_user_object( + self, + user_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> Optional[dict]: + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client + + _user_id_rate_limits = await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=self.internal_usage_cache.dual_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=None, + ) + + if _user_id_rate_limits is None: + return None + + return _user_id_rate_limits.model_dump() diff --git a/litellm/tests/test_parallel_request_limiter.py b/litellm/tests/test_parallel_request_limiter.py index 93e86404e52..dee7fefbf02 100644 --- a/litellm/tests/test_parallel_request_limiter.py +++ b/litellm/tests/test_parallel_request_limiter.py @@ -295,7 +295,13 @@ async def test_pre_call_hook_user_tpm_limits(): local_cache = DualCache() # create user with tpm/rpm limits user_id = "test-user" - user_obj = {"tpm_limit": 9, "rpm_limit": 10} + user_obj = { + "tpm_limit": 9, + "rpm_limit": 10, + "user_id": user_id, + "user_email": "user_email", + "max_budget": None, + } local_cache.set_cache(key=user_id, value=user_obj) @@ -331,6 +337,7 @@ async def test_pre_call_hook_user_tpm_limits(): ## Expected cache val: {"current_requests": 0, "current_tpm": 0, "current_rpm": 1} try: + print("cache=local_cache", local_cache.in_memory_cache.cache_dict) await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache,