zhuchaowe/maxkb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support aws proxy_url

Browse Source
This commit is contained in:
wxg0103 2024-12-24 18:40:59 +08:00 committed by wxg
parent f1cca66012
commit 0cae0c865b
2 changed files with 21 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
apps/setting/models_provider/impl/aws_bedrock_model_provider/credential/llm.py
View File

@ -79,6 +79,7 @@ class BedrockLLMModelCredential(BaseForm, BaseModelCredential):
region_name = forms.TextInputField('Region Name', required=True) region_name = forms.TextInputField('Region Name', required=True)
access_key_id = forms.TextInputField('Access Key ID', required=True) access_key_id = forms.TextInputField('Access Key ID', required=True)
secret_access_key = forms.PasswordInputField('Secret Access Key', required=True) secret_access_key = forms.PasswordInputField('Secret Access Key', required=True)
base_url = forms.TextInputField('Proxy URL', required=False)
def get_model_params_setting_form(self, model_name): def get_model_params_setting_form(self, model_name):
return BedrockLLMModelParams() return BedrockLLMModelParams()

23
apps/setting/models_provider/impl/aws_bedrock_model_provider/model/llm.py
View File

@ -1,4 +1,6 @@
from typing import List, Dict from typing import List, Dict
from botocore.config import Config
from langchain_community.chat_models import BedrockChat from langchain_community.chat_models import BedrockChat
from setting.models_provider.base_model_provider import MaxKBBaseModel from setting.models_provider.base_model_provider import MaxKBBaseModel
@ -33,19 +35,34 @@ class BedrockModel(MaxKBBaseModel, BedrockChat):
return False return False
def __init__(self, model_id: str, region_name: str, credentials_profile_name: str, def __init__(self, model_id: str, region_name: str, credentials_profile_name: str,
streaming: bool = False, **kwargs): streaming: bool = False, config: Config = None, **kwargs):
super().__init__(model_id=model_id, region_name=region_name, super().__init__(model_id=model_id, region_name=region_name,
credentials_profile_name=credentials_profile_name, streaming=streaming, **kwargs) credentials_profile_name=credentials_profile_name, streaming=streaming, config=config,
**kwargs)
@classmethod @classmethod
def new_instance(cls, model_type: str, model_name: str, model_credential: Dict[str, str], def new_instance(cls, model_type: str, model_name: str, model_credential: Dict[str, str],
**model_kwargs) -> 'BedrockModel': **model_kwargs) -> 'BedrockModel':
optional_params = MaxKBBaseModel.filter_optional_params(model_kwargs) optional_params = MaxKBBaseModel.filter_optional_params(model_kwargs)
config = {}
# 判断model_kwargs是否包含 base_url 且不为空
if 'base_url' in model_credential and model_credential['base_url']:
proxy_url = model_credential['base_url']
config = Config(
proxies={
'http': proxy_url,
'https': proxy_url
},
connect_timeout=60,
read_timeout=60
)
return cls( return cls(
model_id=model_name, model_id=model_name,
region_name=model_credential['region_name'], region_name=model_credential['region_name'],
credentials_profile_name=model_credential['credentials_profile_name'], credentials_profile_name=model_credential['credentials_profile_name'],
streaming=model_kwargs.pop('streaming', True), streaming=model_kwargs.pop('streaming', True),
model_kwargs=optional_params model_kwargs=optional_params,
config=config
) )