scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge pull request #4546 from BerriAI/litellm_cohere_tool_calling_fix

Browse source 
Cohere tool calling fix
This commit is contained in:
Krish Dholakia 2024-07-04 15:25:22 -07:00 • committed by GitHub
commit 20b0a63f56
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 520 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
litellm/llms/cohere_chat.py
View file

@ -1,13 +1,19 @@
import os, types
import json
import os
import time
import traceback
import types
from enum import Enum
import requests # type: ignore
import time, traceback
from typing import Callable, Optional
from litellm.utils import ModelResponse, Choices, Message, Usage
import litellm
import httpx # type: ignore
from .prompt_templates.factory import cohere_message_pt
import requests # type: ignore
import litellm
from litellm.types.llms.cohere import ToolResultObject
from litellm.utils import Choices, Message, ModelResponse, Usage
from .prompt_templates.factory import cohere_message_pt, cohere_messages_pt_v2
class CohereError(Exception):
@ -112,7 +118,7 @@ class CohereChatConfig:
def validate_environment(api_key):
headers = {
"Request-Source":"unspecified:litellm",
"Request-Source": "unspecified:litellm",
"accept": "application/json",
"content-type": "application/json",
}
@ -196,17 +202,17 @@ def completion(
api_base: str,
model_response: ModelResponse,
print_verbose: Callable,
optional_params: dict,
encoding,
api_key,
logging_obj,
optional_params=None,
litellm_params=None,
logger_fn=None,
):
headers = validate_environment(api_key)
completion_url = api_base
model = model
prompt, tool_results = cohere_message_pt(messages=messages)
most_recent_message, chat_history = cohere_messages_pt_v2(messages=messages)
## Load Config
config = litellm.CohereConfig.get_config()
@ -221,18 +227,18 @@ 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
if isinstance(most_recent_message, dict):
optional_params["tool_results"] = [most_recent_message]
elif isinstance(most_recent_message, str):
optional_params["message"] = most_recent_message
data = {
"model": model,
"message": prompt,
**optional_params,
}
## LOGGING
logging_obj.pre_call(
input=prompt,
input=most_recent_message,
api_key=api_key,
additional_args={
"complete_input_dict": data,
@ -256,7 +262,7 @@ def completion(
else:
## LOGGING
logging_obj.post_call(
input=prompt,
input=most_recent_message,
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},

297
litellm/llms/prompt_templates/factory.py
View file

@ -1415,16 +1415,37 @@ def convert_to_documents(
return documents
def convert_openai_message_to_cohere_tool_result(message):
from litellm.types.llms.cohere import (
CallObject,
ChatHistory,
ChatHistoryChatBot,
ChatHistorySystem,
ChatHistoryToolResult,
ChatHistoryUser,
ToolCallObject,
ToolResultObject,
)
def convert_openai_message_to_cohere_tool_result(
message, tool_calls: List
) -> ToolResultObject:
"""
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"},
},
"""
"""
OpenAI message with a function call looks like:
{
"role": "function",
"name": "get_current_weather",
"content": "function result goes here",
}
"""
"""
Cohere tool_results look like:
@ -1434,7 +1455,6 @@ def convert_openai_message_to_cohere_tool_result(message):
"parameters": {
"day": "2023-09-29"
},
"generation_id": "4807c924-9003-4d6b-8069-eda03962c465"
},
"outputs": [
{
@ -1444,30 +1464,255 @@ def convert_openai_message_to_cohere_tool_result(message):
]
},
"""
content_str: str = message.get("content", "")
if len(content_str) > 0:
try:
content = json.loads(content_str)
except json.JSONDecodeError:
content = {"result": content_str}
else:
content = {}
name = ""
arguments = {}
# Recover name from last message with tool calls
if len(tool_calls) > 0:
tools = tool_calls
msg_tool_call_id = message.get("tool_call_id", None)
for tool in tools:
prev_tool_call_id = tool.get("id", None)
if (
msg_tool_call_id
and prev_tool_call_id
and msg_tool_call_id == prev_tool_call_id
):
name = tool.get("function", {}).get("name", "")
arguments_str = tool.get("function", {}).get("arguments", "")
if arguments_str is not None and len(arguments_str) > 0:
arguments = json.loads(arguments_str)
tool_call_id = message.get("tool_call_id")
name = message.get("name")
content = message.get("content")
if message["role"] == "function":
name = message.get("name")
cohere_tool_result: ToolResultObject = {
"call": CallObject(name=name, parameters=arguments),
"outputs": [content],
}
return cohere_tool_result
else:
# We can't determine from openai message format whether it's a successful or
# error call result so default to the successful result template
# Create the Cohere tool_result dictionary
cohere_tool_result = {
"call": {
"name": name,
"parameters": {"location": "San Francisco, CA"},
"generation_id": tool_call_id,
},
"outputs": convert_to_documents(content),
cohere_tool_result = {
"call": CallObject(name=name, parameters=arguments),
"outputs": [content],
}
return cohere_tool_result
def get_all_tool_calls(messages: List) -> List:
"""
Returns extracted list of `tool_calls`.
Done to handle openai no longer returning tool call 'name' in tool results.
"""
tool_calls: List = []
for m in messages:
if m.get("tool_calls", None) is not None:
if isinstance(m["tool_calls"], list):
tool_calls.extend(m["tool_calls"])
return tool_calls
def convert_to_cohere_tool_invoke(tool_calls: list) -> List[ToolCallObject]:
"""
OpenAI tool invokes:
{
"role": "assistant",
"content": null,
"tool_calls": [
{
"id": "call_abc123",
"type": "function",
"function": {
"name": "get_current_weather",
"arguments": "{\n\"location\": \"Boston, MA\"\n}"
}
}
]
},
"""
"""
Cohere tool invokes:
{
"role": "CHATBOT",
"tool_calls": [{"name": "get_weather", "parameters": {"location": "San Francisco, CA"}}]
}
return cohere_tool_result
"""
cohere_tool_invoke: List[ToolCallObject] = [
{
"name": get_attribute_or_key(
get_attribute_or_key(tool, "function"), "name"
),
"parameters": json.loads(
get_attribute_or_key(
get_attribute_or_key(tool, "function"), "arguments"
)
),
}
for tool in tool_calls
if get_attribute_or_key(tool, "type") == "function"
]
return cohere_tool_invoke
def cohere_messages_pt_v2(
messages: List,
) -> Tuple[Union[str, ToolResultObject], ChatHistory]:
"""
Returns a tuple(Union[tool_result, message], chat_history)
- if last message is tool result -> return 'tool_result'
- if last message is text -> return message (str)
- return preceding messages as 'chat_history'
Note:
- cannot specify message if the last entry in chat history contains tool results
- message must be at least 1 token long or tool results must be specified.
"""
tool_calls: List = get_all_tool_calls(messages=messages)
## GET MOST RECENT MESSAGE
most_recent_message = messages.pop(-1)
returned_message: Union[ToolResultObject, str] = ""
if (
most_recent_message.get("role", "") is not None
and most_recent_message["role"] == "tool"
):
# tool result
returned_message = convert_openai_message_to_cohere_tool_result(
most_recent_message, tool_calls
)
else:
content: Union[str, List] = most_recent_message.get("content")
if isinstance(content, str):
returned_message = content
else:
for chunk in content:
if chunk.get("type") == "text":
returned_message += chunk.get("text")
## CREATE CHAT HISTORY
user_message_types = {"user"}
tool_message_types = {"tool", "function"}
# reformat messages to ensure user/assistant are alternating, if there's either 2 consecutive 'user' messages or 2 consecutive 'assistant' message, merge them.
new_messages: ChatHistory = []
msg_i = 0
while msg_i < len(messages):
user_content: str = ""
init_msg_i = msg_i
## MERGE CONSECUTIVE USER CONTENT ##
while msg_i < len(messages) and messages[msg_i]["role"] in user_message_types:
if isinstance(messages[msg_i]["content"], list):
for m in messages[msg_i]["content"]:
if m.get("type", "") == "text":
user_content += m["text"]
else:
user_content += messages[msg_i]["content"]
msg_i += 1
if len(user_content) > 0:
new_messages.append(ChatHistoryUser(role="USER", message=user_content))
system_content: str = ""
## MERGE CONSECUTIVE SYSTEM CONTENT ##
while msg_i < len(messages) and messages[msg_i]["role"] == "system":
if isinstance(messages[msg_i]["content"], list):
for m in messages[msg_i]["content"]:
if m.get("type", "") == "text":
system_content += m["text"]
else:
system_content += messages[msg_i]["content"]
msg_i += 1
if len(system_content) > 0:
new_messages.append(
ChatHistorySystem(role="SYSTEM", message=system_content)
)
assistant_content: str = ""
assistant_tool_calls: List[ToolCallObject] = []
## MERGE CONSECUTIVE ASSISTANT CONTENT ##
while msg_i < len(messages) and messages[msg_i]["role"] == "assistant":
assistant_text = (
messages[msg_i].get("content") or ""
) # either string or none
if assistant_text:
assistant_content += assistant_text
if messages[msg_i].get(
"tool_calls", []
): # support assistant tool invoke conversion
assistant_tool_calls.extend(
convert_to_cohere_tool_invoke(messages[msg_i]["tool_calls"])
)
if messages[msg_i].get("function_call"):
assistant_tool_calls.extend(
convert_to_cohere_tool_invoke(messages[msg_i]["function_call"])
)
msg_i += 1
if len(assistant_content) > 0:
new_messages.append(
ChatHistoryChatBot(
role="CHATBOT",
message=assistant_content,
tool_calls=assistant_tool_calls,
)
)
## MERGE CONSECUTIVE TOOL RESULTS
tool_results: List[ToolResultObject] = []
while msg_i < len(messages) and messages[msg_i]["role"] in tool_message_types:
tool_results.append(
convert_openai_message_to_cohere_tool_result(
messages[msg_i], tool_calls
)
)
msg_i += 1
if len(tool_results) > 0:
new_messages.append(
ChatHistoryToolResult(role="TOOL", tool_results=tool_results)
)
if msg_i == init_msg_i: # prevent infinite loops
raise Exception(
"Invalid Message passed in - {}. File an issue https://github.com/BerriAI/litellm/issues".format(
messages[msg_i]
)
)
return returned_message, new_messages
def cohere_message_pt(messages: list):
tool_calls: List = get_all_tool_calls(messages=messages)
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_result = convert_openai_message_to_cohere_tool_result(
message, tool_calls=tool_calls
)
tool_results.append(tool_result)
elif message.get("content"):
prompt += message["content"] + "\n\n"
@ -1658,6 +1903,26 @@ def azure_text_pt(messages: list):
return prompt
###### AZURE AI #######
def stringify_json_tool_call_content(messages: List) -> List:
"""
- Check 'content' in tool role -> convert to dict (if not) -> stringify
Done for azure_ai/cohere calls to handle results of a tool call
"""
for m in messages:
if m["role"] == "tool" and isinstance(m["content"], str):
# check if content is a valid json object
try:
json.loads(m["content"])
except json.JSONDecodeError:
m["content"] = json.dumps({"result": m["content"]})
return messages
###### AMAZON BEDROCK #######
from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock

68
litellm/main.py
View file

@ -113,6 +113,7 @@ from .llms.prompt_templates.factory import (
function_call_prompt,
map_system_message_pt,
prompt_factory,
stringify_json_tool_call_content,
)
from .llms.text_completion_codestral import CodestralTextCompletion
from .llms.triton import TritonChatCompletion
@ -1114,6 +1115,73 @@ def completion(
"api_base": api_base,
},
)
elif custom_llm_provider == "azure_ai":
api_base = (
api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there
or litellm.api_base
or get_secret("AZURE_AI_API_BASE")
)
# set API KEY
api_key = (
api_key
or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there
or litellm.openai_key
or get_secret("AZURE_AI_API_KEY")
)
headers = headers or litellm.headers
## LOAD CONFIG - if set
config = litellm.OpenAIConfig.get_config()
for k, v in config.items():
if (
k not in optional_params
): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v
## FOR COHERE
if "command-r" in model: # make sure tool call in messages are str
messages = stringify_json_tool_call_content(messages=messages)
## COMPLETION CALL
try:
response = openai_chat_completions.completion(
model=model,
messages=messages,
headers=headers,
model_response=model_response,
print_verbose=print_verbose,
api_key=api_key,
api_base=api_base,
acompletion=acompletion,
logging_obj=logging,
optional_params=optional_params,
litellm_params=litellm_params,
logger_fn=logger_fn,
timeout=timeout, # type: ignore
custom_prompt_dict=custom_prompt_dict,
client=client, # pass AsyncOpenAI, OpenAI client
organization=organization,
custom_llm_provider=custom_llm_provider,
)
except Exception as e:
## LOGGING - log the original exception returned
logging.post_call(
input=messages,
api_key=api_key,
original_response=str(e),
additional_args={"headers": headers},
)
raise e
if optional_params.get("stream", False):
## LOGGING
logging.post_call(
input=messages,
api_key=api_key,
original_response=response,
additional_args={"headers": headers},
)
elif (
custom_llm_provider == "text-completion-openai"
or "ft:babbage-002" in model

11
litellm/tests/test_amazing_vertex_completion.py
View file

@ -1121,7 +1121,7 @@ async def test_gemini_pro_httpx_custom_api_base(provider):
assert "hello" in mock_call.call_args.kwargs["headers"]
@pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call")
# @pytest.mark.skip(reason="exhausted vertex quota. need to refactor to mock the call")
@pytest.mark.parametrize("sync_mode", [True])
@pytest.mark.parametrize("provider", ["vertex_ai"])
@pytest.mark.asyncio
@ -1159,7 +1159,7 @@ async def test_gemini_pro_function_calling(provider, sync_mode):
# The result of the tool call is added to the history
{
"role": "tool",
"tool_call_id": "call_123",
"tool_call_id": "call_123",
"content": "27 degrees celsius and clear in San Francisco, CA",
},
# Now the assistant can reply with the result of the tool call.
@ -1381,6 +1381,7 @@ async def test_vertexai_aembedding():
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
def test_tool_name_conversion():
messages = [
@ -1424,7 +1425,8 @@ def test_tool_name_conversion():
# assert that the last tool response has the corresponding tool name
assert (
translated_messages[-1]["parts"][0]["function_response"]["name"] == "get_weather"
translated_messages[-1]["parts"][0]["function_response"]["name"]
== "get_weather"
)
@ -1585,6 +1587,7 @@ def test_prompt_factory():
print(f"\n\ntranslated_messages: {translated_messages}\ntranslated_messages")
def test_prompt_factory_nested():
messages = [
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
@ -1606,4 +1609,4 @@ def test_prompt_factory_nested():
assert "text" in message["parts"][0], "Missing 'text' from 'parts'"
assert isinstance(
message["parts"][0]["text"], str
), "'text' value not a string."
), "'text' value not a string."

97
litellm/tests/test_completion.py
View file

@ -408,6 +408,103 @@ def test_completion_claude_3_function_call(model):
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("sync_mode", [True])
@pytest.mark.parametrize(
"model, api_key, api_base",
[
("gpt-3.5-turbo", None, None),
("claude-3-opus-20240229", None, None),
("command-r", None, None),
("anthropic.claude-3-sonnet-20240229-v1:0", None, None),
(
"azure_ai/command-r-plus",
os.getenv("AZURE_COHERE_API_KEY"),
os.getenv("AZURE_COHERE_API_BASE"),
),
],
)
@pytest.mark.asyncio
async def test_model_function_invoke(model, sync_mode, api_key, api_base):
try:
litellm.set_verbose = True
messages = [
{
"role": "system",
"content": "Your name is Litellm Bot, you are a helpful assistant",
},
# User asks for their name and weather in San Francisco
{
"role": "user",
"content": "Hello, what is your name and can you tell me the weather?",
},
# Assistant replies with a tool call
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"index": 0,
"function": {
"name": "get_weather",
"arguments": '{"location": "San Francisco, CA"}',
},
}
],
},
# The result of the tool call is added to the history
{
"role": "tool",
"tool_call_id": "call_123",
"content": "27 degrees celsius and clear in San Francisco, CA",
},
# Now the assistant can reply with the result of the tool call.
]
tools = [
{
"type": "function",
"function": {
"name": "get_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",
}
},
"required": ["location"],
},
},
}
]
data = {
"model": model,
"messages": messages,
"tools": tools,
"api_key": api_key,
"api_base": api_base,
}
if sync_mode:
response = litellm.completion(**data)
else:
response = await litellm.acompletion(**data)
print(f"response: {response}")
except litellm.RateLimitError as e:
pass
except Exception as e:
if "429 Quota exceeded" in str(e):
pass
else:
pytest.fail("An unexpected exception occurred - {}".format(str(e)))
@pytest.mark.asyncio
async def test_anthropic_no_content_error():
"""

46
litellm/types/llms/cohere.py Normal file
View file

@ -0,0 +1,46 @@
from typing import Iterable, List, Optional, Union
from typing_extensions import Literal, Required, TypedDict
class CallObject(TypedDict):
name: str
parameters: dict
class ToolResultObject(TypedDict):
call: CallObject
outputs: List[dict]
class ChatHistoryToolResult(TypedDict, total=False):
role: Required[Literal["TOOL"]]
tool_results: List[ToolResultObject]
class ToolCallObject(TypedDict):
name: str
parameters: dict
class ChatHistoryUser(TypedDict, total=False):
role: Required[Literal["USER"]]
message: str
tool_calls: List[ToolCallObject]
class ChatHistorySystem(TypedDict, total=False):
role: Required[Literal["SYSTEM"]]
message: str
tool_calls: List[ToolCallObject]
class ChatHistoryChatBot(TypedDict, total=False):
role: Required[Literal["CHATBOT"]]
message: str
tool_calls: List[ToolCallObject]
ChatHistory = List[
Union[ChatHistorySystem, ChatHistoryChatBot, ChatHistoryUser, ChatHistoryToolResult]
]