Merge pull request #2474 from BerriAI/litellm_support_command_r

[New-Model] Cohere/command-r

litellm/tests/test_completion.py

@@ -1984,6 +1984,50 @@ def test_completion_cohere():
         pytest.fail(f"Error occurred: {e}")
 
 
+# FYI - cohere_chat looks quite unstable, even when testing locally
+def test_chat_completion_cohere():
+    try:
+        litellm.set_verbose = True
+        messages = [
+            {"role": "system", "content": "You're a good bot"},
+            {
+                "role": "user",
+                "content": "Hey",
+            },
+        ]
+        response = completion(
+            model="cohere_chat/command-r",
+            messages=messages,
+            max_tokens=10,
+        )
+        print(response)
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
+def test_chat_completion_cohere_stream():
+    try:
+        litellm.set_verbose = False
+        messages = [
+            {"role": "system", "content": "You're a good bot"},
+            {
+                "role": "user",
+                "content": "Hey",
+            },
+        ]
+        response = completion(
+            model="cohere_chat/command-r",
+            messages=messages,
+            max_tokens=10,
+            stream=True,
+        )
+        print(response)
+        for chunk in response:
+            print(chunk)
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_azure_cloudflare_api():
     litellm.set_verbose = True
     try: