Merge pull request #3626 from BerriAI/litellm_reset_spend_per_team_api_key

feat - reset spend per team, api_key [Only Master Key]

docs/my-website/docs/proxy/cost_tracking.md

@@ -126,6 +126,31 @@ Output from script
 </Tabs>
 
 
+## Reset Team, API Key Spend - MASTER KEY ONLY
+
+Use `/global/spend/reset` if you want to:
+- Reset the Spend for all API Keys, Teams. The `spend` for ALL Teams and Keys in `LiteLLM_TeamTable` and `LiteLLM_VerificationToken` will be set to `spend=0`
+
+- LiteLLM will maintain all the logs in `LiteLLMSpendLogs` for Auditing Purposes
+
+### Request 
+Only the `LITELLM_MASTER_KEY` you set can access this route
+```shell
+curl -X POST \
+  'http://localhost:4000/global/spend/reset' \
+  -H 'Authorization: Bearer sk-1234' \
+  -H 'Content-Type: application/json'
+```
+
+### Expected Responses
+
+```shell
+{"message":"Spend for all API Keys and Teams reset successfully","status":"success"}
+```
+
+
+
+
 ## Spend Tracking for Azure
 
 Set base model for cost tracking azure image-gen call

litellm/proxy/_types.py

@@ -89,6 +89,11 @@ class LiteLLMRoutes(enum.Enum):
         "/v1/models",
     ]
 
+    # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend
+    master_key_only_routes: List = [
+        "/global/spend/reset",
+    ]
+
     info_routes: List = [
         "/key/info",
         "/team/info",

litellm/proxy/proxy_server.py

@@ -589,6 +589,15 @@ async def user_api_key_auth(
             )
 
             return _user_api_key_obj
+
+        ## IF it's not a master key
+        ## Route should not be in master_key_only_routes
+        if route in LiteLLMRoutes.master_key_only_routes.value:
+            raise Exception(
+                f"Tried to access route={route}, which is only for MASTER KEY"
+            )
+
+        ## Check DB
         if isinstance(
             api_key, str
         ):  # if generated token, make sure it starts with sk-.
@@ -5922,6 +5931,42 @@ async def view_spend_logs(
         )
 
 
+@router.post(
+    "/global/spend/reset",
+    tags=["Budget & Spend Tracking"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def global_spend_reset():
+    """
+    ADMIN ONLY / MASTER KEY Only Endpoint
+
+    Globally reset spend for All API Keys and Teams, maintain LiteLLM_SpendLogs
+
+    1. LiteLLM_SpendLogs will maintain the logs on spend, no data gets deleted from there
+    2. LiteLLM_VerificationTokens spend will be set = 0
+    3. LiteLLM_TeamTable spend will be set = 0
+
+    """
+    global prisma_client
+    if prisma_client is None:
+        raise ProxyException(
+            message="Prisma Client is not initialized",
+            type="internal_error",
+            param="None",
+            code=status.HTTP_401_UNAUTHORIZED,
+        )
+
+    await prisma_client.db.litellm_verificationtoken.update_many(
+        data={"spend": 0.0}, where={}
+    )
+    await prisma_client.db.litellm_teamtable.update_many(data={"spend": 0.0}, where={})
+
+    return {
+        "message": "Spend for all API Keys and Teams reset successfully",
+        "status": "success",
+    }
+
+
 @router.get(
     "/global/spend/logs",
     tags=["Budget & Spend Tracking"],

litellm/tests/test_key_generate_prisma.py

@@ -2013,3 +2013,74 @@ async def test_master_key_hashing(prisma_client):
     except Exception as e:
         print("Got Exception", e)
         pytest.fail(f"Got exception {e}")
+
+
+@pytest.mark.asyncio
+async def test_reset_spend_authentication(prisma_client):
+    """
+    1. Test master key can access this route  -> ONLY MASTER KEY SHOULD BE ABLE TO RESET SPEND
+    2. Test that non-master key gets rejected
+    3. Test that non-master key with role == "proxy_admin" or admin gets rejected
+    """
+
+    print("prisma client=", prisma_client)
+
+    master_key = "sk-1234"
+
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", master_key)
+
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    from litellm.proxy.proxy_server import user_api_key_cache
+
+    bearer_token = "Bearer " + master_key
+
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/global/spend/reset")
+
+    # Test 1 - Master Key
+    result: UserAPIKeyAuth = await user_api_key_auth(
+        request=request, api_key=bearer_token
+    )
+
+    print("result from user auth with Master key", result)
+    assert result.token is not None
+
+    # Test 2 - Non-Master Key
+    _response = await new_user(
+        data=NewUserRequest(
+            tpm_limit=20,
+        )
+    )
+
+    generate_key = "Bearer " + _response.key
+
+    try:
+        await user_api_key_auth(request=request, api_key=generate_key)
+        pytest.fail(f"This should have failed!. IT's an expired key")
+    except Exception as e:
+        print("Got Exception", e)
+        assert (
+            "Tried to access route=/global/spend/reset, which is only for MASTER KEY"
+            in e.message
+        )
+
+    # Test 3 - Non-Master Key with role == "proxy_admin" or admin
+    _response = await new_user(
+        data=NewUserRequest(
+            user_role="proxy_admin",
+            tpm_limit=20,
+        )
+    )
+
+    generate_key = "Bearer " + _response.key
+
+    try:
+        await user_api_key_auth(request=request, api_key=generate_key)
+        pytest.fail(f"This should have failed!. IT's an expired key")
+    except Exception as e:
+        print("Got Exception", e)
+        assert (
+            "Tried to access route=/global/spend/reset, which is only for MASTER KEY"
+            in e.message
+        )