Add llmobs_enabled.

Author: purple4reina

datadog_lambda/config.py

@@ -59,6 +59,7 @@ class Config:
     trace_extractor = _get_env("DD_TRACE_EXTRACTOR")
     capture_payload_max_depth = _get_env("DD_CAPTURE_PAYLOAD_MAX_DEPTH", 10, int)
     profiling_enabled = _get_env("DD_PROFILING_ENABLED", "false", as_bool)
+    llmobs_enabled = _get_env("DD_LLMOBS_ENABLED", "false", as_bool)
 
     @property
     def fips_mode_enabled(self):

datadog_lambda/wrapper.py

@@ -48,8 +48,7 @@
 if config.profiling_enabled:
     from ddtrace.profiling import profiler
 
-llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1")
-if llmobs_env_var:
+if config.llmobs_enabled:
     from ddtrace.llmobs import LLMObs
 
 exception_replay_env_var = os.environ.get(
@@ -203,7 +202,7 @@ def __init__(self, func):
             patch_all()
 
             # Enable LLM Observability
-            if llmobs_env_var:
+            if config.llmobs_enabled:
                 LLMObs.enable()
 
             # Enable Exception Replay
@@ -378,7 +377,7 @@ def _after(self, event, context):
                 # logs api
                 flush_extension()
 
-            if llmobs_env_var:
+            if config.llmobs_enabled:
                 LLMObs.flush()
 
             if self.encode_authorizer_context and is_authorizer_response(self.response):

tests/test_config.py

@@ -146,6 +146,15 @@ def set_env(key, value):
     ("DD_PROFILING_ENABLED", "profiling_enabled", "1", True),  # CHANGED
     ("DD_PROFILING_ENABLED", "profiling_enabled", "0", False),
     ("DD_PROFILING_ENABLED", "profiling_enabled", "purple", False),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", None, False),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "", False),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "true", True),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "TRUE", True),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "false", False),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "FALSE", False),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "1", True),  # CHANGED
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "0", False),
+    ("DD_LLMOBS_ENABLED", "llmobs_enabled", "purple", False),
 )