From 2dbc95653e6cd4f13593e57ae7eea8eeaed6f7ff Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 12 Mar 2024 13:19:17 -0700
Subject: [PATCH] (feat) cohere tool calling

---
 litellm/llms/cohere_chat.py             | 75 +++++++++++++++++++++++++
 litellm/tests/test_cohere_completion.py |  4 +-
 2 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/cohere_chat.py b/litellm/llms/cohere_chat.py
index 9027572e6a..ecdb6ffb25 100644
--- a/litellm/llms/cohere_chat.py
+++ b/litellm/llms/cohere_chat.py
@@ -116,6 +116,75 @@ def validate_environment(api_key):
     return headers
 
 
+def translate_openai_tool_to_cohere(openai_tool):
+    # cohere tools look like this
+    """
+    {
+       "name": "query_daily_sales_report",
+       "description": "Connects to a database to retrieve overall sales volumes and sales information for a given day.",
+       "parameter_definitions": {
+           "day": {
+               "description": "Retrieves sales data for this day, formatted as YYYY-MM-DD.",
+               "type": "str",
+               "required": True
+           }
+       }
+    }
+    """
+
+    # OpenAI tools look like this
+    """
+    {
+        "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"],
+            },
+        },
+    }
+    """
+    cohere_tool = {
+        "name": openai_tool["function"]["name"],
+        "description": openai_tool["function"]["description"],
+        "parameter_definitions": {},
+    }
+
+    for param_name, param_def in openai_tool["function"]["parameters"][
+        "properties"
+    ].items():
+        required_params = (
+            openai_tool.get("function", {}).get("parameters", {}).get("required", [])
+        )
+        cohere_param_def = {
+            "description": param_def.get("description", ""),
+            "type": param_def.get("type", ""),
+            "required": param_name in required_params,
+        }
+        cohere_tool["parameter_definitions"][param_name] = cohere_param_def
+
+    return cohere_tool
+
+
+def construct_cohere_tool(tools=None):
+    if tools is None:
+        tools = []
+    cohere_tools = []
+    for tool in tools:
+        cohere_tool = translate_openai_tool_to_cohere(tool)
+        cohere_tools.append(cohere_tool)
+    return cohere_tools
+
+
 def completion(
     model: str,
     messages: list,
@@ -142,6 +211,12 @@ 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
+        cohere_tools = construct_cohere_tool(tools=optional_params["tools"])
+        optional_params["tools"] = cohere_tools
+
     data = {
         "model": model,
         "message": prompt,
diff --git a/litellm/tests/test_cohere_completion.py b/litellm/tests/test_cohere_completion.py
index 683f97eeea..932a243245 100644
--- a/litellm/tests/test_cohere_completion.py
+++ b/litellm/tests/test_cohere_completion.py
@@ -64,16 +64,14 @@ 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",
+                "content": "What is the weather like in Boston?",
             },
         ]
         response = completion(
             model="cohere_chat/command-r",
             messages=messages,
-            max_tokens=10,
             tools=[
                 {
                     "type": "function",