[ref] use one method to get boto client for aws bedrock

api/core/model_runtime/model_providers/bedrock/get_bedrock_client.py

@@ -0,0 +1,21 @@
+import boto3
+from botocore.config import Config
+
+
+def get_bedrock_client(service_name, credentials=None):
+    client_config = Config(region_name=credentials["aws_region"])
+    aws_access_key_id = (credentials["aws_access_key_id"],)
+    aws_secret_access_key = credentials["aws_secret_access_key"]
+    if aws_access_key_id and aws_secret_access_key:
+        # 使用 AKSK 方式
+        client = boto3.client(
+            service_name=service_name,
+            config=client_config,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+    else:
+        # 使用 IAM 角色方式
+        client = boto3.client(service_name=service_name, config=client_config)
+
+    return client

api/core/model_runtime/model_providers/bedrock/llm/llm.py

@@ -7,6 +7,7 @@
 
 # 3rd import
 import boto3
+from api.core.model_runtime.model_providers.bedrock.get_bedrock_client import get_bedrock_client
 from botocore.config import Config
 from botocore.exceptions import (
     ClientError,
@@ -173,13 +174,7 @@ def _generate_with_converse(
         :param stream: is stream response
         :return: full response or stream response chunk generator result
         """
-        bedrock_client = boto3.client(
-            service_name="bedrock-runtime",
-            aws_access_key_id=credentials.get("aws_access_key_id"),
-            aws_secret_access_key=credentials.get("aws_secret_access_key"),
-            region_name=credentials["aws_region"],
-        )
-
+        bedrock_client = get_bedrock_client("bedrock-runtime", credentials)
         system, prompt_message_dicts = self._convert_converse_prompt_messages(prompt_messages)
         inference_config, additional_model_fields = self._convert_converse_api_model_parameters(model_parameters, stop)
 

api/core/model_runtime/model_providers/bedrock/rerank/rerank.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
-import boto3
-from botocore.config import Config
+from api.core.model_runtime.model_providers.bedrock.get_bedrock_client import get_bedrock_client
 
 from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
 from core.model_runtime.errors.invoke import (
@@ -48,13 +47,7 @@ def _invoke(
             return RerankResult(model=model, docs=docs)
 
         # initialize client
-        client_config = Config(region_name=credentials["aws_region"])
-        bedrock_runtime = boto3.client(
-            service_name="bedrock-agent-runtime",
-            config=client_config,
-            aws_access_key_id=credentials.get("aws_access_key_id", ""),
-            aws_secret_access_key=credentials.get("aws_secret_access_key"),
-        )
+        bedrock_runtime = get_bedrock_client("bedrock-agent-runtime", credentials)
         queries = [{"type": "TEXT", "textQuery": {"text": query}}]
         text_sources = []
         for text in docs:

api/core/model_runtime/model_providers/bedrock/text_embedding/text_embedding.py

@@ -3,8 +3,7 @@
 import time
 from typing import Optional
 
-import boto3
-from botocore.config import Config
+from api.core.model_runtime.model_providers.bedrock.get_bedrock_client import get_bedrock_client
 from botocore.exceptions import (
     ClientError,
     EndpointConnectionError,
@@ -48,14 +47,7 @@ def _invoke(
         :param input_type: input type
         :return: embeddings result
         """
-        client_config = Config(region_name=credentials["aws_region"])
-
-        bedrock_runtime = boto3.client(
-            service_name="bedrock-runtime",
-            config=client_config,
-            aws_access_key_id=credentials.get("aws_access_key_id"),
-            aws_secret_access_key=credentials.get("aws_secret_access_key"),
-        )
+        bedrock_runtime = get_bedrock_client("bedrock-runtime", credentials)
 
         embeddings = []
         token_usage = 0