Skip to content

Commit c824105

Browse files
Hailong-amHailong-am
authored
Merge branch 'main' into chatbot_invalidpaylod
2 parents 14bcb07 + 247a88c commit c824105

File tree

79 files changed

+9742
-15
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

79 files changed

+9742
-15
lines changed

‎client/src/main/java/org/opensearch/ml/client/MachineLearningClient.java‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.opensearch.ml.common.input.Input;
2525
import org.opensearch.ml.common.input.MLInput;
2626
import org.opensearch.ml.common.output.MLOutput;
27+
import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
2728
import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
2829
import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
2930
import org.opensearch.ml.common.transport.connector.MLCreateConnectorResponse;
@@ -463,6 +464,24 @@ default ActionFuture<MLRegisterAgentResponse> registerAgent(MLAgent mlAgent) {
463464
*/
464465
void registerAgent(MLAgent mlAgent, ActionListener<MLRegisterAgentResponse> listener);
465466

467+
/**
468+
* Get agent
469+
* @param agentId The id of the agent to get
470+
* @return the result future
471+
*/
472+
default ActionFuture<MLAgentGetResponse> getAgent(String agentId) {
473+
PlainActionFuture<MLAgentGetResponse> actionFuture = PlainActionFuture.newFuture();
474+
getAgent(agentId, actionFuture);
475+
return actionFuture;
476+
}
477+
478+
/**
479+
* Get agent
480+
* @param agentId The id of the agent to get
481+
* @param listener a listener to be notified of the result
482+
*/
483+
void getAgent(String agentId, ActionListener<MLAgentGetResponse> listener);
484+
466485
/**
467486
* Delete agent
468487
* @param agentId The id of the agent to delete

‎client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,9 @@
3636
import org.opensearch.ml.common.transport.MLTaskResponse;
3737
import org.opensearch.ml.common.transport.agent.MLAgentDeleteAction;
3838
import org.opensearch.ml.common.transport.agent.MLAgentDeleteRequest;
39+
import org.opensearch.ml.common.transport.agent.MLAgentGetAction;
40+
import org.opensearch.ml.common.transport.agent.MLAgentGetRequest;
41+
import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
3942
import org.opensearch.ml.common.transport.agent.MLRegisterAgentAction;
4043
import org.opensearch.ml.common.transport.agent.MLRegisterAgentRequest;
4144
import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
@@ -291,6 +294,13 @@ public void registerAgent(MLAgent mlAgent, ActionListener<MLRegisterAgentRespons
291294
client.execute(MLRegisterAgentAction.INSTANCE, mlRegisterAgentRequest, getMLRegisterAgentResponseActionListener(listener));
292295
}
293296

297+
@Override
298+
public void getAgent(String agentId, ActionListener<MLAgentGetResponse> listener) {
299+
MLAgentGetRequest mlAgentGetRequest = MLAgentGetRequest.builder().agentId(agentId).build();
300+
301+
client.execute(MLAgentGetAction.INSTANCE, mlAgentGetRequest, getMlGetAgentResponseActionListener(listener));
302+
}
303+
294304
@Override
295305
public void deleteAgent(String agentId, String tenantId, ActionListener<DeleteResponse> listener) {
296306
MLAgentDeleteRequest agentDeleteRequest = new MLAgentDeleteRequest(agentId, tenantId);
@@ -345,6 +355,10 @@ private ActionListener<MLRegisterAgentResponse> getMLRegisterAgentResponseAction
345355
return wrapActionListener(listener, MLRegisterAgentResponse::fromActionResponse);
346356
}
347357

358+
private ActionListener<MLAgentGetResponse> getMlGetAgentResponseActionListener(ActionListener<MLAgentGetResponse> listener) {
359+
return wrapActionListener(listener, MLAgentGetResponse::fromActionResponse);
360+
}
361+
348362
private ActionListener<MLExecuteTaskResponse> getMLExecuteResponseActionListener(ActionListener<MLExecuteTaskResponse> listener) {
349363
return wrapActionListener(listener, MLExecuteTaskResponse::fromActionResponse);
350364
}

‎client/src/test/java/org/opensearch/ml/client/MachineLearningClientTest.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
4949
import org.opensearch.ml.common.output.MLOutput;
5050
import org.opensearch.ml.common.output.MLTrainingOutput;
51+
import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
5152
import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
5253
import org.opensearch.ml.common.transport.config.MLConfigGetResponse;
5354
import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
@@ -106,6 +107,9 @@ public class MachineLearningClientTest {
106107
@Mock
107108
MLRegisterAgentResponse registerAgentResponse;
108109

110+
@Mock
111+
MLAgentGetResponse getAgentResponse;
112+
109113
@Mock
110114
MLConfigGetResponse configGetResponse;
111115

@@ -247,6 +251,11 @@ public void registerAgent(MLAgent mlAgent, ActionListener<MLRegisterAgentRespons
247251
listener.onResponse(registerAgentResponse);
248252
}
249253

254+
@Override
255+
public void getAgent(String agentId, ActionListener<MLAgentGetResponse> listener) {
256+
listener.onResponse(getAgentResponse);
257+
}
258+
250259
@Override
251260
public void deleteAgent(String agentId, String tenantId, ActionListener<DeleteResponse> listener) {
252261
listener.onResponse(deleteResponse);
@@ -535,6 +544,11 @@ public void testRegisterAgent() {
535544
assertEquals(registerAgentResponse, machineLearningClient.registerAgent(mlAgent).actionGet());
536545
}
537546

547+
@Test
548+
public void testGetAgent() {
549+
assertEquals(getAgentResponse, machineLearningClient.getAgent("agentId").actionGet());
550+
}
551+
538552
@Test
539553
public void deleteAgent() {
540554
assertEquals(deleteResponse, machineLearningClient.deleteAgent("agentId").actionGet());

‎client/src/test/java/org/opensearch/ml/client/MachineLearningNodeClientTest.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,9 @@
8787
import org.opensearch.ml.common.transport.MLTaskResponse;
8888
import org.opensearch.ml.common.transport.agent.MLAgentDeleteAction;
8989
import org.opensearch.ml.common.transport.agent.MLAgentDeleteRequest;
90+
import org.opensearch.ml.common.transport.agent.MLAgentGetAction;
91+
import org.opensearch.ml.common.transport.agent.MLAgentGetRequest;
92+
import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
9093
import org.opensearch.ml.common.transport.agent.MLRegisterAgentAction;
9194
import org.opensearch.ml.common.transport.agent.MLRegisterAgentRequest;
9295
import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
@@ -207,6 +210,9 @@ public class MachineLearningNodeClientTest {
207210
@Mock
208211
ActionListener<MLRegisterAgentResponse> registerAgentResponseActionListener;
209212

213+
@Mock
214+
ActionListener<MLAgentGetResponse> getAgentResponseActionListener;
215+
210216
@Mock
211217
ActionListener<DeleteResponse> deleteAgentActionListener;
212218

@@ -1167,6 +1173,27 @@ public void testRegisterAgent() {
11671173
assertEquals(agentId, (argumentCaptor.getValue()).getAgentId());
11681174
}
11691175

1176+
@Test
1177+
public void testGetAgent() {
1178+
String agentId = "agentId";
1179+
MLAgent mlAgent = MLAgent.builder().name("Agent name").type(MLAgentType.FLOW.name()).build();
1180+
1181+
doAnswer(invocation -> {
1182+
ActionListener<MLAgentGetResponse> actionListener = invocation.getArgument(2);
1183+
MLAgentGetResponse output = new MLAgentGetResponse(mlAgent);
1184+
actionListener.onResponse(output);
1185+
return null;
1186+
}).when(client).execute(eq(MLAgentGetAction.INSTANCE), any(), any());
1187+
1188+
ArgumentCaptor<MLAgentGetResponse> argumentCaptor = ArgumentCaptor.forClass(MLAgentGetResponse.class);
1189+
1190+
machineLearningNodeClient.getAgent(agentId, getAgentResponseActionListener);
1191+
1192+
verify(client).execute(eq(MLAgentGetAction.INSTANCE), isA(MLAgentGetRequest.class), any());
1193+
verify(getAgentResponseActionListener).onResponse(argumentCaptor.capture());
1194+
assertEquals(mlAgent, (argumentCaptor.getValue()).getMlAgent());
1195+
}
1196+
11701197
@Test
11711198
public void deleteAgent() {
11721199

‎common/src/main/java/org/opensearch/ml/common/CommonValue.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ public class CommonValue {
3333
public static String HOT_BOX_TYPE = "hot";
3434
// warm node
3535
public static String WARM_BOX_TYPE = "warm";
36+
public static final String ML_INDEX_INSIGHT_CONFIG_INDEX = ".plugins-ml-index-insight-config";
37+
public static final String ML_INDEX_INSIGHT_STORAGE_INDEX = ".plugins-ml-index-insight-storage";
38+
3639
public static final String ML_MODEL_GROUP_INDEX = ".plugins-ml-model-group";
3740
public static final String ML_MODEL_INDEX = ".plugins-ml-model";
3841
public static final String ML_TASK_INDEX = ".plugins-ml-task";
@@ -68,6 +71,8 @@ public class CommonValue {
6871
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
6972
public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
7073
public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";
74+
public static final String ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_config.json";
75+
public static final String ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_storage.json";
7176

7277
// Calculate Versions independently of OpenSearch core version
7378
public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");
@@ -111,4 +116,9 @@ public class CommonValue {
111116

112117
// TOOL Constants
113118
public static final String TOOL_INPUT_SCHEMA_FIELD = "input_schema";
119+
120+
public static final String INDEX_INSIGHT_AGENT_NAME = "os_index_insight_agent";
121+
public static final long INDEX_INSIGHT_GENERATING_TIMEOUT = 3 * 60 * 1000;
122+
public static final long INDEX_INSIGHT_UPDATE_INTERVAL = 24 * 60 * 60 * 1000;
123+
114124
}

‎common/src/main/java/org/opensearch/ml/common/MLIndex.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX_MAPPING_PATH;
1616
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1717
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
18+
import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_CONFIG_INDEX;
19+
import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH;
20+
import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_STORAGE_INDEX;
21+
import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH;
1822
import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX;
1923
import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX_MAPPING_PATH;
2024
import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
@@ -50,7 +54,9 @@ public enum MLIndex {
5054
MEMORY_CONTAINER(ML_MEMORY_CONTAINER_INDEX, false, ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH),
5155
MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
5256
MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH),
53-
JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH);
57+
JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH),
58+
INDEX_INSIGHT_CONFIG(ML_INDEX_INSIGHT_CONFIG_INDEX, false, ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH),
59+
INDEX_INSIGHT_STORAGE(ML_INDEX_INSIGHT_STORAGE_INDEX, false, ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH);
5460

5561
private final String indexName;
5662
// whether we use an alias for the index

0 commit comments

Comments
 (0)