[Fix]: /messages - allow using dynamic AWS params (#10769)

* fix: dynamic AWS params added for messages routes

* Update tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff 2025-05-12 14:09:17 -07:00 committed by GitHub
parent 6de80a29cb
commit 51930c07c5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 68 additions and 1 deletions

View file

@ -1118,7 +1118,9 @@ class BaseLLMHTTPHandler:
headers, signed_json_body = anthropic_messages_provider_config.sign_request(
headers=headers,
optional_params=anthropic_messages_optional_request_params,
optional_params=dict(
litellm_params
), # dynamic aws_* params are passed under litellm_params
request_data=request_body,
api_base=request_url,
stream=stream,

View file

@ -23,6 +23,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.router import Router
import importlib
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
# Load environment variables
load_dotenv()
@ -612,3 +613,67 @@ async def test_anthropic_messages_with_thinking():
assert response == mock_response.json.return_value
return response
@pytest.mark.asyncio
async def test_anthropic_messages_bedrock_credentials_passthrough():
"""
Test that AWS credentials are correctly passed through to BaseAWSLLM.get_credentials
when using anthropic.messages.acreate with a bedrock model
"""
# Mock the get_credentials method
with unittest.mock.patch.object(BaseAWSLLM, 'get_credentials') as mock_get_credentials:
# Create a proper mock for credentials with the necessary attributes
mock_credentials = unittest.mock.MagicMock()
mock_credentials.access_key = "mock_access_key"
mock_credentials.secret_key = "mock_secret_key"
mock_credentials.token = "mock_session_token"
mock_get_credentials.return_value = mock_credentials
# We also need to mock the actual AWS request signing to avoid real API calls
with unittest.mock.patch('botocore.auth.SigV4Auth.add_auth'):
# Set up mock for AsyncHTTPHandler.post to avoid actual API calls
with unittest.mock.patch('litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post') as mock_post:
# Configure mock response
mock_response = unittest.mock.MagicMock()
mock_response.raise_for_status = unittest.mock.MagicMock()
mock_response.json.return_value = {
"id": "msg_bedrock_123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "This is a mock response"}],
"model": "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
mock_post.return_value = mock_response
# Test all AWS credentials parameters
test_params = {
"aws_access_key_id": "test_access_key",
"aws_secret_access_key": "test_secret_key",
"aws_session_token": "test_session_token",
"aws_region_name": "us-west-2",
"aws_role_name": "test_role_name",
"aws_session_name": "test_session_name",
"aws_profile_name": "test_profile",
"aws_web_identity_token": "test_web_identity_token",
"aws_sts_endpoint": "https://sts.test-region.amazonaws.com",
}
# Call the function with AWS credentials
await litellm.anthropic.messages.acreate(
messages=[{"role": "user", "content": "Hello, test credentials"}],
model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
max_tokens=100,
**test_params
)
# Verify get_credentials was called with the correct parameters
mock_get_credentials.assert_called_once()
call_args = mock_get_credentials.call_args[1]
# Assert that our test credentials were passed correctly
for param_name, param_value in test_params.items():
assert call_args[param_name] == param_value, f"Parameter {param_name} was not passed correctly"