From 82a943a358492fa9f4ad7cece39384448bf12147 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Mon, 10 Mar 2025 13:25:05 -0400 Subject: [PATCH 1/3] enable llmobs agentless with parsed api_key --- datadog_lambda/wrapper.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 2632d22e..3c97f706 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -56,10 +56,16 @@ if profiling_env_var: from ddtrace.profiling import profiler +llmobs_api_key = None llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1") if llmobs_env_var: + from datadog_lambda.api import init_api + from datadog import api from ddtrace.llmobs import LLMObs + init_api() + llmobs_api_key = api._api_key + logger = logging.getLogger(__name__) DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG" @@ -229,7 +235,10 @@ def __init__(self, func): # Enable LLM Observability if llmobs_env_var: - LLMObs.enable() + LLMObs.enable( + agentless_enabled=True, + api_key=llmobs_api_key, + ) logger.debug("datadog_lambda_wrapper initialized") except Exception as e: From b2cc1ba781a39fa4ebfb6e29704d42d956107815 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Mon, 10 Mar 2025 17:06:10 -0400 Subject: [PATCH 2/3] extract getting api key to its own function --- datadog_lambda/api.py | 59 ++++++++++++++++++++++++--------------- datadog_lambda/wrapper.py | 6 ++-- 2 files changed, 39 insertions(+), 26 deletions(-) diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index b5414fd9..0e5d95e8 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -4,6 +4,7 @@ logger = logging.getLogger(__name__) KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName" +api_key = None def decrypt_kms_api_key(kms_client, ciphertext): @@ -46,6 +47,41 @@ def decrypt_kms_api_key(kms_client, ciphertext): return plaintext +def get_api_key() -> str: + """ + Gets the Datadog API key from the environment variables or secrets manager. + Extracts the result to a global value to avoid repeated calls to the secrets manager from different products. + """ + global api_key + if api_key: + return api_key + + import boto3 + + DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "") + DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "") + DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "") + DD_API_KEY = os.environ.get( + "DD_API_KEY", os.environ.get("DATADOG_API_KEY", "") + ) + + if DD_API_KEY_SECRET_ARN: + api_key = boto3.client("secretsmanager").get_secret_value( + SecretId=DD_API_KEY_SECRET_ARN + )["SecretString"] + elif DD_API_KEY_SSM_NAME: + api_key = boto3.client("ssm").get_parameter( + Name=DD_API_KEY_SSM_NAME, WithDecryption=True + )["Parameter"]["Value"] + elif DD_KMS_API_KEY: + kms_client = boto3.client("kms") + api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY) + else: + api_key = DD_API_KEY + + return api_key + + def init_api(): if not os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true": # Make sure that this package would always be lazy-loaded/outside from the critical path @@ -54,28 +90,7 @@ def init_api(): from datadog import api if not api._api_key: - import boto3 - - DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "") - DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "") - DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "") - DD_API_KEY = os.environ.get( - "DD_API_KEY", os.environ.get("DATADOG_API_KEY", "") - ) - - if DD_API_KEY_SECRET_ARN: - api._api_key = boto3.client("secretsmanager").get_secret_value( - SecretId=DD_API_KEY_SECRET_ARN - )["SecretString"] - elif DD_API_KEY_SSM_NAME: - api._api_key = boto3.client("ssm").get_parameter( - Name=DD_API_KEY_SSM_NAME, WithDecryption=True - )["Parameter"]["Value"] - elif DD_KMS_API_KEY: - kms_client = boto3.client("kms") - api._api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY) - else: - api._api_key = DD_API_KEY + api._api_key = get_api_key() logger.debug("Setting DATADOG_API_KEY of length %d", len(api._api_key)) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 3c97f706..6afa9a07 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -59,12 +59,10 @@ llmobs_api_key = None llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1") if llmobs_env_var: - from datadog_lambda.api import init_api - from datadog import api + from datadog_lambda.api import get_api_key from ddtrace.llmobs import LLMObs - init_api() - llmobs_api_key = api._api_key + llmobs_api_key = get_api_key() logger = logging.getLogger(__name__) From 06e4596dd25187435bd77c61021d94822729b7a6 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Mon, 10 Mar 2025 19:41:13 -0400 Subject: [PATCH 3/3] lint --- datadog_lambda/api.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index 0e5d95e8..a114fe8f 100644 --- a/datadog_lambda/api.py +++ b/datadog_lambda/api.py @@ -50,7 +50,8 @@ def decrypt_kms_api_key(kms_client, ciphertext): def get_api_key() -> str: """ Gets the Datadog API key from the environment variables or secrets manager. - Extracts the result to a global value to avoid repeated calls to the secrets manager from different products. + Extracts the result to a global value to avoid repeated calls to the + secrets manager from different products. """ global api_key if api_key: @@ -61,9 +62,7 @@ def get_api_key() -> str: DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "") DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "") DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "") - DD_API_KEY = os.environ.get( - "DD_API_KEY", os.environ.get("DATADOG_API_KEY", "") - ) + DD_API_KEY = os.environ.get("DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")) if DD_API_KEY_SECRET_ARN: api_key = boto3.client("secretsmanager").get_secret_value(