split off cohere fix

getsentry · Sep 19, 2024 · 9f0b114 · 9f0b114
1 parent 9f410a2
commit 9f0b114
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py
@@ -26,19 +26,14 @@
     from cohere import (
         ChatStreamEndEvent,
         NonStreamedChatResponse,
+        StreamedChatResponse_StreamEnd,
     )
 
     if TYPE_CHECKING:
         from cohere import StreamedChatResponse
 except ImportError:
     raise DidNotEnable("Cohere not installed")
 
-try:
-    # cohere 5.9.3+
-    from cohere import StreamEndStreamedChatResponse
-except ImportError:
-    from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse
-
 
 COLLECTED_CHAT_PARAMS = {
     "model": SPANDATA.AI_MODEL_ID,
@@ -194,7 +189,7 @@ def new_iterator():
                     with capture_internal_exceptions():
                         for x in old_iterator:
                             if isinstance(x, ChatStreamEndEvent) or isinstance(
-                                x, StreamEndStreamedChatResponse
+                                x, StreamedChatResponse_StreamEnd
                             ):
                                 collect_chat_response_fields(
                                     span,