feature: allow choosing js payload by alias in private method

Author: keerthanvasist

src/sagemaker/jumpstart/payload_utils.py

@@ -62,6 +62,7 @@ def _construct_payload(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    alias: Optional[str] = None,
 ) -> Optional[JumpStartSerializablePayload]:
     """Returns example payload from prompt.
 
@@ -102,7 +103,9 @@ def _construct_payload(
     if payloads is None or len(payloads) == 0:
         return None
 
-    payload_to_use: JumpStartSerializablePayload = list(payloads.values())[0]
+    payload_to_use: JumpStartSerializablePayload = (
+        payloads[alias] if alias else list(payloads.values())[0]
+    )
 
     prompt_key: Optional[str] = payload_to_use.prompt_key
     if prompt_key is None:

tests/unit/sagemaker/jumpstart/test_payload_utils.py

@@ -32,10 +32,36 @@ def test_construct_payload(self, patched_get_model_specs):
         region = "us-west-2"
 
         constructed_payload_body = _construct_payload(
-            prompt="kobebryant",
-            model_id=model_id,
-            model_version="*",
-            region=region,
+            prompt="kobebryant", model_id=model_id, model_version="*", region=region
+        ).body
+
+        self.assertEqual(
+            {
+                "hello": {"prompt": "kobebryant"},
+                "seed": 43,
+            },
+            constructed_payload_body,
+        )
+
+        # Unsupported model
+        self.assertIsNone(
+            _construct_payload(
+                prompt="blah",
+                model_id="default_payloads",
+                model_version="*",
+                region=region,
+            )
+        )
+
+    @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    def test_construct_payload_with_specific_alias(self, patched_get_model_specs):
+        patched_get_model_specs.side_effect = get_special_model_spec
+
+        model_id = "prompt-key"
+        region = "us-west-2"
+
+        constructed_payload_body = _construct_payload(
+            prompt="kobebryant", model_id=model_id, model_version="*", region=region, alias="Dog"
         ).body
 
         self.assertEqual(