diff --git a/memgpt/local_llm/llm_chat_completion_wrappers/dolphin.py b/memgpt/local_llm/llm_chat_completion_wrappers/dolphin.py new file mode 100644 index 00000000..d1ae5f22 --- /dev/null +++ b/memgpt/local_llm/llm_chat_completion_wrappers/dolphin.py @@ -0,0 +1,243 @@ +import json + +from .wrapper_base import LLMChatCompletionWrapper + + +class Dolphin21MistralWrapper(LLMChatCompletionWrapper): + """Wrapper for Dolphin 2.1 Mistral 7b: https://huggingface.co/ehartford/dolphin-2.1-mistral-7b + + Note: this wrapper formats a prompt that only generates JSON, no inner thoughts + """ + + def __init__( + self, + simplify_json_content=True, + clean_function_args=True, + include_assistant_prefix=True, + include_opening_brace_in_prefix=True, + include_section_separators=False, + ): + self.simplify_json_content = simplify_json_content + self.clean_func_args = clean_function_args + self.include_assistant_prefix = include_assistant_prefix + self.include_opening_brance_in_prefix = include_opening_brace_in_prefix + self.include_section_separators = include_section_separators + + def chat_completion_to_prompt(self, messages, functions): + """Example for airoboros: https://huggingface.co/jondurbin/airoboros-l2-70b-2.1#prompt-format + + <|im_start|>system + You are Dolphin, a helpful AI assistant.<|im_end|> + <|im_start|>user + {prompt}<|im_end|> + <|im_start|>assistant + + Do function spec Airoboros style inside the system message: + Functions support: https://huggingface.co/jondurbin/airoboros-l2-70b-2.1#agentfunction-calling + + As an AI assistant, please select the most suitable function and parameters from the list of available functions below, based on the user's input. Provide your response in JSON format. + + Input: I want to know how many times 'Python' is mentioned in my text file. + + Available functions: + file_analytics: + description: This tool performs various operations on a text file. + params: + action: The operation we want to perform on the data, such as "count_occurrences", "find_line", etc. + filters: + keyword: The word or phrase we want to search for. + + OpenAI functions schema style: + + { + "name": "send_message", + "description": "Sends a message to the human user", + "parameters": { + "type": "object", + "properties": { + # https://json-schema.org/understanding-json-schema/reference/array.html + "message": { + "type": "string", + "description": "Message contents. All unicode (including emojis) are supported.", + }, + }, + "required": ["message"], + } + }, + """ + prompt = "" + + # <|im_start|>system + # You are Dolphin, a helpful AI assistant.<|im_end|> + + IM_START_TOKEN = "<|im_start|>" + IM_END_TOKEN = "<|im_end|>" + + # System instructions go first + assert messages[0]["role"] == "system" + prompt += f"{IM_START_TOKEN}system" + prompt += f"\n{messages[0]['content']}" + + # Next is the functions preamble + def create_function_description(schema): + # airorobos style + func_str = "" + func_str += f"{schema['name']}:" + func_str += f"\n description: {schema['description']}" + func_str += f"\n params:" + for param_k, param_v in schema["parameters"]["properties"].items(): + # TODO we're ignoring type + func_str += f"\n {param_k}: {param_v['description']}" + # TODO we're ignoring schema['parameters']['required'] + return func_str + + # prompt += f"\nPlease select the most suitable function and parameters from the list of available functions below, based on the user's input. Provide your response in JSON format." + prompt += f"\nPlease select the most suitable function and parameters from the list of available functions below, based on the ongoing conversation. Provide your response in JSON format." + prompt += f"\nAvailable functions:" + for function_dict in functions: + prompt += f"\n{create_function_description(function_dict)}" + + # Put functions INSIDE system message (TODO experiment with this) + prompt += IM_END_TOKEN + + def create_function_call(function_call): + """Go from ChatCompletion to Airoboros style function trace (in prompt) + + ChatCompletion data (inside message['function_call']): + "function_call": { + "name": ... + "arguments": { + "arg1": val1, + ... + } + + Airoboros output: + { + "function": "send_message", + "params": { + "message": "Hello there! I am Sam, an AI developed by Liminal Corp. How can I assist you today?" + } + } + """ + airo_func_call = { + "function": function_call["name"], + "params": json.loads(function_call["arguments"]), + } + return json.dumps(airo_func_call, indent=2) + + # option (1): from HF README: + # <|im_start|>user + # {prompt}<|im_end|> + # <|im_start|>assistant + # {assistant reply} + # {function output (if function)} + + # option (2): take liberties + # <|im_start|>user + # {prompt}<|im_end|> + # <|im_start|>assistant + # or + # <|im_start|>function + + # Add a sep for the conversation + # if self.include_section_separators: + # prompt += "\n### INPUT" + + # Last are the user/assistant messages + for message in messages[1:]: + assert message["role"] in ["user", "assistant", "function"], message + + if message["role"] == "user": + if self.simplify_json_content: + try: + content_json = json.loads(message["content"]) + content_simple = content_json["message"] + prompt += ( + f"\n{IM_START_TOKEN}user\n{content_simple}{IM_END_TOKEN}" + ) + # prompt += f"\nUSER: {content_simple}" + except: + prompt += f"\n{IM_START_TOKEN}user\n{message['content']}{IM_END_TOKEN}" + # prompt += f"\nUSER: {message['content']}" + elif message["role"] == "assistant": + prompt += f"\n{IM_START_TOKEN}assistant" + if message["content"] is not None: + prompt += f"\n{message['content']}" + # prompt += f"\nASSISTANT: {message['content']}" + # need to add the function call if there was one + if message["function_call"]: + prompt += f"\n{create_function_call(message['function_call'])}" + prompt += f"{IM_END_TOKEN}" + elif message["role"] == "function": + # TODO find a good way to add this + # prompt += f"\nASSISTANT: (function return) {message['content']}" + prompt += f"\n{IM_START_TOKEN}assistant" + prompt += f"\nFUNCTION RETURN: {message['content']}" + # prompt += f"\nFUNCTION RETURN: {message['content']}" + continue + else: + raise ValueError(message) + + # Add a sep for the response + # if self.include_section_separators: + # prompt += "\n### RESPONSE" + + if self.include_assistant_prefix: + # prompt += f"\nASSISTANT:" + prompt += f"\n{IM_START_TOKEN}assistant" + if self.include_opening_brance_in_prefix: + prompt += "\n{" + + print(prompt) + return prompt + + def clean_function_args(self, function_name, function_args): + """Some basic MemGPT-specific cleaning of function args""" + cleaned_function_name = function_name + cleaned_function_args = function_args.copy() + + if function_name == "send_message": + # strip request_heartbeat + cleaned_function_args.pop("request_heartbeat", None) + + # TODO more cleaning to fix errors LLM makes + return cleaned_function_name, cleaned_function_args + + def output_to_chat_completion_response(self, raw_llm_output): + """Turn raw LLM output into a ChatCompletion style response with: + "message" = { + "role": "assistant", + "content": ..., + "function_call": { + "name": ... + "arguments": { + "arg1": val1, + ... + } + } + } + """ + if self.include_opening_brance_in_prefix and raw_llm_output[0] != "{": + raw_llm_output = "{" + raw_llm_output + + try: + function_json_output = json.loads(raw_llm_output) + except Exception as e: + raise Exception(f"Failed to decode JSON from LLM output:\n{raw_llm_output}") + function_name = function_json_output["function"] + function_parameters = function_json_output["params"] + + if self.clean_func_args: + function_name, function_parameters = self.clean_function_args( + function_name, function_parameters + ) + + message = { + "role": "assistant", + "content": None, + "function_call": { + "name": function_name, + "arguments": json.dumps(function_parameters), + }, + } + return message