From 06efff365dd082888653e393508e7f6c2e4f1831 Mon Sep 17 00:00:00 2001 From: colin-sentry <161344340+colin-sentry@users.noreply.github.com> Date: Wed, 15 May 2024 12:48:33 -0400 Subject: [PATCH] Add tags + data passing functionality to @ai_track (#3071) Co-authored-by: Anton Pirker --- sentry_sdk/ai/monitoring.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/sentry_sdk/ai/monitoring.py b/sentry_sdk/ai/monitoring.py index f5f9cd7aad..bd48ffa053 100644 --- a/sentry_sdk/ai/monitoring.py +++ b/sentry_sdk/ai/monitoring.py @@ -31,7 +31,12 @@ def wrapped(*args, **kwargs): # type: (Any, Any) -> Any curr_pipeline = _ai_pipeline_name.get() op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline") + with start_span(description=description, op=op, **span_kwargs) as span: + for k, v in kwargs.pop("sentry_tags", {}).items(): + span.set_tag(k, v) + for k, v in kwargs.pop("sentry_data", {}).items(): + span.set_data(k, v) if curr_pipeline: span.set_data("ai.pipeline.name", curr_pipeline) return f(*args, **kwargs)