From 39cf2c0aac1e845e8f6a4521fd54293c9e3f8b49 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Fri, 28 Jun 2024 10:07:58 -0400 Subject: [PATCH 1/4] add force flush --- datadog_lambda/wrapper.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index ba31f2be..d844d16a 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" @@ -330,6 +334,9 @@ def _after(self, event, context): if should_trace_cold_start: trace_ctx = tracer.current_trace_context() + if llmobs_env_var: + LLMObs.flush() + if self.span: if dd_capture_lambda_payload_enabled: tag_object.tag_object(self.span, "function.request", event) From 4144819dd29e6b9c1a10e2363b2697d27f88f21e Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Fri, 28 Jun 2024 15:41:39 -0400 Subject: [PATCH 2/4] add enable to init --- datadog_lambda/wrapper.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index d844d16a..a7432ce0 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -225,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)) From 18f3be8808eb6f0443b24015116308d856313d74 Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Sat, 29 Jun 2024 09:26:31 -0400 Subject: [PATCH 3/4] move flush to after parent aws span finishes --- datadog_lambda/wrapper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index a7432ce0..6965c9a4 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -338,9 +338,6 @@ def _after(self, event, context): if should_trace_cold_start: trace_ctx = tracer.current_trace_context() - if llmobs_env_var: - LLMObs.flush() - if self.span: if dd_capture_lambda_payload_enabled: tag_object.tag_object(self.span, "function.request", event) @@ -350,6 +347,9 @@ def _after(self, event, context): self.span.set_tag("http.status_code", status_code) self.span.finish() + if llmobs_env_var: + LLMObs.flush() + if self.inferred_span: if status_code: self.inferred_span.set_tag("http.status_code", status_code) From 9c52f42c17b662dd89f0ccc39d968ed9135c22ab Mon Sep 17 00:00:00 2001 From: Sam Brenner Date: Mon, 1 Jul 2024 13:55:10 -0400 Subject: [PATCH 4/4] move flushing to be with other flushes --- datadog_lambda/wrapper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 6965c9a4..35c98496 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -347,9 +347,6 @@ def _after(self, event, context): self.span.set_tag("http.status_code", status_code) self.span.finish() - if llmobs_env_var: - LLMObs.flush() - if self.inferred_span: if status_code: self.inferred_span.set_tag("http.status_code", status_code) @@ -384,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")