add bedrock distribution code (#358)

* add bedrock distribution code

* fix linter error

* add bedrock shields support

* linter fixes

* working bedrock safety

* change to return only one violation

* remove env var reading

* refereshable boto credentials

* remove env vars

* address raghu's feedback

* fix session_ttl passing

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>

llama_stack/providers/adapters/inference/bedrock/bedrock.py

@@ -6,9 +6,7 @@
 
 from typing import *  # noqa: F403
 
-import boto3
 from botocore.client import BaseClient
-from botocore.config import Config
 
 from llama_models.llama3.api.chat_format import ChatFormat
 from llama_models.llama3.api.tokenizer import Tokenizer
@@ -16,7 +14,9 @@ from llama_models.llama3.api.tokenizer import Tokenizer
 from llama_stack.providers.utils.inference.model_registry import ModelRegistryHelper
 
 from llama_stack.apis.inference import *  # noqa: F403
+
 from llama_stack.providers.adapters.inference.bedrock.config import BedrockConfig
+from llama_stack.providers.utils.bedrock.client import create_bedrock_client
 
 
 BEDROCK_SUPPORTED_MODELS = {
@@ -34,7 +34,7 @@ class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
         )
         self._config = config
 
-        self._client = _create_bedrock_client(config)
+        self._client = create_bedrock_client(config)
         self.formatter = ChatFormat(Tokenizer.get_instance())
 
     @property
@@ -437,43 +437,3 @@ class BedrockInferenceAdapter(ModelRegistryHelper, Inference):
         contents: List[InterleavedTextMedia],
     ) -> EmbeddingsResponse:
         raise NotImplementedError()
-
-
-def _create_bedrock_client(config: BedrockConfig) -> BaseClient:
-    retries_config = {
-        k: v
-        for k, v in dict(
-            total_max_attempts=config.total_max_attempts,
-            mode=config.retry_mode,
-        ).items()
-        if v is not None
-    }
-
-    config_args = {
-        k: v
-        for k, v in dict(
-            region_name=config.region_name,
-            retries=retries_config if retries_config else None,
-            connect_timeout=config.connect_timeout,
-            read_timeout=config.read_timeout,
-        ).items()
-        if v is not None
-    }
-
-    boto3_config = Config(**config_args)
-
-    session_args = {
-        k: v
-        for k, v in dict(
-            aws_access_key_id=config.aws_access_key_id,
-            aws_secret_access_key=config.aws_secret_access_key,
-            aws_session_token=config.aws_session_token,
-            region_name=config.region_name,
-            profile_name=config.profile_name,
-        ).items()
-        if v is not None
-    }
-
-    boto3_session = boto3.session.Session(**session_args)
-
-    return boto3_session.client("bedrock-runtime", config=boto3_config)

llama_stack/providers/adapters/inference/bedrock/config.py

@@ -3,53 +3,12 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-from typing import *  # noqa: F403
 
 from llama_models.schema_utils import json_schema_type
-from pydantic import BaseModel, Field
+
+from llama_stack.providers.utils.bedrock.config import BedrockBaseConfig
 
 
 @json_schema_type
-class BedrockConfig(BaseModel):
-    aws_access_key_id: Optional[str] = Field(
-        default=None,
-        description="The AWS access key to use. Default use environment variable: AWS_ACCESS_KEY_ID",
-    )
-    aws_secret_access_key: Optional[str] = Field(
-        default=None,
-        description="The AWS secret access key to use. Default use environment variable: AWS_SECRET_ACCESS_KEY",
-    )
-    aws_session_token: Optional[str] = Field(
-        default=None,
-        description="The AWS session token to use. Default use environment variable: AWS_SESSION_TOKEN",
-    )
-    region_name: Optional[str] = Field(
-        default=None,
-        description="The default AWS Region to use, for example, us-west-1 or us-west-2."
-        "Default use environment variable: AWS_DEFAULT_REGION",
-    )
-    profile_name: Optional[str] = Field(
-        default=None,
-        description="The profile name that contains credentials to use."
-        "Default use environment variable: AWS_PROFILE",
-    )
-    total_max_attempts: Optional[int] = Field(
-        default=None,
-        description="An integer representing the maximum number of attempts that will be made for a single request, "
-        "including the initial attempt. Default use environment variable: AWS_MAX_ATTEMPTS",
-    )
-    retry_mode: Optional[str] = Field(
-        default=None,
-        description="A string representing the type of retries Boto3 will perform."
-        "Default use environment variable: AWS_RETRY_MODE",
-    )
-    connect_timeout: Optional[float] = Field(
-        default=60,
-        description="The time in seconds till a timeout exception is thrown when attempting to make a connection. "
-        "The default is 60 seconds.",
-    )
-    read_timeout: Optional[float] = Field(
-        default=60,
-        description="The time in seconds till a timeout exception is thrown when attempting to read from a connection."
-        "The default is 60 seconds.",
-    )
+class BedrockConfig(BedrockBaseConfig):
+    pass