Skip to content

Commit

Permalink
[MLOB-1009] Auto-Instrument LLM Observability (#500)
Browse files Browse the repository at this point in the history
* add force flush

* add enable to init

* move flush to after parent aws span finishes

* move flushing to be with other flushes
  • Loading branch information
sabrenner authored Jul 1, 2024
1 parent 82970d1 commit 4dd3f62
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions datadog_lambda/wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 4dd3f62

Please sign in to comment.