Add support for SageMaker Inference Components in sagemaker chat

litellm/llms/openai/chat/gpt_transformation.py

@@ -383,11 +383,16 @@ def transform_request(
             dict: The transformed request. Sent as the body of the API call.
         """
         messages = self._transform_messages(messages=messages, model=model)
-        return {
+        model_id = optional_params.get("model_id", None)
+        request_data = {
             "model": model,
             "messages": messages,
             **optional_params,
         }
+        if model_id:
+            del request_data["model"]
+
+        return request_data
 
     async def async_transform_request(
         self,
@@ -401,11 +406,16 @@ async def async_transform_request(
             messages=messages, model=model, is_async=True
         )
 
-        return {
+        model_id = optional_params.get("model_id", None)
+        request_data = {
             "model": model,
             "messages": transformed_messages,
             **optional_params,
         }
+        if model_id:
+            del request_data["model"]
+
+        return request_data
 
     def _passed_in_tools(self, optional_params: dict) -> bool:
         return optional_params.get("tools", None) is not None

litellm/llms/sagemaker/chat/transformation.py

@@ -104,9 +104,6 @@ def sign_request(
         stream: Optional[bool] = None,
         fake_stream: Optional[bool] = None,
     ) -> Tuple[dict, Optional[bytes]]:
-        model_id = optional_params.get("model_id", None)
-        if model_id:
-            del request_data["model"]
         return self._sign_request(
             service_name="sagemaker",
             headers=headers,

tests/litellm/llms/sagemaker/test_sagemaker_common_utils.py

@@ -153,25 +153,13 @@ def test_inference_component_header(self):
 
     def test_inference_component_model_not_in_request(self):
         """Test that `model` is not part of the request body"""
-        test_params = {"model_id": "test"}
 
-        with patch(
-            "litellm.llms.sagemaker.chat.transformation.SagemakerChatConfig._sign_request"
-        ) as mock_sign_request:
-            self.config.sign_request(
-                headers={"X-Amzn-SageMaker-Inference-Component": "test"},
-                optional_params=test_params,
-                request_data={"model": self.model},
-                api_base="",
-            )
-
-            mock_sign_request.assert_called_once_with(
-                service_name="sagemaker",
-                headers={"X-Amzn-SageMaker-Inference-Component": "test"},
-                optional_params=test_params,
-                request_data={},
-                api_base="",
-                model=None,
-                stream=None,
-                fake_stream=None,
-            )
+        result = self.config.transform_request(
+            model=self.model,
+            messages=self.messages,
+            optional_params=self.optional_params,
+            litellm_params=None,
+            headers=None,
+        )
+
+        assert "model" not in result