fix(integrations): Auto-wrap root gen_ai spans for openai, cohere, langgraph, huggingface_hub

Author: constantinius

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,

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,

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,

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,

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):

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,14 +2268,15 @@ 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(
                 model="some-model", messages=[{"role": "system", "content": "hello"}]
             )
 
         (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,14 +2885,15 @@ 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(
                 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()
 
@@ -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])