mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-08 22:25:25 +00:00
Add Add per model group header forwarding for Bedrock Invoke API (#16042)
This commit is contained in:
@@ -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():
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user