From f9a5b692a7b4ed7046ba5035cb1b592aa0cd9f0b Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 5 Dec 2023 14:22:46 -0800 Subject: [PATCH] fix: update test assumption for streaming endpoint of chat-bison@001 PiperOrigin-RevId: 588188368 --- tests/system/aiplatform/test_language_models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/system/aiplatform/test_language_models.py b/tests/system/aiplatform/test_language_models.py index a9533d735b..1403f9eefe 100644 --- a/tests/system/aiplatform/test_language_models.py +++ b/tests/system/aiplatform/test_language_models.py @@ -35,6 +35,7 @@ ChatModel, InputOutputTextPair, TextGenerationModel, + TextGenerationResponse, TextEmbeddingModel, ) @@ -256,7 +257,7 @@ def test_chat_model_send_message_streaming(self): message1 = "Are my favorite movies based on a book series?" for response in chat.send_message_streaming(message1): - assert response.text + assert isinstance(response, TextGenerationResponse) assert len(chat.message_history) == 2 assert chat.message_history[0].author == chat.USER_AUTHOR assert chat.message_history[0].content == message1 @@ -267,7 +268,7 @@ def test_chat_model_send_message_streaming(self): message2, temperature=0.1, ): - assert response2.text + assert isinstance(response2, TextGenerationResponse) assert len(chat.message_history) == 4 assert chat.message_history[2].author == chat.USER_AUTHOR assert chat.message_history[2].content == message2