chore: LLM - Fixed the test_code_chat_model_send_message_streaming …

…test PiperOrigin-RevId: 565637309
googleapis · Sep 15, 2023 · 6fb30bc · 6fb30bc
1 parent 5b2492c
commit 6fb30bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/system/aiplatform/test_language_models.py b/tests/system/aiplatform/test_language_models.py
@@ -306,7 +306,7 @@ def test_code_generation_streaming(self):
     def test_code_chat_model_send_message_streaming(self):
         aiplatform.init(project=e2e_base._PROJECT, location=e2e_base._LOCATION)
 
-        chat_model = language_models.ChatModel.from_pretrained("codechat-bison@001")
+        chat_model = language_models.CodeChatModel.from_pretrained("codechat-bison@001")
         chat = chat_model.start_chat()
 
         message1 = "Please help write a function to calculate the max of two numbers"