add framework name to UserAgent header in AWS Bedrock API call (#13159)

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()

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