From 372fcb6e4cbfc53dffd483b5c63cacab3195942f Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 5 Sep 2024 12:35:04 -0700
Subject: [PATCH] allow internal user to view their own spend

---
 .../spend_management_endpoints.py             | 58 ++++++++++++++++++-
 1 file changed, 57 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 5f9d99c1bd..75dd9e280c 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -1910,6 +1910,52 @@ async def global_spend():
         )
 
 
+async def global_spend_key_internal_user(
+    user_api_key_dict: UserAPIKeyAuth, limit: int = 10
+):
+    from litellm.proxy.proxy_server import prisma_client
+
+    if prisma_client is None:
+        raise HTTPException(status_code=500, detail={"error": "No db connected"})
+
+    user_id = user_api_key_dict.user_id
+    if user_id is None:
+        raise HTTPException(status_code=500, detail={"error": "No user_id found"})
+
+    sql_query = """
+            WITH top_api_keys AS (
+            SELECT 
+                api_key,
+                SUM(spend) as total_spend
+            FROM 
+                "LiteLLM_SpendLogs"
+            WHERE 
+                "user" = $1
+            GROUP BY 
+                api_key
+            ORDER BY 
+                total_spend DESC
+            LIMIT $2  -- Adjust this number to get more or fewer top keys
+        )
+        SELECT 
+            t.api_key,
+            t.total_spend,
+            v.key_alias,
+            v.key_name
+        FROM 
+            top_api_keys t
+        LEFT JOIN 
+            "LiteLLM_VerificationToken" v ON t.api_key = v.token
+        ORDER BY 
+            t.total_spend DESC;
+    
+    """
+
+    response = await prisma_client.db.query_raw(sql_query, user_id, limit)
+
+    return response
+
+
 @router.get(
     "/global/spend/keys",
     tags=["Budget & Spend Tracking"],
@@ -1920,7 +1966,8 @@ async def global_spend_keys(
     limit: int = fastapi.Query(
         default=None,
         description="Number of keys to get. Will return Top 'n' keys.",
-    )
+    ),
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
     [BETA] This is a beta endpoint. It will change.
@@ -1929,6 +1976,15 @@ async def global_spend_keys(
     """
     from litellm.proxy.proxy_server import prisma_client
 
+    if (
+        user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER
+        or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
+    ):
+        response = await global_spend_key_internal_user(
+            user_api_key_dict=user_api_key_dict
+        )
+
+        return response
     if prisma_client is None:
         raise HTTPException(status_code=500, detail={"error": "No db connected"})
     sql_query = f"""SELECT * FROM "Last30dKeysBySpend" LIMIT {limit};"""