another fix

cohere-ai · Aug 30, 2024 · ce9a274 · ce9a274
1 parent 2ec0181
commit ce9a274
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/tests/test_client_v2.py b/tests/test_client_v2.py
@@ -1,10 +1,9 @@
-import json
 import os
 import typing
 import unittest
 
 import cohere
-from cohere import TextContent, ChatDocument, ToolMessage2, UserMessage, AssistantMessage
+from cohere import ToolMessage2, UserMessage, AssistantMessage
 
 co = cohere.ClientV2(timeout=10000)
 
@@ -43,13 +42,10 @@ def test_chat_documents(self) -> None:
             {"title": "widget sales 2020", "text": "2 million"},
             {"title": "widget sales 2021", "text": "4 million"},
         ]
-        docs: typing.List[ChatDocument] = []
-        for doc in documents:
-            docs.append(ChatDocument(id="1", document=doc))
         response = co.chat(
             messages=cohere.v2.UserMessage(
                 content=cohere.v2.TextContent(text="how many widges were sold in 2020?"),
-                documents=docs,
+                documents=documents,
             ),
         )