Add Add per model group header forwarding for Bedrock Invoke API (#16042)

This commit is contained in:
Sameer Kankute
2025-10-30 20:10:17 -07:00
committed by GitHub
parent bb79a75a89
commit c1369a07ba
2 changed files with 89 additions and 1 deletions
@@ -1774,11 +1774,16 @@ class BaseLLMHTTPHandler:
provider_specific_header=provider_specific_header,
custom_llm_provider=custom_llm_provider,
)
forwarded_headers = kwargs.get("headers", None)
if forwarded_headers and extra_headers:
merged_headers = {**forwarded_headers, **extra_headers}
else:
merged_headers = forwarded_headers or extra_headers
(
headers,
api_base,
) = anthropic_messages_provider_config.validate_anthropic_messages_environment(
headers=extra_headers or {},
headers=merged_headers or {},
model=model,
messages=messages,
optional_params=anthropic_messages_optional_request_params,
@@ -581,6 +581,89 @@ async def test_anthropic_messages_with_extra_headers():
return response
@pytest.mark.asyncio
async def test_bedrock_messages_api_header_forwarding():
"""
Test that headers from kwargs (set by proxy's add_headers_to_llm_call_by_model_group)
are correctly passed to validate_anthropic_messages_environment for Bedrock Invoke API.
This verifies that forward_client_headers_to_llm_api works for Bedrock Invoke API (Messages API).
Issue: When calling Anthropic models via the Messages API, LiteLLM makes a call to
Bedrock's Invoke API, and custom headers were not being forwarded, even though
they worked correctly for Chat Completions API with Bedrock's Converse API.
"""
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.router import GenericLiteLLMParams
handler = BaseLLMHTTPHandler()
# Headers that would be set by the proxy when forward_client_headers_to_llm_api is configured
custom_headers = {
"X-Custom-Header": "CustomValue",
"X-Request-ID": "req-123",
}
# Mock the provider config
mock_provider_config = MagicMock()
# We'll check what headers are passed to this method
mock_provider_config.validate_anthropic_messages_environment.return_value = (
{"Authorization": "Bearer test"},
"https://bedrock-runtime.us-east-1.amazonaws.com/invoke"
)
mock_provider_config.transform_anthropic_messages_request.return_value = {"model": "test"}
mock_provider_config.get_complete_url.return_value = "https://test.com"
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.transform_anthropic_messages_response.return_value = {"id": "test"}
# Mock HTTP client to prevent actual network calls
with unittest.mock.patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") as mock_get_client:
mock_http_client = AsyncMock()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"id": "test", "content": []}
mock_response.text = "{}"
mock_http_client.post.return_value = mock_response
mock_get_client.return_value = mock_http_client
# Mock logging object
mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
mock_logging_obj.model_call_details = {}
# Call the handler with headers in kwargs
try:
await handler.async_anthropic_messages_handler(
model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
messages=[{"role": "user", "content": "Hello"}],
anthropic_messages_provider_config=mock_provider_config,
anthropic_messages_optional_request_params={"max_tokens": 100},
custom_llm_provider="bedrock",
litellm_params=GenericLiteLLMParams(
api_key="test-key",
aws_region_name="us-east-1"
),
logging_obj=mock_logging_obj,
api_key="test-key",
stream=False,
kwargs={"headers": custom_headers} # Headers set by proxy
)
except Exception:
pass # Ignore errors, we're only checking if headers were passed
# Verify that validate_anthropic_messages_environment was called
assert mock_provider_config.validate_anthropic_messages_environment.called
# Get the headers that were passed
call_args = mock_provider_config.validate_anthropic_messages_environment.call_args
passed_headers = call_args[1]["headers"]
# The custom headers from kwargs should be in the passed headers
assert "X-Custom-Header" in passed_headers or "x-custom-header" in passed_headers
assert "X-Request-ID" in passed_headers or "x-request-id" in passed_headers
@pytest.mark.asyncio
async def test_anthropic_messages_with_thinking():
"""