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 71

Merge 2a12aa17eb into 52023601a8

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

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

@ -55,7 +55,9 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile(
r"(?:^|\.)sts(?:-fips)?\.([a-z0-9-]+)\.(?:amazonaws\.com(?:\.cn)?|vpce\.amazonaws\.com)"
)
SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"})
SIGV4_COMPUTED_HEADERS: Final = frozenset(
{"authorization", "x-amz-date", "x-amz-security-token", "date", "x-amz-content-sha256"}
)
_AWS_SESSION_TAGS_ADAPTER: Final[TypeAdapter[tuple[AwsSessionTag, ...]]] = TypeAdapter(tuple[AwsSessionTag, ...])
@ -1699,6 +1701,8 @@ class BaseAWSLLM(SignsRequestsWithAWS):
if header_value is None:
continue
header_lower = header_name.lower()
if header_lower in SIGV4_COMPUTED_HEADERS:
continue
if header_lower in aws_headers or header_lower.startswith("x-amz-") or header_lower.startswith("x-amzn-"):
aws_signature_headers[header_name] = header_value

4
tests/unit/llms/bedrock/chat/test_invoke_handler.py
View file

@ -1697,11 +1697,11 @@ def test_filter_headers_for_aws_signature():
expected_aws_headers = {
"Content-Type": "application/json",
"Host": "bedrock-runtime.us-east-1.amazonaws.com",
"x-amz-date": "20240101T120000Z",
"x-amz-security-token": "test-token",
}
assert filtered_headers == expected_aws_headers, f"Expected {expected_aws_headers}, got {filtered_headers}"
excluded_headers = [
"x-amz-date",
"x-amz-security-token",
"x-custom-header",
"x-litellm-user-id",
"x-forwarded-for",

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

@ -16,7 +16,7 @@ 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
@ -3077,7 +3077,7 @@ def test_filter_headers_skips_none_values():
filtered = llm._filter_headers_for_aws_signature(headers)
assert filtered["Content-Type"] == "application/json"
assert filtered["x-amz-date"] == "20240101T000000Z"
assert "x-amz-date" not in filtered
assert "x-amz-security-token" not in filtered
assert "x-amzn-bedrock-kb-session-id" not in filtered
assert "host" not in filtered
@ -4221,6 +4221,97 @@ def test_shared_json_signer_preserves_body_and_signs_for_the_requested_service()
assert "/us-west-2/s3/aws4_request" in request.headers["Authorization"]
def _request_with_payload_hash(
entrypoint: Literal["invoke", "prepared", "mantle"],
headers: dict[str, str],
request_data: dict[str, str],
url: str,
) -> tuple[dict[str, str], str | bytes]:
from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig
optional_params: Final = {
"aws_access_key_id": "payload-hash-access",
"aws_secret_access_key": "payload-hash-secret",
"aws_region_name": "us-east-1",
}
if entrypoint == "invoke":
invoke_headers, invoke_body = BaseAWSLLM()._sign_request(
service_name="bedrock",
headers=headers,
optional_params=optional_params,
request_data=request_data,
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=headers,
optional_params=optional_params,
request_data=request_data,
api_base=url,
)
assert mantle_body is not None
return mantle_headers, mantle_body
prepared: Final = BaseAWSLLM().get_request_headers(
credentials=Credentials("payload-hash-access", "payload-hash-secret"),
aws_region_name="us-east-1",
extra_headers=None,
endpoint_url=url,
data=json.dumps(request_data),
headers=headers,
supports_bearer_token=False,
)
assert isinstance(prepared.body, (str, bytes))
return dict(prepared.headers), prepared.body
@pytest.mark.parametrize("hash_header", (None, "x-amz-content-sha256", "X-Amz-Content-SHA256"))
@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"),
),
)
def test_sigv4_resign_uses_current_body_hash_and_preserves_forwarded_headers(
entrypoint: Literal["invoke", "prepared", "mantle"], url: str, hash_header: str | None
) -> None:
import hashlib
original: Final = {"input": "before retry"}
original_hash: Final = hashlib.sha256(json.dumps(original).encode()).hexdigest()
input_headers: Final = {
"Content-Type": "application/json",
"anthropic-version": "2023-06-01",
**({hash_header: original_hash} if hash_header is not None else {}),
}
first_headers, first_body = _request_with_payload_hash(entrypoint, input_headers, original, url)
assert first_headers["Authorization"].split("Signature=")[1] == _recomputed_sigv4_signature(
url=url,
secret_key="payload-hash-secret",
authorization=first_headers["Authorization"],
headers=first_headers,
body=first_body,
)
previous_headers: Final = dict(first_headers)
retry_headers, retry_body = _request_with_payload_hash(entrypoint, first_headers, {"input": "after retry: é"}, url)
retry_bytes: Final = retry_body if isinstance(retry_body, bytes) else retry_body.encode()
expected_hash: Final = hashlib.sha256(retry_bytes).hexdigest()
normalized: Final = {name.lower(): value for name, value in retry_headers.items()}
assert normalized.get("x-amz-content-sha256", expected_hash) == expected_hash
assert retry_headers["Authorization"].split("Signature=")[1] == _recomputed_sigv4_signature(
url=url,
secret_key="payload-hash-secret",
authorization=retry_headers["Authorization"],
headers=retry_headers,
body=retry_body,
)
assert retry_headers["anthropic-version"] == "2023-06-01"
assert first_headers == previous_headers
_CONVERSE_SIGNED_RESPONSE: Final = {
"output": {"message": {"role": "assistant", "content": [{"text": "signed"}]}},
"stopReason": "end_turn",

122
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -2981,26 +2981,25 @@ async def test_alist_input_items_surfaces_upstream_error_status():
@pytest.mark.asyncio
async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_request(monkeypatch):
"""Regression: after Bedrock rejects a replayed thinking block (400 invalid signature),
the strip-and-retry re-sign must not inherit attempt 1's SigV4 Authorization/X-Amz-Date;
reusing them over the new stripped body makes AWS return 403 SignatureDoesNotMatch."""
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeMessagesConfig,
)
@pytest.mark.parametrize("hash_header", (None, "x-amz-content-sha256", "X-Amz-Content-SHA256"))
@pytest.mark.parametrize("stream", (False, True))
async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_request(
monkeypatch: pytest.MonkeyPatch, hash_header: str | None, stream: bool
) -> None:
import hashlib
for env_var in ("AWS_BEARER_TOKEN_BEDROCK", "AWS_SESSION_TOKEN", "AWS_PROFILE"):
monkeypatch.delenv(env_var, raising=False)
handler = BaseLLMHTTPHandler()
provider_config = AmazonAnthropicClaudeMessagesConfig()
litellm_params = GenericLiteLLMParams(
aws_access_key_id="AKIAIOSFODNN7EXAMPLE",
aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
handler: Final = BaseLLMHTTPHandler()
provider_config: Final = AmazonAnthropicClaudeMessagesConfig()
litellm_params: Final = GenericLiteLLMParams(
aws_access_key_id="payload-hash-access",
aws_secret_access_key="payload-hash-secret",
aws_region_name="us-east-1",
)
request_url = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke"
request_body = {
request_url: Final = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test-model/invoke"
request_body: Final = {
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 100,
"messages": [
@ -3015,60 +3014,71 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques
{"role": "user", "content": "continue"},
],
}
first_attempt_headers, signed_json_body = provider_config.sign_request(
headers={"Content-Type": "application/json"},
original_hash: Final = hashlib.sha256(json.dumps(request_body).encode()).hexdigest()
first_headers, signed_body = provider_config.sign_request(
headers={
"Content-Type": "application/json",
**({hash_header: original_hash} if hash_header is not None else {}),
},
optional_params=dict(litellm_params),
request_data=request_body,
api_base=request_url,
api_key=None,
stream=False,
stream=stream,
fake_stream=False,
model="test-model",
)
posts: Final[list[httpx.Request]] = [] # mutable-ok: the injected transport records request order
posts: list = []
invalid_signature_response = httpx.Response(
400,
text='{"message": "messages.1.content.0: Invalid `signature` in `thinking` block"}',
request=httpx.Request("POST", request_url),
)
ok_response = httpx.Response(200, json={"id": "msg_1"}, request=httpx.Request("POST", request_url))
def bedrock_peer(request: httpx.Request) -> httpx.Response:
posts.append(request)
body_hash: Final = hashlib.sha256(request.content).hexdigest()
if request.headers.get("x-amz-content-sha256", body_hash) != body_hash:
return httpx.Response(
403,
headers={"x-amzn-errortype": "SignatureDoesNotMatch"},
json={"message": "SignatureDoesNotMatch: payload hash does not match request bytes"},
)
if len(posts) == 1:
return httpx.Response(
400,
json={"message": "messages.1.content.0: Invalid `signature` in `thinking` block"},
)
return httpx.Response(200, json={"id": "msg_1"})
class FakeAsyncClient:
async def post(self, url, headers, data, stream=False, logging_obj=None, timeout=None):
posts.append({"headers": dict(headers), "data": data})
return invalid_signature_response if len(posts) == 1 else ok_response
logging_obj: Final = Mock(baseline_cache_context=None)
logging_obj.model_call_details = {}
response = await handler._async_post_anthropic_messages_with_http_error_retry(
async_httpx_client=FakeAsyncClient(),
request_url=request_url,
headers=dict(first_attempt_headers),
signed_json_body=signed_json_body,
request_body=request_body,
stream=False,
logging_obj=logging_obj,
provider_config=provider_config,
litellm_params=litellm_params,
api_key=None,
model="test-model",
)
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(bedrock_peer))
logging_obj: Final = Mock(baseline_cache_context=None, model_call_details={})
try:
response: Final = await handler._async_post_anthropic_messages_with_http_error_retry(
async_httpx_client=client,
request_url=request_url,
headers=first_headers,
signed_json_body=signed_body,
request_body=request_body,
stream=stream,
logging_obj=logging_obj,
provider_config=provider_config,
litellm_params=litellm_params,
api_key=None,
model="test-model",
)
await response.aclose()
finally:
await client.close()
assert response.status_code == 200
assert len(posts) == 2
retry_payload = json.loads(posts[1]["data"])
retry_blocks = [
block
for message in retry_payload["messages"]
if isinstance(message.get("content"), list)
for block in message["content"]
]
assert retry_blocks and all(block["type"] != "thinking" for block in retry_blocks)
retry_authorization = posts[1]["headers"]["Authorization"]
assert retry_authorization.startswith("AWS4-HMAC-SHA256")
assert retry_authorization != first_attempt_headers["Authorization"]
assert posts[0].content != posts[1].content
assert json.loads(posts[1].content) == {
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 100,
"messages": [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": [{"type": "text", "text": "ok"}]},
{"role": "user", "content": "continue"},
],
}
assert posts[1].headers["authorization"] != posts[0].headers["authorization"]
def test_aws_signing_overrides_only_fills_missing_credentials():