Merge branch 'main' into main

Author: dhrubo-os

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

@@ -24,6 +24,7 @@
 import org.opensearch.ml.common.input.Input;
 import org.opensearch.ml.common.input.MLInput;
 import org.opensearch.ml.common.output.MLOutput;
+import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
 import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
 import org.opensearch.ml.common.transport.connector.MLCreateConnectorResponse;
@@ -463,6 +464,24 @@ default ActionFuture<MLRegisterAgentResponse> registerAgent(MLAgent mlAgent) {
      */
     void registerAgent(MLAgent mlAgent, ActionListener<MLRegisterAgentResponse> listener);
 
+    /**
+     * Get agent
+     * @param agentId The id of the agent to get
+     * @return the result future
+     */
+    default ActionFuture<MLAgentGetResponse> getAgent(String agentId) {
+        PlainActionFuture<MLAgentGetResponse> actionFuture = PlainActionFuture.newFuture();
+        getAgent(agentId, actionFuture);
+        return actionFuture;
+    }
+
+    /**
+     * Get agent
+     * @param agentId The id of the agent to get
+     * @param listener a listener to be notified of the result
+     */
+    void getAgent(String agentId, ActionListener<MLAgentGetResponse> listener);
+
     /**
      * Delete agent
      * @param agentId The id of the agent to delete

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

@@ -36,6 +36,9 @@
 import org.opensearch.ml.common.transport.MLTaskResponse;
 import org.opensearch.ml.common.transport.agent.MLAgentDeleteAction;
 import org.opensearch.ml.common.transport.agent.MLAgentDeleteRequest;
+import org.opensearch.ml.common.transport.agent.MLAgentGetAction;
+import org.opensearch.ml.common.transport.agent.MLAgentGetRequest;
+import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentAction;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentRequest;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
@@ -291,6 +294,13 @@ public void registerAgent(MLAgent mlAgent, ActionListener<MLRegisterAgentRespons
         client.execute(MLRegisterAgentAction.INSTANCE, mlRegisterAgentRequest, getMLRegisterAgentResponseActionListener(listener));
     }
 
+    @Override
+    public void getAgent(String agentId, ActionListener<MLAgentGetResponse> listener) {
+        MLAgentGetRequest mlAgentGetRequest = MLAgentGetRequest.builder().agentId(agentId).build();
+
+        client.execute(MLAgentGetAction.INSTANCE, mlAgentGetRequest, getMlGetAgentResponseActionListener(listener));
+    }
+
     @Override
     public void deleteAgent(String agentId, String tenantId, ActionListener<DeleteResponse> listener) {
         MLAgentDeleteRequest agentDeleteRequest = new MLAgentDeleteRequest(agentId, tenantId);
@@ -345,6 +355,10 @@ private ActionListener<MLRegisterAgentResponse> getMLRegisterAgentResponseAction
         return wrapActionListener(listener, MLRegisterAgentResponse::fromActionResponse);
     }
 
+    private ActionListener<MLAgentGetResponse> getMlGetAgentResponseActionListener(ActionListener<MLAgentGetResponse> listener) {
+        return wrapActionListener(listener, MLAgentGetResponse::fromActionResponse);
+    }
+
     private ActionListener<MLExecuteTaskResponse> getMLExecuteResponseActionListener(ActionListener<MLExecuteTaskResponse> listener) {
         return wrapActionListener(listener, MLExecuteTaskResponse::fromActionResponse);
     }

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

@@ -48,6 +48,7 @@
 import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
 import org.opensearch.ml.common.output.MLOutput;
 import org.opensearch.ml.common.output.MLTrainingOutput;
+import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
 import org.opensearch.ml.common.transport.config.MLConfigGetResponse;
 import org.opensearch.ml.common.transport.connector.MLCreateConnectorInput;
@@ -106,6 +107,9 @@ public class MachineLearningClientTest {
     @Mock
     MLRegisterAgentResponse registerAgentResponse;
 
+    @Mock
+    MLAgentGetResponse getAgentResponse;
+
     @Mock
     MLConfigGetResponse configGetResponse;
 
@@ -247,6 +251,11 @@ public void registerAgent(MLAgent mlAgent, ActionListener<MLRegisterAgentRespons
                 listener.onResponse(registerAgentResponse);
             }
 
+            @Override
+            public void getAgent(String agentId, ActionListener<MLAgentGetResponse> listener) {
+                listener.onResponse(getAgentResponse);
+            }
+
             @Override
             public void deleteAgent(String agentId, String tenantId, ActionListener<DeleteResponse> listener) {
                 listener.onResponse(deleteResponse);
@@ -535,6 +544,11 @@ public void testRegisterAgent() {
         assertEquals(registerAgentResponse, machineLearningClient.registerAgent(mlAgent).actionGet());
     }
 
+    @Test
+    public void testGetAgent() {
+        assertEquals(getAgentResponse, machineLearningClient.getAgent("agentId").actionGet());
+    }
+
     @Test
     public void deleteAgent() {
         assertEquals(deleteResponse, machineLearningClient.deleteAgent("agentId").actionGet());

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

@@ -87,6 +87,9 @@
 import org.opensearch.ml.common.transport.MLTaskResponse;
 import org.opensearch.ml.common.transport.agent.MLAgentDeleteAction;
 import org.opensearch.ml.common.transport.agent.MLAgentDeleteRequest;
+import org.opensearch.ml.common.transport.agent.MLAgentGetAction;
+import org.opensearch.ml.common.transport.agent.MLAgentGetRequest;
+import org.opensearch.ml.common.transport.agent.MLAgentGetResponse;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentAction;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentRequest;
 import org.opensearch.ml.common.transport.agent.MLRegisterAgentResponse;
@@ -207,6 +210,9 @@ public class MachineLearningNodeClientTest {
     @Mock
     ActionListener<MLRegisterAgentResponse> registerAgentResponseActionListener;
 
+    @Mock
+    ActionListener<MLAgentGetResponse> getAgentResponseActionListener;
+
     @Mock
     ActionListener<DeleteResponse> deleteAgentActionListener;
 
@@ -1167,6 +1173,27 @@ public void testRegisterAgent() {
         assertEquals(agentId, (argumentCaptor.getValue()).getAgentId());
     }
 
+    @Test
+    public void testGetAgent() {
+        String agentId = "agentId";
+        MLAgent mlAgent = MLAgent.builder().name("Agent name").type(MLAgentType.FLOW.name()).build();
+
+        doAnswer(invocation -> {
+            ActionListener<MLAgentGetResponse> actionListener = invocation.getArgument(2);
+            MLAgentGetResponse output = new MLAgentGetResponse(mlAgent);
+            actionListener.onResponse(output);
+            return null;
+        }).when(client).execute(eq(MLAgentGetAction.INSTANCE), any(), any());
+
+        ArgumentCaptor<MLAgentGetResponse> argumentCaptor = ArgumentCaptor.forClass(MLAgentGetResponse.class);
+
+        machineLearningNodeClient.getAgent(agentId, getAgentResponseActionListener);
+
+        verify(client).execute(eq(MLAgentGetAction.INSTANCE), isA(MLAgentGetRequest.class), any());
+        verify(getAgentResponseActionListener).onResponse(argumentCaptor.capture());
+        assertEquals(mlAgent, (argumentCaptor.getValue()).getMlAgent());
+    }
+
     @Test
     public void deleteAgent() {
 

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

@@ -33,6 +33,9 @@ public class CommonValue {
     public static String HOT_BOX_TYPE = "hot";
     // warm node
     public static String WARM_BOX_TYPE = "warm";
+    public static final String ML_INDEX_INSIGHT_CONFIG_INDEX = ".plugins-ml-index-insight-config";
+    public static final String ML_INDEX_INSIGHT_STORAGE_INDEX = ".plugins-ml-index-insight-storage";
+
     public static final String ML_MODEL_GROUP_INDEX = ".plugins-ml-model-group";
     public static final String ML_MODEL_INDEX = ".plugins-ml-model";
     public static final String ML_TASK_INDEX = ".plugins-ml-task";
@@ -68,6 +71,8 @@ public class CommonValue {
     public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
     public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
     public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";
+    public static final String ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_config.json";
+    public static final String ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_storage.json";
 
     // Calculate Versions independently of OpenSearch core version
     public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");
@@ -111,4 +116,9 @@ public class CommonValue {
 
     // TOOL Constants
     public static final String TOOL_INPUT_SCHEMA_FIELD = "input_schema";
+
+    public static final String INDEX_INSIGHT_AGENT_NAME = "os_index_insight_agent";
+    public static final long INDEX_INSIGHT_GENERATING_TIMEOUT = 3 * 60 * 1000;
+    public static final long INDEX_INSIGHT_UPDATE_INTERVAL = 24 * 60 * 60 * 1000;
+
 }

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

@@ -15,6 +15,10 @@
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
+import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_CONFIG_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH;
+import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_STORAGE_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
@@ -50,7 +54,9 @@ public enum MLIndex {
     MEMORY_CONTAINER(ML_MEMORY_CONTAINER_INDEX, false, ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH),
     MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
     MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH),
-    JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH);
+    JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH),
+    INDEX_INSIGHT_CONFIG(ML_INDEX_INSIGHT_CONFIG_INDEX, false, ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH),
+    INDEX_INSIGHT_STORAGE(ML_INDEX_INSIGHT_STORAGE_INDEX, false, ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH);
 
     private final String indexName;
     // whether we use an alias for the index