From d136238f6fb385a12cd8188276133f1c404181d2 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 12 Mar 2024 12:35:52 -0700 Subject: [PATCH 1/4] (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" ): From 2dbc95653e6cd4f13593e57ae7eea8eeaed6f7ff Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 12 Mar 2024 13:19:17 -0700 Subject: [PATCH 2/4] (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 9027572e6a6..ecdb6ffb250 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 683f97eeea2..932a2432456 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", From 5b0b251d423eb9abe773695ab24ebc4da1754b4b Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 12 Mar 2024 14:24:48 -0700 Subject: [PATCH 3/4] (feat) support tool_calling on cohere command-r --- litellm/llms/cohere_chat.py | 29 ++++- litellm/llms/prompt_templates/factory.py | 59 +++++++++++ litellm/tests/test_cohere_completion.py | 129 +++++++++++++++++++++++ 3 files changed, 216 insertions(+), 1 deletion(-) diff --git a/litellm/llms/cohere_chat.py b/litellm/llms/cohere_chat.py index ecdb6ffb250..c51ef8deda5 100644 --- a/litellm/llms/cohere_chat.py +++ b/litellm/llms/cohere_chat.py @@ -7,6 +7,7 @@ from typing import Callable, Optional from litellm.utils import ModelResponse, Choices, Message, Usage import litellm import httpx +from .prompt_templates.factory import cohere_message_pt class CohereError(Exception): @@ -201,7 +202,7 @@ def completion( headers = validate_environment(api_key) completion_url = api_base model = model - prompt = " ".join(message["content"] for message in messages) + prompt, tool_results = cohere_message_pt(messages=messages) ## Load Config config = litellm.CohereConfig.get_config() @@ -216,6 +217,8 @@ def completion( _is_function_call = True cohere_tools = construct_cohere_tool(tools=optional_params["tools"]) optional_params["tools"] = cohere_tools + if len(tool_results) > 0: + optional_params["tool_results"] = tool_results data = { "model": model, @@ -262,6 +265,30 @@ def completion( except Exception as e: raise CohereError(message=response.text, status_code=response.status_code) + ## Tool calling response + cohere_tools_response = completion_response.get("tool_calls", None) + if cohere_tools_response is not None and cohere_tools_response is not []: + # convert cohere_tools_response to OpenAI response format + tool_calls = [] + for tool in cohere_tools_response: + function_name = tool.get("name", "") + generation_id = tool.get("generation_id", "") + parameters = tool.get("parameters", {}) + tool_call = { + "id": f"call_{generation_id}", + "type": "function", + "function": { + "name": function_name, + "arguments": json.dumps(parameters), + }, + } + tool_calls.append(tool_call) + _message = litellm.Message( + tool_calls=tool_calls, + content=None, + ) + model_response.choices[0].message = _message # type: ignore + ## CALCULATING USAGE - use cohere `billed_units` for returning usage billed_units = completion_response.get("meta", {}).get("billed_units", {}) diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index ae12d954a89..97caa9389c5 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -652,6 +652,65 @@ def parse_xml_params(xml_content): ### +def convert_openai_message_to_cohere_tool_result(message): + """ + OpenAI message with a tool result looks like: + { + "tool_call_id": "tool_1", + "role": "tool", + "name": "get_current_weather", + "content": {"location": "San Francisco, CA", "unit": "fahrenheit", "temperature": "72"}, + }, + """ + + """ + Cohere tool_results look like: + { + "call": { + "name": "query_daily_sales_report", + "parameters": { + "day": "2023-09-29" + }, + "generation_id": "4807c924-9003-4d6b-8069-eda03962c465" + }, + "outputs": [ + { + "date": "2023-09-29", + "summary": "Total Sales Amount: 10000, Total Units Sold: 250" + } + ] + }, + """ + + tool_call_id = message.get("tool_call_id") + name = message.get("name") + content = message.get("content") + + # Create the Cohere tool_result dictionary + cohere_tool_result = { + "call": { + "name": name, + "parameters": {"location": "San Francisco, CA"}, + "generation_id": tool_call_id, + }, + "outputs": [content], + } + return cohere_tool_result + + +def cohere_message_pt(messages: list): + prompt = "" + tool_results = [] + for message in messages: + # check if this is a tool_call result + if message["role"] == "tool": + tool_result = convert_openai_message_to_cohere_tool_result(message) + tool_results.append(tool_result) + else: + prompt += message["content"] + return prompt, tool_results + + def amazon_titan_pt( messages: list, ): # format - https://github.com/BerriAI/litellm/issues/1896 diff --git a/litellm/tests/test_cohere_completion.py b/litellm/tests/test_cohere_completion.py index 932a2432456..9c3c9bf93c3 100644 --- a/litellm/tests/test_cohere_completion.py +++ b/litellm/tests/test_cohere_completion.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import embedding, completion, completion_cost, Timeout from litellm import RateLimitError +import json litellm.num_retries = 3 @@ -99,3 +100,131 @@ def test_chat_completion_cohere_tool_calling(): print(response) except Exception as e: pytest.fail(f"Error occurred: {e}") + + # def get_current_weather(location, unit="fahrenheit"): + # """Get the current weather in a given location""" + # if "tokyo" in location.lower(): + # return json.dumps({"location": "Tokyo", "temperature": "10", "unit": unit}) + # elif "san francisco" in location.lower(): + # return json.dumps({"location": "San Francisco", "temperature": "72", "unit": unit}) + # elif "paris" in location.lower(): + # return json.dumps({"location": "Paris", "temperature": "22", "unit": unit}) + # else: + # return json.dumps({"location": location, "temperature": "unknown"}) + + # def test_chat_completion_cohere_tool_with_result_calling(): + # # end to end cohere command-r with tool calling + # # Step 1 - Send available tools + # # Step 2 - Execute results + # # Step 3 - Send results to command-r + # try: + # litellm.set_verbose = True + # import json + + # # Step 1 - Send available tools + # 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"], + # }, + # }, + # } + # ] + + # messages = [ + # { + # "role": "user", + # "content": "What is the weather like in Boston?", + # }, + # ] + # response = completion( + # model="cohere_chat/command-r", + # messages=messages, + # tools=tools, + # ) + # print("Response with tools to call", response) + # print(response) + + # # step 2 - Execute results + # tool_calls = response.tool_calls + + # available_functions = { + # "get_current_weather": get_current_weather, + # } # only one function in this example, but you can have multiple + + # for tool_call in tool_calls: + # function_name = tool_call.function.name + # function_to_call = available_functions[function_name] + # function_args = json.loads(tool_call.function.arguments) + # function_response = function_to_call( + # location=function_args.get("location"), + # unit=function_args.get("unit"), + # ) + # messages.append( + # { + # "tool_call_id": tool_call.id, + # "role": "tool", + # "name": function_name, + # "content": function_response, + # } + # ) # extend conversation with function response + + # print("messages with tool call results", messages) + + # messages = [ + # { + # "role": "user", + # "content": "What is the weather like in Boston?", + # }, + # { + # "tool_call_id": "tool_1", + # "role": "tool", + # "name": "get_current_weather", + # "content": {"location": "San Francisco, CA", "unit": "fahrenheit", "temperature": "72"}, + # }, + # ] + # respone = completion( + # model="cohere_chat/command-r", + # messages=messages, + # 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(respone) + except Exception as e: + pytest.fail(f"Error occurred: {e}") From b9bfc7c36c2d38ba90caa594cb535d741dc587f2 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 12 Mar 2024 14:31:43 -0700 Subject: [PATCH 4/4] (fix) use cohere_chat optional params --- litellm/tests/test_cohere_completion.py | 2 -- litellm/utils.py | 38 +++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_cohere_completion.py b/litellm/tests/test_cohere_completion.py index 9c3c9bf93c3..372c87b4000 100644 --- a/litellm/tests/test_cohere_completion.py +++ b/litellm/tests/test_cohere_completion.py @@ -22,7 +22,6 @@ def test_chat_completion_cohere(): try: litellm.set_verbose = True messages = [ - {"role": "system", "content": "You're a good bot"}, { "role": "user", "content": "Hey", @@ -42,7 +41,6 @@ def test_chat_completion_cohere_stream(): try: litellm.set_verbose = False messages = [ - {"role": "system", "content": "You're a good bot"}, { "role": "user", "content": "Hey", diff --git a/litellm/utils.py b/litellm/utils.py index 1fd8434338a..6abba4afa0d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4401,6 +4401,31 @@ def get_optional_params( optional_params["presence_penalty"] = presence_penalty if stop is not None: optional_params["stop_sequences"] = stop + elif custom_llm_provider == "cohere_chat": + ## check if unsupported param passed in + supported_params = get_supported_openai_params( + model=model, custom_llm_provider=custom_llm_provider + ) + _check_valid_arg(supported_params=supported_params) + # handle cohere params + if stream: + optional_params["stream"] = stream + if temperature is not None: + optional_params["temperature"] = temperature + if max_tokens is not None: + optional_params["max_tokens"] = max_tokens + if n is not None: + optional_params["num_generations"] = n + if top_p is not None: + optional_params["p"] = top_p + if frequency_penalty is not None: + optional_params["frequency_penalty"] = frequency_penalty + if presence_penalty is not None: + optional_params["presence_penalty"] = presence_penalty + if stop is not None: + optional_params["stop_sequences"] = stop + if tools is not None: + optional_params["tools"] = tools elif custom_llm_provider == "maritalk": ## check if unsupported param passed in supported_params = get_supported_openai_params( @@ -5084,6 +5109,19 @@ def get_supported_openai_params(model: str, custom_llm_provider: str): "stop", "n", ] + elif custom_llm_provider == "cohere_chat": + return [ + "stream", + "temperature", + "max_tokens", + "top_p", + "frequency_penalty", + "presence_penalty", + "stop", + "n", + "tools", + "tool_choice", + ] elif custom_llm_provider == "maritalk": return [ "stream",