diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 6fd2086db9..1b6f9067ee 100644 --- 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,