From b09c283cc4a49bbd70ab32efa66d69ee64a42751 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 29 Jun 2024 11:33:02 -0700
Subject: [PATCH] feat - add spend report grouped by api key

---
 .../spend_management_endpoints.py             | 46 ++++++++++++++++++-
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 1fbd95b3cf..aafd14c623 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -817,9 +817,9 @@ async def get_global_spend_report(
         default=None,
         description="Time till which to view spend",
     ),
-    group_by: Optional[Literal["team", "customer"]] = fastapi.Query(
+    group_by: Optional[Literal["team", "customer", "api_key"]] = fastapi.Query(
         default="team",
-        description="Group spend by internal team or customer",
+        description="Group spend by internal team or customer or api_key",
     ),
 ):
     """
@@ -992,6 +992,48 @@ async def get_global_spend_report(
                 return []
 
             return db_response
+        elif group_by == "api_key":
+            sql_query = """
+                WITH SpendByModelApiKey AS (
+                    SELECT
+                        sl.api_key,
+                        sl.model,
+                        SUM(sl.spend) AS model_cost,
+                        SUM(sl.prompt_tokens) AS model_input_tokens,
+                        SUM(sl.completion_tokens) AS model_output_tokens
+                    FROM
+                        "LiteLLM_SpendLogs" sl
+                    WHERE
+                        sl."startTime" BETWEEN $1::date AND $2::date
+                    GROUP BY
+                        sl.api_key,
+                        sl.model
+                )
+                SELECT
+                    api_key,
+                    SUM(model_cost) AS total_cost,
+                    SUM(model_input_tokens) AS total_input_tokens,
+                    SUM(model_output_tokens) AS total_output_tokens,
+                    jsonb_agg(jsonb_build_object(
+                        'model', model,
+                        'total_cost', model_cost,
+                        'total_input_tokens', model_input_tokens,
+                        'total_output_tokens', model_output_tokens
+                    )) AS model_details
+                FROM
+                    SpendByModelApiKey
+                GROUP BY
+                    api_key
+                ORDER BY
+                    total_cost DESC;
+            """
+            db_response = await prisma_client.db.query_raw(
+                sql_query, start_date_obj, end_date_obj
+            )
+            if db_response is None:
+                return []
+
+            return db_response
 
     except Exception as e:
         raise HTTPException(