Skip to content

Commit 0b65cc9

Browse files
rcho93rcho93
committed
spotless applied
Signed-off-by: seungwon cho <[email protected]>
1 parent 9fdc5cd commit 0b65cc9

File tree

10 files changed

+105
-96
lines changed

10 files changed

+105
-96
lines changed

‎common/src/main/java/org/opensearch/ml/common/transport/prompt/MLImportPromptInput.java

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,15 @@ public class MLImportPromptInput implements ToXContentObject, Writeable {
3939
private String tenantId;
4040

4141
@Builder(toBuilder = true)
42-
public MLImportPromptInput(String name, String tag, @NonNull String promptManagementType, @NonNull String publicKey, @NonNull String accessKey, String limit, String tenantId) {
42+
public MLImportPromptInput(
43+
String name,
44+
String tag,
45+
@NonNull String promptManagementType,
46+
@NonNull String publicKey,
47+
@NonNull String accessKey,
48+
String limit,
49+
String tenantId
50+
) {
4351
Objects.requireNonNull(promptManagementType, "must specify prompt management type");
4452
Objects.requireNonNull(publicKey, "public key can not be null");
4553
Objects.requireNonNull(accessKey, "access key can not be null");
@@ -113,15 +121,15 @@ public static MLImportPromptInput parse(XContentParser parser) throws IOExceptio
113121
}
114122
}
115123
return MLImportPromptInput
116-
.builder()
117-
.name(name)
118-
.tag(tag)
119-
.promptManagementType(promptManagementType)
120-
.publicKey(publicKey)
121-
.accessKey(accessKey)
122-
.limit(limit)
123-
.tenantId(tenantId)
124-
.build();
124+
.builder()
125+
.name(name)
126+
.tag(tag)
127+
.promptManagementType(promptManagementType)
128+
.publicKey(publicKey)
129+
.accessKey(accessKey)
130+
.limit(limit)
131+
.tenantId(tenantId)
132+
.build();
125133
}
126134

127135
@Override

‎common/src/main/java/org/opensearch/ml/common/transport/prompt/MLUpdatePromptInput.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@
2121
import org.opensearch.core.xcontent.ToXContentObject;
2222
import org.opensearch.core.xcontent.XContentBuilder;
2323
import org.opensearch.core.xcontent.XContentParser;
24+
import org.opensearch.ml.common.prompt.PromptExtraConfig;
2425

2526
import lombok.Builder;
2627
import lombok.Data;
2728
import lombok.Setter;
28-
import org.opensearch.ml.common.prompt.PromptExtraConfig;
2929

3030
/**
3131
* MLUpdatePromptInput is the input class for MLUpdatePromptAction. It contains the parameters

‎plugin/src/main/java/org/opensearch/ml/action/prompt/ImportPromptTransportAction.java

Lines changed: 25 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -89,30 +89,27 @@ protected void doExecute(Task task, MLImportPromptRequest mlImportPromptRequest,
8989
String accessKey = mlImportPromptInput.getAccessKey();
9090
String tenantId = mlImportPromptInput.getTenantId();
9191

92-
// TODO: pass in PromptManagementType inside request body
9392
try {
9493
AbstractPromptManagement promptManagement = init(
9594
promptManagementType,
9695
PromptExtraConfig.builder().publicKey(publicKey).accessKey(accessKey).build()
9796
);
9897
List<MLPrompt> mlPromptList = promptManagement.importPrompts(mlImportPromptInput);
9998
Map<String, String> responseBody = new HashMap<>();
99+
if (mlPromptList.isEmpty()) {
100+
listener.onResponse(new MLImportPromptResponse(responseBody));
101+
return;
102+
}
100103
AtomicInteger remainingMLPrompts = new AtomicInteger(mlPromptList.size());
101104
for (MLPrompt mlPrompt : mlPromptList) {
102105
mlPrompt.encrypt(promptManagementType, mlEngine::encrypt, tenantId);
103-
handleDuplicateName(
104-
mlPrompt,
105-
tenantId,
106-
ActionListener.wrap(
107-
promptId ->
108-
{
109-
if (promptId == null) {
110-
indexPrompt(mlPrompt, responseBody, remainingMLPrompts, listener);
111-
} else {
112-
updateImportResponseBody(promptId, mlPrompt.getName(), responseBody, remainingMLPrompts, listener);
113-
}
114-
}, listener::onFailure
115-
));
106+
handleDuplicateName(mlPrompt, tenantId, ActionListener.wrap(promptId -> {
107+
if (promptId == null) {
108+
indexPrompt(mlPrompt, responseBody, remainingMLPrompts, listener);
109+
} else {
110+
updateImportResponseBody(promptId, mlPrompt.getName(), responseBody, remainingMLPrompts, listener);
111+
}
112+
}, listener::onFailure));
116113
}
117114
} catch (Exception e) {
118115
handleFailure(e, null, listener, "Failed to import " + promptManagementType + " Prompts into System Index");
@@ -184,17 +181,17 @@ private void updateImportResponseBody(
184181

185182
private void handleDuplicateName(MLPrompt importingPrompt, String tenantId, ActionListener<String> wrappedListener) throws IOException {
186183
String name = importingPrompt.getName();
187-
SearchResponse searchResponse = mlPromptManager.validateUniquePromptName(name, tenantId);
184+
SearchResponse searchResponse = mlPromptManager.searchPromptByName(name, tenantId);
188185
if (searchResponse != null
189-
&& searchResponse.getHits().getTotalHits() != null
190-
&& searchResponse.getHits().getTotalHits().value() != 0) {
186+
&& searchResponse.getHits().getTotalHits() != null
187+
&& searchResponse.getHits().getTotalHits().value() != 0) {
191188
String promptId = searchResponse.getHits().getAt(0).getId();
192189
GetDataObjectRequest getDataObjectRequest = GetDataObjectRequest
193-
.builder()
194-
.index(ML_PROMPT_INDEX)
195-
.id(promptId)
196-
.tenantId(tenantId)
197-
.build();
190+
.builder()
191+
.index(ML_PROMPT_INDEX)
192+
.id(promptId)
193+
.tenantId(tenantId)
194+
.build();
198195
MLPrompt existingMLPrompt = mlPromptManager.getPrompt(getDataObjectRequest);
199196

200197
// check the prompt management type
@@ -203,18 +200,17 @@ private void handleDuplicateName(MLPrompt importingPrompt, String tenantId, Acti
203200
throw new IllegalArgumentException("Provided name: " + name + " is already being used by ML Prompt with id: " + promptId);
204201
} else if (promptManagementType.equalsIgnoreCase(LANGFUSE)) {
205202
// update the existing langfuse prompt with new content if the version des not match
206-
UpdateDataObjectRequest updateDataObjectRequest = UpdateDataObjectRequest.builder()
203+
UpdateDataObjectRequest updateDataObjectRequest = UpdateDataObjectRequest
204+
.builder()
207205
.index(ML_PROMPT_INDEX)
208206
.id(promptId)
209207
.tenantId(tenantId)
210208
.dataObject(importingPrompt)
211209
.build();
212-
mlPromptManager.updatePromptIndex(updateDataObjectRequest, promptId, ActionListener.wrap(
213-
updateResponse -> {
214-
log.info("{} Prompt with promptId: {} updated successfully", promptManagementType, promptId);
215-
wrappedListener.onResponse(promptId);
216-
}, wrappedListener::onFailure
217-
));
210+
mlPromptManager.updatePromptIndex(updateDataObjectRequest, promptId, ActionListener.wrap(updateResponse -> {
211+
log.info("{} Prompt with promptId: {} updated successfully", promptManagementType, promptId);
212+
wrappedListener.onResponse(promptId);
213+
}, wrappedListener::onFailure));
218214
}
219215
} else {
220216
// provided name is unique, good to be imported

‎plugin/src/main/java/org/opensearch/ml/action/prompt/UpdatePromptTransportAction.java

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,38 +7,31 @@
77

88
import static org.opensearch.ml.common.CommonValue.ML_PROMPT_INDEX;
99
import static org.opensearch.ml.prompt.AbstractPromptManagement.init;
10-
import static org.opensearch.ml.prompt.MLPromptManager.MLPromptNameAlreadyExists;
1110
import static org.opensearch.ml.prompt.MLPromptManager.TAG_RESTRICTION_ERR_MESSAGE;
1211
import static org.opensearch.ml.prompt.MLPromptManager.UNIQUE_NAME_ERR_MESSAGE;
1312
import static org.opensearch.ml.prompt.MLPromptManager.handleFailure;
1413
import static org.opensearch.ml.prompt.MLPromptManager.validateTags;
1514

16-
import java.time.Instant;
17-
1815
import org.apache.commons.lang3.StringUtils;
1916
import org.opensearch.action.search.SearchResponse;
2017
import org.opensearch.action.support.ActionFilters;
2118
import org.opensearch.action.support.HandledTransportAction;
2219
import org.opensearch.action.update.UpdateResponse;
2320
import org.opensearch.common.inject.Inject;
24-
import org.opensearch.common.util.concurrent.ThreadContext;
2521
import org.opensearch.core.action.ActionListener;
2622
import org.opensearch.ml.common.prompt.MLPrompt;
2723
import org.opensearch.ml.common.prompt.PromptExtraConfig;
2824
import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
2925
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptAction;
3026
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptInput;
3127
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptRequest;
32-
import org.opensearch.ml.engine.encryptor.Encryptor;
3328
import org.opensearch.ml.engine.encryptor.EncryptorImpl;
3429
import org.opensearch.ml.prompt.AbstractPromptManagement;
3530
import org.opensearch.ml.prompt.MLPromptManager;
3631
import org.opensearch.ml.utils.TenantAwareHelper;
3732
import org.opensearch.remote.metadata.client.GetDataObjectRequest;
3833
import org.opensearch.remote.metadata.client.SdkClient;
3934
import org.opensearch.remote.metadata.client.UpdateDataObjectRequest;
40-
import org.opensearch.remote.metadata.client.UpdateDataObjectResponse;
41-
import org.opensearch.remote.metadata.common.SdkClientUtils;
4235
import org.opensearch.search.SearchHit;
4336
import org.opensearch.tasks.Task;
4437
import org.opensearch.transport.TransportService;
@@ -110,10 +103,10 @@ protected void doExecute(Task task, MLUpdatePromptRequest mlUpdatePromptRequest,
110103
try {
111104
if (StringUtils.isNotBlank(mlUpdatePromptInput.getName())) {
112105
SearchResponse searchResponse = mlPromptManager
113-
.validateUniquePromptName(mlUpdatePromptInput.getName(), mlUpdatePromptInput.getTenantId());
106+
.searchPromptByName(mlUpdatePromptInput.getName(), mlUpdatePromptInput.getTenantId());
114107
if (searchResponse != null
115-
&& searchResponse.getHits().getTotalHits() != null
116-
&& searchResponse.getHits().getTotalHits().value() != 0) {
108+
&& searchResponse.getHits().getTotalHits() != null
109+
&& searchResponse.getHits().getTotalHits().value() != 0) {
117110
SearchHit hit = searchResponse.getHits().getAt(0);
118111
String id = hit.getId();
119112
actionListener.onFailure(new IllegalArgumentException(UNIQUE_NAME_ERR_MESSAGE + id));
@@ -129,10 +122,11 @@ protected void doExecute(Task task, MLUpdatePromptRequest mlUpdatePromptRequest,
129122
.getPromptAsync(
130123
getDataObjectRequest,
131124
promptId,
132-
ActionListener.wrap(
133-
mlPrompt -> handleGetPrompt(mlPrompt, mlUpdatePromptInput, promptId, tenantId, actionListener),
134-
e -> handleFailure(e, promptId, actionListener, "Failed to get ML Prompt {}")
135-
)
125+
ActionListener
126+
.wrap(
127+
mlPrompt -> handleGetPrompt(mlPrompt, mlUpdatePromptInput, promptId, tenantId, actionListener),
128+
e -> handleFailure(e, promptId, actionListener, "Failed to get ML Prompt {}")
129+
)
136130
);
137131
} catch (Exception exception) {
138132
handleFailure(exception, promptId, actionListener, "Failed to search ML Prompt Index");

‎plugin/src/main/java/org/opensearch/ml/prompt/AbstractPromptManagement.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@
1515
import org.opensearch.ml.common.prompt.PromptExtraConfig;
1616
import org.opensearch.ml.common.transport.prompt.MLCreatePromptInput;
1717
import org.opensearch.ml.common.transport.prompt.MLImportPromptInput;
18-
19-
import lombok.Getter;
2018
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptInput;
2119
import org.opensearch.remote.metadata.client.UpdateDataObjectRequest;
2220

21+
import lombok.Getter;
22+
2323
@Getter
2424
public abstract class AbstractPromptManagement implements ToXContentObject {
2525
public static AbstractPromptManagement init(String promptManagementType, PromptExtraConfig extraConfig) {

‎plugin/src/main/java/org/opensearch/ml/prompt/LangfusePromptManagement.java

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
1010
import static org.opensearch.ml.common.CommonValue.ML_PROMPT_INDEX;
1111
import static org.opensearch.ml.common.prompt.MLPrompt.LANGFUSE;
12-
import static org.opensearch.ml.common.prompt.MLPrompt.MLPROMPT;
1312
import static org.opensearch.ml.prompt.MLPromptManagement.INITIAL_VERSION;
1413

1514
import java.io.IOException;
@@ -21,7 +20,6 @@
2120
import java.util.Map;
2221

2322
import org.opensearch.OpenSearchStatusException;
24-
import org.opensearch.action.search.SearchResponse;
2523
import org.opensearch.common.xcontent.LoggingDeprecationHandler;
2624
import org.opensearch.core.rest.RestStatus;
2725
import org.opensearch.core.xcontent.NamedXContentRegistry;
@@ -31,6 +29,8 @@
3129
import org.opensearch.ml.common.prompt.PromptExtraConfig;
3230
import org.opensearch.ml.common.transport.prompt.MLCreatePromptInput;
3331
import org.opensearch.ml.common.transport.prompt.MLImportPromptInput;
32+
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptInput;
33+
import org.opensearch.remote.metadata.client.UpdateDataObjectRequest;
3434

3535
import com.langfuse.client.LangfuseClient;
3636
import com.langfuse.client.core.LangfuseClientApiException;
@@ -48,9 +48,6 @@
4848

4949
import lombok.Getter;
5050
import lombok.extern.log4j.Log4j2;
51-
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptInput;
52-
import org.opensearch.remote.metadata.client.GetDataObjectRequest;
53-
import org.opensearch.remote.metadata.client.UpdateDataObjectRequest;
5451

5552
@Log4j2
5653
@Getter
@@ -295,8 +292,8 @@ public List<MLPrompt> importPrompts(MLImportPromptInput mlImportPromptInput) {
295292
}
296293

297294
PromptMetaListResponse promptMetaListResponse = langfuseClient
298-
.prompts()
299-
.list(ListPromptsMetaRequest.builder().tag(tag).limit(Integer.parseInt(limit)).build());
295+
.prompts()
296+
.list(ListPromptsMetaRequest.builder().tag(tag).limit(Integer.parseInt(limit)).build());
300297

301298
List<PromptMeta> promptMetas = promptMetaListResponse.getData();
302299
List<MLPrompt> mlPromptList = new ArrayList<>();
@@ -335,12 +332,12 @@ public List<MLPrompt> importPrompts(MLImportPromptInput mlImportPromptInput) {
335332
public UpdateDataObjectRequest updatePrompt(MLUpdatePromptInput mlUpdatePromptInput, MLPrompt mlPrompt) {
336333
getPrompt(mlPrompt);
337334
MLCreatePromptInput updateContent = MLCreatePromptInput
338-
.builder()
339-
.name(mlPrompt.getName())
340-
.tags(mlPrompt.getTags())
341-
.prompt(mlPrompt.getPrompt())
342-
.promptExtraConfig(mlPrompt.getPromptExtraConfig())
343-
.build();
335+
.builder()
336+
.name(mlPrompt.getName())
337+
.tags(mlPrompt.getTags())
338+
.prompt(mlPrompt.getPrompt())
339+
.promptExtraConfig(mlPrompt.getPromptExtraConfig())
340+
.build();
344341

345342
// Langfuse does not allow users to change prompt's name
346343
if (mlUpdatePromptInput.getTags() != null) {

‎plugin/src/main/java/org/opensearch/ml/prompt/MLPromptManagement.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
package org.opensearch.ml.prompt;
77

8+
import static org.opensearch.ml.common.CommonValue.ML_PROMPT_INDEX;
9+
810
import java.io.IOException;
911
import java.time.Instant;
1012
import java.util.List;
@@ -15,12 +17,10 @@
1517
import org.opensearch.ml.common.prompt.MLPrompt;
1618
import org.opensearch.ml.common.transport.prompt.MLCreatePromptInput;
1719
import org.opensearch.ml.common.transport.prompt.MLImportPromptInput;
18-
19-
import lombok.extern.log4j.Log4j2;
2020
import org.opensearch.ml.common.transport.prompt.MLUpdatePromptInput;
2121
import org.opensearch.remote.metadata.client.UpdateDataObjectRequest;
2222

23-
import static org.opensearch.ml.common.CommonValue.ML_PROMPT_INDEX;
23+
import lombok.extern.log4j.Log4j2;
2424

2525
@Log4j2
2626
public class MLPromptManagement extends AbstractPromptManagement {

0 commit comments

Comments
 (0)