diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py index 4d6a4a244c..388b86f1e0 100644 --- a/sentry_sdk/integrations/cohere.py +++ b/sentry_sdk/integrations/cohere.py @@ -26,6 +26,7 @@ from cohere import ( ChatStreamEndEvent, NonStreamedChatResponse, + StreamedChatResponse_StreamEnd, ) if TYPE_CHECKING: @@ -33,12 +34,6 @@ 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,