Skip to content

Commit f9a5b69

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: update test assumption for streaming endpoint of chat-bison@001
PiperOrigin-RevId: 588188368
1 parent f58baf3 commit f9a5b69

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

tests/system/aiplatform/test_language_models.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
ChatModel,
3636
InputOutputTextPair,
3737
TextGenerationModel,
38+
TextGenerationResponse,
3839
TextEmbeddingModel,
3940
)
4041

@@ -256,7 +257,7 @@ def test_chat_model_send_message_streaming(self):
256257

257258
message1 = "Are my favorite movies based on a book series?"
258259
for response in chat.send_message_streaming(message1):
259-
assert response.text
260+
assert isinstance(response, TextGenerationResponse)
260261
assert len(chat.message_history) == 2
261262
assert chat.message_history[0].author == chat.USER_AUTHOR
262263
assert chat.message_history[0].content == message1
@@ -267,7 +268,7 @@ def test_chat_model_send_message_streaming(self):
267268
message2,
268269
temperature=0.1,
269270
):
270-
assert response2.text
271+
assert isinstance(response2, TextGenerationResponse)
271272
assert len(chat.message_history) == 4
272273
assert chat.message_history[2].author == chat.USER_AUTHOR
273274
assert chat.message_history[2].content == message2

0 commit comments

Comments
 (0)