use get_start_span_function

Author: alexander-alderman-webb

sentry_sdk/integrations/google_genai/__init__.py

@@ -73,7 +73,7 @@ def new_generate_content_stream(
 
         _model, contents, model_name = prepare_generate_content_args(args, kwargs)
 
-        chat_span = sentry_sdk.start_span(
+        chat_span = get_start_span_function()(
             op=OP.GEN_AI_CHAT,
             name=f"chat {model_name}",
             origin=ORIGIN,
@@ -132,7 +132,7 @@ async def new_async_generate_content_stream(
 
         _model, contents, model_name = prepare_generate_content_args(args, kwargs)
 
-        chat_span = sentry_sdk.start_span(
+        chat_span = get_start_span_function()(
             op=OP.GEN_AI_CHAT,
             name=f"chat {model_name}",
             origin=ORIGIN,
@@ -187,7 +187,7 @@ def new_generate_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
         model, contents, model_name = prepare_generate_content_args(args, kwargs)
 
-        with sentry_sdk.start_span(
+        with get_start_span_function()(
             op=OP.GEN_AI_CHAT,
             name=f"chat {model_name}",
             origin=ORIGIN,
@@ -225,7 +225,7 @@ async def new_async_generate_content(
 
         model, contents, model_name = prepare_generate_content_args(args, kwargs)
 
-        with sentry_sdk.start_span(
+        with get_start_span_function()(
             op=OP.GEN_AI_CHAT,
             name=f"chat {model_name}",
             origin=ORIGIN,
@@ -259,7 +259,7 @@ def new_embed_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
         model_name, contents = prepare_embed_content_args(args, kwargs)
 
-        with sentry_sdk.start_span(
+        with get_start_span_function()(
             op=OP.GEN_AI_EMBEDDINGS,
             name=f"embeddings {model_name}",
             origin=ORIGIN,
@@ -294,7 +294,7 @@ async def new_async_embed_content(
 
         model_name, contents = prepare_embed_content_args(args, kwargs)
 
-        with sentry_sdk.start_span(
+        with get_start_span_function()(
             op=OP.GEN_AI_EMBEDDINGS,
             name=f"embeddings {model_name}",
             origin=ORIGIN,