Enable LLM Observability with agentless_enabled=True by default with a parsed API key (#572)

* enable llmobs agentless with parsed api_key

* extract getting api key to its own function

* lint

Author: sabrenner

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,40 @@ 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 +89,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))
 

datadog_lambda/wrapper.py

@@ -56,10 +56,14 @@
 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 get_api_key
     from ddtrace.llmobs import LLMObs
 
+    llmobs_api_key = get_api_key()
+
 logger = logging.getLogger(__name__)
 
 DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG"
@@ -229,7 +233,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: