apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-07-05 15:52:12 +02:00
Code Issues Projects Releases Packages Wiki Activity

Update ArchFunctionHandler

Browse source
This commit is contained in:
Shuguang Chen 2024-12-08 16:41:45 -08:00
parent 95e167c2f6
commit e0d4ee7357
1 changed files with 62 additions and 78 deletions
Download patch file Download diff file Expand all files Collapse all files

140
model_server/src/core/function_calling.py
View file

@ -4,7 +4,7 @@ import builtins
import textwrap import textwrap
from openai import OpenAI from openai import OpenAI
from typing import Any, Dict, List, Tuple, Union from typing import Any, Dict, List
from overrides import override from overrides import override
from src.core.model_utils import ( from src.core.model_utils import (
Message, Message,
@ -299,9 +299,7 @@ class ArchFunctionHandler(ArchBaseHandler):
# Attempt to parse the corrected string to ensure it’s valid JSON # Attempt to parse the corrected string to ensure it’s valid JSON
return fixed_str.replace("'", '"') return fixed_str.replace("'", '"')
def _extract_tool_calls( def _extract_tool_calls(self, content: str) -> Dict[str, any]:
self, content: str
) -> Tuple[List[Dict[str, Any]], bool, Union[str, Exception]]:
""" """
Extracts tool call information from a given string. Extracts tool call information from a given string.
@ -309,10 +307,10 @@ class ArchFunctionHandler(ArchBaseHandler):
content (str): The content string containing potential tool call information. content (str): The content string containing potential tool call information.
Returns: Returns:
Tuple[List[Dict[str, Any]], bool, Union[str, Exception]]: Dict: A dictionary of extraction, including:
- A list of tool call dictionaries. - "result": A list of tool call dictionaries.
- A boolean indicating if the extraction was valid. - "status": A boolean indicating if the extraction was valid.
- An error message or exception if extraction failed. - "message": An error message or exception if extraction failed.
""" """
tool_calls, is_valid, error_message = [], True, "" tool_calls, is_valid, error_message = [], True, ""
@ -348,11 +346,11 @@ class ArchFunctionHandler(ArchBaseHandler):
flag = False flag = False
return tool_calls, is_valid, error_message return {"result": tool_calls, "status": is_valid, "message": "error_message"}
def _verify_tool_calls( def _verify_tool_calls(
self, tools: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]] self, tools: List[Dict[str, Any]], tool_calls: List[Dict[str, Any]]
) -> Tuple[bool, Union[Dict[str, Any], None], str]: ) -> Dict[str, any]:
""" """
Verifies the validity of extracted tool calls against the provided tools. Verifies the validity of extracted tool calls against the provided tools.
@ -361,13 +359,13 @@ class ArchFunctionHandler(ArchBaseHandler):
tool_calls (List[Dict[str, Any]]): A list of tool calls to verify. tool_calls (List[Dict[str, Any]]): A list of tool calls to verify.
Returns: Returns:
Tuple[bool, Union[Dict[str, Any], None], str]: Dict: A dictionary of verification, including:
- A boolean indicating if the tool calls are valid. - "status": A boolean indicating if the tool calls are valid.
- The invalid tool call dictionary if any. - "invalid_tool_call": A dictionary of the invalid tool call if any.
- An error message. - "message": An error message.
""" """
is_valid, error_tool_call, error_message = True, None, "" is_valid, invalid_tool_call, error_message = True, None, ""
functions = {} functions = {}
for tool in tools: for tool in tools:
@ -390,9 +388,9 @@ class ArchFunctionHandler(ArchBaseHandler):
for required_param in functions[func_name].get("required", []): for required_param in functions[func_name].get("required", []):
if required_param not in func_args: if required_param not in func_args:
is_valid = False is_valid = False
error_tool_call = tool_call invalid_tool_call = tool_call
error_message = f"`{required_param}` is requiried by the function `{func_name}` but not found in the tool call!" error_message = f"`{required_param}` is requiried by the function `{func_name}` but not found in the tool call!"
return is_valid, error_tool_call, error_message return is_valid, invalid_tool_call, error_message
# Verify the data type of each parameter in the tool calls # Verify the data type of each parameter in the tool calls
for param_name, param_value in func_args: for param_name, param_value in func_args:
@ -403,46 +401,36 @@ class ArchFunctionHandler(ArchBaseHandler):
param_value, self.support_data_types[data_type] param_value, self.support_data_types[data_type]
): ):
is_valid = False is_valid = False
error_tool_call = tool_call invalid_tool_call = tool_call
error_message = f"Parameter `{param_name}` is expected to have the data type `{self.support_data_types[data_type]}`, but got `{type(param_value)}`." error_message = f"Parameter `{param_name}` is expected to have the data type `{self.support_data_types[data_type]}`, but got `{type(param_value)}`."
return is_valid, error_tool_call, error_message return is_valid, invalid_tool_call, error_message
return is_valid, error_tool_call, error_message return {
"status": is_valid,
"invalid_tool_call": invalid_tool_call,
"message": error_message,
}
def _prefill_response(self, messages: List[Dict[str, str]]): def _add_prefill_message(self, messages: List[Dict[str, str]]):
""" """
Prefills the response with the tool call prefix. Update messages and generation params for prompt prefilling
Args: Args:
messages (List[Dict[str, str]]): A list of messages. messages (List[Dict[str, str]]): A list of messages.
tools (List[Dict[str, Any]]): A list of tools.
Returns: Returns:
List[Dict[str, str]]: A list of messages with the prefill prefix. prefill_messages (List[Dict[str, str]]): A list of messages.
""" """
messages.append( return messages + [
{ {
"role": "assistant", "role": "assistant",
"content": random.choice(self.prefill_prefix), "content": random.choice(self.prefill_prefix),
} }
) ]
prefill_response = self.client.chat.completions.create(
messages=messages,
model=self.model_name,
stream=False,
extra_body={
**self.generation_params,
**self.prefill_params,
},
)
return prefill_response
@override @override
async def chat_completion( async def chat_completion(self, req: ChatMessage) -> ChatCompletionResponse:
self, req: ChatMessage, enable_prefilling=True
) -> ChatCompletionResponse:
""" """
Generates a chat completion response for a given request. Generates a chat completion response for a given request.
@ -459,60 +447,56 @@ class ArchFunctionHandler(ArchBaseHandler):
messages = self._process_messages(req.messages, req.tools) messages = self._process_messages(req.messages, req.tools)
# Retrieve the first token, handling the Stream object carefully # always enable `stream=True` to collect model responses
response = self.client.chat.completions.create( response = self.client.chat.completions.create(
messages=messages, messages=messages,
model=self.model_name, model=self.model_name,
stream=enable_prefilling, stream=True,
extra_body=self.generation_params, extra_body=self.generation_params,
) )
model_response = "" model_response, has_tool_call = "", None
if enable_prefilling: for token in response:
has_tool_call = None token_content = token.choices[0].delta.content.strip()
model_response = "" if token_content:
for token in response: if has_tool_call is None and token_content != "<tool_call>":
token_content = token.choices[0].delta.content.strip() has_tool_call = False
response.close()
break
else:
has_tool_call = True
if token_content: if has_tool_call is True:
if has_tool_call is None and token_content != "<tool_call>": model_response += token_content
has_tool_call = False
response.close()
break
else:
has_tool_call = True
if has_tool_call is True: # start parameter gathering if the model is not generating tool calls
model_response += token_content if has_tool_call is False:
prefill_response = self.client.chat.completions.create(
messages=self._add_prefill_message(messages),
model=self.model_name,
extra_body={
**self.generation_params,
**self.prefill_params,
},
)
model_response = prefill_response.choices[0].message.content
# start parameter gathering if the model is not generating a tool call # Extract tool calls from model response
if has_tool_call is False: extracted = self._extract_tool_calls(model_response)
prefill_response = self._prefill_response(messages)
model_response = prefill_response.choices[0].message.content
else:
model_response = response.choices[0].message.content
( if extracted["tool_calls"]:
tool_calls,
extraction_status,
extraction_error_message,
) = self._extract_tool_calls(model_response)
if tool_calls:
# [TODO] Review: define the behavior in the case that tool call extraction fails # [TODO] Review: define the behavior in the case that tool call extraction fails
# if not extraction_status: # if not extracted["status"]:
( verified = self._verify_tool_calls(
verification_status, tools=req.tools, tool_calls=extracted["tool_calls"]
invalid_tool_call, )
verification_error_message,
) = self._verify_tool_calls(tools=req.tools, tool_calls=tool_calls)
# [TODO] Review: In the case that tool calls are invalid, define the protocol to collect debugging output and the behavior to handle it appropriately # [TODO] Review: In the case that tool calls are invalid, define the protocol to collect debugging output and the behavior to handle it appropriately
if verification_status: if verified["status"]:
model_response = Message(content="", tool_calls=tool_calls) model_response = Message(content="", tool_calls=extracted["tool_calls"])
# else: # else:
else: else: