scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

Merge abe07dda76 into 52023601a8

Browse source
This commit is contained in:
Micah Mills 2026-10-18 00:54:00 +00:00 • committed by GitHub
commit 497ea8fca6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 185 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -39,7 +39,7 @@ from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams, AwsSe
if TYPE_CHECKING:
from botocore.awsrequest import AWSPreparedRequest
from botocore.credentials import Credentials
from botocore.credentials import Credentials, ReadOnlyCredentials
else:
Credentials = Any
AWSPreparedRequest = Any
@ -1617,7 +1617,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
@tracer.wrap()
def get_request_headers(
self,
credentials: Credentials | None,
credentials: "Credentials | ReadOnlyCredentials | None",
aws_region_name: str,
extra_headers: dict | None,
endpoint_url: str,
@ -1652,7 +1652,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
# Filter headers for AWS signature calculation
# AWS SigV4 only includes specific headers in signature calculation
aws_signature_headers: Final = self._filter_headers_for_aws_signature(headers)
sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
sigv4: Final = SigV4Auth(_freeze_aws_credentials(credentials), "bedrock", aws_region_name)
request = AWSRequest(
method="POST",
url=endpoint_url,
@ -1746,7 +1746,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model=model)
credentials: Final[Credentials] = self.resolve_credentials(auth_params, aws_region_name)
sigv4: Final = SigV4Auth(credentials, service_name, aws_region_name)
sigv4: Final = SigV4Auth(_freeze_aws_credentials(credentials), service_name, aws_region_name)
headers = headers or {}
if not any(header_name.lower() == "content-type" for header_name in headers):
headers = {"Content-Type": "application/json", **headers}
@ -1778,8 +1778,16 @@ class BaseAWSLLM(SignsRequestsWithAWS):
return request_headers_dict, request.body
def _freeze_aws_credentials(credentials: "Credentials | ReadOnlyCredentials") -> "ReadOnlyCredentials":
from botocore.credentials import Credentials
if isinstance(credentials, Credentials):
return credentials.get_frozen_credentials()
return credentials
def sign_aws_json_post(
get_credentials: Callable[[], Credentials],
get_credentials: Callable[[], "Credentials | ReadOnlyCredentials"],
service_name: str,
aws_region_name: str | None,
url: str,
@ -1791,7 +1799,7 @@ def sign_aws_json_post(
from botocore.awsrequest import AWSRequest
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
SigV4Auth(_freeze_aws_credentials(get_credentials()), service_name, aws_region_name).add_auth(aws_request)
return aws_request.prepare()

204
tests/unit/llms/bedrock/test_base_aws_llm.py
View file

@ -16,12 +16,12 @@ from fastapi.testclient import TestClient
from collections.abc import Callable
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Final, Optional
from typing import Any, Dict, Final, Literal, Optional
from unittest.mock import MagicMock, Mock, patch
from botocore.awsrequest import AWSPreparedRequest, AWSRequest
from botocore.auth import SigV4Auth
from botocore.credentials import Credentials
from botocore.credentials import Credentials, ReadOnlyCredentials
from botocore.exceptions import ClientError, NoCredentialsError
import litellm
@ -768,38 +768,32 @@ def test_get_request_headers_preserves_bearer_payload(from_environment):
assert result.method == "POST"
def test_get_request_headers_with_sigv4():
# Setup
llm = BaseAWSLLM()
credentials = Credentials("test_key", "test_secret", "test_token")
headers = {"Content-Type": "application/json"}
def test_get_request_headers_with_sigv4() -> None:
credentials: Final = Credentials("test_key", "test_secret", "test_token")
result: Final = BaseAWSLLM().get_request_headers(
credentials=credentials,
aws_region_name="us-west-2",
extra_headers=None,
endpoint_url="https://api.example.com",
data='{"prompt": "test"}',
headers={"Content-Type": "application/json"},
supports_bearer_token=False,
)
# Create mock request and SigV4 instance
mock_request = MagicMock(spec=AWSRequest)
mock_request.headers = headers.copy()
mock_request.prepare.return_value = MagicMock(spec=AWSPreparedRequest)
mock_sigv4 = MagicMock()
# Test without bearer token (should use SigV4)
with (
patch.dict(os.environ, {}, clear=True),
patch("botocore.auth.SigV4Auth", return_value=mock_sigv4) as mock_sigv4_class,
patch("botocore.awsrequest.AWSRequest", return_value=mock_request),
):
result = llm.get_request_headers(
credentials=credentials,
aws_region_name="us-west-2",
extra_headers=None,
endpoint_url="https://api.example.com",
data='{"prompt": "test"}',
headers=headers,
)
# Verify SigV4 authentication and result
mock_sigv4_class.assert_called_once_with(credentials, "bedrock", "us-west-2")
mock_sigv4.add_auth.assert_called_once_with(mock_request)
assert result == mock_request.prepare.return_value
assert result.method == "POST"
assert result.url == "https://api.example.com"
assert result.body == '{"prompt": "test"}'
assert result.headers["X-Amz-Security-Token"] == "test_token"
authorization: Final = result.headers["Authorization"]
assert authorization.split("Credential=")[1].split("/")[0] == "test_key"
assert "/us-west-2/bedrock/aws4_request" in authorization
assert authorization.split("Signature=")[1] == _recomputed_sigv4_signature(
url=result.url,
secret_key="test_secret",
authorization=authorization,
headers=dict(result.headers),
body=result.body,
)
def test_get_request_headers_without_credentials_or_bearer_token_raises_no_credentials():
@ -4221,6 +4215,150 @@ def test_shared_json_signer_preserves_body_and_signs_for_the_requested_service()
assert "/us-west-2/s3/aws4_request" in request.headers["Authorization"]
@pytest.fixture
def rotating_sigv4_credentials(monkeypatch: pytest.MonkeyPatch) -> Credentials:
from itertools import count
from types import SimpleNamespace
import boto3
from botocore.credentials import DeferredRefreshableCredentials
instant: Final = datetime(2030, 1, 1, tzinfo=timezone.utc)
generations: Final = count(1)
def refresh() -> dict[str, str]:
generation: Final = next(generations)
return {
"access_key": f"rotation-access-{generation}",
"secret_key": f"rotation-secret-{generation}",
"token": f"rotation-token-{generation}",
"expiry_time": (instant + timedelta(seconds=60)).isoformat(),
}
credentials: Final = DeferredRefreshableCredentials(
refresh_using=refresh,
method="assume-role-with-web-identity",
time_fetcher=lambda: instant,
)
credentials._advisory_refresh_timeout = 120
credentials._mandatory_refresh_timeout = 90
session: Final = SimpleNamespace(get_credentials=lambda: credentials)
monkeypatch.setattr(boto3, "Session", lambda: session)
monkeypatch.setattr("botocore.auth.get_current_datetime", lambda: instant)
return credentials
def _request_with_rotating_credentials(
entrypoint: Literal["invoke", "prepared", "mantle", "json_post"],
credentials: Credentials | ReadOnlyCredentials,
url: str,
) -> tuple[dict[str, str], str | bytes]:
from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post
from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig
body: Final = {"input": "credential rotation"}
if entrypoint == "invoke":
invoke_headers, invoke_body = BaseAWSLLM()._sign_request(
service_name="bedrock",
headers={"Content-Type": "application/json"},
optional_params={"aws_region_name": "us-east-1"},
request_data=body,
api_base=url,
)
assert invoke_body is not None
return invoke_headers, invoke_body
if entrypoint == "mantle":
mantle_headers, mantle_body = BedrockMantleAnthropicMessagesConfig().sign_request(
headers={"Content-Type": "application/json"},
optional_params={"aws_region_name": "us-east-1"},
request_data=body,
api_base=url,
)
assert mantle_body is not None
return mantle_headers, mantle_body
if entrypoint == "json_post":
prepared: Final = sign_aws_json_post(
get_credentials=lambda: credentials,
service_name="bedrock",
aws_region_name="us-east-1",
url=url,
body=json.dumps(body),
headers={"Content-Type": "application/json"},
)
assert isinstance(prepared.body, (str, bytes))
return dict(prepared.headers), prepared.body
prepared_request: Final = BaseAWSLLM().get_request_headers(
credentials=credentials,
aws_region_name="us-east-1",
extra_headers=None,
endpoint_url=url,
data=json.dumps(body),
headers={"Content-Type": "application/json"},
supports_bearer_token=False,
)
assert isinstance(prepared_request.body, (str, bytes))
return dict(prepared_request.headers), prepared_request.body
def _verified_credential_generation(url: str, headers: dict[str, str], body: str | bytes) -> str:
authorization: Final = headers["Authorization"]
access_key: Final = authorization.split("Credential=")[1].split("/")[0]
generation: Final = access_key.removeprefix("rotation-access-")
assert headers["X-Amz-Security-Token"] == f"rotation-token-{generation}", headers
assert authorization.split("Signature=")[1] == _recomputed_sigv4_signature(
url=url,
secret_key=f"rotation-secret-{generation}",
authorization=authorization,
headers=headers,
body=body,
), headers
return generation
@pytest.mark.parametrize(
("entrypoint", "url"),
(
("invoke", "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke"),
("prepared", "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/converse"),
("mantle", "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages"),
("json_post", "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages"),
),
)
def test_sigv4_rotation_uses_one_generation_per_request_and_refreshes_between_requests(
entrypoint: Literal["invoke", "prepared", "mantle", "json_post"],
url: str,
rotating_sigv4_credentials: Credentials,
) -> None:
first_headers, first_body = _request_with_rotating_credentials(entrypoint, rotating_sigv4_credentials, url)
first_generation: Final = _verified_credential_generation(url, first_headers, first_body)
second_headers, second_body = _request_with_rotating_credentials(entrypoint, rotating_sigv4_credentials, url)
second_generation: Final = _verified_credential_generation(url, second_headers, second_body)
assert second_generation != first_generation
@pytest.mark.parametrize("entrypoint", ("prepared", "json_post"))
@pytest.mark.parametrize("frozen", (False, True))
@pytest.mark.parametrize("token", (None, "rotation-token-static"))
def test_sigv4_signs_with_static_or_frozen_credentials(
entrypoint: Literal["prepared", "json_post"], frozen: bool, token: str | None
) -> None:
credentials: Final = Credentials("rotation-access-static", "rotation-secret-static", token)
signing_credentials: Final = credentials.get_frozen_credentials() if frozen else credentials
url: Final = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke"
headers, body = _request_with_rotating_credentials(entrypoint, signing_credentials, url)
assert headers.get("X-Amz-Security-Token") == token
authorization: Final = headers["Authorization"]
assert "Credential=rotation-access-static/" in authorization
assert authorization.split("Signature=")[1] == _recomputed_sigv4_signature(
url=url,
secret_key="rotation-secret-static",
authorization=authorization,
headers=headers,
body=body,
)
_CONVERSE_SIGNED_RESPONSE: Final = {
"output": {"message": {"role": "assistant", "content": [{"text": "signed"}]}},
"stopReason": "end_turn",