AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-14 20:55:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: bedrock get_num_tokens prompt_messages parameter name err (#4932)

Browse Source
This commit is contained in:
takatost 2024-06-05 01:53:05 +08:00 committed by GitHub
parent 21ac2afb3a
commit f44d1e62d2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
api/core/model_runtime/model_providers/bedrock/llm/llm.py
View File

@ -358,27 +358,26 @@ class BedrockLargeLanguageModel(LargeLanguageModel):
return message_dict return message_dict
def get_num_tokens(self, model: str, credentials: dict, messages: list[PromptMessage] | str, def get_num_tokens(self, model: str, credentials: dict, prompt_messages: list[PromptMessage] | str,
tools: Optional[list[PromptMessageTool]] = None) -> int: tools: Optional[list[PromptMessageTool]] = None) -> int:
""" """
Get number of tokens for given prompt messages Get number of tokens for given prompt messages
:param model: model name :param model: model name
:param credentials: model credentials :param credentials: model credentials
:param messages: prompt messages or message string :param prompt_messages: prompt messages or message string
:param tools: tools for tool calling :param tools: tools for tool calling
:return:md = genai.GenerativeModel(model) :return:md = genai.GenerativeModel(model)
""" """
prefix = model.split('.')[0] prefix = model.split('.')[0]
model_name = model.split('.')[1] model_name = model.split('.')[1]
if isinstance(messages, str): if isinstance(prompt_messages, str):
prompt = messages prompt = prompt_messages
else: else:
prompt = self._convert_messages_to_prompt(messages, prefix, model_name) prompt = self._convert_messages_to_prompt(prompt_messages, prefix, model_name)
return self._get_num_tokens_by_gpt2(prompt) return self._get_num_tokens_by_gpt2(prompt)
def validate_credentials(self, model: str, credentials: dict) -> None: def validate_credentials(self, model: str, credentials: dict) -> None:
""" """
Validate model credentials Validate model credentials