change agent model searcher map to set

Signed-off-by: xinyual <[email protected]>
xinyual · Nov 19, 2024 · de464e9 · de464e9
1 parent ba587fd
commit de464e9
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/AgentModelsSearcher.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/AgentModelsSearcher.java
@@ -4,7 +4,6 @@
 import static org.opensearch.ml.common.CommonValue.TOOL_MODEL_RELATED_FIELD_PREFIX;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -17,14 +16,13 @@
 import org.opensearch.search.builder.SearchSourceBuilder;
 
 public class AgentModelsSearcher {
-    private Map<String, List<String>> relatedModelIdMap;
+    private Set<String> relatedModelIdSet;
 
     public AgentModelsSearcher(Map<String, Tool.Factory> toolFactories) {
-        relatedModelIdMap = new HashMap<>();
+        relatedModelIdSet = new HashSet<>();
         for (Map.Entry<String, Tool.Factory> entry : toolFactories.entrySet()) {
-            String toolType = entry.getKey();
             Tool.Factory toolFactory = entry.getValue();
-            relatedModelIdMap.put(toolType, toolFactory.getAllModelKeys());
+            relatedModelIdSet.addAll(toolFactory.getAllModelKeys());
         }
     }
 
@@ -41,8 +39,8 @@ public SearchRequest constructQueryRequest(String candidateModelId) {
 
     private List<String> collectAllKeys() {
         Set<String> keys = new HashSet<>();
-        for (Map.Entry<String, List<String>> entry : relatedModelIdMap.entrySet()) {
-            keys.addAll(entry.getValue());
+        for (String modelKey : relatedModelIdSet) {
+            keys.add(modelKey);
         }
         return new ArrayList<>(keys);
     }