feat(cohere): Add span streaming support (#6479)

Enable the Cohere integration to use span streaming when configured. Add
support for both StreamedSpan and traditional span lifecycle patterns.

Fixes PY-2315
Fixes #6013

Author: ericapisani

sentry_sdk/integrations/cohere.py

@@ -6,9 +6,11 @@
 from sentry_sdk.ai.monitoring import record_token_usage
 from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized
 from sentry_sdk.consts import SPANDATA
+from sentry_sdk.traces import StreamedSpan
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Iterator
+    from typing import Any, Callable, Iterator, Union
 
     from sentry_sdk.tracing import Span
 
@@ -90,9 +92,18 @@ def _capture_exception(exc: "Any") -> None:
     sentry_sdk.capture_event(event, hint=hint)
 
 
+def _end_span(span: "Any") -> None:
+    if isinstance(span, StreamedSpan):
+        span.end()
+    else:
+        span.__exit__(None, None, None)
+
+
 def _wrap_chat(f: "Callable[..., Any]", streaming: bool) -> "Callable[..., Any]":
     def collect_chat_response_fields(
-        span: "Span", res: "NonStreamedChatResponse", include_pii: bool
+        span: "Union[Span, StreamedSpan]",
+        res: "NonStreamedChatResponse",
+        include_pii: bool,
     ) -> None:
         if include_pii:
             if hasattr(res, "text"):
@@ -129,6 +140,9 @@ def collect_chat_response_fields(
     @wraps(f)
     def new_chat(*args: "Any", **kwargs: "Any") -> "Any":
         integration = sentry_sdk.get_client().get_integration(CohereIntegration)
+        is_span_streaming_enabled = has_span_streaming_enabled(
+            sentry_sdk.get_client().options
+        )
 
         if (
             integration is None
@@ -139,19 +153,29 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any":
 
         message = kwargs.get("message")
 
-        span = get_start_span_function()(
-            op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE,
-            name="cohere.client.Chat",
-            origin=CohereIntegration.origin,
-        )
-        span.__enter__()
+        if is_span_streaming_enabled:
+            span = sentry_sdk.traces.start_span(
+                name="cohere.client.Chat",
+                attributes={
+                    "sentry.op": consts.OP.COHERE_CHAT_COMPLETIONS_CREATE,
+                    "sentry.origin": CohereIntegration.origin,
+                },
+            )
+        else:
+            span = get_start_span_function()(
+                op=consts.OP.COHERE_CHAT_COMPLETIONS_CREATE,
+                name="cohere.client.Chat",
+                origin=CohereIntegration.origin,
+            )
+            span.__enter__()
         try:
             res = f(*args, **kwargs)
         except Exception as e:
             exc_info = sys.exc_info()
             with capture_internal_exceptions():
                 _capture_exception(e)
                 span.__exit__(*exc_info)
+
             reraise(*exc_info)
 
         with capture_internal_exceptions():
@@ -195,8 +219,7 @@ def new_iterator() -> "Iterator[StreamedChatResponse]":
                                     and integration.include_prompts,
                                 )
                             yield x
-
-                    span.__exit__(None, None, None)
+                    _end_span(span)
 
                 return new_iterator()
             elif isinstance(res, NonStreamedChatResponse):
@@ -206,11 +229,11 @@ def new_iterator() -> "Iterator[StreamedChatResponse]":
                     include_pii=should_send_default_pii()
                     and integration.include_prompts,
                 )
-                span.__exit__(None, None, None)
+                _end_span(span)
             else:
                 set_data_normalized(span, "unknown_response", True)
-                span.__exit__(None, None, None)
-            return res
+                _end_span(span)
+        return res
 
     return new_chat
 
@@ -222,11 +245,26 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
         if integration is None:
             return f(*args, **kwargs)
 
-        with get_start_span_function()(
-            op=consts.OP.COHERE_EMBEDDINGS_CREATE,
-            name="Cohere Embedding Creation",
-            origin=CohereIntegration.origin,
-        ) as span:
+        is_span_streaming_enabled = has_span_streaming_enabled(
+            sentry_sdk.get_client().options
+        )
+
+        if is_span_streaming_enabled:
+            span_ctx = sentry_sdk.traces.start_span(
+                name="Cohere Embedding Creation",
+                attributes={
+                    "sentry.op": consts.OP.COHERE_EMBEDDINGS_CREATE,
+                    "sentry.origin": CohereIntegration.origin,
+                },
+            )
+        else:
+            span_ctx = get_start_span_function()(
+                op=consts.OP.COHERE_EMBEDDINGS_CREATE,
+                name="Cohere Embedding Creation",
+                origin=CohereIntegration.origin,
+            )
+
+        with span_ctx as span:
             if "texts" in kwargs and (
                 should_send_default_pii() and integration.include_prompts
             ):