diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 920de3cc822..62e5eba0132 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -206,11 +206,9 @@ async def get_end_user_object( if end_user_id is None: return None - + _key = "end_user_id:{}".format(end_user_id) # check if in cache - cached_user_obj = user_api_key_cache.async_get_cache( - key="end_user_id:{}".format(end_user_id) - ) + cached_user_obj = await user_api_key_cache.async_get_cache(key=_key) if cached_user_obj is not None: if isinstance(cached_user_obj, dict): return LiteLLM_EndUserTable(**cached_user_obj) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f50fddf7fbc..1ba29dcc699 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1834,6 +1834,9 @@ async def update_cache( ) async def _update_end_user_cache(): + if end_user_id is None or response_cost is None: + return + _id = "end_user_id:{}".format(end_user_id) try: # Fetch the existing cost for the given user @@ -1846,7 +1849,7 @@ async def update_cache( if litellm.max_end_user_budget is not None: max_end_user_budget = litellm.max_end_user_budget existing_spend_obj = LiteLLM_EndUserTable( - user_id=_id, + user_id=end_user_id, spend=0, blocked=False, litellm_budget_table=LiteLLM_BudgetTable( @@ -1874,7 +1877,7 @@ async def update_cache( existing_spend_obj.spend = new_spend user_api_key_cache.set_cache(key=_id, value=existing_spend_obj.json()) except Exception as e: - verbose_proxy_logger.debug( + verbose_proxy_logger.error( f"An error occurred updating end user cache: {str(e)}\n\n{traceback.format_exc()}" ) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 08618c98891..e6f2437e7d7 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -418,9 +418,16 @@ def test_call_with_user_over_budget(prisma_client): print(vars(e)) +def test_end_user_cache_write_unit_test(): + """ + assert end user object is being written to cache as expected + """ + pass + + def test_call_with_end_user_over_budget(prisma_client): # Test if a user passed to /chat/completions is tracked & fails when they cross their budget - # we only check this when litellm.max_user_budget is set + # we only check this when litellm.max_end_user_budget is set import random setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)