diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index cc205e62dc9..f9581e494e8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -678,6 +678,12 @@ class BaseAWSLLM: headers: dict, api_key: Optional[str] = None, ) -> AWSPreparedRequest: + # update User-Agent field in headers to include framework name + if "User-Agent" in headers: + headers["User-Agent"] += " x-client-framework:litellm" + else: + headers["User-Agent"] = "x-client-framework:litellm" + if api_key is not None: aws_bearer_token: Optional[str] = api_key else: @@ -733,6 +739,13 @@ class BaseAWSLLM: Returns: Tuple[dict, Optional[str]]: A tuple containing the headers and the json str body of the request """ + # update User-Agent field in headers to include framework name + headers = headers or {} + if "User-Agent" in headers: + headers["User-Agent"] += " x-client-framework:litellm" + else: + headers["User-Agent"] = "x-client-framework:litellm" + if api_key is not None: aws_bearer_token: Optional[str] = api_key else: @@ -740,7 +753,6 @@ class BaseAWSLLM: # If aws bearer token is set, use it directly in the header if aws_bearer_token: - headers = headers or {} headers["Content-Type"] = "application/json" headers["Authorization"] = f"Bearer {aws_bearer_token}" return headers, json.dumps(request_data).encode() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index a2bbaa620b2..3ccace1abf5 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -205,6 +205,8 @@ def test_sign_request_with_env_var_bearer_token(): assert result_headers["Authorization"] == "Bearer test_token" assert result_headers["Content-Type"] == "application/json" assert result_headers["Custom-Header"] == "test" + assert "User-Agent" in result_headers + assert "x-client-framework:litellm" in result_headers["User-Agent"] assert result_body == json.dumps(request_data).encode() @@ -218,6 +220,7 @@ def test_sign_request_with_sigv4(): mock_request.headers = { "Authorization": "AWS4-HMAC-SHA256 Credential=test", "Content-Type": "application/json", + "User-Agent": "x-client-framework:litellm" } mock_request.body = b'{"prompt": "test"}' @@ -252,6 +255,8 @@ def test_sign_request_with_sigv4(): assert "Authorization" in result_headers assert result_headers["Authorization"] != "Bearer test_token" assert result_headers["Content-Type"] == "application/json" + assert "User-Agent" in result_headers + assert "x-client-framework:litellm" in result_headers["User-Agent"] assert result_body == mock_request.body @@ -283,6 +288,8 @@ def test_sign_request_with_api_key_bearer_token(): assert result_headers["Authorization"] == f"Bearer {api_key}" assert result_headers["Content-Type"] == "application/json" assert result_headers["Custom-Header"] == "test" + assert "User-Agent" in result_headers + assert "x-client-framework:litellm" in result_headers["User-Agent"] assert result_body == json.dumps(request_data).encode() @@ -318,6 +325,7 @@ def test_get_request_headers_with_env_var_bearer_token(): # Assert assert mock_request.headers["Authorization"] == "Bearer test_token" + assert mock_request.headers["User-Agent"] == "x-client-framework:litellm" assert result == mock_prepared_request @@ -333,13 +341,16 @@ def test_get_request_headers_with_sigv4(): mock_request.prepare.return_value = MagicMock(spec=AWSPreparedRequest) mock_sigv4 = MagicMock() + + def mock_aws_request_init(method, url, data, headers): + mock_request.headers.update(headers) + return mock_request # Test without bearer token (should use SigV4) - with patch.dict(os.environ, {}, clear=True), patch( - "botocore.auth.SigV4Auth", return_value=mock_sigv4 - ) as mock_sigv4_class, patch( - "botocore.awsrequest.AWSRequest", return_value=mock_request - ): + with patch.dict(os.environ, {}, clear=True), \ + patch('botocore.auth.SigV4Auth', return_value=mock_sigv4) as mock_sigv4_class, \ + patch('botocore.awsrequest.AWSRequest', side_effect=mock_aws_request_init): + result = llm.get_request_headers( credentials=credentials, aws_region_name="us-west-2", @@ -352,6 +363,8 @@ def test_get_request_headers_with_sigv4(): # Verify SigV4 authentication and result mock_sigv4_class.assert_called_once_with(credentials, "bedrock", "us-west-2") mock_sigv4.add_auth.assert_called_once_with(mock_request) + assert "User-Agent" in mock_request.headers + assert "x-client-framework:litellm" in mock_request.headers["User-Agent"] assert result == mock_request.prepare.return_value @@ -392,6 +405,7 @@ def test_get_request_headers_with_api_key_bearer_token(): # Assert assert mock_request.headers["Authorization"] == f"Bearer {api_key}" + assert mock_request.headers["User-Agent"] == "x-client-framework:litellm" assert result == mock_prepared_request