From d136238f6fb385a12cd8188276133f1c404181d2 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 12 Mar 2024 12:35:52 -0700
Subject: [PATCH] (v0) tool calling

---
 litellm/llms/cohere.py                  |  14 ++++
 litellm/tests/test_cohere_completion.py | 103 ++++++++++++++++++++++++
 litellm/utils.py                        |   1 +
 3 files changed, 118 insertions(+)
 create mode 100644 litellm/tests/test_cohere_completion.py

diff --git a/litellm/llms/cohere.py b/litellm/llms/cohere.py
index 40b65439b21..960dc66d37d 100644
--- a/litellm/llms/cohere.py
+++ b/litellm/llms/cohere.py
@@ -22,6 +22,12 @@ class CohereError(Exception):
         )  # Call the base class constructor with the parameters it needs
 
 
+def construct_cohere_tool(tools=None):
+    if tools is None:
+        tools = []
+    return {"tools": tools}
+
+
 class CohereConfig:
     """
     Reference: https://docs.cohere.com/reference/generate
@@ -145,6 +151,14 @@ def completion(
         ):  # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in
             optional_params[k] = v
 
+    ## Handle Tool Calling
+    if "tools" in optional_params:
+        _is_function_call = True
+        tool_calling_system_prompt = construct_cohere_tool(
+            tools=optional_params["tools"]
+        )
+        optional_params["tools"] = tool_calling_system_prompt
+
     data = {
         "model": model,
         "prompt": prompt,
diff --git a/litellm/tests/test_cohere_completion.py b/litellm/tests/test_cohere_completion.py
new file mode 100644
index 00000000000..683f97eeea2
--- /dev/null
+++ b/litellm/tests/test_cohere_completion.py
@@ -0,0 +1,103 @@
+import sys, os
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import os, io
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest
+import litellm
+from litellm import embedding, completion, completion_cost, Timeout
+from litellm import RateLimitError
+
+litellm.num_retries = 3
+
+
+# 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_chat_completion_cohere_tool_calling():
+    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,
+            tools=[
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "get_current_weather",
+                        "description": "Get the current weather in a given location",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {
+                                "location": {
+                                    "type": "string",
+                                    "description": "The city and state, e.g. San Francisco, CA",
+                                },
+                                "unit": {
+                                    "type": "string",
+                                    "enum": ["celsius", "fahrenheit"],
+                                },
+                            },
+                            "required": ["location"],
+                        },
+                    },
+                }
+            ],
+        )
+        print(response)
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
diff --git a/litellm/utils.py b/litellm/utils.py
index 3c1bf989a7d..1fd8434338a 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -4269,6 +4269,7 @@ def get_optional_params(
             and custom_llm_provider != "together_ai"
             and custom_llm_provider != "mistral"
             and custom_llm_provider != "anthropic"
+            and custom_llm_provider != "cohere_chat"
             and custom_llm_provider != "bedrock"
             and custom_llm_provider != "ollama_chat"
         ):