From 02cf18be83c7efb06dafe8feb7a5205a056158ea Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 4 Nov 2024 11:11:09 -0800
Subject: [PATCH] StandardLoggingBudgetMetadata

---
 .../proxy/hooks/proxy_track_cost_callback.py  |  4 +--
 litellm/types/utils.py                        | 25 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 6d5b441cd..2de2b0673 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -73,12 +73,10 @@ async def _PROXY_track_cost_callback(
         key_alias = standard_logging_payload.get("metadata", {}).get(
             "user_api_key_alias"
         )
-        response_cost: Optional[float] = standard_logging_payload.get("response_cost")
-
-        # end_user_max_budget = metadata.get("user_api_end_user_max_budget", None)
         end_user_max_budget = standard_logging_payload.get("metadata", {}).get(
             "user_api_end_user_max_budget"
         )
+        response_cost: Optional[float] = standard_logging_payload.get("response_cost")
 
         if response_cost is not None:
             if user_api_key is not None or user_id is not None or team_id is not None:
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 2d0e262fe..2b5a1cdfd 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -1416,7 +1416,28 @@ class AdapterCompletionStreamWrapper:
             raise StopAsyncIteration
 
 
+class StandardLoggingBudgetMetadata(TypedDict, total=False):
+    """
+    Store Budget related metadata for Team, Internal User, End User etc
+    """
+
+    user_api_end_user_max_budget: Optional[float]
+
+
 class StandardLoggingUserAPIKeyMetadata(TypedDict):
+    """
+    Store User API Key related metadata to identify the request
+
+    Example:
+    user_api_key_hash: "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
+    user_api_key_alias: "litellm-key-123"
+    user_api_key_org_id: "123"
+    user_api_key_team_id: "456"
+    user_api_key_user_id: "789"
+    user_api_key_team_alias: "litellm-team-123"
+
+    """
+
     user_api_key_hash: Optional[str]  # hash of the litellm virtual key used
     user_api_key_alias: Optional[str]
     user_api_key_org_id: Optional[str]
@@ -1425,7 +1446,9 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict):
     user_api_key_team_alias: Optional[str]
 
 
-class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
+class StandardLoggingMetadata(
+    StandardLoggingUserAPIKeyMetadata, StandardLoggingBudgetMetadata
+):
     """
     Specific metadata k,v pairs logged to integration for easier cost tracking
     """