added user_api_key_spend

Author: mubashir1osmani

litellm/integrations/datadog/datadog_llm_obs.py

@@ -148,23 +148,16 @@ async def async_send_batch(self):
                     ),
                 ),
             }
-            # serialize datetime objects - for budget reset time in spend metrics
-            import json
-            from datetime import datetime, date
 
-            def custom_json_encoder(obj):
-                if isinstance(obj, (datetime, date)):
-                    return obj.isoformat()
-                raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
+            # serialize datetime objects - for budget reset time in spend metrics
+            from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
-            # Serialize payload with custom encoder for debugging
             try:
-                verbose_logger.debug("payload %s", json.dumps(payload, indent=4, default=custom_json_encoder))
+                verbose_logger.debug("payload %s", safe_dumps(payload))
             except Exception as debug_error:
                 verbose_logger.debug("payload serialization failed: %s", str(debug_error))
 
-            # Convert payload to JSON string with custom encoder for HTTP request
-            json_payload = json.dumps(payload, default=custom_json_encoder)
+            json_payload = safe_dumps(payload)
 
             response = await self.async_client.post(
                 url=self.intake_url,
@@ -338,7 +331,6 @@ def _get_response_messages(
         if isinstance(response_obj, str):
             try:
                 import ast
-
                 response_obj = ast.literal_eval(response_obj)
             except (ValueError, SyntaxError):
                 try:
@@ -573,48 +565,55 @@ def _get_spend_metrics(
         Get the spend metrics from the standard logging payload
         """
         spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics()
-        
-        # Get response cost for litellm_spend_metric
-        response_cost = standard_logging_payload.get("response_cost", 0.0)
-        if response_cost > 0:
-            spend_metrics["litellm_spend_metric"] = response_cost
+
+        # send response cost
+        spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0)
 
         # Get budget information from metadata
         metadata = standard_logging_payload.get("metadata", {})
 
         # API key max budget
         user_api_key_max_budget = metadata.get("user_api_key_max_budget")
         if user_api_key_max_budget is not None:
-            # type casting to make sure its a float value
+            spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)
+
+        # API key spend
+        user_api_key_spend = metadata.get("user_api_key_spend")
+        if user_api_key_spend is not None:
             try:
-                if isinstance(user_api_key_max_budget, (int, float)):
-                    spend_metrics["litellm_api_key_max_budget_metric"] = float(user_api_key_max_budget)
-                elif isinstance(user_api_key_max_budget, str):
-                    spend_metrics["litellm_api_key_max_budget_metric"] = float(user_api_key_max_budget)
+                spend_metrics["user_api_key_spend"] = float(user_api_key_spend)
             except (ValueError, TypeError):
-                verbose_logger.debug(f"Invalid user_api_key_max_budget value: {user_api_key_max_budget}")
+                verbose_logger.debug(f"Invalid user_api_key_spend value: {user_api_key_spend}")
 
-        # API key budget remaining hours
+        # API key budget reset datetime
         user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at")
         if user_api_key_budget_reset_at is not None:
             try:
-                from datetime import datetime
-                budget_reset_at: datetime
+                from datetime import datetime, timezone
+
+                budget_reset_at = None
                 if isinstance(user_api_key_budget_reset_at, str):
-                    # Parse ISO string if it's a string
-                    budget_reset_at = datetime.fromisoformat(user_api_key_budget_reset_at.replace('Z', '+00:00'))
+                    # Handle ISO format strings that might have 'Z' suffix
+                    iso_string = user_api_key_budget_reset_at.replace('Z', '+00:00')
+                    budget_reset_at = datetime.fromisoformat(iso_string)
                 elif isinstance(user_api_key_budget_reset_at, datetime):
                     budget_reset_at = user_api_key_budget_reset_at
-                else:
-                    verbose_logger.debug(f"Invalid user_api_key_budget_reset_at type: {type(user_api_key_budget_reset_at)}")
-                    return spend_metrics
-                
-                remaining_hours = (
-                    budget_reset_at - datetime.now(budget_reset_at.tzinfo)
-                ).total_seconds() / 3600
-                spend_metrics["litellm_api_key_budget_remaining_hours_metric"] = max(0, remaining_hours)
+
+                if budget_reset_at is not None:
+                    # Preserve timezone info if already present
+                    if budget_reset_at.tzinfo is None:
+                        budget_reset_at = budget_reset_at.replace(tzinfo=timezone.utc)
+
+                    # Convert to ISO string format for JSON serialization
+                    # This prevents circular reference issues and ensures proper timezone representation
+                    iso_string = budget_reset_at.isoformat()
+                    spend_metrics["user_api_key_budget_reset_at"] = iso_string
+                    
+                    # Debug logging to verify the conversion
+                    verbose_logger.debug(f"Converted budget_reset_at to ISO format: {iso_string}")
             except Exception as e:
-                verbose_logger.debug(f"Error calculating remaining hours for budget reset: {e}")
+                verbose_logger.debug(f"Error processing budget reset datetime: {e}")
+                verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}")
 
         return spend_metrics
 

litellm/proxy/litellm_pre_call_utils.py

@@ -546,6 +546,7 @@ def get_sanitized_user_information_from_key(
             user_api_key_end_user_id=user_api_key_dict.end_user_id,
             user_api_key_user_email=user_api_key_dict.user_email,
             user_api_key_request_route=user_api_key_dict.request_route,
+            user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at,
         )
         return user_api_key_logged_metadata
 

litellm/types/integrations/datadog_llm_obs.py

@@ -83,6 +83,7 @@ class DDLLMObsLatencyMetrics(TypedDict, total=False):
     guardrail_overhead_time_ms: float
 
 class DDLLMObsSpendMetrics(TypedDict, total=False):
-    litellm_spend_metric: float
-    litellm_api_key_max_budget_metric: float
-    litellm_api_key_budget_remaining_hours_metric: float
+    response_cost: float
+    user_api_key_spend: float
+    user_api_key_max_budget: float
+    user_api_key_budget_reset_at: str

litellm/types/utils.py

@@ -1807,6 +1807,9 @@ async def __anext__(self):
 class StandardLoggingUserAPIKeyMetadata(TypedDict):
     user_api_key_hash: Optional[str]  # hash of the litellm virtual key used
     user_api_key_alias: Optional[str]
+    user_api_key_spend: Optional[float]
+    user_api_key_max_budget: Optional[float] = None
+    user_api_key_budget_reset_at: Optional[str] = None
     user_api_key_org_id: Optional[str]
     user_api_key_team_id: Optional[str]
     user_api_key_user_id: Optional[str]