fix: aws computer tokens error

--bug=1055094 --user=王孝刚 【应用】aws模型在应用编排中使用时,回答内容会带一串报错信息,并且消耗tokens显示为0 https://www.tapd.cn/57709429/s/1689448
This commit is contained in:
wxg0103 2025-04-23 10:41:18 +08:00
parent f24337d5f3
commit e58f95832b

View File

@ -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")