fix(sagemaker): include payload additional_args

Copy SageMakerAIPayloadSchema's additional_args into request payloads
where provided - previously these were being ignored. Includes unit
tests.

Author: athewsey

src/strands/models/sagemaker.py

@@ -238,6 +238,10 @@ def format_request(
             },
         }
 
+        payload_additional_args = self.payload_config.get("additional_args")
+        if payload_additional_args:
+            payload.update(payload_additional_args)
+
         # Remove tools and tool_choice if tools = []
         if not payload["tools"]:
             payload.pop("tools")

tests/strands/models/test_sagemaker.py

@@ -112,12 +112,13 @@ def test_init_with_all_params(self, boto_session):
             "endpoint_name": "test-endpoint",
             "inference_component_name": "test-component",
             "region_name": "us-west-2",
-            "additional_args": {"test_arg_name": "test_arg_value"},
+            "additional_args": {"test_req_arg_name": "test_req_arg_value"},
         }
         payload_config = {
             "stream": False,
             "max_tokens": 1024,
             "temperature": 0.7,
+            "additional_args": {"test_payload_arg_name": "test_payload_arg_value"},
         }
         client_config = BotocoreConfig(user_agent_extra="test-agent")
 
@@ -130,10 +131,11 @@ def test_init_with_all_params(self, boto_session):
 
         assert model.endpoint_config["endpoint_name"] == "test-endpoint"
         assert model.endpoint_config["inference_component_name"] == "test-component"
-        assert model.endpoint_config["additional_args"]["test_arg_name"] == "test_arg_value"
+        assert model.endpoint_config["additional_args"]["test_req_arg_name"] == "test_req_arg_value"
         assert model.payload_config["stream"] is False
         assert model.payload_config["max_tokens"] == 1024
         assert model.payload_config["temperature"] == 0.7
+        assert model.payload_config["additional_args"]["test_payload_arg_name"] == "test_payload_arg_value"
 
         boto_session.client.assert_called_once_with(
             service_name="sagemaker-runtime",
@@ -246,16 +248,24 @@ def test_format_request_with_additional_args(self, boto_session, endpoint_config
         endpoint_config_ext = {
             **endpoint_config,
             "additional_args": {
-                "extra_key": "extra_value",
+                "extra_request_key": "extra_request_value",
+            },
+        }
+        payload_config_ext = {
+            **payload_config,
+            "additional_args": {
+                "extra_payload_key": "extra_payload_value",
             },
         }
         model = SageMakerAIModel(
             boto_session=boto_session,
             endpoint_config=endpoint_config_ext,
-            payload_config=payload_config,
+            payload_config=payload_config_ext,
         )
         request = model.format_request(messages)
-        assert request.get("extra_key") == "extra_value"
+        assert request.get("extra_request_key") == "extra_request_value"
+        payload = json.loads(request["Body"])
+        assert payload.get("extra_payload_key") == "extra_payload_value"
 
     @pytest.mark.asyncio
     async def test_stream_with_streaming_enabled(self, sagemaker_client, model, messages):