From 1a2523fd158282df89ba875ddc84193c3e81515e Mon Sep 17 00:00:00 2001
From: MaFee921 <31881301+2284730142@users.noreply.github.com>
Date: Wed, 5 Feb 2025 12:24:24 +0800
Subject: [PATCH] feat: bedrock_endpoint_url (#12838)

---
 .../model_runtime/model_providers/bedrock/bedrock.yaml   | 9 +++++++++
 .../model_providers/bedrock/get_bedrock_client.py        | 2 ++
 2 files changed, 11 insertions(+)

diff --git a/api/core/model_runtime/model_providers/bedrock/bedrock.yaml b/api/core/model_runtime/model_providers/bedrock/bedrock.yaml
index 63d6640348..358c942c3f 100644
--- a/api/core/model_runtime/model_providers/bedrock/bedrock.yaml
+++ b/api/core/model_runtime/model_providers/bedrock/bedrock.yaml
@@ -123,6 +123,15 @@ provider_credential_schema:
             en_US: AWS GovCloud (US-West)
             zh_Hans: AWS GovCloud (US-West)
             ja_JP: AWS GovCloud (米国西部)
+    - variable: bedrock_endpoint_url
+      label:
+        zh_Hans: Bedrock Endpoint URL
+        en_US: Bedrock Endpoint URL
+      type: text-input
+      required: false
+      placeholder:
+        zh_Hans: 在此输入您的 Bedrock Endpoint URL, 如：https://123456.cloudfront.net
+        en_US: Enter your Bedrock Endpoint URL, e.g. https://123456.cloudfront.net
     - variable: model_for_validation
       required: false
       label:
diff --git a/api/core/model_runtime/model_providers/bedrock/get_bedrock_client.py b/api/core/model_runtime/model_providers/bedrock/get_bedrock_client.py
index 2ad37cef3b..61cd72889e 100644
--- a/api/core/model_runtime/model_providers/bedrock/get_bedrock_client.py
+++ b/api/core/model_runtime/model_providers/bedrock/get_bedrock_client.py
@@ -13,6 +13,7 @@ def get_bedrock_client(service_name: str, credentials: Mapping[str, str]):
     client_config = Config(region_name=region_name)
     aws_access_key_id = credentials.get("aws_access_key_id")
     aws_secret_access_key = credentials.get("aws_secret_access_key")
+    bedrock_endpoint_url = credentials.get("bedrock_endpoint_url")
 
     if aws_access_key_id and aws_secret_access_key:
         # use aksk to call bedrock
@@ -21,6 +22,7 @@ def get_bedrock_client(service_name: str, credentials: Mapping[str, str]):
             config=client_config,
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            **({"endpoint_url": bedrock_endpoint_url} if bedrock_endpoint_url else {}),
         )
     else:
         # use iam without aksk to call