From 3686789c362943e11b4ac1f2119053c0e6be378f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 13 May 2024 18:23:23 -0700 Subject: [PATCH 1/5] feat - reset spend per team, api_key --- enterprise/utils.py | 4 +-- litellm/proxy/proxy_server.py | 57 +++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/enterprise/utils.py b/enterprise/utils.py index 1b5036e96..90b14314c 100644 --- a/enterprise/utils.py +++ b/enterprise/utils.py @@ -44,8 +44,8 @@ async def ui_get_spend_by_tags(start_date: str, end_date: str, prisma_client): # print("tags - spend") # print(response) # Bar Chart 1 - Spend per tag - Top 10 tags by spend - total_spend_per_tag = collections.defaultdict(float) - total_requests_per_tag = collections.defaultdict(int) + total_spend_per_tag: collections.defaultdict = collections.defaultdict(float) + total_requests_per_tag: collections.defaultdict = collections.defaultdict(int) for row in response: tag_name = row["individual_request_tag"] tag_spend = row["total_spend"] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6388bd5fd..9b2e01bc8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5922,6 +5922,63 @@ async def view_spend_logs( ) +@router.get( + "/global/spend/reset", + tags=["Budget & Spend Tracking"], + dependencies=[Depends(user_api_key_auth)], +) +async def global_spend_reset( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + 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, master_key + if prisma_client is None: + raise ProxyException( + message="Prisma Client is not initialized", + type="internal_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) + + if master_key is None: + raise ProxyException( + message="Master key is not initialized, please set LITELLM_MASTER_KEY in .env", + type="internal_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) + + if user_api_key_dict.api_key is None: + raise ProxyException( + message="no api_key passed", + type="auth_error", + param="master_key", + code=status.HTTP_401_UNAUTHORIZED, + ) + + if not secrets.compare_digest(master_key, user_api_key_dict.api_key): + raise ProxyException( + message="/global/spend/reset Route only allowed for master key", + type="auth_error", + param="master_key", + 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={}) + + @router.get( "/global/spend/logs", tags=["Budget & Spend Tracking"], From 7e56e2722658766fb93c24d4dcd34ea0ab89efc6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 14 May 2024 11:04:50 -0700 Subject: [PATCH 2/5] fix security for global_spend_reset --- litellm/proxy/_types.py | 5 +++++ litellm/proxy/proxy_server.py | 39 ++++++++++------------------------- 2 files changed, 16 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f91b768e6..5dbdd3a71 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -79,6 +79,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 9b2e01bc8..30f0275b6 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-. @@ -5927,9 +5936,7 @@ async def view_spend_logs( tags=["Budget & Spend Tracking"], dependencies=[Depends(user_api_key_auth)], ) -async def global_spend_reset( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +async def global_spend_reset(): """ ADMIN ONLY / MASTER KEY Only Endpoint @@ -5940,7 +5947,7 @@ async def global_spend_reset( 3. LiteLLM_TeamTable spend will be set = 0 """ - global prisma_client, master_key + global prisma_client if prisma_client is None: raise ProxyException( message="Prisma Client is not initialized", @@ -5949,30 +5956,6 @@ async def global_spend_reset( code=status.HTTP_401_UNAUTHORIZED, ) - if master_key is None: - raise ProxyException( - message="Master key is not initialized, please set LITELLM_MASTER_KEY in .env", - type="internal_error", - param="None", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if user_api_key_dict.api_key is None: - raise ProxyException( - message="no api_key passed", - type="auth_error", - param="master_key", - code=status.HTTP_401_UNAUTHORIZED, - ) - - if not secrets.compare_digest(master_key, user_api_key_dict.api_key): - raise ProxyException( - message="/global/spend/reset Route only allowed for master key", - type="auth_error", - param="master_key", - code=status.HTTP_401_UNAUTHORIZED, - ) - await prisma_client.db.litellm_verificationtoken.update_many( data={"spend": 0.0}, where={} ) From ca41e6590e3ac4a353a2eb11c936d989c220999c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 14 May 2024 11:28:35 -0700 Subject: [PATCH 3/5] test - auth on /reset/spend --- litellm/tests/test_key_generate_prisma.py | 71 +++++++++++++++++++++++ 1 file changed, 71 insertions(+) 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 + ) From 787c02c8dbf5be6d5e30de1d90e8c7e21a245972 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 14 May 2024 11:38:39 -0700 Subject: [PATCH 4/5] fix - return success spend reset --- litellm/proxy/proxy_server.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 30f0275b6..720033419 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5931,7 +5931,7 @@ async def view_spend_logs( ) -@router.get( +@router.post( "/global/spend/reset", tags=["Budget & Spend Tracking"], dependencies=[Depends(user_api_key_auth)], @@ -5961,6 +5961,11 @@ async def global_spend_reset(): ) 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", From fe1f5369ab446760050ca5072dbaf1f95b27f1b8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 14 May 2024 11:47:43 -0700 Subject: [PATCH 5/5] docs - global/spend/reset --- docs/my-website/docs/proxy/cost_tracking.md | 25 +++++++++++++++++++++ 1 file changed, 25 insertions(+) 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