diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index ba31f2be..35c98496 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -54,6 +54,10 @@ if profiling_env_var: from ddtrace.profiling import profiler +llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1") +if llmobs_env_var: + from ddtrace.llmobs import LLMObs + logger = logging.getLogger(__name__) DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG" @@ -221,6 +225,10 @@ def __init__(self, func): # Patch third-party libraries for tracing patch_all() + # Enable LLM Observability + if llmobs_env_var: + LLMObs.enable() + logger.debug("datadog_lambda_wrapper initialized") except Exception as e: logger.error(format_err_with_traceback(e)) @@ -373,6 +381,9 @@ def _after(self, event, context): # logs api flush_extension() + if llmobs_env_var: + LLMObs.flush() + if self.encode_authorizer_context and is_authorizer_response(self.response): self._inject_authorizer_span_headers( event.get("requestContext", {}).get("requestId")