diff --git a/docs/my-website/docs/observability/datadog.md b/docs/my-website/docs/observability/datadog.md index 6f785be1013..9385b0020cf 100644 --- a/docs/my-website/docs/observability/datadog.md +++ b/docs/my-website/docs/observability/datadog.md @@ -253,3 +253,12 @@ LiteLLM supports customizing the following Datadog environment variables \* **Required when using Direct API** (default): `DD_API_KEY` and `DD_SITE` are required \* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required for **Datadog Logs**. (**Note: `DD_API_KEY` IS REQUIRED for Datadog LLM Observability**) +## Automatic Tags + +LiteLLM automatically adds the following tags to your Datadog logs and metrics if the information is available in the request: + +| Tag | Description | Source | +|-----|-------------|--------| +| `team` | The team alias or ID associated with the API Key | `user_api_key_team_alias`, `team_alias`, `user_api_key_team_id`, or `team_id` in metadata | +| `request_tag` | Custom tags passed in the request | `request_tags` in logging payload | + diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 2eb94b59dd8..a961d4f9244 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -93,7 +93,9 @@ def _aggregate_costs( Aggregates costs by Provider, Model, and Date. Returns a list of DatadogFOCUSCostEntry. """ - aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {} + aggregator: Dict[ + Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry + ] = {} for log in logs: try: @@ -167,10 +169,20 @@ def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]: metadata = log.get("metadata", {}) if metadata: # Add user info - if "user_api_key_alias" in metadata: + # Add user info + if metadata.get("user_api_key_alias"): tags["user"] = str(metadata["user_api_key_alias"]) - if "user_api_key_team_alias" in metadata: - tags["team"] = str(metadata["user_api_key_team_alias"]) + + # Add Team Tag + team_tag = ( + metadata.get("user_api_key_team_alias") + or metadata.get("team_alias") # type: ignore + or metadata.get("user_api_key_team_id") + or metadata.get("team_id") # type: ignore + ) + + if team_tag: + tags["team"] = str(team_tag) # model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get() model_group = metadata.get("model_group") # type: ignore[misc] if model_group: diff --git a/litellm/integrations/datadog/datadog_handler.py b/litellm/integrations/datadog/datadog_handler.py index e2f30f2f614..0406f1e5d20 100644 --- a/litellm/integrations/datadog/datadog_handler.py +++ b/litellm/integrations/datadog/datadog_handler.py @@ -55,4 +55,15 @@ def get_datadog_tags( request_tags = standard_logging_object.get("request_tags", []) or [] tags.extend(f"request_tag:{tag}" for tag in request_tags) + # Add Team Tag + metadata = standard_logging_object.get("metadata", {}) or {} + team_tag = ( + metadata.get("user_api_key_team_alias") + or metadata.get("team_alias") + or metadata.get("user_api_key_team_id") + or metadata.get("team_id") + ) + if team_tag: + tags.append(f"team:{team_tag}") + return ",".join(tags) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index bdbbc7579b7..6a14e42c485 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1335,7 +1335,11 @@ def set_cost_breakdown( ) # Store additional costs if provided (free-form dict for extensibility) - if additional_costs and isinstance(additional_costs, dict) and len(additional_costs) > 0: + if ( + additional_costs + and isinstance(additional_costs, dict) + and len(additional_costs) > 0 + ): self.cost_breakdown["additional_costs"] = additional_costs # Store discount information if provided @@ -4519,13 +4523,19 @@ def get_standard_logging_metadata( requester_custom_headers=None, cold_storage_object_key=None, user_api_key_auth_metadata=None, + team_alias=None, + team_id=None, ) if isinstance(metadata, dict): for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS: clean_metadata[key] = metadata[key] # type: ignore user_api_key = metadata.get("user_api_key") - if user_api_key and isinstance(user_api_key, str) and is_valid_sha256_hash(user_api_key): + if ( + user_api_key + and isinstance(user_api_key, str) + and is_valid_sha256_hash(user_api_key) + ): clean_metadata["user_api_key_hash"] = user_api_key _potential_requester_metadata = metadata.get( "metadata", None @@ -5279,6 +5289,8 @@ def get_standard_logging_metadata( user_api_key_request_route=None, cold_storage_object_key=None, user_api_key_auth_metadata=None, + team_alias=None, + team_id=None, ) if isinstance(metadata, dict): # Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5f8798c7712..8eb13f333f3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2532,6 +2532,8 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata): cold_storage_object_key: Optional[ str ] # S3/GCS object key for cold storage retrieval + team_alias: Optional[str] + team_id: Optional[str] class StandardLoggingAdditionalHeaders(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py b/tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py new file mode 100644 index 00000000000..3f1d2be4137 --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_tags_regression.py @@ -0,0 +1,91 @@ +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../")) + +from litellm.integrations.datadog.datadog_handler import get_datadog_tags +from litellm.integrations.datadog.datadog_cost_management import ( + DatadogCostManagementLogger, +) +from litellm.types.utils import StandardLoggingPayload, StandardLoggingMetadata + + +class TestDatadogTagsRegression: + @pytest.fixture + def mock_env_vars(self): + """Mock environment variables to isolate environment.""" + with patch.dict( + os.environ, + { + "DD_ENV": "test-env", + "DD_SERVICE": "test-service", + "DD_VERSION": "1.0.0", + "HOSTNAME": "test-host", + "POD_NAME": "test-pod", + "DD_API_KEY": "mock-api-key", + "DD_APP_KEY": "mock-app-key", + }, + ): + yield + + def test_get_datadog_tags_regression(self, mock_env_vars): + """ + Regression Test: Ensure that get_datadog_tags still produces basic tags correctly + AND now includes the new team tag when provided. + """ + # Case 1: Legacy behavior (no team info) + payload_legacy = StandardLoggingPayload(metadata={}) + tags_legacy = get_datadog_tags(payload_legacy) + + # Verify base tags exist (legacy requirement) + assert "env:test-env" in tags_legacy + assert "service:test-service" in tags_legacy + # Verify NO team tag (should not invent one) + assert "team:" not in tags_legacy + + # Case 2: New feature (team info provided) + payload_with_team = StandardLoggingPayload( + metadata=StandardLoggingMetadata(user_api_key_team_alias="regression-team") + ) + tags_with_team = get_datadog_tags(payload_with_team) + + # Verify base tags STILL exist + assert "env:test-env" in tags_with_team + assert "service:test-service" in tags_with_team + # Verify NEW team tag is added + assert "team:regression-team" in tags_with_team + + @pytest.mark.asyncio + async def test_datadog_cost_management_tags_regression(self, mock_env_vars): + """ + Regression Test: Ensure DatadogCostManagementLogger extracts tags correctly, + preserving existing behavior while adding the team tag capability. + """ + logger = DatadogCostManagementLogger() + + # Case 1: Legacy metadata (user alias only) + payload_legacy = StandardLoggingPayload( + metadata=StandardLoggingMetadata(user_api_key_alias="legacy-user") + ) + + tags_legacy = logger._extract_tags(payload_legacy) + + assert tags_legacy["env"] == "test-env" + assert tags_legacy["user"] == "legacy-user" + assert "team" not in tags_legacy # Should not exist + + # Case 2: New metadata (team alias) + payload_new = StandardLoggingPayload( + metadata=StandardLoggingMetadata( + user_api_key_alias="new-user", user_api_key_team_alias="new-team-alias" + ) + ) + + tags_new = logger._extract_tags(payload_new) + + assert tags_new["env"] == "test-env" + assert tags_new["user"] == "new-user" + assert tags_new["team"] == "new-team-alias" # New feature verified