From d2ec9e4328d64aba4a57d7f053bd1b041b1f2f28 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 19 Apr 2025 08:04:56 -0700
Subject: [PATCH] fix: fix linting error

---
 .../common_daily_activity.py                  | 23 +++++++++++--------
 .../internal_user_endpoints.py                |  2 --
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py
index d5de85311b..94cff2489d 100644
--- a/litellm/proxy/management_endpoints/common_daily_activity.py
+++ b/litellm/proxy/management_endpoints/common_daily_activity.py
@@ -101,6 +101,19 @@ def update_breakdown_metrics(
     return breakdown
 
 
+async def get_api_key_metadata(
+    prisma_client: PrismaClient,
+    api_keys: set[str],
+) -> Dict[str, Dict[str, Any]]:
+    """Update api key metadata for a single record."""
+    key_records = await prisma_client.db.litellm_verificationtoken.find_many(
+        where={"token": {"in": list(api_keys)}}
+    )
+    return {
+        k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records
+    }
+
+
 async def get_daily_activity(
     prisma_client: Optional[PrismaClient],
     table_name: str,
@@ -177,15 +190,7 @@ async def get_daily_activity(
         model_metadata: Dict[str, Dict[str, Any]] = {}
         provider_metadata: Dict[str, Dict[str, Any]] = {}
         if api_keys:
-            key_records = await prisma_client.db.litellm_verificationtoken.find_many(
-                where={"token": {"in": list(api_keys)}}
-            )
-            api_key_metadata.update(
-                {
-                    k.token: {"key_alias": k.key_alias, "team_id": k.team_id}
-                    for k in key_records
-                }
-            )
+            api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
 
         # Process results
         results = []
diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index e2f1077095..42a5dd7e81 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -36,8 +36,6 @@ from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
 from litellm.proxy.utils import handle_exception_on_proxy
 from litellm.types.proxy.management_endpoints.common_daily_activity import (
     BreakdownMetrics,
-    DailySpendData,
-    DailySpendMetadata,
     KeyMetadata,
     KeyMetricWithMetadata,
     LiteLLM_DailyUserSpend,