diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 35b47a6243..a88f16ffb1 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -4,7 +4,7 @@ from sentry_sdk import consts from sentry_sdk.ai.monitoring import record_token_usage -from sentry_sdk.ai.utils import set_data_normalized +from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized from sentry_sdk.consts import SPANDATA from sentry_sdk.tracing_utils import set_span_errored @@ -142,7 +142,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any": message = kwargs.get("message") - span = sentry_sdk.start_span( + span = get_start_span_function()( op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE, name="cohere.client.Chat", origin=CohereIntegration.origin, @@ -225,7 +225,7 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any": if integration is None: return f(*args, **kwargs) - with sentry_sdk.start_span( + with get_start_span_function()( op=consts.OP.COHERE_EMBEDDINGS_CREATE, name="Cohere Embedding Creation", origin=CohereIntegration.origin, diff --git a/sentry_sdk/integrations/huggingface_hub.py b/sentry_sdk/integrations/huggingface_hub.py index a388baa136..835acc7279 100644 --- a/sentry_sdk/integrations/huggingface_hub.py +++ b/sentry_sdk/integrations/huggingface_hub.py @@ -5,7 +5,11 @@ import sentry_sdk from sentry_sdk.ai.monitoring import record_token_usage -from sentry_sdk.ai.utils import _set_span_data_attribute, set_data_normalized +from sentry_sdk.ai.utils import ( + _set_span_data_attribute, + get_start_span_function, + set_data_normalized, +) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration from sentry_sdk.scope import should_send_default_pii @@ -97,7 +101,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any": }, ) else: - span = sentry_sdk.start_span( + span = get_start_span_function()( op=op, name=f"{operation_name} {model}", origin=HuggingfaceHubIntegration.origin, diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 1334769685..acf630add2 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -3,6 +3,7 @@ import sentry_sdk from sentry_sdk.ai.utils import ( + get_start_span_function, normalize_message_roles, set_data_normalized, truncate_and_annotate_messages, @@ -159,7 +160,7 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" ) - with sentry_sdk.start_span( + with get_start_span_function()( op=OP.GEN_AI_INVOKE_AGENT, name=span_name, origin=LanggraphIntegration.origin, @@ -219,7 +220,7 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" ) - with sentry_sdk.start_span( + with get_start_span_function()( op=OP.GEN_AI_INVOKE_AGENT, name=span_name, origin=LanggraphIntegration.origin, diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index f0f95eef58..295ab96d34 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -25,6 +25,7 @@ ) from sentry_sdk.ai.monitoring import record_token_usage from sentry_sdk.ai.utils import ( + get_start_span_function, normalize_message_roles, set_data_normalized, truncate_and_annotate_embedding_inputs, @@ -713,7 +714,7 @@ def _new_sync_chat_completion(f: "Any", *args: "Any", **kwargs: "Any") -> "Any": model = kwargs.get("model") - span = sentry_sdk.start_span( + span = get_start_span_function()( op=consts.OP.GEN_AI_CHAT, name=f"chat {model}", origin=OpenAIIntegration.origin, @@ -781,7 +782,7 @@ async def _new_async_chat_completion(f: "Any", *args: "Any", **kwargs: "Any") -> model = kwargs.get("model") - span = sentry_sdk.start_span( + span = get_start_span_function()( op=consts.OP.GEN_AI_CHAT, name=f"chat {model}", origin=OpenAIIntegration.origin, @@ -1177,7 +1178,7 @@ def _new_sync_embeddings_create(f: "Any", *args: "Any", **kwargs: "Any") -> "Any model = kwargs.get("model") - with sentry_sdk.start_span( + with get_start_span_function()( op=consts.OP.GEN_AI_EMBEDDINGS, name=f"embeddings {model}", origin=OpenAIIntegration.origin, @@ -1209,7 +1210,7 @@ async def _new_async_embeddings_create( model = kwargs.get("model") - with sentry_sdk.start_span( + with get_start_span_function()( op=consts.OP.GEN_AI_EMBEDDINGS, name=f"embeddings {model}", origin=OpenAIIntegration.origin, @@ -1263,7 +1264,7 @@ def _new_sync_responses_create(f: "Any", *args: "Any", **kwargs: "Any") -> "Any" model = kwargs.get("model") - span = sentry_sdk.start_span( + span = get_start_span_function()( op=consts.OP.GEN_AI_RESPONSES, name=f"responses {model}", origin=OpenAIIntegration.origin, @@ -1321,7 +1322,7 @@ async def _new_async_responses_create(f: "Any", *args: "Any", **kwargs: "Any") - model = kwargs.get("model") - span = sentry_sdk.start_span( + span = get_start_span_function()( op=consts.OP.GEN_AI_RESPONSES, name=f"responses {model}", origin=OpenAIIntegration.origin, diff --git a/tests/integrations/cohere/test_cohere.py b/tests/integrations/cohere/test_cohere.py index 9528eae6b7..8c26ac0a91 100644 --- a/tests/integrations/cohere/test_cohere.py +++ b/tests/integrations/cohere/test_cohere.py @@ -162,8 +162,9 @@ def test_bad_chat(sentry_init, capture_events): with pytest.raises(httpx.HTTPError): client.chat(model="some-model", message="hello") - (event,) = events + (event, transaction) = events assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" def test_span_status_error(sentry_init, capture_events): diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index 2a7b8834b0..e6dc003ccc 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -2164,7 +2164,7 @@ def test_bad_chat_completion( ) if stream_gen_ai_spans: - items = capture_items("event") + items = capture_items("event", "transaction") client = OpenAI(api_key="z") client.chat.completions._post = mock.Mock( @@ -2177,6 +2177,7 @@ def test_bad_chat_completion( ) (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") else: events = capture_events() @@ -2190,9 +2191,10 @@ def test_bad_chat_completion( messages=[{"role": "system", "content": "hello"}], ) - (event,) = events + (event, transaction) = events assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @@ -2266,7 +2268,7 @@ async def test_bad_chat_completion_async( side_effect=OpenAIError("API rate limit reached") ) if stream_gen_ai_spans: - items = capture_items("event") + items = capture_items("event", "transaction") with pytest.raises(OpenAIError): await client.chat.completions.create( @@ -2274,6 +2276,7 @@ async def test_bad_chat_completion_async( ) (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") else: events = capture_events() @@ -2282,9 +2285,10 @@ async def test_bad_chat_completion_async( model="some-model", messages=[{"role": "system", "content": "hello"}] ) - (event,) = events + (event, transaction) = events assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @@ -2834,21 +2838,23 @@ def test_embeddings_create_raises_error( ) if stream_gen_ai_spans: - items = capture_items("event") + items = capture_items("event", "transaction") with pytest.raises(OpenAIError): client.embeddings.create(input="hello", model="text-embedding-3-large") (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") else: events = capture_events() with pytest.raises(OpenAIError): client.embeddings.create(input="hello", model="text-embedding-3-large") - (event,) = events + (event, transaction) = events assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @@ -2879,7 +2885,7 @@ async def test_embeddings_create_raises_error_async( ) if stream_gen_ai_spans: - items = capture_items("event") + items = capture_items("event", "transaction") with pytest.raises(OpenAIError): await client.embeddings.create( @@ -2887,6 +2893,7 @@ async def test_embeddings_create_raises_error_async( ) (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") else: events = capture_events() @@ -2895,9 +2902,10 @@ async def test_embeddings_create_raises_error_async( input="hello", model="text-embedding-3-large" ) - (event,) = events + (event, transaction) = events assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])