Skip to content

Commit 6bb7c63

Browse files
xinyualxinyual
committed
change name and reformat logic
Signed-off-by: xinyual <xinyual@amazon.com>
1 parent cde3934 commit 6bb7c63

File tree

5 files changed

+80
-121
lines changed

5 files changed

+80
-121
lines changed

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ public class CommonValue {
8181
public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
8282
public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
8383
public static final Integer ML_MEMORY_MESSAGE_INDEX_SCHEMA_VERSION = 1;
84-
public static final String TOOL_MODEL_RELATED_FIELD_PREFIX = "tools.parameters.";
84+
public static final String TOOL_PARAMETERS_PREFIX = "tools.parameters.";
8585
public static final String USER_FIELD_MAPPING = " \""
8686
+ CommonValue.USER
8787
+ "\": {\n"

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

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
package org.opensearch.ml.engine.tools;
1+
package org.opensearch.ml.engine.utils;
22

33
import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
4-
import static org.opensearch.ml.common.CommonValue.TOOL_MODEL_RELATED_FIELD_PREFIX;
4+
import static org.opensearch.ml.common.CommonValue.TOOL_PARAMETERS_PREFIX;
55

66
import java.util.HashSet;
77
import java.util.Map;
@@ -28,7 +28,7 @@ public SearchRequest constructQueryRequest(String candidateModelId) {
2828
SearchRequest searchRequest = new SearchRequest(ML_AGENT_INDEX);
2929
BoolQueryBuilder shouldQuery = QueryBuilders.boolQuery();
3030
for (String keyField : relatedModelIdSet) {
31-
shouldQuery.should(QueryBuilders.termsQuery(TOOL_MODEL_RELATED_FIELD_PREFIX + keyField, candidateModelId));
31+
shouldQuery.should(QueryBuilders.termsQuery(TOOL_PARAMETERS_PREFIX + keyField, candidateModelId));
3232
}
3333
searchRequest.source(new SearchSourceBuilder().query(shouldQuery));
3434
return searchRequest;

‎plugin/src/main/java/org/opensearch/ml/action/models/DeleteModelTransportAction.java

+41-48
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,13 @@
2828
import java.util.concurrent.CountDownLatch;
2929
import java.util.concurrent.atomic.AtomicBoolean;
3030
import java.util.function.Function;
31+
import java.util.function.Supplier;
3132

3233
import org.apache.commons.lang3.tuple.Pair;
3334
import org.opensearch.OpenSearchStatusException;
3435
import org.opensearch.ResourceNotFoundException;
3536
import org.opensearch.action.ActionRequest;
37+
import org.opensearch.action.ActionType;
3638
import org.opensearch.action.delete.DeleteRequest;
3739
import org.opensearch.action.delete.DeleteResponse;
3840
import org.opensearch.action.get.GetRequest;
@@ -66,7 +68,8 @@
6668
import org.opensearch.ml.common.transport.model.MLModelDeleteAction;
6769
import org.opensearch.ml.common.transport.model.MLModelDeleteRequest;
6870
import org.opensearch.ml.common.transport.model.MLModelGetRequest;
69-
import org.opensearch.ml.engine.tools.AgentModelsSearcher;
71+
import org.opensearch.ml.common.utils.StringUtils;
72+
import org.opensearch.ml.engine.utils.AgentModelsSearcher;
7073
import org.opensearch.ml.helper.ModelAccessControlHelper;
7174
import org.opensearch.ml.utils.RestActionUtils;
7275
import org.opensearch.search.SearchHit;
@@ -284,48 +287,27 @@ private void checkAgentBeforeDeleteModel(String modelId, ActionListener<Boolean>
284287
}
285288

286289
private void checkIngestPipelineBeforeDeleteModel(String modelId, ActionListener<Boolean> actionListener) {
287-
GetPipelineRequest getPipelineRequest = new GetPipelineRequest();
288-
client.execute(GetPipelineAction.INSTANCE, getPipelineRequest, ActionListener.wrap(ingestPipelineResponse -> {
289-
List<String> allDependentPipelineIds = findDependentPipelines(
290-
ingestPipelineResponse.pipelines(),
291-
modelId,
292-
org.opensearch.ingest.PipelineConfiguration::getConfigAsMap,
293-
org.opensearch.ingest.PipelineConfiguration::getId
294-
);
295-
if (allDependentPipelineIds.isEmpty()) {
296-
actionListener.onResponse(true);
297-
} else {
298-
actionListener
299-
.onFailure(
300-
new OpenSearchStatusException(
301-
String
302-
.format(
303-
Locale.ROOT,
304-
"%d ingest pipelines are still using this model, please delete or update the pipelines first: %s",
305-
allDependentPipelineIds.size(),
306-
Arrays.toString(allDependentPipelineIds.toArray(new String[0]))
307-
),
308-
RestStatus.CONFLICT
309-
)
310-
);
311-
}
312-
}, e -> {
313-
log.error("Failed to delete ML Model: " + modelId, e);
314-
actionListener.onFailure(e);
315-
316-
}));
290+
checkPipelineBeforeDeleteModel(modelId, actionListener, "ingest", GetPipelineRequest::new, GetPipelineAction.INSTANCE);
317291

318292
}
319293

320294
private void checkSearchPipelineBeforeDeleteModel(String modelId, ActionListener<Boolean> actionListener) {
321-
GetSearchPipelineRequest getSearchPipelineRequest = new GetSearchPipelineRequest();
322-
client.execute(GetSearchPipelineAction.INSTANCE, getSearchPipelineRequest, ActionListener.wrap(searchPipelineResponse -> {
323-
List<String> allDependentPipelineIds = findDependentPipelines(
324-
searchPipelineResponse.pipelines(),
325-
modelId,
326-
org.opensearch.search.pipeline.PipelineConfiguration::getConfigAsMap,
327-
org.opensearch.search.pipeline.PipelineConfiguration::getId
328-
);
295+
checkPipelineBeforeDeleteModel(modelId, actionListener, "search", GetSearchPipelineRequest::new, GetSearchPipelineAction.INSTANCE);
296+
297+
}
298+
299+
private void checkPipelineBeforeDeleteModel(
300+
String modelId,
301+
ActionListener<Boolean> actionListener,
302+
String pipelineType,
303+
Supplier<ActionRequest> requestSupplier,
304+
ActionType actionType
305+
) {
306+
ActionRequest request = requestSupplier.get();
307+
client.execute(actionType, request, ActionListener.wrap(pipelineResponse -> {
308+
String responseString = pipelineResponse.toString();
309+
Map<String, Object> allConfigMap = StringUtils.fromJson(pipelineResponse.toString(), "");
310+
List<String> allDependentPipelineIds = findDependentPipelinesEasy(allConfigMap, modelId);
329311
if (allDependentPipelineIds.isEmpty()) {
330312
actionListener.onResponse(true);
331313
} else {
@@ -335,8 +317,9 @@ private void checkSearchPipelineBeforeDeleteModel(String modelId, ActionListener
335317
String
336318
.format(
337319
Locale.ROOT,
338-
"%d search pipelines are still using this model, please delete or update the pipelines first: %s",
320+
"%d %s pipelines are still using this model, please delete or update the pipelines first: %s",
339321
allDependentPipelineIds.size(),
322+
pipelineType,
340323
Arrays.toString(allDependentPipelineIds.toArray(new String[0]))
341324
),
342325
RestStatus.CONFLICT
@@ -479,6 +462,18 @@ private Boolean isModelNotDeployed(MLModelState mlModelState) {
479462
&& !mlModelState.equals(MLModelState.PARTIALLY_DEPLOYED);
480463
}
481464

465+
private List<String> findDependentPipelinesEasy(Map<String, Object> allConfigMap, String candidateModelId) {
466+
List<String> dependentPipelineConfigurations = new ArrayList<>();
467+
for (Map.Entry<String, Object> entry : allConfigMap.entrySet()) {
468+
String id = entry.getKey();
469+
Map<String, Object> config = (Map<String, Object>) entry.getValue();
470+
if (searchThroughConfig(config, candidateModelId)) {
471+
dependentPipelineConfigurations.add(id);
472+
}
473+
}
474+
return dependentPipelineConfigurations;
475+
}
476+
482477
private <T> List<String> findDependentPipelines(
483478
List<T> pipelineConfigurations,
484479
String candidateModelId,
@@ -533,24 +528,22 @@ private Boolean searchThroughConfig(Object searchCandidate, String candidateId)
533528
}
534529

535530
private String formatAgentErrorMessage(SearchHit[] hits) {
536-
boolean isHidden = false;
537531
List<String> agentIds = new ArrayList<>();
538532
for (SearchHit hit : hits) {
539533
Map<String, Object> sourceAsMap = hit.getSourceAsMap();
540-
isHidden = isHidden || Boolean.parseBoolean((String) sourceAsMap.getOrDefault(MLAgent.IS_HIDDEN_FIELD, false));
541-
agentIds.add(hit.getId());
542-
}
543-
if (isHidden) {
544-
return String
545-
.format(Locale.ROOT, "%d agents are still using this model, please delete or update the agents first", hits.length);
534+
Boolean isHidden = (Boolean) sourceAsMap.getOrDefault(MLAgent.IS_HIDDEN_FIELD, false);
535+
if (!isHidden) {
536+
agentIds.add(hit.getId());
537+
}
546538
}
547539
return String
548540
.format(
549541
Locale.ROOT,
550-
"%d agents are still using this model, please delete or update the agents first: %s",
542+
"%d agents are still using this model, please delete or update the agents first, all visible agents are: %s",
551543
hits.length,
552544
Arrays.toString(agentIds.toArray(new String[0]))
553545
);
546+
554547
}
555548

556549
// this method is only to stub static method.

‎plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -178,14 +178,14 @@
178178
import org.opensearch.ml.engine.indices.MLInputDatasetHandler;
179179
import org.opensearch.ml.engine.memory.ConversationIndexMemory;
180180
import org.opensearch.ml.engine.memory.MLMemoryManager;
181-
import org.opensearch.ml.engine.tools.AgentModelsSearcher;
182181
import org.opensearch.ml.engine.tools.AgentTool;
183182
import org.opensearch.ml.engine.tools.CatIndexTool;
184183
import org.opensearch.ml.engine.tools.ConnectorTool;
185184
import org.opensearch.ml.engine.tools.IndexMappingTool;
186185
import org.opensearch.ml.engine.tools.MLModelTool;
187186
import org.opensearch.ml.engine.tools.SearchIndexTool;
188187
import org.opensearch.ml.engine.tools.VisualizationsTool;
188+
import org.opensearch.ml.engine.utils.AgentModelsSearcher;
189189
import org.opensearch.ml.helper.ConnectorAccessControlHelper;
190190
import org.opensearch.ml.helper.ModelAccessControlHelper;
191191
import org.opensearch.ml.memory.ConversationalMemoryHandler;

0 commit comments

Comments
 (0)