wehub-resource-sync
320 lines
12 KiB
Python
320 lines
12 KiB
Python
import ast
|
|
import json
|
|
import os
|
|
import time
|
|
from typing import Any
|
|
|
|
import cohere
|
|
from bfcl_eval.model_handler.base_handler import BaseHandler
|
|
from bfcl_eval.constants.type_mappings import GORILLA_TO_OPENAPI
|
|
from bfcl_eval.constants.enums import ModelStyle
|
|
from bfcl_eval.model_handler.utils import (
|
|
convert_to_tool,
|
|
retry_with_backoff,
|
|
extract_system_prompt,
|
|
)
|
|
from tenacity.stop import stop_after_attempt
|
|
|
|
|
|
class CohereHandler(BaseHandler):
|
|
client: cohere.ClientV2
|
|
|
|
def __init__(
|
|
self,
|
|
model_name,
|
|
temperature,
|
|
registry_name,
|
|
is_fc_model,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__(model_name, temperature, registry_name, is_fc_model, **kwargs)
|
|
self.model_style = ModelStyle.COHERE
|
|
self.is_fc_model = True
|
|
self.client = cohere.ClientV2(api_key=os.getenv("COHERE_API_KEY"))
|
|
|
|
def decode_ast(self, result, language, has_tool_call_tag):
|
|
decoded_output = []
|
|
if isinstance(result, list):
|
|
for tool_call in result:
|
|
name = tool_call["tool_name"]
|
|
params = tool_call["parameters"]
|
|
decoded_output.append({name: params})
|
|
return decoded_output
|
|
|
|
def decode_execute(self, result, has_tool_call_tag):
|
|
execution_list = []
|
|
if isinstance(result, list):
|
|
for tool_call in result:
|
|
parameter_key_value_list = []
|
|
for parameter_name, parameter_value in tool_call["parameters"].items():
|
|
parameter_key_value_list.append(
|
|
"{}={}".format(parameter_name, repr(parameter_value))
|
|
)
|
|
execution_list.append(
|
|
"{}({})".format(
|
|
tool_call["tool_name"], ",".join(parameter_key_value_list)
|
|
)
|
|
)
|
|
return execution_list
|
|
|
|
#### FC methods ####
|
|
|
|
def _query_FC(self, inference_data: dict):
|
|
if system_message := inference_data.get("system_message"):
|
|
system_turn = [
|
|
cohere.SystemChatMessageV2(
|
|
role="system",
|
|
content=system_message,
|
|
)
|
|
]
|
|
else:
|
|
system_turn = []
|
|
all_chat_turns = system_turn + inference_data["chat_turns"]
|
|
|
|
response, latency = self.generate_with_backoff(
|
|
messages=all_chat_turns,
|
|
tools=[load_cohere_tool(tool=tool) for tool in inference_data["tools"]],
|
|
)
|
|
|
|
model_tool_calls = []
|
|
chat_turn_to_append = cohere.AssistantChatMessageV2(role="assistant")
|
|
response_message = response.message
|
|
if response_message.tool_calls:
|
|
chat_turn_to_append.tool_calls = response_message.tool_calls
|
|
for tool_call in response_message.tool_calls:
|
|
model_tool_calls.append(
|
|
{
|
|
"tool_name": tool_call.function.name,
|
|
"parameters": json.loads(tool_call.function.arguments),
|
|
}
|
|
)
|
|
if response_message.tool_plan:
|
|
chat_turn_to_append.tool_plan = response_message.tool_plan
|
|
if response_message.content:
|
|
if chat_turn_to_append.content is None:
|
|
chat_turn_to_append.content = []
|
|
|
|
for msg in response_message.content:
|
|
if hasattr(msg, "thinking"):
|
|
chat_turn_to_append.content.append(
|
|
cohere.ThinkingAssistantMessageV2ContentItem(thinking=msg.thinking)
|
|
)
|
|
else:
|
|
chat_turn_to_append.content.append(
|
|
cohere.TextAssistantMessageV2ContentItem(text=msg.text)
|
|
)
|
|
if response_message.citations:
|
|
chat_turn_to_append.citations = response_message.citations
|
|
inference_data["chat_turns"].append(chat_turn_to_append)
|
|
|
|
input_token: float | None = None
|
|
output_token: float | None = None
|
|
if response.usage and response.usage.billed_units:
|
|
input_token = response.usage.billed_units.input_tokens
|
|
output_token = response.usage.billed_units.output_tokens
|
|
|
|
metadata = {
|
|
"model_responses": (
|
|
chat_turn_to_append.content if chat_turn_to_append.content else None
|
|
),
|
|
"tool_calls": model_tool_calls,
|
|
"chat_history": [],
|
|
"input_token": input_token or 0,
|
|
"output_token": output_token or 0,
|
|
}
|
|
return metadata, latency
|
|
|
|
@retry_with_backoff(error_type=Exception, stop=stop_after_attempt(5), reraise=True)
|
|
def generate_with_backoff(
|
|
self,
|
|
messages: list,
|
|
tools: list[cohere.types.ToolV2]
|
|
) -> tuple[cohere.v2.types.V2ChatResponse, float]:
|
|
start_time = time.time()
|
|
api_response = self.client.chat(
|
|
model=self.model_name.replace("-FC", ""),
|
|
messages=messages,
|
|
tools=tools,
|
|
citation_options=cohere.CitationOptions(mode="OFF"),
|
|
temperature=self.temperature,
|
|
)
|
|
end_time = time.time()
|
|
|
|
return api_response, end_time - start_time
|
|
|
|
def _pre_query_processing_FC(self, inference_data: dict, test_entry: dict) -> dict:
|
|
turns = []
|
|
for turn_idx, turn in enumerate(test_entry["question"]):
|
|
if turn_idx == 0: # we only extract system message from the first turn
|
|
system_message = extract_system_prompt(turn)
|
|
if system_message:
|
|
inference_data["system_message"] = (
|
|
system_message # we log system message if necessary
|
|
)
|
|
if len(turn) > 0:
|
|
turns.append(preprocess_chat_turns(turn))
|
|
else:
|
|
turns.append(
|
|
[]
|
|
) # for miss_func categories, the turn to supplement function will be empty
|
|
assert len(turns) == len(test_entry["question"])
|
|
test_entry["question"] = turns
|
|
return inference_data
|
|
|
|
def _compile_tools(self, inference_data: dict, test_entry: dict) -> dict:
|
|
functions: list = test_entry["function"]
|
|
|
|
tools = convert_to_tool(functions, GORILLA_TO_OPENAPI, self.model_style)
|
|
inference_data["tools"] = tools
|
|
|
|
return inference_data
|
|
|
|
def _parse_query_response_FC(self, api_response: Any) -> dict:
|
|
reasoning_content = ""
|
|
|
|
if len(api_response["tool_calls"]) > 0: # non empty tool call list
|
|
model_responses = api_response[
|
|
"tool_calls"
|
|
] # list: {"tool_name": , "parameters"}
|
|
else:
|
|
if isinstance(api_response["model_responses"], list):
|
|
model_responses = []
|
|
for item in api_response["model_responses"]:
|
|
if isinstance(item, cohere.types.ThinkingAssistantMessageV2ContentItem):
|
|
reasoning_content += item.thinking
|
|
continue
|
|
elif isinstance(item, cohere.types.TextAssistantMessageV2ContentItem):
|
|
model_responses.append(item.text)
|
|
else:
|
|
model_responses.append(item)
|
|
model_responses = "\n".join(model_responses)
|
|
else:
|
|
model_responses = api_response["model_responses"]
|
|
|
|
return {
|
|
"model_responses": model_responses,
|
|
"reasoning_content": reasoning_content or None,
|
|
"tool_calls": api_response["tool_calls"],
|
|
"chat_history": api_response["chat_history"],
|
|
"input_token": api_response["input_token"],
|
|
"output_token": api_response["output_token"],
|
|
}
|
|
|
|
def add_first_turn_message_FC(
|
|
self, inference_data: dict, first_turn_message: list[dict]
|
|
) -> dict:
|
|
chat_turns = []
|
|
for message in first_turn_message:
|
|
message_role = message["role"]
|
|
assert message_role in [
|
|
"user",
|
|
"assistant",
|
|
], "message role must be in ['user', 'assistant']"
|
|
if message_role == "user":
|
|
chat_turns.append(
|
|
cohere.UserChatMessageV2(role="user", content=message["content"])
|
|
)
|
|
else:
|
|
chat_turns.append(
|
|
cohere.AssistantChatMessageV2(
|
|
role="assistant",
|
|
content=message["content"],
|
|
)
|
|
)
|
|
inference_data["chat_turns"] = chat_turns
|
|
inference_data["raw_prompt"] = []
|
|
inference_data["raw_completion"] = []
|
|
return inference_data
|
|
|
|
def _add_next_turn_user_message_FC(
|
|
self, inference_data: dict, user_message: list[dict]
|
|
) -> dict:
|
|
assert "chat_turns" in inference_data, "expected chat_turns to be present"
|
|
for message in user_message:
|
|
message_role = message["role"]
|
|
if message_role == "user":
|
|
inference_data["chat_turns"].append(
|
|
cohere.UserChatMessageV2(role="user", content=message["content"])
|
|
)
|
|
elif message_role == "assistant":
|
|
inference_data["chat_turns"].append(
|
|
cohere.AssistantChatMessageV2(
|
|
role="assistant",
|
|
content=message["content"],
|
|
)
|
|
)
|
|
else:
|
|
raise Exception(f"Role {message_role} is undefined!")
|
|
if inference_data["chat_turns"][-1].role != "user":
|
|
# if last turn is not user turn - we suffixing a user turn at the end of the conversation history
|
|
inference_data["chat_turns"].append(
|
|
cohere.UserChatMessageV2(role="user", content="")
|
|
)
|
|
return inference_data
|
|
|
|
def _add_assistant_message_FC(
|
|
self, inference_data: dict, model_response_data: dict
|
|
) -> dict:
|
|
# Cohere has all the messages in the chat history already, so no need to add anything here
|
|
return inference_data
|
|
|
|
def _add_execution_results_FC(
|
|
self, inference_data: dict, execution_results: list[str], model_response_data: dict
|
|
) -> dict:
|
|
if execution_results:
|
|
# non-empty execution_results, the last turn of inference_data["chat_turns"] must be a tool use turn
|
|
# otherwise, do nothing
|
|
assert (
|
|
inference_data["chat_turns"][-1].role == "assistant"
|
|
), "last turn must be tool use turn and from the assistant"
|
|
assert inference_data["chat_turns"][
|
|
-1
|
|
].tool_calls, "last turn must have tool calls"
|
|
assert len(inference_data["chat_turns"][-1].tool_calls) == len(
|
|
execution_results
|
|
), "Number of execution result must match number of tool calls from last turn!"
|
|
tool_call_messages = []
|
|
for tool_call, execution_result in zip(
|
|
inference_data["chat_turns"][-1].tool_calls, execution_results
|
|
):
|
|
tool_call_id = tool_call.id
|
|
try:
|
|
tool_execution_result = ast.literal_eval(execution_result)
|
|
except:
|
|
tool_execution_result = execution_result
|
|
if isinstance(tool_execution_result, dict):
|
|
if "id" in tool_execution_result:
|
|
tool_execution_result["ID"] = tool_execution_result["id"]
|
|
del tool_execution_result["id"]
|
|
result_to_render = json.dumps(tool_execution_result)
|
|
else:
|
|
result_to_render = execution_result
|
|
one_tool_call_output = cohere.ToolChatMessageV2(
|
|
tool_call_id=tool_call_id,
|
|
content=[
|
|
cohere.TextToolContent(type="text", text=result_to_render),
|
|
],
|
|
)
|
|
tool_call_messages.append(one_tool_call_output)
|
|
inference_data["chat_turns"].extend(tool_call_messages)
|
|
return inference_data
|
|
|
|
|
|
def preprocess_chat_turns(all_messages: list[dict]) -> list[dict]:
|
|
processed_messages: list[dict] = []
|
|
for message in all_messages:
|
|
processed_messages.append(message)
|
|
return processed_messages
|
|
|
|
|
|
def load_cohere_tool(tool: dict):
|
|
function = tool["function"]
|
|
return cohere.ToolV2(
|
|
type="function",
|
|
function=cohere.ToolV2Function(
|
|
name=function["name"],
|
|
description=function["description"],
|
|
parameters=function["parameters"],
|
|
),
|
|
)
|