From 17251cdf4ba4f9d1b5861f63bb4db93f7d693a48 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Tue, 28 Jan 2025 14:00:57 -0800
Subject: [PATCH] Add deepseek as a trusted endpoint. (#3440) (#3446)

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
(cherry picked from commit 06a6021ffe5bc68d27db576895561831b7dc7aaa)

Co-authored-by: Nathalie Jonathan <143617992+nathaliellenaa@users.noreply.github.com>
---
 .../ml/settings/MLCommonsSettings.java        |  1 +
 .../TransportCreateConnectorActionTests.java  | 60 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java b/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java
index c7a464a8ef..072878c8b2 100644
--- a/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java
+++ b/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java
@@ -165,6 +165,7 @@ private MLCommonsSettings() {}
                     "^https://api\\.sagemaker\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
                     "^https://api\\.openai\\.com/.*$",
                     "^https://api\\.cohere\\.ai/.*$",
+                    "^https://api\\.deepseek\\.com/.*$",
                     "^https://bedrock-runtime\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
                     "^https://bedrock-agent-runtime\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
                     "^https://bedrock\\..*[a-z0-9-]\\.amazonaws\\.com/.*$",
diff --git a/plugin/src/test/java/org/opensearch/ml/action/connector/TransportCreateConnectorActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/connector/TransportCreateConnectorActionTests.java
index 030f91fece..d321b3578a 100644
--- a/plugin/src/test/java/org/opensearch/ml/action/connector/TransportCreateConnectorActionTests.java
+++ b/plugin/src/test/java/org/opensearch/ml/action/connector/TransportCreateConnectorActionTests.java
@@ -118,7 +118,12 @@ public class TransportCreateConnectorActionTests extends OpenSearchTestCase {
     private ArgumentCaptor<PutDataObjectRequest> putDataObjectRequestArgumentCaptor;
 
     private static final List<String> TRUSTED_CONNECTOR_ENDPOINTS_REGEXES = ImmutableList
-        .of("^https://runtime\\.sagemaker\\..*\\.amazonaws\\.com/.*$", "^https://api\\.openai\\.com/.*$", "^https://api\\.cohere\\.ai/.*$");
+        .of(
+            "^https://runtime\\.sagemaker\\..*\\.amazonaws\\.com/.*$",
+            "^https://api\\.openai\\.com/.*$",
+            "^https://api\\.cohere\\.ai/.*$",
+            "^https://api\\.deepseek\\.com/.*$"
+        );
 
     @Before
     public void setup() {
@@ -539,4 +544,57 @@ public void test_execute_URL_notMatchingExpression_exception() {
             argumentCaptor.getValue().getMessage()
         );
     }
+
+    public void test_connector_creation_success_deepseek() {
+        TransportCreateConnectorAction action = new TransportCreateConnectorAction(
+            transportService,
+            actionFilters,
+            mlIndicesHandler,
+            client,
+            sdkClient,
+            mlEngine,
+            connectorAccessControlHelper,
+            settings,
+            clusterService,
+            mlModelManager,
+            mlFeatureEnabledSetting
+        );
+        doAnswer(invocation -> {
+            ActionListener<Boolean> listener = invocation.getArgument(0);
+            listener.onResponse(true);
+            return null;
+        }).when(mlIndicesHandler).initMLConnectorIndex(isA(ActionListener.class));
+
+        doAnswer(invocation -> {
+            ActionListener<IndexResponse> listener = invocation.getArgument(1);
+            listener.onResponse(indexResponse);
+            return null;
+        }).when(client).index(any(IndexRequest.class), isA(ActionListener.class));
+
+        List<ConnectorAction> actions = new ArrayList<>();
+        actions
+            .add(
+                ConnectorAction
+                    .builder()
+                    .actionType(ConnectorAction.ActionType.PREDICT)
+                    .method("POST")
+                    .url("https://api.deepseek.com/v1/chat/completions")
+                    .build()
+            );
+
+        Map<String, String> credential = ImmutableMap.of("access_key", "mockKey", "secret_key", "mockSecret");
+        MLCreateConnectorInput mlCreateConnectorInput = MLCreateConnectorInput
+            .builder()
+            .name(randomAlphaOfLength(5))
+            .description(randomAlphaOfLength(10))
+            .version("1")
+            .protocol(ConnectorProtocols.HTTP)
+            .credential(credential)
+            .actions(actions)
+            .build();
+
+        MLCreateConnectorRequest request = new MLCreateConnectorRequest(mlCreateConnectorInput);
+        action.doExecute(task, request, actionListener);
+        verify(actionListener).onResponse(any(MLCreateConnectorResponse.class));
+    }
 }