Fixes

cohere-ai · Sep 19, 2024 · f41396c · f41396c
1 parent 4cc09ff
commit f41396c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_client_v2.py b/tests/test_client_v2.py
@@ -14,12 +14,12 @@
 class TestClientV2(unittest.TestCase):
 
     def test_chat(self) -> None:
-        response = co.chat(model="command-r-plus", messages=[cohere.UserMessage(message="hello world!")])
+        response = co.chat(model="command-r-plus", messages=[cohere.UserChatMessageV2(content="hello world!")])
 
         print(response.message)
 
     def test_chat_stream(self) -> None:
-        stream = co.chat_stream(model="command-r-plus", messages=[cohere.UserMessage(message="hello world!")])
+        stream = co.chat_stream(model="command-r-plus", messages=[cohere.UserChatMessageV2(content="hello world!")])
 
         events = set()