mirror of
https://github.com/1Panel-dev/MaxKB.git
synced 2025-12-25 17:22:55 +00:00
fix: aws computer tokens error
--bug=1055094 --user=王孝刚 【应用】aws模型在应用编排中使用时,回答内容会带一串报错信息,并且消耗tokens显示为0 https://www.tapd.cn/57709429/s/1689448
This commit is contained in:
parent
f24337d5f3
commit
e58f95832b
|
|
@ -1,10 +1,12 @@
|
|||
import os
|
||||
import re
|
||||
from typing import Dict
|
||||
from typing import Dict, List
|
||||
|
||||
from botocore.config import Config
|
||||
from langchain_community.chat_models import BedrockChat
|
||||
from langchain_core.messages import BaseMessage, get_buffer_string
|
||||
|
||||
from common.config.tokenizer_manage_config import TokenizerManage
|
||||
from setting.models_provider.base_model_provider import MaxKBBaseModel
|
||||
|
||||
|
||||
|
|
@ -72,6 +74,20 @@ class BedrockModel(MaxKBBaseModel, BedrockChat):
|
|||
config=config
|
||||
)
|
||||
|
||||
def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
|
||||
try:
|
||||
return super().get_num_tokens_from_messages(messages)
|
||||
except Exception as e:
|
||||
tokenizer = TokenizerManage.get_tokenizer()
|
||||
return sum([len(tokenizer.encode(get_buffer_string([m]))) for m in messages])
|
||||
|
||||
def get_num_tokens(self, text: str) -> int:
|
||||
try:
|
||||
return super().get_num_tokens(text)
|
||||
except Exception as e:
|
||||
tokenizer = TokenizerManage.get_tokenizer()
|
||||
return len(tokenizer.encode(text))
|
||||
|
||||
|
||||
def _update_aws_credentials(profile_name, access_key_id, secret_access_key):
|
||||
credentials_path = os.path.join(os.path.expanduser("~"), ".aws", "credentials")
|
||||
|
|
|
|||
Loading…
Reference in New Issue