From c7997aa420075cab7cbcf2dee49417217fe1dba4 Mon Sep 17 00:00:00 2001
From: Amir Sadoughi <amir@2ndset.ai>
Date: Wed, 18 Feb 2026 21:58:28 -0500
Subject: [PATCH] Add model_id (Inference Component) support to sagemaker_chat

Extract model_id from optional_params in SagemakerChatConfig.sign_request()
and inject it as the X-Amzn-SageMaker-Inference-Component header, enabling
users to target specific inference components via the Messages API path.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/llms/sagemaker/chat/transformation.py |  4 +
 .../test_sagemaker_chat_transformation.py     | 82 +++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100644 tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py

diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py
index 2b458fbc438..e7196fb63b4 100644
--- a/litellm/llms/sagemaker/chat/transformation.py
+++ b/litellm/llms/sagemaker/chat/transformation.py
@@ -98,6 +98,10 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
         stream: Optional[bool] = None,
         fake_stream: Optional[bool] = None,
     ) -> Tuple[dict, Optional[bytes]]:
+        model_id = optional_params.get("model_id", None)
+        if model_id is not None:
+            headers = headers or {}
+            headers["X-Amzn-SageMaker-Inference-Component"] = model_id
         return self._sign_request(
             service_name="sagemaker",
             headers=headers,
diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py
new file mode 100644
index 00000000000..061fd475f2a
--- /dev/null
+++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py
@@ -0,0 +1,82 @@
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig
+
+
+class TestSagemakerChatSignRequest:
+    def setup_method(self):
+        self.config = SagemakerChatConfig()
+
+    @patch.object(SagemakerChatConfig, "_sign_request")
+    def test_sign_request_injects_model_id_header(self, mock_sign_request):
+        """
+        Test that model_id in optional_params is injected as the
+        X-Amzn-SageMaker-Inference-Component header before signing.
+        """
+        mock_sign_request.return_value = ({"Authorization": "signed"}, b'{}')
+
+        headers = {"Content-Type": "application/json"}
+        optional_params = {"model_id": "my-inference-component"}
+
+        self.config.sign_request(
+            headers=headers,
+            optional_params=optional_params,
+            request_data={"messages": []},
+            api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/my-endpoint/invocations",
+            model="my-endpoint",
+        )
+
+        # Verify _sign_request was called with the inference component header
+        call_kwargs = mock_sign_request.call_args
+        signed_headers = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers")
+        assert signed_headers["X-Amzn-SageMaker-Inference-Component"] == "my-inference-component"
+
+    @patch.object(SagemakerChatConfig, "_sign_request")
+    def test_sign_request_no_model_id_no_header(self, mock_sign_request):
+        """
+        Test that when model_id is not provided, the inference component
+        header is not added.
+        """
+        mock_sign_request.return_value = ({"Authorization": "signed"}, b'{}')
+
+        headers = {"Content-Type": "application/json"}
+        optional_params = {}
+
+        self.config.sign_request(
+            headers=headers,
+            optional_params=optional_params,
+            request_data={"messages": []},
+            api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/my-endpoint/invocations",
+            model="my-endpoint",
+        )
+
+        call_kwargs = mock_sign_request.call_args
+        signed_headers = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers")
+        assert "X-Amzn-SageMaker-Inference-Component" not in signed_headers
+
+    @patch.object(SagemakerChatConfig, "_sign_request")
+    def test_sign_request_model_id_with_none_headers(self, mock_sign_request):
+        """
+        Test that model_id injection works even when headers is initially None.
+        """
+        mock_sign_request.return_value = ({"Authorization": "signed"}, b'{}')
+
+        optional_params = {"model_id": "component-abc"}
+
+        self.config.sign_request(
+            headers=None,
+            optional_params=optional_params,
+            request_data={"messages": []},
+            api_base="https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/my-endpoint/invocations",
+            model="my-endpoint",
+        )
+
+        call_kwargs = mock_sign_request.call_args
+        signed_headers = call_kwargs.kwargs.get("headers") or call_kwargs[1].get("headers")
+        assert signed_headers is not None
+        assert signed_headers["X-Amzn-SageMaker-Inference-Component"] == "component-abc"