From e85fe570038ab7f18eec22d962f6e7ce646434bf Mon Sep 17 00:00:00 2001 From: Billy Trend Date: Thu, 19 Sep 2024 14:16:37 +0100 Subject: [PATCH] Fix type --- tests/test_client_v2.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_client_v2.py b/tests/test_client_v2.py index 946d7a958..6630a1b58 100644 --- a/tests/test_client_v2.py +++ b/tests/test_client_v2.py @@ -43,8 +43,8 @@ def test_chat_documents(self) -> None: {"title": "widget sales 2021", "text": "4 million"}, ] response = co.chat( - messages=cohere.UserMessage( - message=cohere.TextContent(text="how many widges were sold in 2020?"), + messages=cohere.UserV2ChatMessage( + content=cohere.TextContent(text="how many widges were sold in 2020?"), documents=documents, ), ) @@ -68,8 +68,8 @@ def test_chat_tools(self) -> None: }, } tools = [cohere.V2Tool(type="function", function=get_weather_tool)] - messages: typing.List[typing.Union[UserMessage, AssistantMessage, None, cohere.V2ToolMessage]] = [ - UserMessage(message="what is the weather in Toronto?") + messages: cohere.ChatMessages = [ + cohere.UserV2ChatMessage(content="what is the weather in Toronto?") ] res = co.chat(model="command-r-plus", tools=tools, messages=messages) @@ -77,7 +77,7 @@ def test_chat_tools(self) -> None: tool_result = {"temperature": "30C"} tool_content = [cohere.Content(output=tool_result, text="The weather in Toronto is 30C")] messages.append(res.message) - messages.append(cohere.V2ToolMessage(tool_call_id=res.message.tool_calls[0].id, tool_content=tool_content)) + messages.append(cohere.ToolV2ChatMessage(tool_call_id=res.message.tool_calls[0].id, tool_content=tool_content)) res = co.chat(tools=tools, messages=messages) print(res.message)