From 6fd1795d25d746791d4cd152883e92442d0bf33a Mon Sep 17 00:00:00 2001 From: Su Yang Date: Mon, 18 Mar 2024 00:44:09 +0800 Subject: [PATCH] feat: Allow users to specify AWS Bedrock validation models (#2857) --- .../model_runtime/model_providers/bedrock/bedrock.py | 5 ++--- .../model_runtime/model_providers/bedrock/bedrock.yaml | 9 +++++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/api/core/model_runtime/model_providers/bedrock/bedrock.py b/api/core/model_runtime/model_providers/bedrock/bedrock.py index aa322fc664..96cb90280e 100644 --- a/api/core/model_runtime/model_providers/bedrock/bedrock.py +++ b/api/core/model_runtime/model_providers/bedrock/bedrock.py @@ -17,10 +17,9 @@ class BedrockProvider(ModelProvider): """ try: model_instance = self.get_model_instance(ModelType.LLM) - - # Use `gemini-pro` model for validate, + bedrock_validate_model_name = credentials.get('model_for_validation', 'amazon.titan-text-lite-v1') model_instance.validate_credentials( - model='amazon.titan-text-lite-v1', + model=bedrock_validate_model_name, credentials=credentials ) except CredentialsValidateFailedError as ex: diff --git a/api/core/model_runtime/model_providers/bedrock/bedrock.yaml b/api/core/model_runtime/model_providers/bedrock/bedrock.yaml index 1458b830cd..05cd402d4e 100644 --- a/api/core/model_runtime/model_providers/bedrock/bedrock.yaml +++ b/api/core/model_runtime/model_providers/bedrock/bedrock.yaml @@ -69,3 +69,12 @@ provider_credential_schema: label: en_US: AWS GovCloud (US-West) zh_Hans: AWS GovCloud (US-West) + - variable: model_for_validation + required: false + label: + en_US: Available Model Name + zh_Hans: 可用模型名称 + type: text-input + placeholder: + en_US: A model you have access to (e.g. amazon.titan-text-lite-v1) for validation. + zh_Hans: 为了进行验证，请输入一个您可用的模型名称 (例如：amazon.titan-text-lite-v1)