added tests

Author: mubashir1osmani

litellm/integrations/datadog/datadog_llm_obs.py

@@ -148,17 +148,19 @@ async def async_send_batch(self):
                     ),
                 ),
             }
-            
+
             # serialize datetime objects - for budget reset time in spend metrics
             from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-            
+
             try:
                 verbose_logger.debug("payload %s", safe_dumps(payload))
             except Exception as debug_error:
-                verbose_logger.debug("payload serialization failed: %s", str(debug_error))
-            
+                verbose_logger.debug(
+                    "payload serialization failed: %s", str(debug_error)
+                )
+
             json_payload = safe_dumps(payload)
-            
+
             response = await self.async_client.post(
                 url=self.intake_url,
                 content=json_payload,
@@ -331,6 +333,7 @@ def _get_response_messages(
         if isinstance(response_obj, str):
             try:
                 import ast
+
                 response_obj = ast.literal_eval(response_obj)
             except (ValueError, SyntaxError):
                 try:
@@ -557,21 +560,23 @@ def _get_latency_metrics(
                 )
 
         return latency_metrics
-    
+
     def _get_spend_metrics(
-            self, standard_logging_payload: StandardLoggingPayload
+        self, standard_logging_payload: StandardLoggingPayload
     ) -> DDLLMObsSpendMetrics:
         """
         Get the spend metrics from the standard logging payload
         """
         spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics()
 
         # send response cost
-        spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0)
+        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:
@@ -583,7 +588,9 @@ def _get_spend_metrics(
             try:
                 spend_metrics["user_api_key_spend"] = float(user_api_key_spend)
             except (ValueError, TypeError):
-                verbose_logger.debug(f"Invalid user_api_key_spend value: {user_api_key_spend}")
+                verbose_logger.debug(
+                    f"Invalid user_api_key_spend value: {user_api_key_spend}"
+                )
 
         # API key budget reset datetime
         user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at")
@@ -594,7 +601,7 @@ def _get_spend_metrics(
                 budget_reset_at = None
                 if isinstance(user_api_key_budget_reset_at, str):
                     # Handle ISO format strings that might have 'Z' suffix
-                    iso_string = user_api_key_budget_reset_at.replace('Z', '+00:00')
+                    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
@@ -608,9 +615,11 @@ def _get_spend_metrics(
                     # 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}")
+                    verbose_logger.debug(
+                        f"Converted budget_reset_at to ISO format: {iso_string}"
+                    )
             except Exception as e:
                 verbose_logger.debug(f"Error processing budget reset datetime: {e}")
                 verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}")

litellm/types/integrations/datadog_llm_obs.py

@@ -82,8 +82,9 @@ class DDLLMObsLatencyMetrics(TypedDict, total=False):
     litellm_overhead_time_ms: float
     guardrail_overhead_time_ms: float
 
+
 class DDLLMObsSpendMetrics(TypedDict, total=False):
     response_cost: float
     user_api_key_spend: float
     user_api_key_max_budget: float
-    user_api_key_budget_reset_at: str
+    user_api_key_budget_reset_at: str

tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py

@@ -1,7 +1,7 @@
 import asyncio
 import os
 import sys
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Optional
 from unittest.mock import Mock, patch, MagicMock
 
@@ -901,89 +901,12 @@ def test_tool_call_response_handling(self, mock_env_vars):
             output_function_info = output_tool_calls[0].get("function", {})
             assert output_function_info.get("name") == "format_response"
 
-
-def create_standard_logging_payload() -> StandardLoggingPayload:
-    """Create a standard logging payload for testing"""
-    return {
-        "id": "test_id",
-        "trace_id": "test_trace_id",
-        "call_type": "completion",
-        "stream": False,
-        "response_cost": 0.1,
-        "response_cost_failure_debug_info": None,
-        "status": "success",
-        "custom_llm_provider": None,
-        "total_tokens": 30,
-        "prompt_tokens": 20,
-        "completion_tokens": 10,
-        "startTime": 1234567890.0,
-        "endTime": 1234567891.0,
-        "completionStartTime": 1234567890.5,
-        "response_time": 1.0,
-        "model_map_information": {
-            "model_map_key": "gpt-3.5-turbo",
-            "model_map_value": None
-        },
-        "model": "gpt-3.5-turbo",
-        "model_id": "model-123",
-        "model_group": "openai-gpt",
-        "api_base": "https://api.openai.com",
-        "metadata": {
-            "user_api_key_hash": "test_hash",
-            "user_api_key_org_id": None,
-            "user_api_key_alias": "test_alias",
-            "user_api_key_team_id": "test_team",
-            "user_api_key_user_id": "test_user",
-            "user_api_key_team_alias": "test_team_alias",
-            "user_api_key_end_user_id": None,
-            "user_api_key_request_route": None,
-            "user_api_key_max_budget": None,
-            "user_api_key_budget_reset_at": None,
-            "user_api_key_user_email": None,
-            "spend_logs_metadata": None,
-            "requester_ip_address": "127.0.0.1",
-            "requester_metadata": None,
-            "requester_custom_headers": None,
-            "prompt_management_metadata": None,
-            "mcp_tool_call_metadata": None,
-            "vector_store_request_metadata": None,
-            "applied_guardrails": None,
-            "usage_object": None,
-            "cold_storage_object_key": None,
-        },
-        "cache_hit": False,
-        "cache_key": None,
-        "saved_cache_cost": 0.0,
-        "request_tags": [],
-        "end_user": None,
-        "requester_ip_address": "127.0.0.1",
-        "messages": [{"role": "user", "content": "Hello, world!"}],
-        "response": {"choices": [{"message": {"content": "Hi there!"}}]},
-        "error_str": None,
-        "model_parameters": {"stream": True},
-        "hidden_params": {
-            "model_id": "model-123",
-            "cache_key": None,
-            "api_base": "https://api.openai.com",
-            "response_cost": "0.1",
-            "additional_headers": None,
-            "litellm_overhead_time_ms": None,
-            "batch_models": None,
-            "litellm_model_name": None,
-            "usage_object": None,
-        },
-        "error_information": None,
-        "guardrail_information": None,
-        "standard_built_in_tools_params": None,
-    }  # type: ignore
-
-
 def create_standard_logging_payload_with_spend_metrics() -> StandardLoggingPayload:
     """Create a StandardLoggingPayload object with spend metrics for testing"""
     from datetime import datetime, timezone
 
-    # Create a budget reset time 24 hours from now
-    budget_reset_at = datetime.now(timezone.utc) + timedelta(hours=24)
+    # Create a budget reset time 10 days from now (using "10d" format)
+    budget_reset_at = datetime.now(timezone.utc) + timedelta(days=10)
 
     return {
         "id": "test-request-id-spend",
@@ -1019,8 +942,9 @@ def create_standard_logging_payload_with_spend_metrics() -> StandardLoggingPaylo
             "user_api_key_user_email": None,
             "user_api_key_end_user_id": None,
             "user_api_key_request_route": None,
+            "user_api_key_spend": 0.67,
             "user_api_key_max_budget": 10.0,  # $10 max budget
-            "user_api_key_budget_reset_at": budget_reset_at.isoformat(),
+            "user_api_key_budget_reset_at": budget_reset_at.isoformat(),  # ISO format: 2025-09-26T...
             "spend_logs_metadata": None,
             "requester_ip_address": "127.0.0.1",
             "requester_metadata": None,
@@ -1064,23 +988,32 @@ async def test_datadog_llm_obs_spend_metrics(mock_env_vars):
     """Test that budget metrics are properly extracted and logged"""
     datadog_llm_obs_logger = DataDogLLMObsLogger()
 
-    # Create a standard logging payload with budget metadata
-    payload = create_standard_logging_payload()
+    # Create a standard logging payload with spend metrics
+    payload = create_standard_logging_payload_with_spend_metrics()
 
-    # Add budget information to metadata
-    payload["metadata"]["user_api_key_max_budget"] = 10.0
-    payload["metadata"]["user_api_key_budget_reset_at"] = "2025-09-15T00:00:00+00:00"
+    # Show the budget reset time in ISO format
+    budget_reset_iso = payload["metadata"]["user_api_key_budget_reset_at"]
+    print(f"Budget reset time (ISO format): {budget_reset_iso}")
+    from datetime import datetime, timezone
+    print(f"Current time: {datetime.now(timezone.utc).isoformat()}")
 
     # Test the _get_spend_metrics method
     spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload)
 
     # Verify budget metrics are present
-    assert "litellm_api_key_max_budget_metric" in spend_metrics
-    assert spend_metrics["litellm_api_key_max_budget_metric"] == 10.0
-
-    assert "litellm_api_key_budget_remaining_hours_metric" in spend_metrics
-    # The remaining hours should be calculated based on the reset time
-    assert spend_metrics["litellm_api_key_budget_remaining_hours_metric"] >= 0
+    assert "user_api_key_max_budget" in spend_metrics
+    assert spend_metrics["user_api_key_max_budget"] == 10.0
+
+    assert "user_api_key_budget_reset_at" in spend_metrics
+    # The budget reset should be a datetime string in ISO format
+    budget_reset = spend_metrics["user_api_key_budget_reset_at"]
+    assert isinstance(budget_reset, str)
+    print(f"Budget reset datetime: {budget_reset}")
+    # Should be close to 10 days from now
+    budget_reset_dt = datetime.fromisoformat(budget_reset.replace('Z', '+00:00'))
+    now = datetime.now(timezone.utc)
+    time_diff = (budget_reset_dt - now).total_seconds() / 86400  # days
+    assert 9.5 <= time_diff <= 10.5  # Should be close to 10 days
 
     print(f"Spend metrics: {spend_metrics}")
 
@@ -1091,25 +1024,30 @@ async def test_datadog_llm_obs_spend_metrics_no_budget(mock_env_vars):
     datadog_llm_obs_logger = DataDogLLMObsLogger()
 
     # Create a standard logging payload without budget metadata
-    payload = create_standard_logging_payload()
+    payload = create_standard_logging_payload_with_spend_metrics()
+
+    # Remove budget-related metadata to test no-budget scenario
+    payload["metadata"].pop("user_api_key_max_budget", None)
+    payload["metadata"].pop("user_api_key_budget_reset_at", None)
 
     # Test the _get_spend_metrics method
     spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload)
 
     # Verify only response cost is present
-    assert "litellm_spend_metric" in spend_metrics
-    assert spend_metrics["litellm_spend_metric"] == 0.1
+    assert "response_cost" in spend_metrics
+    assert spend_metrics["response_cost"] == 0.15
 
     # Budget metrics should not be present
-    assert "litellm_api_key_max_budget_metric" not in spend_metrics
-    assert "litellm_api_key_budget_remaining_hours_metric" not in spend_metrics
+    assert "user_api_key_max_budget" not in spend_metrics
+    assert "user_api_key_budget_reset_at" not in spend_metrics
 
     print(f"Spend metrics (no budget): {spend_metrics}")
 
 
 @pytest.mark.asyncio
 async def test_spend_metrics_in_datadog_payload(mock_env_vars):
     """Test that spend metrics are correctly included in DataDog LLM Observability payloads"""
+    from datetime import datetime
     datadog_llm_obs_logger = DataDogLLMObsLogger()
 
     standard_payload = create_standard_logging_payload_with_spend_metrics()
@@ -1138,17 +1076,28 @@ async def test_spend_metrics_in_datadog_payload(mock_env_vars):
     spend_metrics = metadata.get("spend_metrics", {})
     assert spend_metrics, "Spend metrics should exist in metadata"
 
-    # Check that all three spend metrics are present
-    assert "litellm_spend_metric" in spend_metrics
-    assert "litellm_api_key_max_budget_metric" in spend_metrics
-    assert "litellm_api_key_budget_remaining_hours_metric" in spend_metrics
+    # Check that all metrics are present
+    assert "response_cost" in spend_metrics
+    assert "user_api_key_spend" in spend_metrics
+    assert "user_api_key_max_budget" in spend_metrics
+    assert "user_api_key_budget_reset_at" in spend_metrics
 
     # Verify the values are correct
-    assert spend_metrics["litellm_spend_metric"] == 0.15  # response_cost
-    assert spend_metrics["litellm_api_key_max_budget_metric"] == 10.0  # max budget
-
-    # Verify remaining hours is a reasonable value (should be close to 24 since we set it to 24 hours from now)
-    remaining_hours = spend_metrics["litellm_api_key_budget_remaining_hours_metric"]
-    assert isinstance(remaining_hours, (int, float))
-    assert 20 <= remaining_hours <= 25  # Should be close to 24 hours
-
+    assert spend_metrics["response_cost"] == 0.15  # response_cost
+    assert spend_metrics["user_api_key_spend"] == 0.67  # lol
+    assert spend_metrics["user_api_key_max_budget"] == 10.0  # max budget
+
+    # Verify budget reset is a datetime string in ISO format
+    budget_reset = spend_metrics["user_api_key_budget_reset_at"]
+    assert isinstance(budget_reset, str)
+    print(f"Budget reset in payload: {budget_reset}")    # In StandardLoggingUserAPIKeyMetadata
+    user_api_key_budget_reset_at: Optional[str] = None
+    
+    # In DDLLMObsSpendMetrics  
+    user_api_key_budget_reset_at: str
+    # Should be close to 10 days from now
+    from datetime import datetime, timezone
+    budget_reset_dt = datetime.fromisoformat(budget_reset.replace('Z', '+00:00'))
+    now = datetime.now(timezone.utc)
+    time_diff = (budget_reset_dt - now).total_seconds() / 86400  # days
+    assert 9.5 <= time_diff <= 10.5  # Should be close to 10 days