forked from phoenix-oss/llama-stack-mirror
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>
This commit is contained in:
parent
6ebd553da5
commit
093c9f1987
16 changed files with 429 additions and 135 deletions
76
llama_stack/providers/utils/bedrock/client.py
Normal file
76
llama_stack/providers/utils/bedrock/client.py
Normal file
|
@ -0,0 +1,76 @@
|
|||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the terms described in the LICENSE file in
|
||||
# the root directory of this source tree.
|
||||
|
||||
|
||||
import boto3
|
||||
from botocore.client import BaseClient
|
||||
from botocore.config import Config
|
||||
|
||||
from llama_stack.providers.utils.bedrock.config import BedrockBaseConfig
|
||||
from llama_stack.providers.utils.bedrock.refreshable_boto_session import (
|
||||
RefreshableBotoSession,
|
||||
)
|
||||
|
||||
|
||||
def create_bedrock_client(
|
||||
config: BedrockBaseConfig, service_name: str = "bedrock-runtime"
|
||||
) -> BaseClient:
|
||||
"""Creates a boto3 client for Bedrock services with the given configuration.
|
||||
|
||||
Args:
|
||||
config: The Bedrock configuration containing AWS credentials and settings
|
||||
service_name: The AWS service name to create client for (default: "bedrock-runtime")
|
||||
|
||||
Returns:
|
||||
A configured boto3 client
|
||||
"""
|
||||
if config.aws_access_key_id and config.aws_secret_access_key:
|
||||
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 = {
|
||||
"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,
|
||||
"session_ttl": config.session_ttl,
|
||||
}
|
||||
|
||||
# Remove None values
|
||||
session_args = {k: v for k, v in session_args.items() if v is not None}
|
||||
|
||||
boto3_session = boto3.session.Session(**session_args)
|
||||
return boto3_session.client(service_name, config=boto3_config)
|
||||
else:
|
||||
return (
|
||||
RefreshableBotoSession(
|
||||
region_name=config.region_name,
|
||||
profile_name=config.profile_name,
|
||||
session_ttl=config.session_ttl,
|
||||
)
|
||||
.refreshable_session()
|
||||
.client(service_name)
|
||||
)
|
Loading…
Add table
Add a link
Reference in a new issue