feat: Allow users to specify AWS Bedrock validation models (#2857)

langgenius · Mar 17, 2024 · 6fd1795 · 6fd1795
1 parent f770232
commit 6fd1795
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/api/core/model_runtime/model_providers/bedrock/bedrock.py b/api/core/model_runtime/model_providers/bedrock/bedrock.py
@@ -17,10 +17,9 @@ def validate_provider_credentials(self, credentials: dict) -> None:
         """
         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
@@ -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)