scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

Add support for AWS assume_role with a session token

Browse source
This commit is contained in:
Steven Anton 2025-08-23 21:52:18 -07:00
parent 827d1def02
commit d95578464e
3 changed files with 17 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/llms/bedrock/base_aws_llm.py
View file

@ -202,6 +202,7 @@ class BaseAWSLLM:
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
)
@ -554,6 +555,7 @@ class BaseAWSLLM:
self,
aws_access_key_id: Optional[str],
aws_secret_access_key: Optional[str],
aws_session_token: Optional[str],
aws_role_name: str,
aws_session_name: str,
) -> Tuple[Credentials, Optional[int]]:
@ -614,6 +616,7 @@ class BaseAWSLLM:
"sts",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
)
sts_response = sts_client.assume_role(

1
tests/llm_translation/test_aws_base_llm.py
View file

@ -88,6 +88,7 @@ def test_auth_with_aws_role(mock_boto3_client, base_aws_llm):
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="test_access",
aws_secret_access_key="test_secret",
aws_session_token="test_token",
aws_role_name="test_role",
aws_session_name="test_session",
)

14
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -576,6 +576,7 @@ def test_eks_irsa_ambient_credentials_used():
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
@ -630,6 +631,7 @@ def test_explicit_credentials_used_when_provided():
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
aws_session_token="assumed-session-token",
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
@ -639,6 +641,7 @@ def test_explicit_credentials_used_when_provided():
"sts",
aws_access_key_id="explicit-access-key",
aws_secret_access_key="explicit-secret-key",
aws_session_token="assumed-session-token",
)
# Should call assume_role
@ -687,6 +690,7 @@ def test_partial_credentials_still_use_ambient():
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id="AKIAEXAMPLE",
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="test-session"
)
@ -695,7 +699,8 @@ def test_partial_credentials_still_use_ambient():
mock_boto3_client.assert_called_once_with(
"sts",
aws_access_key_id="AKIAEXAMPLE",
aws_secret_access_key=None
aws_secret_access_key=None,
aws_session_token=None,
)
# Should still call assume_role
@ -737,6 +742,7 @@ def test_cross_account_role_assumption():
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole",
aws_session_name="cross-account-session"
)
@ -789,6 +795,7 @@ def test_role_assumption_with_custom_session_name():
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="evals-bedrock-session"
)
@ -832,6 +839,7 @@ def test_role_assumption_ttl_calculation():
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="ttl-test-session"
)
@ -858,6 +866,7 @@ def test_role_assumption_error_handling():
base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole",
aws_session_name="error-test-session"
)
@ -911,6 +920,7 @@ def test_multiple_role_assumptions_in_sequence():
credentials1, ttl1 = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole",
aws_session_name="session-1"
)
@ -919,6 +929,7 @@ def test_multiple_role_assumptions_in_sequence():
credentials2, ttl2 = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole",
aws_session_name="session-2"
)
@ -980,6 +991,7 @@ def test_auth_with_aws_role_irsa_environment():
creds, ttl = base_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name='arn:aws:iam::222222222222:role/target-role',
aws_session_name='test-session'
)