Skip to content

Commit

Permalink
(feat) dd logger - set tags according to the values set by those env …
Browse files Browse the repository at this point in the history
…vars (#6933)

* dd logger, inherit from .envs

* test_datadog_payload_environment_variables

* fix _get_datadog_service
  • Loading branch information
ishaan-jaff authored Nov 27, 2024
1 parent fe151db commit a6da3de
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 13 deletions.
44 changes: 32 additions & 12 deletions litellm/integrations/datadog/datadog.py
Original file line number Diff line number Diff line change
Expand Up @@ -279,11 +279,11 @@ def create_datadog_logging_payload(
verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload)

dd_payload = DatadogPayload(
ddsource=os.getenv("DD_SOURCE", "litellm"),
ddtags="",
hostname="",
ddsource=self._get_datadog_source(),
ddtags=self._get_datadog_tags(),
hostname=self._get_datadog_hostname(),
message=json_payload,
service="litellm-server",
service=self._get_datadog_service(),
status=status,
)
return dd_payload
Expand Down Expand Up @@ -387,11 +387,11 @@ async def async_post_call_failure_hook(
json_payload = json.dumps(_exception_payload)
verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload)
dd_payload = DatadogPayload(
ddsource=os.getenv("DD_SOURCE", "litellm"),
ddtags="",
hostname="",
ddsource=self._get_datadog_source(),
ddtags=self._get_datadog_tags(),
hostname=self._get_datadog_hostname(),
message=json_payload,
service="litellm-server",
service=self._get_datadog_service(),
status=DataDogStatus.ERROR,
)

Expand Down Expand Up @@ -473,11 +473,31 @@ def _create_v0_logging_payload(
verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload)

dd_payload = DatadogPayload(
ddsource=os.getenv("DD_SOURCE", "litellm"),
ddtags="",
hostname="",
ddsource=self._get_datadog_source(),
ddtags=self._get_datadog_tags(),
hostname=self._get_datadog_hostname(),
message=json_payload,
service="litellm-server",
service=self._get_datadog_service(),
status=DataDogStatus.INFO,
)
return dd_payload

@staticmethod
def _get_datadog_tags():
return f"env:{os.getenv('DD_ENV', 'unknown')},service:{os.getenv('DD_SERVICE', 'litellm')},version:{os.getenv('DD_VERSION', 'unknown')}"

@staticmethod
def _get_datadog_source():
return os.getenv("DD_SOURCE", "litellm")

@staticmethod
def _get_datadog_service():
return os.getenv("DD_SERVICE", "litellm-server")

@staticmethod
def _get_datadog_hostname():
return ""

@staticmethod
def _get_datadog_env():
return os.getenv("DD_ENV", "unknown")
2 changes: 1 addition & 1 deletion litellm/proxy/proxy_config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@ include:
- model_config.yaml

litellm_settings:
callbacks: ["prometheus"]
callbacks: ["datadog"]
44 changes: 44 additions & 0 deletions tests/logging_callback_tests/test_datadog.py
Original file line number Diff line number Diff line change
Expand Up @@ -422,3 +422,47 @@ def __init__(self):

except Exception as e:
pytest.fail(f"Test failed with exception: {str(e)}")


@pytest.mark.asyncio
async def test_datadog_payload_environment_variables():
"""Test that DataDog payload correctly includes environment variables in the payload structure"""
try:
# Set test environment variables
test_env = {
"DD_ENV": "test-env",
"DD_SERVICE": "test-service",
"DD_VERSION": "1.0.0",
"DD_SOURCE": "test-source",
"DD_API_KEY": "fake-key",
"DD_SITE": "datadoghq.com",
}

with patch.dict(os.environ, test_env):
dd_logger = DataDogLogger()
standard_payload = create_standard_logging_payload()

# Create the payload
dd_payload = dd_logger.create_datadog_logging_payload(
kwargs={"standard_logging_object": standard_payload},
response_obj=None,
start_time=datetime.now(),
end_time=datetime.now(),
)

print("dd payload=", json.dumps(dd_payload, indent=2))

# Verify payload structure and environment variables
assert (
dd_payload["ddsource"] == "test-source"
), "Incorrect source in payload"
assert (
dd_payload["service"] == "test-service"
), "Incorrect service in payload"
assert (
dd_payload["ddtags"]
== "env:test-env,service:test-service,version:1.0.0"
), "Incorrect tags in payload"

except Exception as e:
pytest.fail(f"Test failed with exception: {str(e)}")

0 comments on commit a6da3de

Please sign in to comment.