scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

(bedrock_httpx.py) - Add support for custom STS endpoints, e.g. for FIPS.

Browse source
This commit is contained in:
David Manouchehri 2024-07-31 16:23:06 +00:00
parent 7db85fbdb7
commit ffbe1845ef
No known key found for this signature in database
2 changed files with 52 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/llms/bedrock_httpx.py
View file

@ -383,6 +383,7 @@ class BedrockLLM(BaseLLM):
aws_profile_name: Optional[str] = None,
aws_role_name: Optional[str] = None,
aws_web_identity_token: Optional[str] = None,
aws_sts_endpoint: Optional[str] = None,
):
"""
Return a boto3.Credentials object
@ -403,6 +404,7 @@ class BedrockLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
]
# Iterate over parameters and update if needed
@ -421,6 +423,7 @@ class BedrockLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
) = params_to_check
### CHECK STS ###
@ -432,12 +435,19 @@ class BedrockLLM(BaseLLM):
print_verbose(
f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}"
)
if aws_sts_endpoint is None:
sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com"
else:
sts_endpoint = aws_sts_endpoint
iam_creds_cache_key = json.dumps(
{
"aws_web_identity_token": aws_web_identity_token,
"aws_role_name": aws_role_name,
"aws_session_name": aws_session_name,
"aws_region_name": aws_region_name,
"aws_sts_endpoint": sts_endpoint,
}
)
@ -454,7 +464,7 @@ class BedrockLLM(BaseLLM):
sts_client = boto3.client(
"sts",
region_name=aws_region_name,
endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com",
endpoint_url=sts_endpoint,
)
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
@ -849,6 +859,7 @@ class BedrockLLM(BaseLLM):
"aws_bedrock_runtime_endpoint", None
) # https://bedrock-runtime.{region_name}.amazonaws.com
aws_web_identity_token = optional_params.pop("aws_web_identity_token", None)
aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None)
### SET REGION NAME ###
if aws_region_name is None:
@ -878,6 +889,7 @@ class BedrockLLM(BaseLLM):
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
)
### SET RUNTIME ENDPOINT ###
@ -1536,6 +1548,7 @@ class BedrockConverseLLM(BaseLLM):
aws_profile_name: Optional[str] = None,
aws_role_name: Optional[str] = None,
aws_web_identity_token: Optional[str] = None,
aws_sts_endpoint: Optional[str] = None,
):
"""
Return a boto3.Credentials object
@ -1552,6 +1565,7 @@ class BedrockConverseLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
]
# Iterate over parameters and update if needed
@ -1570,6 +1584,7 @@ class BedrockConverseLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
) = params_to_check
### CHECK STS ###
@ -1578,12 +1593,22 @@ class BedrockConverseLLM(BaseLLM):
and aws_role_name is not None
and aws_session_name is not None
):
print_verbose(
f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}"
)
if aws_sts_endpoint is None:
sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com"
else:
sts_endpoint = aws_sts_endpoint
iam_creds_cache_key = json.dumps(
{
"aws_web_identity_token": aws_web_identity_token,
"aws_role_name": aws_role_name,
"aws_session_name": aws_session_name,
"aws_region_name": aws_region_name,
"aws_sts_endpoint": sts_endpoint,
}
)
@ -1600,7 +1625,7 @@ class BedrockConverseLLM(BaseLLM):
sts_client = boto3.client(
"sts",
region_name=aws_region_name,
endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com",
endpoint_url=sts_endpoint,
)
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
@ -1815,6 +1840,7 @@ class BedrockConverseLLM(BaseLLM):
"aws_bedrock_runtime_endpoint", None
) # https://bedrock-runtime.{region_name}.amazonaws.com
aws_web_identity_token = optional_params.pop("aws_web_identity_token", None)
aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None)
### SET REGION NAME ###
if aws_region_name is None:
@ -1844,6 +1870,7 @@ class BedrockConverseLLM(BaseLLM):
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
)
### SET RUNTIME ENDPOINT ###

24
litellm/tests/test_secret_manager.py
View file

@ -13,7 +13,7 @@ import pytest
from litellm import get_secret
from litellm.proxy.secret_managers.aws_secret_manager import load_aws_secret_manager
from litellm.llms.azure import get_azure_ad_token_from_oidc
from litellm.llms.bedrock_httpx import BedrockLLM
from litellm.llms.bedrock_httpx import BedrockLLM, BedrockConverseLLM
@pytest.mark.skip(reason="AWS Suspended Account")
@ -113,3 +113,25 @@ def test_oidc_circle_v1_with_amazon():
aws_role_name=aws_role_name,
aws_session_name="assume-v1-session",
)
@pytest.mark.skipif(
os.environ.get("CIRCLE_OIDC_TOKEN") is None,
reason="Cannot run without being in CircleCI Runner",
)
def test_oidc_circle_v1_with_amazon_fips():
# The purpose of this test is to validate that we can assume a role in a FIPS region
# TODO: This is using ai.moda's IAM role, we should use LiteLLM's IAM role eventually
aws_role_name = (
"arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci-v1-assume-only"
)
aws_web_identity_token = "oidc/circleci/"
bllm = BedrockConverseLLM()
creds = bllm.get_credentials(
aws_region_name="us-west-1",
aws_web_identity_token=aws_web_identity_token,
aws_role_name=aws_role_name,
aws_session_name="assume-v1-session-fips",
aws_sts_endpoint="https://sts-fips.us-west-1.amazonaws.com",
)