diff --git a/docs/my-website/docs/proxy/cost_tracking.md b/docs/my-website/docs/proxy/cost_tracking.md index d4760d83f..8a43d1c9d 100644 --- a/docs/my-website/docs/proxy/cost_tracking.md +++ b/docs/my-website/docs/proxy/cost_tracking.md @@ -126,6 +126,31 @@ Output from script +## 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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c17376622..7a67d6d79 100644 --- a/litellm/proxy/_types.py +++ b/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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6388bd5fd..720033419 100644 --- a/litellm/proxy/proxy_server.py +++ b/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"], diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index e6f2437e7..2eb693cf4 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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 + )