diff --git a/sentry_sdk/ai/monitoring.py b/sentry_sdk/ai/monitoring.py index d8840ad451..423bb50f9d 100644 --- a/sentry_sdk/ai/monitoring.py +++ b/sentry_sdk/ai/monitoring.py @@ -9,6 +9,7 @@ from sentry_sdk.consts import SPANDATA from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing import Span +from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import ContextVar, capture_internal_exceptions, reraise if TYPE_CHECKING: @@ -30,64 +31,141 @@ def get_ai_pipeline_name() -> "Optional[str]": def ai_track(description: str, **span_kwargs: "Any") -> "Callable[[F], F]": def decorator(f: "F") -> "F": def sync_wrapped(*args: "Any", **kwargs: "Any") -> "Any": + client = sentry_sdk.get_client() + curr_pipeline = _ai_pipeline_name.get() op = span_kwargs.pop("op", "ai.run" if curr_pipeline else "ai.pipeline") - with start_span(name=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(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline) - return f(*args, **kwargs) - else: - _ai_pipeline_name.set(description) - try: - res = f(*args, **kwargs) - except Exception as e: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - event, hint = sentry_sdk.utils.event_from_exception( - e, - client_options=sentry_sdk.get_client().options, - mechanism={"type": "ai_monitoring", "handled": False}, - ) - sentry_sdk.capture_event(event, hint=hint) - reraise(*exc_info) - finally: - _ai_pipeline_name.set(None) - return res + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=description, attributes={"sentry.op": op} + ) as span: + for k, v in kwargs.pop("sentry_tags", {}).items(): + span.set_attribute(k, v) + for k, v in kwargs.pop("sentry_data", {}).items(): + span.set_attribute(k, v) + + if curr_pipeline: + span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline) + return f(*args, **kwargs) + else: + _ai_pipeline_name.set(description) + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + event, hint = sentry_sdk.utils.event_from_exception( + e, + client_options=sentry_sdk.get_client().options, + mechanism={ + "type": "ai_monitoring", + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + reraise(*exc_info) + finally: + _ai_pipeline_name.set(None) + return res + + else: + with start_span(name=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(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline) + return f(*args, **kwargs) + else: + _ai_pipeline_name.set(description) + try: + res = f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + event, hint = sentry_sdk.utils.event_from_exception( + e, + client_options=sentry_sdk.get_client().options, + mechanism={ + "type": "ai_monitoring", + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + reraise(*exc_info) + finally: + _ai_pipeline_name.set(None) + return res async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any": + client = sentry_sdk.get_client() + curr_pipeline = _ai_pipeline_name.get() op = span_kwargs.pop("op", "ai.run" if curr_pipeline else "ai.pipeline") - with start_span(name=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(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline) - return await f(*args, **kwargs) - else: - _ai_pipeline_name.set(description) - try: - res = await f(*args, **kwargs) - except Exception as e: - exc_info = sys.exc_info() - with capture_internal_exceptions(): - event, hint = sentry_sdk.utils.event_from_exception( - e, - client_options=sentry_sdk.get_client().options, - mechanism={"type": "ai_monitoring", "handled": False}, - ) - sentry_sdk.capture_event(event, hint=hint) - reraise(*exc_info) - finally: - _ai_pipeline_name.set(None) - return res + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=description, attributes={"sentry.op": op} + ) as span: + for k, v in kwargs.pop("sentry_tags", {}).items(): + span.set_attribute(k, v) + for k, v in kwargs.pop("sentry_data", {}).items(): + span.set_attribute(k, v) + + if curr_pipeline: + span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline) + return await f(*args, **kwargs) + else: + _ai_pipeline_name.set(description) + try: + res = await f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + event, hint = sentry_sdk.utils.event_from_exception( + e, + client_options=sentry_sdk.get_client().options, + mechanism={ + "type": "ai_monitoring", + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + reraise(*exc_info) + finally: + _ai_pipeline_name.set(None) + return res + else: + with start_span(name=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(SPANDATA.GEN_AI_PIPELINE_NAME, curr_pipeline) + return await f(*args, **kwargs) + else: + _ai_pipeline_name.set(description) + try: + res = await f(*args, **kwargs) + except Exception as e: + exc_info = sys.exc_info() + with capture_internal_exceptions(): + event, hint = sentry_sdk.utils.event_from_exception( + e, + client_options=sentry_sdk.get_client().options, + mechanism={ + "type": "ai_monitoring", + "handled": False, + }, + ) + sentry_sdk.capture_event(event, hint=hint) + reraise(*exc_info) + finally: + _ai_pipeline_name.set(None) + return res if inspect.iscoroutinefunction(f): return wraps(f)(async_wrapped) # type: ignore diff --git a/tests/test_ai_monitoring.py b/tests/test_ai_monitoring.py index 70b4a6bcdd..51a2c67f03 100644 --- a/tests/test_ai_monitoring.py +++ b/tests/test_ai_monitoring.py @@ -52,6 +52,31 @@ def pipeline(): assert ai_run_span["description"] == "my tool" +def test_ai_track_span_streaming(sentry_init, capture_items): + sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream") + items = capture_items("span") + + @ai_track("my tool") + def tool(**kwargs): + pass + + @ai_track("some test pipeline") + def pipeline(): + tool() + + with sentry_sdk.traces.start_span(name="span"): + pipeline() + + sentry_sdk.flush() + + spans = [item.payload for item in items] + assert len(spans) == 3 + ai_run_span, ai_pipeline_span, segment = spans + + assert ai_pipeline_span["name"] == "some test pipeline" + assert ai_run_span["name"] == "my tool" + + def test_ai_track_with_tags(sentry_init, capture_events): sentry_init(traces_sample_rate=1.0) events = capture_events() @@ -76,12 +101,39 @@ def pipeline(): ai_run_span = spans[0] if spans[0]["op"] == "ai.run" else spans[1] assert ai_pipeline_span["description"] == "some test pipeline" - print(ai_pipeline_span) assert ai_pipeline_span["tags"]["user"] == "colin" assert ai_pipeline_span["data"]["some_data"] == "value" assert ai_run_span["description"] == "my tool" +def test_ai_track_with_attributes_span_streaming(sentry_init, capture_items): + sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream") + items = capture_items("span") + + @ai_track("my tool") + def tool(**kwargs): + pass + + @ai_track("some test pipeline") + def pipeline(): + tool() + + with sentry_sdk.traces.start_span(name="span"): + pipeline(sentry_tags={"user": "colin"}, sentry_data={"some_data": "value"}) + + sentry_sdk.flush() + + spans = [item.payload for item in items] + + assert len(spans) == 3 + ai_run_span, ai_pipeline_span, segment = spans + + assert ai_pipeline_span["name"] == "some test pipeline" + assert ai_pipeline_span["attributes"]["user"] == "colin" + assert ai_pipeline_span["attributes"]["some_data"] == "value" + assert ai_run_span["name"] == "my tool" + + @pytest.mark.asyncio async def test_ai_track_async(sentry_init, capture_events): sentry_init(traces_sample_rate=1.0) @@ -110,6 +162,32 @@ async def async_pipeline(): assert ai_run_span["description"] == "my async tool" +@pytest.mark.asyncio +async def test_ai_track_async_span_streaming(sentry_init, capture_items): + sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream") + items = capture_items() + + @ai_track("my async tool") + async def async_tool(**kwargs): + pass + + @ai_track("some async test pipeline") + async def async_pipeline(): + await async_tool() + + with sentry_sdk.traces.start_span(name="span"): + await async_pipeline() + + sentry_sdk.flush() + + spans = [item.payload for item in items] + assert len(spans) == 3 + ai_run_span, ai_pipeline_span, segment = spans + + assert ai_pipeline_span["name"] == "some async test pipeline" + assert ai_run_span["name"] == "my async tool" + + @pytest.mark.asyncio async def test_ai_track_async_with_tags(sentry_init, capture_events): sentry_init(traces_sample_rate=1.0) @@ -142,6 +220,38 @@ async def async_pipeline(): assert ai_run_span["description"] == "my async tool" +@pytest.mark.asyncio +async def test_ai_track_async_with_attributes_span_streaming( + sentry_init, capture_items +): + sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream") + items = capture_items("span") + + @ai_track("my async tool") + async def async_tool(**kwargs): + pass + + @ai_track("some async test pipeline") + async def async_pipeline(): + await async_tool() + + with sentry_sdk.traces.start_span(name="span"): + await async_pipeline( + sentry_tags={"user": "czyber"}, sentry_data={"some_data": "value"} + ) + + sentry_sdk.flush() + + spans = [item.payload for item in items] + assert len(spans) == 3 + ai_run_span, ai_pipeline_span, segment = spans + + assert ai_pipeline_span["name"] == "some async test pipeline" + assert ai_pipeline_span["attributes"]["user"] == "czyber" + assert ai_pipeline_span["attributes"]["some_data"] == "value" + assert ai_run_span["name"] == "my async tool" + + def test_ai_track_with_explicit_op(sentry_init, capture_events): sentry_init(traces_sample_rate=1.0) events = capture_events() @@ -162,6 +272,27 @@ def tool(**kwargs): assert span["op"] == "custom.operation" +def test_ai_track_with_explicit_op_span_streaming(sentry_init, capture_items): + sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream") + items = capture_items("span") + + @ai_track("my tool", op="custom.operation") + def tool(**kwargs): + pass + + with sentry_sdk.traces.start_span(name="span"): + tool() + + sentry_sdk.flush() + + spans = [item.payload for item in items] + assert len(spans) == 2 + ai_run_span, segment = spans + + assert ai_run_span["name"] == "my tool" + assert ai_run_span["attributes"]["sentry.op"] == "custom.operation" + + @pytest.mark.asyncio async def test_ai_track_async_with_explicit_op(sentry_init, capture_events): sentry_init(traces_sample_rate=1.0) @@ -183,6 +314,30 @@ async def async_tool(**kwargs): assert span["op"] == "custom.async.operation" +@pytest.mark.asyncio +async def test_ai_track_async_with_explicit_op_span_streaming( + sentry_init, capture_items +): + sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream") + items = capture_items() + + @ai_track("my async tool", op="custom.async.operation") + async def async_tool(**kwargs): + pass + + with sentry_sdk.traces.start_span(name="span"): + await async_tool() + + sentry_sdk.flush() + + spans = [item.payload for item in items] + assert len(spans) == 2 + ai_run_span, segment = spans + + assert ai_run_span["name"] == "my async tool" + assert ai_run_span["attributes"]["sentry.op"] == "custom.async.operation" + + @pytest.fixture def sample_messages(): """Sample messages similar to what gen_ai integrations would use"""