|
18 | 18 | from sagemaker.ai_registry.evaluator import Evaluator, EvaluatorMethod |
19 | 19 | from sagemaker.ai_registry.air_constants import ( |
20 | 20 | RESPONSE_KEY_HUB_CONTENT_VERSION, RESPONSE_KEY_HUB_CONTENT_ARN, |
21 | | - RESPONSE_KEY_CREATION_TIME, RESPONSE_KEY_LAST_MODIFIED_TIME, REWARD_FUNCTION |
| 21 | + RESPONSE_KEY_CREATION_TIME, RESPONSE_KEY_LAST_MODIFIED_TIME, |
| 22 | + REWARD_FUNCTION, REWARD_PROMPT |
22 | 23 | ) |
23 | 24 |
|
24 | 25 |
|
| 26 | +def _keywords_from_import_call(mock_air_hub): |
| 27 | + """Extract the SearchKeyword strings passed to import_hub_content.""" |
| 28 | + tags = mock_air_hub.import_hub_content.call_args.kwargs["tags"] |
| 29 | + return [f"{tag[0]}:{tag[1]}" for tag in tags] |
| 30 | + |
| 31 | + |
25 | 32 | class TestEvaluator: |
26 | 33 | @patch('sagemaker.ai_registry.evaluator.AIRHub') |
27 | 34 | def test_create_with_lambda_arn(self, mock_air_hub): |
@@ -81,6 +88,59 @@ def test_create_with_byoc(self, mock_air_hub, mock_boto3): |
81 | 88 | call_kwargs = mock_lambda_client.create_function.call_args[1] |
82 | 89 | assert call_kwargs["Handler"] == "lambda_function.lambda_handler" |
83 | 90 |
|
| 91 | + @patch('sagemaker.ai_registry.evaluator.AIRHub') |
| 92 | + def test_create_reward_function_always_sets_identity_keywords(self, mock_air_hub): |
| 93 | + """RewardFunction evaluators must always carry the full identity keyword set.""" |
| 94 | + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} |
| 95 | + mock_air_hub.describe_hub_content.return_value = { |
| 96 | + RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0", |
| 97 | + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", |
| 98 | + RESPONSE_KEY_CREATION_TIME: "2024-01-01", |
| 99 | + RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01" |
| 100 | + } |
| 101 | + |
| 102 | + Evaluator.create( |
| 103 | + name="test-evaluator", |
| 104 | + source="arn:aws:lambda:us-west-2:123456789012:function:test", |
| 105 | + type=REWARD_FUNCTION, |
| 106 | + wait=False |
| 107 | + ) |
| 108 | + |
| 109 | + keywords = _keywords_from_import_call(mock_air_hub) |
| 110 | + assert "@subtype:aws/evaluator" in keywords |
| 111 | + assert "@evaluatortype:rewardfunction" in keywords |
| 112 | + assert "@contenttype:byolambda" in keywords |
| 113 | + assert "method:lambda" in keywords |
| 114 | + |
| 115 | + @patch('sagemaker.ai_registry.evaluator.AIRHub') |
| 116 | + def test_create_reward_prompt_always_sets_identity_keywords(self, mock_air_hub): |
| 117 | + """RewardPrompt evaluators have no method but must still be Studio-visible. |
| 118 | +
|
| 119 | + Regression test for P467494019: identity keywords (@subtype/@evaluatortype/ |
| 120 | + @contenttype) were previously gated on ``method is not None``, so reward prompts |
| 121 | + registered with no keywords and were invisible in Studio. |
| 122 | + """ |
| 123 | + mock_air_hub.upload_to_s3.return_value = "s3://bucket/path/prompt.txt" |
| 124 | + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} |
| 125 | + mock_air_hub.describe_hub_content.return_value = { |
| 126 | + RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0", |
| 127 | + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", |
| 128 | + RESPONSE_KEY_CREATION_TIME: "2024-01-01", |
| 129 | + RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01" |
| 130 | + } |
| 131 | + |
| 132 | + Evaluator.create( |
| 133 | + name="test-prompt-evaluator", |
| 134 | + source="s3://bucket/path/prompt.txt", |
| 135 | + type=REWARD_PROMPT, |
| 136 | + wait=False |
| 137 | + ) |
| 138 | + |
| 139 | + keywords = _keywords_from_import_call(mock_air_hub) |
| 140 | + assert "@subtype:aws/evaluator" in keywords |
| 141 | + assert "@evaluatortype:rewardprompt" in keywords |
| 142 | + assert "@contenttype:byocode" in keywords |
| 143 | + |
84 | 144 | @patch('sagemaker.ai_registry.evaluator.AIRHub') |
85 | 145 | def test_get_all(self, mock_air_hub): |
86 | 146 | mock_air_hub.list_hub_content.return_value = { |
|
0 commit comments