Skip to content

Commit 4dd3f62

Browse files
authored
[MLOB-1009] Auto-Instrument LLM Observability (#500)
* add force flush * add enable to init * move flush to after parent aws span finishes * move flushing to be with other flushes
1 parent 82970d1 commit 4dd3f62

File tree

1 file changed

+11
-0
lines changed

1 file changed

+11
-0
lines changed

datadog_lambda/wrapper.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@
5454
if profiling_env_var:
5555
from ddtrace.profiling import profiler
5656

57+
llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1")
58+
if llmobs_env_var:
59+
from ddtrace.llmobs import LLMObs
60+
5761
logger = logging.getLogger(__name__)
5862

5963
DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG"
@@ -221,6 +225,10 @@ def __init__(self, func):
221225
# Patch third-party libraries for tracing
222226
patch_all()
223227

228+
# Enable LLM Observability
229+
if llmobs_env_var:
230+
LLMObs.enable()
231+
224232
logger.debug("datadog_lambda_wrapper initialized")
225233
except Exception as e:
226234
logger.error(format_err_with_traceback(e))
@@ -373,6 +381,9 @@ def _after(self, event, context):
373381
# logs api
374382
flush_extension()
375383

384+
if llmobs_env_var:
385+
LLMObs.flush()
386+
376387
if self.encode_authorizer_context and is_authorizer_response(self.response):
377388
self._inject_authorizer_span_headers(
378389
event.get("requestContext", {}).get("requestId")

0 commit comments

Comments
 (0)