fix(user_api_key_auth.py): respect team budgets over user budget, if key belongs to team

Closes https://github.com/BerriAI/litellm/issues/5097

litellm/tests/test_user_api_key_auth.py

@@ -116,3 +116,65 @@ def test_returned_user_api_key_auth(user_role):
         assert new_obj.user_role == user_role
     else:
         assert new_obj.user_role == "internal_user"
+
+
+@pytest.mark.parametrize("key_ownership", ["user_key", "team_key"])
+@pytest.mark.asyncio
+async def test_user_personal_budgets(key_ownership):
+    """
+    Set a personal budget on a user
+
+    - have it only apply when key belongs to user -> raises BudgetExceededError
+    - if key belongs to team, have key respect team budget -> allows call to go through
+    """
+    import asyncio
+    import time
+
+    from fastapi import Request
+    from starlette.datastructures import URL
+
+    from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+    from litellm.proxy.proxy_server import hash_token, user_api_key_cache
+
+    _user_id = "1234"
+    user_key = "sk-12345678"
+
+    if key_ownership == "user_key":
+        valid_token = UserAPIKeyAuth(
+            token=hash_token(user_key),
+            last_refreshed_at=time.time(),
+            user_id=_user_id,
+            spend=20,
+        )
+    elif key_ownership == "team_key":
+        valid_token = UserAPIKeyAuth(
+            token=hash_token(user_key),
+            last_refreshed_at=time.time(),
+            user_id=_user_id,
+            team_id="my-special-team",
+            team_max_budget=100,
+            spend=20,
+        )
+    await asyncio.sleep(1)
+    user_obj = LiteLLM_UserTable(
+        user_id=_user_id, spend=11, max_budget=10, user_email=""
+    )
+    user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
+    user_api_key_cache.set_cache(key="{}".format(_user_id), value=user_obj)
+
+    setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")
+
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    try:
+        await user_api_key_auth(request=request, api_key="Bearer " + user_key)
+
+        if key_ownership == "user_key":
+            pytest.fail("Expected this call to fail. User is over limit.")
+    except Exception:
+        if key_ownership == "team_key":
+            pytest.fail("Expected this call to work. Key is below team budget.")