From 5a1c29fd8c94fdd4eab326801985acef75f03b1e Mon Sep 17 00:00:00 2001 From: Su Yang Date: Wed, 20 Mar 2024 16:02:13 +0800 Subject: [PATCH] chore: change Yi model SDK to OpenAI (#2910) --- .../model_providers/yi/llm/llm.py | 97 ++++++++++++++++++- 1 file changed, 93 insertions(+), 4 deletions(-) diff --git a/api/core/model_runtime/model_providers/yi/llm/llm.py b/api/core/model_runtime/model_providers/yi/llm/llm.py index 8ad64625149642..d33f38333be9e7 100644 --- a/api/core/model_runtime/model_providers/yi/llm/llm.py +++ b/api/core/model_runtime/model_providers/yi/llm/llm.py @@ -1,30 +1,119 @@ from collections.abc import Generator from typing import Optional, Union +from urllib.parse import urlparse + +import tiktoken from core.model_runtime.entities.llm_entities import LLMResult from core.model_runtime.entities.message_entities import ( PromptMessage, PromptMessageTool, + SystemPromptMessage, ) -from core.model_runtime.model_providers.openai_api_compatible.llm.llm import OAIAPICompatLargeLanguageModel +from core.model_runtime.model_providers.openai.llm.llm import OpenAILargeLanguageModel -class YiLargeLanguageModel(OAIAPICompatLargeLanguageModel): +class YiLargeLanguageModel(OpenAILargeLanguageModel): def _invoke(self, model: str, credentials: dict, prompt_messages: list[PromptMessage], model_parameters: dict, tools: Optional[list[PromptMessageTool]] = None, stop: Optional[list[str]] = None, stream: bool = True, user: Optional[str] = None) \ -> Union[LLMResult, Generator]: self._add_custom_parameters(credentials) + + # yi-vl-plus not support system prompt yet. + if model == "yi-vl-plus": + prompt_message_except_system: list[PromptMessage] = [] + for message in prompt_messages: + if not isinstance(message, SystemPromptMessage): + prompt_message_except_system.append(message) + return super()._invoke(model, credentials, prompt_message_except_system, model_parameters, tools, stop, stream) + return super()._invoke(model, credentials, prompt_messages, model_parameters, tools, stop, stream) def validate_credentials(self, model: str, credentials: dict) -> None: self._add_custom_parameters(credentials) super().validate_credentials(model, credentials) + # refactored from openai model runtime, use cl100k_base for calculate token number + def _num_tokens_from_string(self, model: str, text: str, + tools: Optional[list[PromptMessageTool]] = None) -> int: + """ + Calculate num tokens for text completion model with tiktoken package. + + :param model: model name + :param text: prompt text + :param tools: tools for tool calling + :return: number of tokens + """ + encoding = tiktoken.get_encoding("cl100k_base") + num_tokens = len(encoding.encode(text)) + + if tools: + num_tokens += self._num_tokens_for_tools(encoding, tools) + + return num_tokens + + # refactored from openai model runtime, use cl100k_base for calculate token number + def _num_tokens_from_messages(self, model: str, messages: list[PromptMessage], + tools: Optional[list[PromptMessageTool]] = None) -> int: + """Calculate num tokens for gpt-3.5-turbo and gpt-4 with tiktoken package. + + Official documentation: https://github.com/openai/openai-cookbook/blob/ + main/examples/How_to_format_inputs_to_ChatGPT_models.ipynb""" + encoding = tiktoken.get_encoding("cl100k_base") + tokens_per_message = 3 + tokens_per_name = 1 + + num_tokens = 0 + messages_dict = [self._convert_prompt_message_to_dict(m) for m in messages] + for message in messages_dict: + num_tokens += tokens_per_message + for key, value in message.items(): + # Cast str(value) in case the message value is not a string + # This occurs with function messages + # TODO: The current token calculation method for the image type is not implemented, + # which need to download the image and then get the resolution for calculation, + # and will increase the request delay + if isinstance(value, list): + text = '' + for item in value: + if isinstance(item, dict) and item['type'] == 'text': + text += item['text'] + + value = text + + if key == "tool_calls": + for tool_call in value: + for t_key, t_value in tool_call.items(): + num_tokens += len(encoding.encode(t_key)) + if t_key == "function": + for f_key, f_value in t_value.items(): + num_tokens += len(encoding.encode(f_key)) + num_tokens += len(encoding.encode(f_value)) + else: + num_tokens += len(encoding.encode(t_key)) + num_tokens += len(encoding.encode(t_value)) + else: + num_tokens += len(encoding.encode(str(value))) + + if key == "name": + num_tokens += tokens_per_name + + # every reply is primed with assistant + num_tokens += 3 + + if tools: + num_tokens += self._num_tokens_for_tools(encoding, tools) + + return num_tokens + @staticmethod def _add_custom_parameters(credentials: dict) -> None: credentials['mode'] = 'chat' - + credentials['openai_api_key']=credentials['api_key'] if 'endpoint_url' not in credentials or credentials['endpoint_url'] == "": - credentials['endpoint_url'] = 'https://api.lingyiwanwu.com/v1' + credentials['openai_api_base']='https://api.lingyiwanwu.com' + else: + parsed_url = urlparse(credentials['endpoint_url']) + credentials['openai_api_base']=f"{parsed_url.scheme}://{parsed_url.netloc}"