Skip to content

Commit 9e20a01

Browse files
akolarkunnuakolarkunnu
committed
[Enhancement] Enhance validation for create connector API
This change will address the second part of validation "pre and post processing function validation". Partially resolves opensearch-project#2993 Signed-off-by: Abdul Muneer Kolarkunnu <muneer.kolarkunnu@netapp.com>
1 parent 7e08f9e commit 9e20a01

File tree

2 files changed

+34
-18
lines changed

2 files changed

+34
-18
lines changed

‎common/src/main/java/org/opensearch/ml/common/connector/ConnectorAction.java

+11-3
Original file line numberDiff line numberDiff line change
@@ -212,9 +212,14 @@ public static ConnectorAction parse(XContentParser parser) throws IOException {
212212
}
213213

214214
public void validatePrePostProcessFunctions(Map<String, String> parameters) {
215-
var substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
216-
var endPoint = substitutor.replace(url);
217-
var remoteServer = getRemoteServerFromURL(endPoint);
215+
StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
216+
String endPoint = substitutor.replace(url);
217+
String remoteServer = getRemoteServerFromURL(endPoint);
218+
validatePreProcessFunctions(remoteServer);
219+
validatePostProcessFunctions(remoteServer);
220+
}
221+
222+
private void validatePreProcessFunctions(String remoteServer) {
218223
if (isInBuiltFunction(preProcessFunction)) {
219224
switch (remoteServer) {
220225
case OPENAI:
@@ -270,6 +275,9 @@ public void validatePrePostProcessFunctions(Map<String, String> parameters) {
270275
}
271276
}
272277
}
278+
}
279+
280+
private void validatePostProcessFunctions(String remoteServer) {
273281
if (isInBuiltFunction(postProcessFunction)) {
274282
switch (remoteServer) {
275283
case OPENAI:

‎common/src/test/java/org/opensearch/ml/common/connector/ConnectorActionTest.java

+23-15
Original file line numberDiff line numberDiff line change
@@ -87,23 +87,31 @@ public void constructor_NullMethod() {
8787

8888
@Test
8989
public void connectorWithNullPreProcessFunction() {
90-
var action = new ConnectorAction(TEST_ACTION_TYPE, TEST_METHOD_HTTP, OPENAI_URL, null, TEST_REQUEST_BODY, null, null);
90+
ConnectorAction action = new ConnectorAction(TEST_ACTION_TYPE, TEST_METHOD_HTTP, OPENAI_URL, null, TEST_REQUEST_BODY, null, null);
9191
action.validatePrePostProcessFunctions(Map.of());
9292
assertNotNull(action);
9393
}
9494

9595
@Test
9696
public void connectorWithCustomPainlessScriptPreProcessFunction() {
97-
var preProcessFunction =
97+
String preProcessFunction =
9898
"\"\\n StringBuilder builder = new StringBuilder();\\n builder.append(\\\"\\\\\\\"\\\");\\n String first = params.text_docs[0];\\n builder.append(first);\\n builder.append(\\\"\\\\\\\"\\\");\\n def parameters = \\\"{\\\" +\\\"\\\\\\\"text_inputs\\\\\\\":\\\" + builder + \\\"}\\\";\\n return \\\"{\\\" +\\\"\\\\\\\"parameters\\\\\\\":\\\" + parameters + \\\"}\\\";\"";
99-
var action = new ConnectorAction(TEST_ACTION_TYPE, TEST_METHOD_HTTP, OPENAI_URL, null, TEST_REQUEST_BODY, preProcessFunction, null);
99+
ConnectorAction action = new ConnectorAction(
100+
TEST_ACTION_TYPE,
101+
TEST_METHOD_HTTP,
102+
OPENAI_URL,
103+
null,
104+
TEST_REQUEST_BODY,
105+
preProcessFunction,
106+
null
107+
);
100108
action.validatePrePostProcessFunctions(null);
101109
assertNotNull(action);
102110
}
103111

104112
@Test
105113
public void openAIConnectorWithCorrectInBuiltPrePostProcessFunction() {
106-
var action = new ConnectorAction(
114+
ConnectorAction action = new ConnectorAction(
107115
TEST_ACTION_TYPE,
108116
TEST_METHOD_HTTP,
109117
"https://${parameters.endpoint}/v1/chat/completions",
@@ -118,7 +126,7 @@ public void openAIConnectorWithCorrectInBuiltPrePostProcessFunction() {
118126

119127
@Test
120128
public void openAIConnectorWithWrongInBuiltPrePostProcessFunction() {
121-
var action1 = new ConnectorAction(
129+
ConnectorAction action1 = new ConnectorAction(
122130
TEST_ACTION_TYPE,
123131
TEST_METHOD_HTTP,
124132
OPENAI_URL,
@@ -132,7 +140,7 @@ public void openAIConnectorWithWrongInBuiltPrePostProcessFunction() {
132140
"LLM service is openai, so PreProcessFunction should be connector.pre_process.openai.embedding",
133141
exception.getMessage()
134142
);
135-
var action2 = new ConnectorAction(
143+
ConnectorAction action2 = new ConnectorAction(
136144
TEST_ACTION_TYPE,
137145
TEST_METHOD_HTTP,
138146
OPENAI_URL,
@@ -150,7 +158,7 @@ public void openAIConnectorWithWrongInBuiltPrePostProcessFunction() {
150158

151159
@Test
152160
public void cohereConnectorWithCorrectInBuiltPrePostProcessFunction() {
153-
var action = new ConnectorAction(
161+
ConnectorAction action = new ConnectorAction(
154162
TEST_ACTION_TYPE,
155163
TEST_METHOD_HTTP,
156164
COHERE_URL,
@@ -187,7 +195,7 @@ public void cohereConnectorWithCorrectInBuiltPrePostProcessFunction() {
187195

188196
@Test
189197
public void cohereConnectorWithWrongInBuiltPrePostProcessFunction() {
190-
var action1 = new ConnectorAction(
198+
ConnectorAction action1 = new ConnectorAction(
191199
TEST_ACTION_TYPE,
192200
TEST_METHOD_HTTP,
193201
COHERE_URL,
@@ -202,7 +210,7 @@ public void cohereConnectorWithWrongInBuiltPrePostProcessFunction() {
202210
+ " or connector.pre_process.cohere.multimodal_embedding or connector.pre_process.cohere.rerank",
203211
exception.getMessage()
204212
);
205-
var action2 = new ConnectorAction(
213+
ConnectorAction action2 = new ConnectorAction(
206214
TEST_ACTION_TYPE,
207215
TEST_METHOD_HTTP,
208216
COHERE_URL,
@@ -221,7 +229,7 @@ public void cohereConnectorWithWrongInBuiltPrePostProcessFunction() {
221229

222230
@Test
223231
public void bedrockConnectorWithCorrectInBuiltPrePostProcessFunction() {
224-
var action = new ConnectorAction(
232+
ConnectorAction action = new ConnectorAction(
225233
TEST_ACTION_TYPE,
226234
TEST_METHOD_HTTP,
227235
BEDROCK_URL,
@@ -258,7 +266,7 @@ public void bedrockConnectorWithCorrectInBuiltPrePostProcessFunction() {
258266

259267
@Test
260268
public void bedrockConnectorWithWrongInBuiltPrePostProcessFunction() {
261-
var action1 = new ConnectorAction(
269+
ConnectorAction action1 = new ConnectorAction(
262270
TEST_ACTION_TYPE,
263271
TEST_METHOD_HTTP,
264272
BEDROCK_URL,
@@ -273,7 +281,7 @@ public void bedrockConnectorWithWrongInBuiltPrePostProcessFunction() {
273281
+ " or connector.pre_process.bedrock.multimodal_embedding or connector.pre_process.bedrock.rerank",
274282
exception.getMessage()
275283
);
276-
var action2 = new ConnectorAction(
284+
ConnectorAction action2 = new ConnectorAction(
277285
TEST_ACTION_TYPE,
278286
TEST_METHOD_HTTP,
279287
BEDROCK_URL,
@@ -292,7 +300,7 @@ public void bedrockConnectorWithWrongInBuiltPrePostProcessFunction() {
292300

293301
@Test
294302
public void sagemakerConnectorWithCorrectInBuiltPrePostProcessFunction() {
295-
var action = new ConnectorAction(
303+
ConnectorAction action = new ConnectorAction(
296304
TEST_ACTION_TYPE,
297305
TEST_METHOD_HTTP,
298306
SAGEMAKER_URL,
@@ -318,7 +326,7 @@ public void sagemakerConnectorWithCorrectInBuiltPrePostProcessFunction() {
318326

319327
@Test
320328
public void sagemakerConnectorWithWrongInBuiltPrePostProcessFunction() {
321-
var action1 = new ConnectorAction(
329+
ConnectorAction action1 = new ConnectorAction(
322330
TEST_ACTION_TYPE,
323331
TEST_METHOD_HTTP,
324332
SAGEMAKER_URL,
@@ -333,7 +341,7 @@ public void sagemakerConnectorWithWrongInBuiltPrePostProcessFunction() {
333341
+ " or connector.pre_process.default.rerank",
334342
exception.getMessage()
335343
);
336-
var action2 = new ConnectorAction(
344+
ConnectorAction action2 = new ConnectorAction(
337345
TEST_ACTION_TYPE,
338346
TEST_METHOD_HTTP,
339347
SAGEMAKER_URL,

0 commit comments

Comments
 (0)