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 108

fix(bedrock): resolve bearer auth CI failures

Browse source
This commit is contained in:
jjj-n 2026-08-28 13:08:18 +08:00
parent d5060b3611
commit 4de35621dd
8 changed files with 115 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -373,22 +373,25 @@ class BedrockConverseLLM(BaseAWSLLM):
# before transforming so whichever path runs emits pre_call once, and
# hand down the credentials, region and endpoint this handler already
# resolved so both paths sign as the same principal.
credential_params: Final = (
()
if credentials is None
else (
("aws_access_key_id", credentials.access_key),
("aws_secret_access_key", credentials.secret_key),
("aws_session_token", credentials.token),
)
)
rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy
**optional_params,
"aws_region_name": aws_region_name,
**(
{ # mutable-ok: merged into its mutable parent above
key: value
for key, value in (
("aws_access_key_id", credentials.access_key),
("aws_secret_access_key", credentials.secret_key),
("aws_session_token", credentials.token),
)
if value is not None
}
if credentials is not None
else {}
),
**{ # mutable-ok: merged into its mutable parent above
key: value
for key, value in (
*credential_params,
("aws_region_name", aws_region_name),
)
if value is not None
},
}
serves_via_rust: Final = rust_chat_completions_accepts(
model=model,

23
tests/llm_translation/test_bedrock_moonshot.py
View file

@ -329,29 +329,6 @@ class TestBedrockMoonshotBasic:
# The model ID in the request body should be stripped
assert transformed["model"] == "moonshot.kimi-k2-thinking"
def test_transform_request_bearer_token_skips_aws_credentials(self):
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
config = AmazonMoonshotConfig()
with patch.object(
config,
"get_credentials",
side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"),
) as get_credentials:
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[{"role": "user", "content": "Hello"}],
optional_params={},
litellm_params={"api_key": "bedrock-bearer-token"},
headers={},
)
get_credentials.assert_not_called()
assert transformed["model"] == "moonshot.kimi-k2-thinking"
class TestBedrockMoonshotReasoningContent:
"""Tests for reasoning content extraction."""

20
tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py Normal file
View file

@ -0,0 +1,20 @@
from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import (
AmazonMoonshotConfig,
)
def test_transform_request_bearer_token_skips_aws_credentials():
config = AmazonMoonshotConfig()
transformed = config.transform_request(
model="bedrock/invoke/moonshot.kimi-k2-thinking",
messages=[{"role": "user", "content": "Hello"}],
optional_params={
"aws_profile_name": "litellm-profile-that-does-not-exist",
"aws_region_name": "us-east-1",
},
litellm_params={"api_key": "bedrock-bearer-token"},
headers={},
)
assert transformed["model"] == "moonshot.kimi-k2-thinking"

16
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -138,14 +138,16 @@ def test_the_core_receives_the_credentials_this_handler_already_resolved():
def test_bearer_token_skips_aws_credential_resolution():
seen = _inject()
with patch.object(
BedrockConverseLLM,
"get_credentials",
side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"),
) as get_credentials:
response = BedrockConverseLLM().completion(**_completion_kwargs(api_key="bedrock-bearer-token"))
response = BedrockConverseLLM().completion(
**_completion_kwargs(
api_key="bedrock-bearer-token",
optional_params={
"maxTokens": 16,
"aws_profile_name": "litellm-profile-that-does-not-exist",
},
)
)
get_credentials.assert_not_called()
assert response.choices[0].message.content == "hello from rust"
params = seen["call"][0]["optional_params"]
assert params["aws_region_name"] == "us-east-1"

30
tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py
View file

@ -1,5 +1,5 @@
import json
from unittest.mock import Mock, patch
from unittest.mock import AsyncMock, Mock, patch
import pytest
@ -239,6 +239,34 @@ class TestBedrockAsyncInvokeEmbedding:
assert status_response["status"] == "InProgress"
assert "invocationArn" in status_response
@pytest.mark.asyncio
async def test_async_invoke_status_uses_sigv4_when_bearer_token_is_configured(
self, monkeypatch
):
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
response = Mock(status_code=200)
response.json.return_value = async_invoke_status_response
client = Mock()
client.get = AsyncMock(return_value=response)
with patch(
"litellm.llms.bedrock.embed.embedding.get_async_httpx_client",
return_value=client,
):
status_response = await BedrockEmbedding()._get_async_invoke_status(
invocation_arn="arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123def456",
aws_region_name="us-east-1",
aws_access_key_id="test-access-key",
aws_secret_access_key="test-secret-key",
aws_session_token="test-session-token",
)
request_headers = client.get.await_args.kwargs["headers"]
assert request_headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert status_response == async_invoke_status_response
def test_async_invoke_error_handling_missing_output_s3_uri(self):
"""Test error handling when output_s3_uri is missing for async invoke."""
from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import (

1
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -65,6 +65,7 @@ def test_bedrock_embedding_with_api_key_bearer_token(model, input_type, embed_re
"client": client,
"aws_region_name": "us-east-1",
"aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-east-1.amazonaws.com",
"aws_profile_name": "litellm-profile-that-does-not-exist",
"api_key": test_api_key,
}

19
tests/test_litellm/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py
View file

@ -3,6 +3,7 @@
import base64
import io
from typing import cast
from unittest.mock import Mock
import httpx
import pytest
@ -72,6 +73,24 @@ def test_get_config_class_stability_unchanged():
assert cls is BedrockStabilityImageEditConfig
def test_prepare_request_bearer_token_skips_aws_credentials():
prepared_request = BedrockImageEdit()._prepare_request(
model="amazon.nova-canvas-v1:0",
image=[io.BytesIO(b"image")],
prompt="replace the background",
optional_params={
"aws_profile_name": "litellm-profile-that-does-not-exist",
"aws_region_name": "us-east-1",
},
api_base="https://bedrock-runtime.us-east-1.amazonaws.com",
extra_headers=None,
logging_obj=Mock(),
api_key="bedrock-bearer-token",
)
assert prepared_request.prepped.headers["Authorization"] == "Bearer bedrock-bearer-token"
def test_provider_config_router_returns_nova_for_canvas():
"""ProviderConfigManager routes Nova Canvas to Nova image-edit config."""
cfg = get_bedrock_image_edit_config_for_model("amazon.nova-canvas-v1:0")

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

@ -743,18 +743,14 @@ def test_bearer_token_skips_boto_credential_resolution(
env_bearer_token: str | None,
):
llm = BaseAWSLLM()
optional_params = {"aws_region_name": "us-east-1"}
optional_params = {
"aws_region_name": "us-east-1",
"aws_profile_name": "litellm-profile-that-does-not-exist",
}
with (
patch(
"litellm.llms.bedrock.base_aws_llm.get_secret_str",
return_value=env_bearer_token,
),
patch.object(
llm,
"get_credentials",
side_effect=AssertionError("bearer-token requests must not resolve AWS credentials"),
) as get_credentials,
with patch(
"litellm.llms.bedrock.base_aws_llm.get_secret_str",
return_value=env_bearer_token,
):
credential_info = llm._get_boto_credentials_from_optional_params(
optional_params,
@ -762,28 +758,30 @@ def test_bearer_token_skips_boto_credential_resolution(
supports_bearer_token=True,
)
get_credentials.assert_not_called()
assert credential_info.credentials is None
assert credential_info.aws_region_name == "us-east-1"
def test_shared_boto_helper_requires_credentials_by_default():
llm = BaseAWSLLM()
credentials = Credentials("test_key", "test_secret", "test_token")
with (
patch(
"litellm.llms.bedrock.base_aws_llm.get_secret_str",
return_value="environment-bearer-token",
),
patch.object(llm, "get_credentials", return_value=credentials) as get_credentials,
with patch(
"litellm.llms.bedrock.base_aws_llm.get_secret_str",
return_value="environment-bearer-token",
):
credential_info = llm._get_boto_credentials_from_optional_params(
{"aws_region_name": "us-east-1"},
{
"aws_region_name": "us-east-1",
"aws_access_key_id": "test_key",
"aws_secret_access_key": "test_secret",
"aws_session_token": "test_token",
},
)
get_credentials.assert_called_once()
assert credential_info.credentials is credentials
assert credential_info.credentials is not None
assert credential_info.credentials.access_key == "test_key"
assert credential_info.credentials.secret_key == "test_secret"
assert credential_info.credentials.token == "test_token"
def test_get_request_headers_with_env_var_bearer_token():