Fix cohere testsuite for new release of cohere. (#3098)

* Check for new class to signal end of stream
getsentry · May 23, 2024 · ec23396 · ec23396
1 parent 38c14e9
commit ec23396
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py
@@ -22,7 +22,11 @@
 try:
     from cohere.client import Client
     from cohere.base_client import BaseCohere
-    from cohere import ChatStreamEndEvent, NonStreamedChatResponse
+    from cohere import (
+        ChatStreamEndEvent,
+        NonStreamedChatResponse,
+        StreamedChatResponse_StreamEnd,
+    )
 
     if TYPE_CHECKING:
         from cohere import StreamedChatResponse
@@ -181,7 +185,9 @@ def new_iterator():
 
                     with capture_internal_exceptions():
                         for x in old_iterator:
-                            if isinstance(x, ChatStreamEndEvent):
+                            if isinstance(x, ChatStreamEndEvent) or isinstance(
+                                x, StreamedChatResponse_StreamEnd
+                            ):
                                 collect_chat_response_fields(
                                     span,
                                     x.response,