mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
[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:
parent
6de80a29cb
commit
51930c07c5
2 changed files with 68 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue