Check before delete

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

@@ -42,6 +42,7 @@ public class CommonValue {
     public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
     public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
     public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
+    public static final String TOOL_PARAMETERS_PREFIX = "tools.parameters.";
 
     // Index mapping paths
     public static final String ML_MODEL_GROUP_INDEX_MAPPING_PATH = "index-mappings/ml-model-group.json";

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/MLModelTool.java

@@ -17,8 +17,8 @@
 import org.opensearch.ml.common.output.model.ModelTensorOutput;
 import org.opensearch.ml.common.output.model.ModelTensors;
 import org.opensearch.ml.common.spi.tools.Parser;
-import org.opensearch.ml.common.spi.tools.Tool;
 import org.opensearch.ml.common.spi.tools.ToolAnnotation;
+import org.opensearch.ml.common.spi.tools.WithModelTool;
 import org.opensearch.ml.common.transport.prediction.MLPredictionTaskAction;
 import org.opensearch.ml.common.transport.prediction.MLPredictionTaskRequest;
 import org.opensearch.ml.common.utils.StringUtils;
@@ -33,7 +33,7 @@
  */
 @Log4j2
 @ToolAnnotation(MLModelTool.TYPE)
-public class MLModelTool implements Tool {
+public class MLModelTool implements WithModelTool {
     public static final String TYPE = "MLModelTool";
     public static final String RESPONSE_FIELD = "response_field";
     public static final String MODEL_ID_FIELD = "model_id";
@@ -127,7 +127,7 @@ public boolean validate(Map<String, String> parameters) {
         return true;
     }
 
-    public static class Factory implements Tool.Factory<MLModelTool> {
+    public static class Factory implements WithModelTool.Factory<MLModelTool> {
         private Client client;
 
         private static Factory INSTANCE;
@@ -172,5 +172,10 @@ public String getDefaultType() {
         public String getDefaultVersion() {
             return null;
         }
+
+        @Override
+        public List<String> getAllModelKeys() {
+            return List.of(MODEL_ID_FIELD);
+        }
     }
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/SearchIndexTool.java

@@ -195,5 +195,6 @@ public String getDefaultType() {
         public String getDefaultVersion() {
             return null;
         }
+
     }
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/utils/AgentModelsSearcher.java

@@ -0,0 +1,46 @@
+package org.opensearch.ml.engine.utils;
+
+import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
+import static org.opensearch.ml.common.CommonValue.TOOL_PARAMETERS_PREFIX;
+
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.ml.common.spi.tools.Tool;
+import org.opensearch.ml.common.spi.tools.WithModelTool;
+import org.opensearch.search.builder.SearchSourceBuilder;
+
+public class AgentModelsSearcher {
+    private final Set<String> relatedModelIdSet;
+
+    public AgentModelsSearcher(Map<String, Tool.Factory> toolFactories) {
+        relatedModelIdSet = new HashSet<>();
+        for (Map.Entry<String, Tool.Factory> entry : toolFactories.entrySet()) {
+            Tool.Factory toolFactory = entry.getValue();
+            if (toolFactory instanceof WithModelTool.Factory withModelTool) {
+                relatedModelIdSet.addAll(withModelTool.getAllModelKeys());
+            }
+        }
+    }
+
+    /**
+     * Construct a should query to search all agent which containing candidate model Id
+
+     @param candidateModelId the candidate model Id
+     @return a should search request towards agent index.
+     */
+    public SearchRequest constructQueryRequestToSearchModelIdInsideAgent(String candidateModelId) {
+        SearchRequest searchRequest = new SearchRequest(ML_AGENT_INDEX);
+        BoolQueryBuilder shouldQuery = QueryBuilders.boolQuery();
+        for (String keyField : relatedModelIdSet) {
+            shouldQuery.should(QueryBuilders.termsQuery(TOOL_PARAMETERS_PREFIX + keyField, candidateModelId));
+        }
+        searchRequest.source(new SearchSourceBuilder().query(shouldQuery));
+        return searchRequest;
+    }
+
+}