tests: Fix cohere API change

getsentry · Sep 19, 2024 · 2bd9188 · 2bd9188
1 parent 4f6ccc4
commit 2bd9188
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py
@@ -26,14 +26,19 @@
     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,
@@ -189,7 +194,7 @@ def new_iterator():
                     with capture_internal_exceptions():
                         for x in old_iterator:
                             if isinstance(x, ChatStreamEndEvent) or isinstance(
-                                x, StreamedChatResponse_StreamEnd
+                                x, StreamEndStreamedChatResponse
                             ):
                                 collect_chat_response_fields(
                                     span,