tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-17 02:17:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4982 from Manouchehri/bedrock-sts-custom-1

Browse Source 
Allow Bedrock to set custom STS endpoint for OIDC flow
This commit is contained in:
Krish Dholakia
2024-08-02 23:42:38 -07:00
committed by GitHub
parent 5f13d2ee64 507529e8df
commit fa82b2a412
3 changed files with 55 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock_httpx.py
+29 -2
View File
@@ -386,6 +386,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
@@ -406,6 +407,7 @@ class BedrockLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
]
# Iterate over parameters and update if needed
@@ -424,6 +426,7 @@ class BedrockLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
) = params_to_check
### CHECK STS ###
@@ -435,12 +438,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,
}
)
@@ -457,7 +467,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
@@ -852,6 +862,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:
@@ -881,6 +892,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 ###
@@ -1539,6 +1551,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
@@ -1555,6 +1568,7 @@ class BedrockConverseLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
]
# Iterate over parameters and update if needed
@@ -1573,6 +1587,7 @@ class BedrockConverseLLM(BaseLLM):
aws_profile_name,
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
) = params_to_check
### CHECK STS ###
@@ -1581,12 +1596,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,
}
)
@@ -1603,7 +1628,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
@@ -1818,6 +1843,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:
@@ -1847,6 +1873,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 ###
litellm/tests/test_bedrock_completion.py
+3 -1
View File
@@ -558,7 +558,7 @@ def test_completion_bedrock_httpx_command_r_sts_oidc_auth():
import os
aws_web_identity_token = "oidc/circleci_v2/"
aws_region_name = os.environ["AWS_REGION_NAME"]
aws_region_name = "us-west-2"
# aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"]
# 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"
@@ -575,6 +575,8 @@ def test_completion_bedrock_httpx_command_r_sts_oidc_auth():
aws_web_identity_token=aws_web_identity_token,
aws_role_name=aws_role_name,
aws_session_name="my-test-session",
aws_sts_endpoint="https://sts-fips.us-west-2.amazonaws.com",
aws_bedrock_runtime_endpoint="https://bedrock-runtime-fips.us-west-2.amazonaws.com",
)
# Add any assertions here to check the response
print(response)
litellm/tests/test_secret_manager.py
+23 -1
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",
)