From 3a5d258e054f528d81bd52849cd96ba806ec70f7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 29 Jun 2024 12:00:19 -0700 Subject: [PATCH] raise error on /spend/report endpoint --- litellm/proxy/_types.py | 2 +- litellm/proxy/spend_tracking/spend_management_endpoints.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 640c7695a0..1f1aaf0eea 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1622,7 +1622,7 @@ class ProxyException(Exception): } -class CommonProxyErrors(enum.Enum): +class CommonProxyErrors(str, enum.Enum): db_not_connected_error = "DB not connected" no_llm_router = "No models configured on proxy" not_allowed_access = "Admin-only endpoint. Not allowed to access this." diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index aafd14c623..11406b162f 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -860,7 +860,7 @@ async def get_global_spend_report( start_date_obj = datetime.strptime(start_date, "%Y-%m-%d") end_date_obj = datetime.strptime(end_date, "%Y-%m-%d") - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client try: if prisma_client is None: @@ -868,6 +868,11 @@ async def get_global_spend_report( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) + if premium_user is not True: + raise ValueError( + "/spend/report endpoint" + CommonProxyErrors.not_premium_user.value + ) + if group_by == "team": # first get data from spend logs -> SpendByModelApiKey # then read data from "SpendByModelApiKey" to format the response obj