Skip to content

Commit 4e8afcd

Browse files
pyek-botpyek-bot
authored
[Backport 2.x] Manual backport for adding AWS Rekognition trust endpoint (#3445)
* Added amazon rekognition as a trust endpoint (#3419) * feat: add rekognition trust endpoint Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> * refactor: consistent regex with previous connectors and looser checks on region Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> * test: add test case to validate connector creation successful Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> * chore: spotless apply Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> * chore: linter fixes Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> --------- Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> * fix: handle conflict due to deepseek changes Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com> --------- Signed-off-by: Pavan Yekbote <mail2pavanyekbote@gmail.com>
1 parent fcd8874 commit 4e8afcd

File tree

2 files changed

+67
-2
lines changed

2 files changed

+67
-2
lines changed

‎plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java

+4-1
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,8 @@ private MLCommonsSettings() {}
156156
public static final Setting<Boolean> ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED = Setting
157157
.boolSetting("plugins.ml_commons.offline_batch_inference_enabled", true, Setting.Property.NodeScope, Setting.Property.Dynamic);
158158

159+
public static final String REKOGNITION_TRUST_ENDPOINT_REGEX = "^https://rekognition(-fips)?\\..*[a-z0-9-]\\.amazonaws\\.com$";
160+
159161
public static final Setting<List<String>> ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX = Setting
160162
.listSetting(
161163
"plugins.ml_commons.trusted_connector_endpoints_regex",
@@ -170,7 +172,8 @@ private MLCommonsSettings() {}
170172
"^https://bedrock-agent-runtime\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
171173
"^https://bedrock\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
172174
"^https://textract\\..*[a-z0-9-]\\.amazonaws\\.com$",
173-
"^https://comprehend\\..*[a-z0-9-]\\.amazonaws\\.com$"
175+
"^https://comprehend\\..*[a-z0-9-]\\.amazonaws\\.com$",
176+
REKOGNITION_TRUST_ENDPOINT_REGEX
174177
),
175178
Function.identity(),
176179
Setting.Property.NodeScope,

‎plugin/src/test/java/org/opensearch/ml/action/connector/TransportCreateConnectorActionTests.java

+63-1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
1414
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_ACCESS_CONTROL_ENABLED;
1515
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX;
16+
import static org.opensearch.ml.settings.MLCommonsSettings.REKOGNITION_TRUST_ENDPOINT_REGEX;
1617
import static org.opensearch.ml.task.MLPredictTaskRunnerTests.USER_STRING;
1718
import static org.opensearch.ml.utils.TestHelper.clusterSetting;
1819

@@ -122,6 +123,7 @@ public class TransportCreateConnectorActionTests extends OpenSearchTestCase {
122123
"^https://runtime\\.sagemaker\\..*\\.amazonaws\\.com/.*$",
123124
"^https://api\\.openai\\.com/.*$",
124125
"^https://api\\.cohere\\.ai/.*$",
126+
REKOGNITION_TRUST_ENDPOINT_REGEX,
125127
"^https://api\\.deepseek\\.com/.*$"
126128
);
127129

@@ -559,6 +561,56 @@ public void test_connector_creation_success_deepseek() {
559561
mlModelManager,
560562
mlFeatureEnabledSetting
561563
);
564+
doAnswer(invocation -> {
565+
ActionListener<Boolean> listener = invocation.getArgument(0);
566+
listener.onResponse(true);
567+
return null;
568+
}).when(mlIndicesHandler).initMLConnectorIndex(isA(ActionListener.class));
569+
doAnswer(invocation -> {
570+
ActionListener<IndexResponse> listener = invocation.getArgument(1);
571+
listener.onResponse(indexResponse);
572+
return null;
573+
}).when(client).index(any(IndexRequest.class), isA(ActionListener.class));
574+
List<ConnectorAction> actions = new ArrayList<>();
575+
actions
576+
.add(
577+
ConnectorAction
578+
.builder()
579+
.actionType(ConnectorAction.ActionType.PREDICT)
580+
.method("POST")
581+
.url("https://api.deepseek.com/v1/chat/completions")
582+
.build()
583+
);
584+
Map<String, String> credential = ImmutableMap.of("access_key", "mockKey", "secret_key", "mockSecret");
585+
MLCreateConnectorInput mlCreateConnectorInput = MLCreateConnectorInput
586+
.builder()
587+
.name(randomAlphaOfLength(5))
588+
.description(randomAlphaOfLength(10))
589+
.version("1")
590+
.protocol(ConnectorProtocols.HTTP)
591+
.credential(credential)
592+
.actions(actions)
593+
.build();
594+
MLCreateConnectorRequest request = new MLCreateConnectorRequest(mlCreateConnectorInput);
595+
action.doExecute(task, request, actionListener);
596+
verify(actionListener).onResponse(any(MLCreateConnectorResponse.class));
597+
}
598+
599+
public void test_connector_creation_success_rekognition() {
600+
TransportCreateConnectorAction action = new TransportCreateConnectorAction(
601+
transportService,
602+
actionFilters,
603+
mlIndicesHandler,
604+
client,
605+
sdkClient,
606+
mlEngine,
607+
connectorAccessControlHelper,
608+
settings,
609+
clusterService,
610+
mlModelManager,
611+
mlFeatureEnabledSetting
612+
);
613+
562614
doAnswer(invocation -> {
563615
ActionListener<Boolean> listener = invocation.getArgument(0);
564616
listener.onResponse(true);
@@ -578,7 +630,16 @@ public void test_connector_creation_success_deepseek() {
578630
.builder()
579631
.actionType(ConnectorAction.ActionType.PREDICT)
580632
.method("POST")
581-
.url("https://api.deepseek.com/v1/chat/completions")
633+
.url("https://rekognition.test-region-1.amazonaws.com")
634+
.build()
635+
);
636+
actions
637+
.add(
638+
ConnectorAction
639+
.builder()
640+
.actionType(ConnectorAction.ActionType.PREDICT)
641+
.method("POST")
642+
.url("https://rekognition-fips.test-region-1.amazonaws.com")
582643
.build()
583644
);
584645

@@ -594,6 +655,7 @@ public void test_connector_creation_success_deepseek() {
594655
.build();
595656

596657
MLCreateConnectorRequest request = new MLCreateConnectorRequest(mlCreateConnectorInput);
658+
597659
action.doExecute(task, request, actionListener);
598660
verify(actionListener).onResponse(any(MLCreateConnectorResponse.class));
599661
}

0 commit comments

Comments
 (0)