Skip to content

Commit 49dd64d

Browse files
authored
tests: Fix cohere API change (getsentry#3549)
1 parent 4f6ccc4 commit 49dd64d

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

sentry_sdk/integrations/cohere.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,19 @@
2626
from cohere import (
2727
ChatStreamEndEvent,
2828
NonStreamedChatResponse,
29-
StreamedChatResponse_StreamEnd,
3029
)
3130

3231
if TYPE_CHECKING:
3332
from cohere import StreamedChatResponse
3433
except ImportError:
3534
raise DidNotEnable("Cohere not installed")
3635

36+
try:
37+
# cohere 5.9.3+
38+
from cohere import StreamEndStreamedChatResponse
39+
except ImportError:
40+
from cohere import StreamedChatResponse_StreamEnd as StreamEndStreamedChatResponse
41+
3742

3843
COLLECTED_CHAT_PARAMS = {
3944
"model": SPANDATA.AI_MODEL_ID,
@@ -189,7 +194,7 @@ def new_iterator():
189194
with capture_internal_exceptions():
190195
for x in old_iterator:
191196
if isinstance(x, ChatStreamEndEvent) or isinstance(
192-
x, StreamedChatResponse_StreamEnd
197+
x, StreamEndStreamedChatResponse
193198
):
194199
collect_chat_response_fields(
195200
span,

0 commit comments

Comments
 (0)