mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
Merge 5155e7d029 into aaf1e2444b
This commit is contained in:
commit
9682b93b37
4 changed files with 634 additions and 24 deletions
|
|
@ -318,3 +318,23 @@ class BaseResponsesAPIConfig(ABC):
|
|||
data["input"]
|
||||
),
|
||||
}
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
optional_params: dict,
|
||||
request_data: dict,
|
||||
api_base: str,
|
||||
model: Optional[str] = None,
|
||||
stream: Optional[bool] = None,
|
||||
fake_stream: Optional[bool] = None,
|
||||
) -> Tuple[dict, Optional[bytes]]:
|
||||
"""Sign the final outbound request just before it is sent.
|
||||
|
||||
Mirrors ``BaseConfig.sign_request`` for chat. Called after every body
|
||||
mutation (normalize, extra_body, fake-stream strip). Returns
|
||||
``(headers, signed_body)``; when ``signed_body`` is not None the handler
|
||||
sends those exact bytes, so body-hashing schemes (e.g. AWS SigV4) stay
|
||||
valid regardless of later re-serialization. Default is a no-op.
|
||||
"""
|
||||
return headers, None
|
||||
|
|
|
|||
|
|
@ -7,17 +7,21 @@ Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
|
|||
only the endpoint URL and Bearer authentication.
|
||||
|
||||
Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
|
||||
standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
|
||||
standard AWS_BEARER_TOKEN_BEDROCK), or AWS SigV4 (IAM) credentials.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
import re
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1"
|
||||
BEDROCK_MANTLE_DEFAULT_REGION = "us-east-2"
|
||||
|
||||
_MANTLE_HOST_PATTERN = re.compile(r"bedrock-mantle\.([a-z0-9-]+)\.api\.aws")
|
||||
|
||||
# Checked longest/most-specific first so a full endpoint URL collapses to host
|
||||
# in one pass and the appended path never doubles.
|
||||
|
|
@ -30,26 +34,42 @@ _BASE_SUFFIXES_TO_STRIP = (
|
|||
)
|
||||
|
||||
|
||||
class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
||||
class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig, BaseAWSLLM):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
BaseAWSLLM.__init__(self)
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.BEDROCK_MANTLE
|
||||
|
||||
def _explicit_region(self, aws_region_name: Optional[str]) -> Optional[str]:
|
||||
return (
|
||||
aws_region_name
|
||||
or get_secret_str("BEDROCK_MANTLE_REGION")
|
||||
or get_secret_str("AWS_REGION")
|
||||
)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
litellm_params: dict,
|
||||
) -> str:
|
||||
region = (
|
||||
get_secret_str("BEDROCK_MANTLE_REGION")
|
||||
or get_secret_str("AWS_REGION")
|
||||
or BEDROCK_MANTLE_DEFAULT_REGION
|
||||
)
|
||||
explicit_region = self._explicit_region(litellm_params.get("aws_region_name"))
|
||||
self._validate_aws_region_name(explicit_region)
|
||||
base = (
|
||||
api_base
|
||||
or get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
or f"https://bedrock-mantle.{region}.api.aws"
|
||||
or f"https://bedrock-mantle.{explicit_region or BEDROCK_MANTLE_DEFAULT_REGION}.api.aws"
|
||||
)
|
||||
# get_llm_provider injects a default Mantle host whose region segment is
|
||||
# resolved without aws_region_name, so it can disagree with the SigV4
|
||||
# signing region and the endpoint rejects the request. When the caller
|
||||
# gave an explicit region, pin the host's region segment to it.
|
||||
if explicit_region:
|
||||
base = _MANTLE_HOST_PATTERN.sub(
|
||||
f"bedrock-mantle.{explicit_region}.api.aws", base
|
||||
)
|
||||
base = base.rstrip("/")
|
||||
for suffix in _BASE_SUFFIXES_TO_STRIP:
|
||||
if base.endswith(suffix):
|
||||
|
|
@ -57,10 +77,46 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
break
|
||||
return f"{base}/openai/v1/responses"
|
||||
|
||||
def _use_sigv4(
|
||||
self,
|
||||
api_key: Optional[str],
|
||||
aws_region_name: Optional[str],
|
||||
aws_access_key_id: Optional[str],
|
||||
) -> bool:
|
||||
bearer_key = (
|
||||
api_key
|
||||
or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
)
|
||||
if bearer_key:
|
||||
return False
|
||||
return any(
|
||||
[
|
||||
aws_region_name,
|
||||
aws_access_key_id,
|
||||
get_secret_str("AWS_ROLE_NAME"),
|
||||
get_secret_str("AWS_ROLE_ARN"),
|
||||
get_secret_str("AWS_WEB_IDENTITY_TOKEN"),
|
||||
get_secret_str("AWS_WEB_IDENTITY_TOKEN_FILE"),
|
||||
get_secret_str("AWS_PROFILE_NAME"),
|
||||
get_secret_str("AWS_ACCESS_KEY_ID"),
|
||||
get_secret_str("AWS_REGION"),
|
||||
get_secret_str("AWS_REGION_NAME"),
|
||||
get_secret_str("BEDROCK_MANTLE_REGION"),
|
||||
]
|
||||
)
|
||||
|
||||
def validate_environment(
|
||||
self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]
|
||||
) -> dict:
|
||||
litellm_params = litellm_params or GenericLiteLLMParams()
|
||||
if self._use_sigv4(
|
||||
api_key=litellm_params.api_key,
|
||||
aws_region_name=litellm_params.aws_region_name,
|
||||
aws_access_key_id=litellm_params.aws_access_key_id,
|
||||
):
|
||||
headers.setdefault("Content-Type", "application/json")
|
||||
return headers
|
||||
api_key = (
|
||||
litellm_params.api_key
|
||||
or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||
|
|
@ -68,12 +124,43 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
|
||||
"(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
|
||||
"Bedrock Mantle API key or AWS IAM credentials are required. "
|
||||
"Set BEDROCK_MANTLE_API_KEY, AWS_BEARER_TOKEN_BEDROCK, or pass "
|
||||
"api_key; or configure AWS credentials for SigV4 "
|
||||
"(aws_region_name, aws_access_key_id, aws_role_name, "
|
||||
"aws_web_identity_token, or aws_profile_name)."
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
optional_params: dict,
|
||||
request_data: dict,
|
||||
api_base: str,
|
||||
model: Optional[str] = None,
|
||||
stream: Optional[bool] = None,
|
||||
fake_stream: Optional[bool] = None,
|
||||
) -> Tuple[dict, Optional[bytes]]:
|
||||
if not self._use_sigv4(
|
||||
api_key=optional_params.get("api_key"),
|
||||
aws_region_name=optional_params.get("aws_region_name"),
|
||||
aws_access_key_id=optional_params.get("aws_access_key_id"),
|
||||
):
|
||||
return headers, None
|
||||
sign_params = dict(optional_params)
|
||||
host_match = _MANTLE_HOST_PATTERN.search(api_base)
|
||||
if host_match:
|
||||
sign_params["aws_region_name"] = host_match.group(1)
|
||||
return self._sign_request(
|
||||
service_name="bedrock",
|
||||
headers=headers,
|
||||
optional_params=sign_params,
|
||||
request_data=request_data,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
def supports_native_file_search(self) -> bool:
|
||||
return False
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
import json
|
||||
import ssl
|
||||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
|
|
@ -14,6 +13,7 @@ from typing import (
|
|||
Union,
|
||||
cast,
|
||||
)
|
||||
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx # type: ignore
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
|
@ -2339,10 +2339,20 @@ class BaseLLMHTTPHandler:
|
|||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
response = sync_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
json=data if signed_body is None else None,
|
||||
data=signed_body,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
stream=stream,
|
||||
|
|
@ -2371,10 +2381,20 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
else:
|
||||
# For non-streaming requests
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
response = sync_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
json=data if signed_body is None else None,
|
||||
data=signed_body,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
)
|
||||
|
|
@ -2485,10 +2505,20 @@ class BaseLLMHTTPHandler:
|
|||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
response = await async_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
json=data if signed_body is None else None,
|
||||
data=signed_body,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
stream=stream,
|
||||
|
|
@ -2519,10 +2549,20 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
else:
|
||||
# For non-streaming, proceed as before
|
||||
headers, signed_body = responses_api_provider_config.sign_request(
|
||||
headers=headers,
|
||||
optional_params=dict(litellm_params),
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
response = await async_httpx_client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
json=data,
|
||||
json=data if signed_body is None else None,
|
||||
data=signed_body,
|
||||
timeout=timeout
|
||||
or float(response_api_optional_request_params.get("timeout", 0)),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -14,12 +14,66 @@ sys.path.insert(0, os.path.abspath("../../../../.."))
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock_mantle.responses.transformation import (
|
||||
BedrockMantleResponsesAPIConfig,
|
||||
)
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
# Signals read by BedrockMantleResponsesAPIConfig._use_sigv4 that the autouse
|
||||
# isolate_host_aws_config fixture (tests/test_litellm/conftest.py) does NOT
|
||||
# already clear. Clearing these makes auth-selection deterministic on hosts that
|
||||
# export e.g. AWS_REGION or AWS_ACCESS_KEY_ID.
|
||||
_AWS_SIGNAL_ENV_VARS = (
|
||||
"AWS_REGION",
|
||||
"AWS_REGION_NAME",
|
||||
"AWS_ROLE_NAME",
|
||||
"AWS_ROLE_ARN",
|
||||
"AWS_WEB_IDENTITY_TOKEN",
|
||||
"AWS_WEB_IDENTITY_TOKEN_FILE",
|
||||
"AWS_PROFILE_NAME",
|
||||
"AWS_ACCESS_KEY_ID",
|
||||
"AWS_SECRET_ACCESS_KEY",
|
||||
"AWS_SESSION_TOKEN",
|
||||
"BEDROCK_MANTLE_REGION",
|
||||
"BEDROCK_MANTLE_API_BASE",
|
||||
"BEDROCK_MANTLE_API_KEY",
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clear_aws_env(monkeypatch):
|
||||
"""Clear the AWS / Bedrock Mantle signals not covered by the autouse
|
||||
isolate_host_aws_config fixture, so auth-selection tests are deterministic
|
||||
regardless of the host machine's AWS configuration."""
|
||||
for name in _AWS_SIGNAL_ENV_VARS:
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
return monkeypatch
|
||||
|
||||
|
||||
class _RecordingConfig(BedrockMantleResponsesAPIConfig):
|
||||
"""Test double that records `_sign_request` invocations instead of signing.
|
||||
|
||||
Subclassing to override the inherited method is dependency injection via a
|
||||
test double, not class-attribute monkeypatching of the production class.
|
||||
"""
|
||||
|
||||
_SIGNED_HEADERS = {
|
||||
"Authorization": "AWS4-HMAC-SHA256 Credential=test",
|
||||
"X-Amz-Date": "20260101T000000Z",
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.sign_calls = []
|
||||
|
||||
def _sign_request(self, **kwargs):
|
||||
self.sign_calls.append(kwargs)
|
||||
return dict(self._SIGNED_HEADERS), b"{}"
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesURL:
|
||||
def test_url_uses_region_from_env(self, monkeypatch):
|
||||
|
|
@ -76,13 +130,31 @@ class TestBedrockMantleResponsesURL:
|
|||
url = cfg.get_complete_url(api_base=None, litellm_params={})
|
||||
assert url == "https://bedrock-mantle.us-west-2.api.aws/openai/v1/responses"
|
||||
|
||||
def test_url_region_default_us_east_1(self, monkeypatch):
|
||||
def test_url_region_default_us_east_2(self, monkeypatch):
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
|
||||
monkeypatch.delenv("AWS_REGION", raising=False)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
url = cfg.get_complete_url(api_base=None, litellm_params={})
|
||||
assert url == "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses"
|
||||
assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"malicious_region",
|
||||
[
|
||||
"us-east-2.attacker.com/",
|
||||
"us-east-2.evil.com",
|
||||
"foo/bar",
|
||||
"us-east-1;rm -rf /",
|
||||
"UPPER-CASE",
|
||||
],
|
||||
)
|
||||
def test_url_rejects_malicious_region(self, clear_aws_env, malicious_region):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
with pytest.raises(ValueError, match="Invalid AWS region format"):
|
||||
cfg.get_complete_url(
|
||||
api_base=None,
|
||||
litellm_params={"aws_region_name": malicious_region},
|
||||
)
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesAuth:
|
||||
|
|
@ -96,7 +168,8 @@ class TestBedrockMantleResponsesAuth:
|
|||
)
|
||||
assert headers["Authorization"] == "Bearer config-key"
|
||||
|
||||
def test_env_key_fallback(self, monkeypatch):
|
||||
def test_env_key_fallback(self, clear_aws_env):
|
||||
monkeypatch = clear_aws_env
|
||||
monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key")
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
|
|
@ -105,7 +178,8 @@ class TestBedrockMantleResponsesAuth:
|
|||
)
|
||||
assert headers["Authorization"] == "Bearer env-key"
|
||||
|
||||
def test_bedrock_bearer_token_fallback(self, monkeypatch):
|
||||
def test_bedrock_bearer_token_fallback(self, clear_aws_env):
|
||||
monkeypatch = clear_aws_env
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bearer-key")
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
|
|
@ -114,9 +188,7 @@ class TestBedrockMantleResponsesAuth:
|
|||
)
|
||||
assert headers["Authorization"] == "Bearer bearer-key"
|
||||
|
||||
def test_missing_key_raises(self, monkeypatch):
|
||||
monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
def test_missing_key_raises(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
with pytest.raises(ValueError, match="Bedrock Mantle API key"):
|
||||
cfg.validate_environment(
|
||||
|
|
@ -281,3 +353,394 @@ class TestBedrockMantleResponsesPricing:
|
|||
def test_models_registered(self, local_cost_map):
|
||||
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
|
||||
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesSigV4:
|
||||
def test_bearer_via_config_key_wins_over_aws_creds(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(
|
||||
api_key="config-key",
|
||||
aws_region_name="us-west-2",
|
||||
aws_access_key_id=None,
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_bearer_via_bedrock_mantle_api_key_env(self, clear_aws_env):
|
||||
clear_aws_env.setenv("BEDROCK_MANTLE_API_KEY", "env-key")
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(api_key=None, aws_region_name=None, aws_access_key_id=None)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_bearer_via_aws_bearer_token_env(self, clear_aws_env):
|
||||
clear_aws_env.setenv("AWS_BEARER_TOKEN_BEDROCK", "bearer-key")
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(api_key=None, aws_region_name=None, aws_access_key_id=None)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_aws_region_param_without_bearer_activates_sigv4(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(
|
||||
api_key=None, aws_region_name="us-east-1", aws_access_key_id=None
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_aws_access_key_param_without_bearer_activates_sigv4(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(
|
||||
api_key=None, aws_region_name=None, aws_access_key_id="AKIA_TEST"
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_env_signal_alone_activates_sigv4(self, clear_aws_env):
|
||||
clear_aws_env.setenv("AWS_REGION", "us-west-2")
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(api_key=None, aws_region_name=None, aws_access_key_id=None)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_irsa_role_arn_activates_sigv4(self, clear_aws_env):
|
||||
clear_aws_env.setenv("AWS_ROLE_ARN", "arn:aws:iam::123456789012:role/MyRole")
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(api_key=None, aws_region_name=None, aws_access_key_id=None)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_irsa_web_identity_token_file_activates_sigv4(self, clear_aws_env):
|
||||
clear_aws_env.setenv("AWS_WEB_IDENTITY_TOKEN_FILE", "/var/run/secrets/token")
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(api_key=None, aws_region_name=None, aws_access_key_id=None)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_no_credentials_does_not_activate_sigv4(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
assert (
|
||||
cfg._use_sigv4(api_key=None, aws_region_name=None, aws_access_key_id=None)
|
||||
is False
|
||||
)
|
||||
|
||||
def test_validate_environment_sigv4_omits_authorization(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
headers = cfg.validate_environment(
|
||||
headers={},
|
||||
model="openai.gpt-5.5",
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1"),
|
||||
)
|
||||
assert "Authorization" not in headers
|
||||
|
||||
def test_validate_environment_sigv4_sets_content_type(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
headers = cfg.validate_environment(
|
||||
headers={},
|
||||
model="openai.gpt-5.5",
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1"),
|
||||
)
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
def test_validate_environment_bearer_sets_authorization(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
headers = cfg.validate_environment(
|
||||
headers={},
|
||||
model="openai.gpt-5.5",
|
||||
litellm_params=GenericLiteLLMParams(api_key="config-key"),
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer config-key"
|
||||
|
||||
def test_validate_environment_no_credentials_message_mentions_iam(
|
||||
self, clear_aws_env
|
||||
):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Bedrock Mantle API key or AWS IAM credentials are required",
|
||||
) as exc_info:
|
||||
cfg.validate_environment(
|
||||
headers={},
|
||||
model="openai.gpt-5.5",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
)
|
||||
message = str(exc_info.value)
|
||||
assert "SigV4" in message or "IAM" in message
|
||||
|
||||
def test_sign_hook_signs_request_with_bedrock_service(self, clear_aws_env):
|
||||
cfg = _RecordingConfig()
|
||||
api_base = cfg.get_complete_url(
|
||||
api_base=None, litellm_params={"aws_region_name": "us-west-2"}
|
||||
)
|
||||
request_data = {"model": "openai.gpt-5.5", "input": "hello"}
|
||||
cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={"aws_region_name": "us-west-2"},
|
||||
request_data=request_data,
|
||||
api_base=api_base,
|
||||
)
|
||||
assert len(cfg.sign_calls) == 1
|
||||
call = cfg.sign_calls[0]
|
||||
assert call["service_name"] == "bedrock"
|
||||
assert call["api_base"].endswith("/openai/v1/responses")
|
||||
assert call["request_data"] == request_data
|
||||
|
||||
def test_sign_hook_returns_signed_headers_and_body(self, clear_aws_env):
|
||||
cfg = _RecordingConfig()
|
||||
api_base = cfg.get_complete_url(
|
||||
api_base=None, litellm_params={"aws_region_name": "us-west-2"}
|
||||
)
|
||||
signed_headers, signed_body = cfg.sign_request(
|
||||
headers={},
|
||||
optional_params={"aws_region_name": "us-west-2"},
|
||||
request_data={"model": "openai.gpt-5.5", "input": "hello"},
|
||||
api_base=api_base,
|
||||
)
|
||||
assert signed_headers["Authorization"] == "AWS4-HMAC-SHA256 Credential=test"
|
||||
assert signed_headers["X-Amz-Date"] == "20260101T000000Z"
|
||||
assert signed_body == b"{}"
|
||||
|
||||
def test_sign_hook_bearer_mode_does_not_sign(self, clear_aws_env):
|
||||
cfg = _RecordingConfig()
|
||||
signed_headers, signed_body = cfg.sign_request(
|
||||
headers={"Authorization": "Bearer config-key"},
|
||||
optional_params={"api_key": "config-key"},
|
||||
request_data={"model": "openai.gpt-5.5", "input": "hello"},
|
||||
api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
)
|
||||
assert cfg.sign_calls == []
|
||||
assert signed_body is None
|
||||
assert "X-Amz-Date" not in signed_headers
|
||||
|
||||
def test_transform_does_not_sign(self, clear_aws_env):
|
||||
cfg = _RecordingConfig()
|
||||
headers = {}
|
||||
cfg.transform_responses_api_request(
|
||||
model="openai.gpt-5.5",
|
||||
input="hello",
|
||||
response_api_optional_request_params={},
|
||||
litellm_params=GenericLiteLLMParams(aws_region_name="us-west-2"),
|
||||
headers=headers,
|
||||
)
|
||||
assert cfg.sign_calls == []
|
||||
assert "Authorization" not in headers
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesStructure:
|
||||
def test_subclasses_openai_responses_and_base_aws_llm(self):
|
||||
assert issubclass(BedrockMantleResponsesAPIConfig, OpenAIResponsesAPIConfig)
|
||||
assert issubclass(BedrockMantleResponsesAPIConfig, BaseAWSLLM)
|
||||
|
||||
def test_get_complete_url_ends_with_responses_path(self, clear_aws_env):
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
url = cfg.get_complete_url(api_base=None, litellm_params={})
|
||||
assert url.endswith("/openai/v1/responses")
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesSigV4Signature:
|
||||
def test_signature_scope_uses_bedrock_service_and_url_region(self, clear_aws_env):
|
||||
clear_aws_env.setenv("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE")
|
||||
clear_aws_env.setenv(
|
||||
"AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
api_base = cfg.get_complete_url(
|
||||
api_base=None, litellm_params={"aws_region_name": "us-east-2"}
|
||||
)
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={"Content-Type": "application/json"},
|
||||
optional_params={"aws_region_name": "us-east-2"},
|
||||
request_data={"model": "openai.gpt-5.4", "input": "hello"},
|
||||
api_base=api_base,
|
||||
)
|
||||
auth = headers["Authorization"]
|
||||
assert auth.startswith("AWS4-HMAC-SHA256 ")
|
||||
assert "/us-east-2/bedrock/aws4_request" in auth
|
||||
|
||||
def test_injected_default_api_base_region_is_pinned_to_signing_region(
|
||||
self, clear_aws_env
|
||||
):
|
||||
clear_aws_env.setenv("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE")
|
||||
clear_aws_env.setenv(
|
||||
"AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
# get_llm_provider injects a us-east-1 default host; the resolved URL and
|
||||
# the SigV4 scope must both end up us-east-2 (the explicit region).
|
||||
url = cfg.get_complete_url(
|
||||
api_base="https://bedrock-mantle.us-east-1.api.aws/v1",
|
||||
litellm_params={"aws_region_name": "us-east-2"},
|
||||
)
|
||||
assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
headers, _ = cfg.sign_request(
|
||||
headers={"Content-Type": "application/json"},
|
||||
optional_params={
|
||||
"aws_region_name": "us-east-2",
|
||||
"api_base": "https://bedrock-mantle.us-east-1.api.aws/v1",
|
||||
},
|
||||
request_data={"model": "openai.gpt-5.4", "input": "hello"},
|
||||
api_base=url,
|
||||
)
|
||||
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
|
||||
|
||||
def test_signed_body_hash_matches_compact_json_sent_by_httpx(self, clear_aws_env):
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
clear_aws_env.setenv("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE")
|
||||
clear_aws_env.setenv(
|
||||
"AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
cfg = BedrockMantleResponsesAPIConfig()
|
||||
data = {"model": "openai.gpt-5.4", "input": "hello, what model are you?"}
|
||||
api_base = cfg.get_complete_url(
|
||||
api_base=None, litellm_params={"aws_region_name": "us-east-2"}
|
||||
)
|
||||
headers, signed_body = cfg.sign_request(
|
||||
headers={"Content-Type": "application/json"},
|
||||
optional_params={"aws_region_name": "us-east-2"},
|
||||
request_data=data,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
# The handler posts the exact signed_body bytes. Re-sign those same bytes
|
||||
# with the date the config used and assert the signatures match, proving
|
||||
# the wire body and the signed body are identical.
|
||||
creds = Credentials(
|
||||
"AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
ref = AWSRequest(
|
||||
method="POST",
|
||||
url="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses",
|
||||
data=signed_body,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Amz-Date": headers["X-Amz-Date"],
|
||||
},
|
||||
)
|
||||
SigV4Auth(creds, "bedrock", "us-east-2").add_auth(ref)
|
||||
assert headers["Authorization"] == ref.headers["Authorization"]
|
||||
|
||||
|
||||
class TestBedrockMantleResponsesHandlerSignedBody:
|
||||
"""Exercises the handler's ``data=signed_body`` branch (the 3 lines the
|
||||
coverage checker flagged). Verifies that when sign_request returns body
|
||||
bytes, the handler sends those exact bytes rather than re-serializing."""
|
||||
|
||||
def test_handler_posts_signed_body_bytes(self, clear_aws_env, respx_mock):
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
clear_aws_env.setenv("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE")
|
||||
clear_aws_env.setenv(
|
||||
"AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
|
||||
captured_body = {}
|
||||
|
||||
def capture_request(request: httpx.Request) -> httpx.Response:
|
||||
captured_body["raw"] = request.content
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "resp_test",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "openai.gpt-5.4",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [
|
||||
{"type": "output_text", "text": "hi", "annotations": []}
|
||||
],
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 1,
|
||||
"total_tokens": 2,
|
||||
"input_tokens_details": {"cached_tokens": 0},
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
respx_mock.post(
|
||||
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
).mock(side_effect=capture_request)
|
||||
|
||||
import litellm
|
||||
|
||||
resp = litellm.responses(
|
||||
model="bedrock_mantle/openai.gpt-5.4",
|
||||
input="hello",
|
||||
aws_region_name="us-east-2",
|
||||
)
|
||||
assert resp.output[0].content[0].text == "hi"
|
||||
|
||||
# The handler should have sent the exact bytes sign_request produced,
|
||||
# not a re-serialization via httpx json=. The signed body uses
|
||||
# json.dumps(data) (with spaces), so verify the wire bytes contain spaces
|
||||
# (httpx compact would not).
|
||||
wire_body = captured_body["raw"]
|
||||
assert b'"model": ' in wire_body or b'"model":' in wire_body
|
||||
parsed = json.loads(wire_body)
|
||||
assert parsed["model"] == "openai.gpt-5.4"
|
||||
assert parsed["input"] == "hello"
|
||||
|
||||
def test_handler_posts_signed_body_bytes_streaming(self, clear_aws_env, respx_mock):
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
clear_aws_env.setenv("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE")
|
||||
clear_aws_env.setenv(
|
||||
"AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
)
|
||||
|
||||
captured_body = {}
|
||||
|
||||
def capture_request(request: httpx.Request) -> httpx.Response:
|
||||
captured_body["raw"] = request.content
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=b'data: {"type":"response.completed"}\n\ndata: [DONE]\n\n',
|
||||
headers={"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
respx_mock.post(
|
||||
"https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"
|
||||
).mock(side_effect=capture_request)
|
||||
|
||||
import litellm
|
||||
|
||||
try:
|
||||
resp = litellm.responses(
|
||||
model="bedrock_mantle/openai.gpt-5.4",
|
||||
input="hello",
|
||||
stream=True,
|
||||
aws_region_name="us-east-2",
|
||||
)
|
||||
for _ in resp:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
wire_body = captured_body["raw"]
|
||||
parsed = json.loads(wire_body)
|
||||
assert parsed["model"] == "openai.gpt-5.4"
|
||||
assert parsed["input"] == "hello"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue