fix(auth_checks.py): fix 'get_end_user_object'

await cache get

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)

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()}"
             )
 

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)