Skip to content

Commit 759b91c

Browse files
xinyualylwu-amzn
xinyual
authored and
ylwu-amzn
committed
fix conflict
Signed-off-by: xinyual <[email protected]>
1 parent a5192ab commit 759b91c

File tree

7 files changed

+27
-208
lines changed

7 files changed

+27
-208
lines changed

‎common/src/main/java/org/opensearch/ml/common/indexInsight/AbstractIndexInsightTask.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -214,9 +214,11 @@ protected void beginGeneration(String tenantId, ActionListener<IndexInsight> lis
214214
.lastUpdatedTime(Instant.now())
215215
.build();
216216

217-
writeIndexInsight(indexInsight, tenantId, ActionListener.wrap(r -> { runWithPrerequisites(tenantId, listener); }, e -> {
218-
saveFailedStatus(tenantId, e, listener);
219-
}));
217+
writeIndexInsight(
218+
indexInsight,
219+
tenantId,
220+
ActionListener.wrap(r -> { runWithPrerequisites(tenantId, listener); }, listener::onFailure)
221+
);
220222
}
221223

222224
protected void runWithPrerequisites(String tenantId, ActionListener<IndexInsight> listener) {
@@ -260,6 +262,7 @@ protected void saveFailedStatus(String tenantId, Exception error, ActionListener
260262
.tenantId(tenantId)
261263
.index(sourceIndex)
262264
.taskType(taskType)
265+
.lastUpdatedTime(Instant.now())
263266
.status(IndexInsightTaskStatus.FAILED)
264267
.build();
265268
writeIndexInsight(

‎common/src/main/java/org/opensearch/ml/common/indexInsight/StatisticalDataTask.java‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,6 @@
2525
import org.opensearch.action.search.SearchResponse;
2626
import org.opensearch.cluster.metadata.MappingMetadata;
2727
import org.opensearch.core.action.ActionListener;
28-
import org.opensearch.core.common.Strings;
29-
import org.opensearch.core.xcontent.MediaTypeRegistry;
3028
import org.opensearch.index.query.QueryBuilders;
3129
import org.opensearch.ml.common.utils.mergeMetaDataUtils.MergeRuleHelper;
3230
import org.opensearch.remote.metadata.client.SdkClient;
@@ -340,20 +338,12 @@ private Map<String, Object> parseSearchResult(
340338

341339
String aggregationType = key.substring(0, prefix.length() - 1);
342340

343-
Map<String, Object> aggregationResult = gson
344-
.fromJson(Strings.toString(MediaTypeRegistry.JSON, aggregation), Map.class);
341+
Map<String, Object> aggregationResult = gson.fromJson(aggregation.toString(), Map.class);
345342
Object targetValue;
346-
String targetKey = key;
347-
for (String aggKey : aggregationResult.keySet()) {
348-
if (aggKey.split("#", 2)[1].equals(key)) {
349-
targetKey = aggKey;
350-
break;
351-
}
352-
}
353343
try {
354344
if (prefix.equals(UNIQUE_TERM_PREFIX)) {
355345
// assuming result.get(key) is a Map containing "buckets" -> List<Map<String, Object>>
356-
Map<String, Object> aggResult = (Map<String, Object>) aggregationResult.get(targetKey);
346+
Map<String, Object> aggResult = (Map<String, Object>) aggregationResult.get(key);
357347
List<Map<String, Object>> buckets = aggResult != null
358348
? (List<Map<String, Object>>) aggResult.get("buckets")
359349
: null;
@@ -363,7 +353,7 @@ private Map<String, Object> parseSearchResult(
363353

364354
targetValue = buckets.stream().filter(bucket -> bucket != null).map(bucket -> bucket.get("key")).toList();
365355
} else {
366-
Map<String, Object> aggResult = (Map<String, Object>) aggregationResult.get(targetKey);
356+
Map<String, Object> aggResult = (Map<String, Object>) aggregationResult.get(key);
367357
if (aggResult.containsKey("value_as_string")) {
368358
targetValue = aggResult.get("value_as_string");
369359
} else {

‎common/src/test/java/org/opensearch/ml/common/indexInsight/AbstractIndexInsightTaskTests.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ public void testHandleExistingDoc_Completed_NeedUpdate() throws IOException {
186186
task.handleExistingDoc(source, "test-tenant", listener);
187187

188188
// Verify prerequisite task execution: 5 threadPool calls indicate both prerequisite and main task ran
189-
verify(client, times(5)).threadPool();
189+
verify(client, times(6)).threadPool();
190190
verify(listener).onResponse(any(IndexInsight.class));
191191
}
192192

@@ -215,7 +215,7 @@ public void testHandleExistingDoc_Generating_Timeout() throws IOException {
215215
ActionListener<IndexInsight> listener = mock(ActionListener.class);
216216
task.handleExistingDoc(source, "test-tenant", listener);
217217

218-
verify(client, times(5)).threadPool();
218+
verify(client, times(6)).threadPool();
219219
verify(listener).onResponse(any(IndexInsight.class));
220220
}
221221

@@ -229,7 +229,7 @@ public void testHandleExistingDoc_Failed_Retry() throws IOException {
229229
ActionListener<IndexInsight> listener = mock(ActionListener.class);
230230
task.handleExistingDoc(source, "test-tenant", listener);
231231

232-
verify(client, times(5)).threadPool();
232+
verify(client, times(6)).threadPool();
233233
verify(listener).onResponse(any(IndexInsight.class));
234234
}
235235

@@ -448,7 +448,7 @@ public void testCallLLMWithAgent_SuccessWithJson() {
448448
return null;
449449
}).when(client).execute(any(), any(MLExecuteTaskRequest.class), any(ActionListener.class));
450450

451-
AbstractIndexInsightTask.callLLMWithAgent(client, agentId, prompt, sourceIndex, listener);
451+
AbstractIndexInsightTask.callLLMWithAgent(client, agentId, prompt, sourceIndex, null, listener);
452452

453453
verify(listener).onResponse("parsed response");
454454
}
@@ -480,7 +480,7 @@ public void testCallLLMWithAgent_SuccessWithPlainText() {
480480
return null;
481481
}).when(client).execute(any(), any(MLExecuteTaskRequest.class), any(ActionListener.class));
482482

483-
AbstractIndexInsightTask.callLLMWithAgent(client, agentId, prompt, sourceIndex, listener);
483+
AbstractIndexInsightTask.callLLMWithAgent(client, agentId, prompt, sourceIndex, null, listener);
484484

485485
verify(listener).onResponse(plainResponse);
486486
}
@@ -500,7 +500,7 @@ public void testCallLLMWithAgent_Failure() {
500500
return null;
501501
}).when(client).execute(any(), any(MLExecuteTaskRequest.class), any(ActionListener.class));
502502

503-
AbstractIndexInsightTask.callLLMWithAgent(client, agentId, prompt, sourceIndex, listener);
503+
AbstractIndexInsightTask.callLLMWithAgent(client, agentId, prompt, sourceIndex, null, listener);
504504

505505
verify(listener).onFailure(expectedException);
506506
}

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1193,8 +1193,6 @@ public List<RestHandler> getRestHandlers(
11931193
restMLRemoveMcpToolsAction,
11941194
restMLListMcpToolsAction,
11951195
restMLMcpToolsUpdateAction,
1196-
restMLPutIndexInsightConfigAction,
1197-
restMLGetIndexInsightConfigAction,
11981196
restMLCreateContextManagementTemplateAction,
11991197
restMLGetContextManagementTemplateAction,
12001198
restMLListContextManagementTemplatesAction,

‎plugin/src/test/java/org/opensearch/ml/action/IndexInsight/GetIndexInsightTransportActionTests.java‎

Lines changed: 12 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
4747
import org.opensearch.ml.common.transport.indexInsight.MLIndexInsightGetRequest;
4848
import org.opensearch.ml.common.transport.indexInsight.MLIndexInsightGetResponse;
49+
import org.opensearch.ml.engine.indices.MLIndicesHandler;
4950
import org.opensearch.remote.metadata.client.GetDataObjectResponse;
5051
import org.opensearch.remote.metadata.client.SdkClient;
5152
import org.opensearch.test.OpenSearchTestCase;
@@ -92,6 +93,9 @@ public class GetIndexInsightTransportActionTests extends OpenSearchTestCase {
9293
@Mock
9394
private MLFeatureEnabledSetting mlFeatureEnabledSetting;
9495

96+
@Mock
97+
private MLIndicesHandler mlIndicesHandler;
98+
9599
GetIndexInsightTransportAction getIndexInsightTransportAction;
96100
MLIndexInsightGetRequest mlIndexInsightGetRequest;
97101
ThreadContext threadContext;
@@ -113,7 +117,8 @@ public void setup() throws IOException {
113117
xContentRegistry,
114118
mlFeatureEnabledSetting,
115119
client,
116-
sdkClient
120+
sdkClient,
121+
mlIndicesHandler
117122
)
118123
);
119124

@@ -132,6 +137,12 @@ public void setup() throws IOException {
132137
when(mlFeatureEnabledSetting.isIndexInsightEnabled()).thenReturn(Boolean.TRUE);
133138
when(getMappingsResponse.getMappings()).thenReturn(Map.of("demo", mappingMetadata));
134139

140+
doAnswer(invocation -> {
141+
ActionListener<Boolean> listener = invocation.getArgument(1);
142+
listener.onResponse(true);
143+
return null;
144+
}).when(mlIndicesHandler).initMLIndexIfAbsent(any(), any());
145+
135146
}
136147

137148
@Test
@@ -161,49 +172,6 @@ public void testGetIndexInsight_Successful() {
161172
assertEquals("test_index", argumentCaptor.getValue().getIndexInsight().getIndex());
162173
}
163174

164-
@Test
165-
public void testGetIndexInsight_FailDueToNotEnabled() {
166-
GetResponse getResponse = mock(GetResponse.class);
167-
when(getResponse.isExists()).thenReturn(true);
168-
when(getResponse.getSourceAsString()).thenReturn("{\"is_enable\":false}");
169-
170-
GetDataObjectResponse sdkResponse = mock(GetDataObjectResponse.class);
171-
when(sdkResponse.getResponse()).thenReturn(getResponse);
172-
173-
CompletableFuture<GetDataObjectResponse> future = CompletableFuture.completedFuture(sdkResponse);
174-
175-
when(sdkClient.getDataObjectAsync(any())).thenReturn(future);
176-
IndexInsightTask indexInsightTask = mock(IndexInsightTask.class);
177-
178-
doReturn(indexInsightTask).when(getIndexInsightTransportAction).createTask(any());
179-
IndexInsight insight = new IndexInsight(
180-
"test_index",
181-
"test content",
182-
IndexInsightTaskStatus.COMPLETED,
183-
STATISTICAL_DATA,
184-
Instant.ofEpochMilli(0),
185-
""
186-
);
187-
doAnswer(invocation -> {
188-
ActionListener<IndexInsight> listener = invocation.getArgument(1);
189-
listener.onResponse(insight);
190-
return null;
191-
}).when(indexInsightTask).execute(any(), any());
192-
193-
SearchResponse searchResponse = mock(SearchResponse.class);
194-
doAnswer(invocation -> {
195-
ActionListener<SearchResponse> listener = invocation.getArgument(1);
196-
listener.onResponse(searchResponse);
197-
return null;
198-
}).when(client).search(any(), any());
199-
200-
getIndexInsightTransportAction.doExecute(null, mlIndexInsightGetRequest, actionListener);
201-
ArgumentCaptor<Exception> argumentCaptor = ArgumentCaptor.forClass(Exception.class);
202-
verify(actionListener).onFailure(argumentCaptor.capture());
203-
assertTrue(argumentCaptor.getValue() instanceof RuntimeException);
204-
assertEquals("You are not enabled to use index insight yet, please firstly enable it.", argumentCaptor.getValue().getMessage());
205-
}
206-
207175
@Test
208176
public void testGetIndexInsight_FailToAccess() {
209177
GetResponse getResponse = mock(GetResponse.class);
@@ -246,48 +214,6 @@ public void testGetIndexInsight_FailToAccess() {
246214
assertEquals("no permissions", argumentCaptor.getValue().getMessage());
247215
}
248216

249-
@Test
250-
public void testGetIndexInsight_ContainerNotInitialized() {
251-
GetResponse getResponse = mock(GetResponse.class);
252-
when(getResponse.isExists()).thenReturn(false);
253-
254-
GetDataObjectResponse sdkResponse = mock(GetDataObjectResponse.class);
255-
when(sdkResponse.getResponse()).thenReturn(getResponse);
256-
257-
CompletableFuture<GetDataObjectResponse> future = CompletableFuture.completedFuture(sdkResponse);
258-
259-
when(sdkClient.getDataObjectAsync(any())).thenReturn(future);
260-
IndexInsightTask indexInsightTask = mock(IndexInsightTask.class);
261-
262-
doReturn(indexInsightTask).when(getIndexInsightTransportAction).createTask(any());
263-
IndexInsight insight = new IndexInsight(
264-
"test_index",
265-
"test content",
266-
IndexInsightTaskStatus.COMPLETED,
267-
STATISTICAL_DATA,
268-
Instant.ofEpochMilli(0),
269-
""
270-
);
271-
doAnswer(invocation -> {
272-
ActionListener<IndexInsight> listener = invocation.getArgument(1);
273-
listener.onResponse(insight);
274-
return null;
275-
}).when(indexInsightTask).execute(any(), any());
276-
277-
SearchResponse searchResponse = mock(SearchResponse.class);
278-
doAnswer(invocation -> {
279-
ActionListener<SearchResponse> listener = invocation.getArgument(1);
280-
listener.onResponse(searchResponse);
281-
return null;
282-
}).when(client).search(any(), any());
283-
284-
getIndexInsightTransportAction.doExecute(null, mlIndexInsightGetRequest, actionListener);
285-
ArgumentCaptor<Exception> argumentCaptor = ArgumentCaptor.forClass(Exception.class);
286-
verify(actionListener).onFailure(argumentCaptor.capture());
287-
assertTrue(argumentCaptor.getValue() instanceof RuntimeException);
288-
assertEquals("You are not enabled to use index insight yet, please firstly enable it.", argumentCaptor.getValue().getMessage());
289-
}
290-
291217
@Test
292218
public void testCreateTask() {
293219
MLIndexInsightGetRequest statisticalRequest = new MLIndexInsightGetRequest("test_index", STATISTICAL_DATA, null);

‎plugin/src/test/java/org/opensearch/ml/rest/RestMLGetIndexInsightConfigIT.java‎

Lines changed: 0 additions & 78 deletions
This file was deleted.

‎plugin/src/test/java/org/opensearch/ml/rest/RestMLGetIndexInsightIT.java‎

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -33,16 +33,6 @@ public void testGetIndexInsight_Success() throws IOException, ParseException, In
3333
"is_enable": true
3434
}
3535
""";
36-
Response response = TestHelper
37-
.makeRequest(
38-
client(),
39-
"PUT",
40-
"/_plugins/_ml/index_insight_config",
41-
null,
42-
TestHelper.toHttpEntity(registerAgentRequestBody),
43-
null
44-
);
45-
assertNotNull(response);
4636
createIndexAndPutDoc();
4737

4838
Response responseIndexInsight = TestHelper
@@ -78,16 +68,6 @@ public void testGetIndexInsightWithPattern_Success() throws IOException, ParseEx
7868
"is_enable": true
7969
}
8070
""";
81-
Response response = TestHelper
82-
.makeRequest(
83-
client(),
84-
"PUT",
85-
"/_plugins/_ml/index_insight_config",
86-
null,
87-
TestHelper.toHttpEntity(registerAgentRequestBody),
88-
null
89-
);
90-
assertNotNull(response);
9171
createIndexAndDocWithPattern();
9272

9373
Response responseIndexInsight = TestHelper

0 commit comments

Comments
 (0)