Skip to content

Commit bb0549e

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 4e5455c commit bb0549e

File tree

2 files changed

+39
-36
lines changed

2 files changed

+39
-36
lines changed

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

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -198,18 +198,42 @@ public static ConnectorAction parse(XContentParser parser) throws IOException {
198198
.build();
199199
}
200200

201+
/**
202+
* Validates whether pre and post process functions corresponding to the same llm service or not.
203+
* There are specific pre and post process functions defined for each end llm services, so if you are
204+
* configuring the pre-built functions it has to be from the corresponding list. This method throws
205+
* IllegalArgumentException exception if it is configured wrongly.
206+
*
207+
* @param parameters - connector parameters
208+
*/
201209
public void validatePrePostProcessFunctions(Map<String, String> parameters) {
202210
StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
203211
String endPoint = substitutor.replace(url);
204212
String remoteServer = getRemoteServerFromURL(endPoint);
205-
validatePreProcessFunctions(remoteServer);
206-
validatePostProcessFunctions(remoteServer);
213+
if (isInBuiltProcessFunction(preProcessFunction)) {
214+
validatePreProcessFunctions(remoteServer);
215+
}
216+
if (isInBuiltProcessFunction(postProcessFunction)) {
217+
validatePostProcessFunctions(remoteServer);
218+
}
219+
}
220+
221+
/**
222+
* To get the remote server name from ULR
223+
*
224+
* @param url - remote server url
225+
* @return - returns the corresponding remote server name for url, if server is not in the pre-defined list,
226+
* it returns null
227+
*/
228+
public static String getRemoteServerFromURL(String url) {
229+
return SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES.stream().filter(url::contains).findFirst().orElse("");
230+
}
231+
232+
private boolean isInBuiltProcessFunction(String processFunction) {
233+
return (processFunction != null && processFunction.startsWith(INBUILT_FUNC_PREFIX));
207234
}
208235

209236
private void validatePreProcessFunctions(String remoteServer) {
210-
if (!isInBuiltProcessFunction(preProcessFunction)) {
211-
return;
212-
}
213237
switch (remoteServer) {
214238
case OPENAI:
215239
if (!preProcessFunction.contains(OPENAI)) {
@@ -242,9 +266,6 @@ private void validatePreProcessFunctions(String remoteServer) {
242266
}
243267

244268
private void validatePostProcessFunctions(String remoteServer) {
245-
if (!isInBuiltProcessFunction(postProcessFunction)) {
246-
return;
247-
}
248269
switch (remoteServer) {
249270
case OPENAI:
250271
if (!postProcessFunction.contains(OPENAI)) {
@@ -274,14 +295,6 @@ private String invalidProcessFuncExcText(String remoteServer, String func) {
274295
return "LLM service is " + remoteServer + ", so " + func + " should be " + remoteServer + " " + func;
275296
}
276297

277-
private boolean isInBuiltProcessFunction(String processFunction) {
278-
return (processFunction != null && processFunction.startsWith(INBUILT_FUNC_PREFIX));
279-
}
280-
281-
public static String getRemoteServerFromURL(String url) {
282-
return SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES.stream().filter(url::contains).findFirst().orElse("");
283-
}
284-
285298
public enum ActionType {
286299
PREDICT,
287300
EXECUTE,

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

Lines changed: 10 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -86,14 +86,13 @@ public void constructor_NullMethod() {
8686
}
8787

8888
@Test
89-
public void connectorWithNullPreProcessFunction() {
89+
public void testValidatePrePostProcessFunctionsWithNullPreProcessFunctionSuccess() {
9090
ConnectorAction action = new ConnectorAction(TEST_ACTION_TYPE, TEST_METHOD_HTTP, OPENAI_URL, null, TEST_REQUEST_BODY, null, null);
9191
action.validatePrePostProcessFunctions(Map.of());
92-
assertNotNull(action);
9392
}
9493

9594
@Test
96-
public void connectorWithCustomPainlessScriptPreProcessFunction() {
95+
public void testValidatePrePostProcessFunctionsWithCustomPainlessScriptPreProcessFunctionSuccess() {
9796
String preProcessFunction =
9897
"\"\\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 + \\\"}\\\";\"";
9998
ConnectorAction action = new ConnectorAction(
@@ -106,11 +105,10 @@ public void connectorWithCustomPainlessScriptPreProcessFunction() {
106105
null
107106
);
108107
action.validatePrePostProcessFunctions(null);
109-
assertNotNull(action);
110108
}
111109

112110
@Test
113-
public void openAIConnectorWithCorrectInBuiltPrePostProcessFunction() {
111+
public void testValidatePrePostProcessFunctionsWithOpenAIConnectorCorrectInBuiltPrePostProcessFunctionAndParamsSuccess() {
114112
ConnectorAction action = new ConnectorAction(
115113
TEST_ACTION_TYPE,
116114
TEST_METHOD_HTTP,
@@ -121,11 +119,10 @@ public void openAIConnectorWithCorrectInBuiltPrePostProcessFunction() {
121119
OPENAI_EMBEDDING
122120
);
123121
action.validatePrePostProcessFunctions(Map.of("endpoint", "api.openai.com"));
124-
assertNotNull(action);
125122
}
126123

127124
@Test
128-
public void openAIConnectorWithWrongInBuiltPrePostProcessFunction() {
125+
public void testValidatePrePostProcessFunctionsWithOpenAIConnectorWrongInBuiltPrePostProcessFunctionThrowsException() {
129126
ConnectorAction action1 = new ConnectorAction(
130127
TEST_ACTION_TYPE,
131128
TEST_METHOD_HTTP,
@@ -151,7 +148,7 @@ public void openAIConnectorWithWrongInBuiltPrePostProcessFunction() {
151148
}
152149

153150
@Test
154-
public void cohereConnectorWithCorrectInBuiltPrePostProcessFunction() {
151+
public void testValidatePrePostProcessFunctionsWithCohereConnectorCorrectInBuiltPrePostProcessFunctionSuccess() {
155152
ConnectorAction action = new ConnectorAction(
156153
TEST_ACTION_TYPE,
157154
TEST_METHOD_HTTP,
@@ -162,7 +159,6 @@ public void cohereConnectorWithCorrectInBuiltPrePostProcessFunction() {
162159
COHERE_EMBEDDING
163160
);
164161
action.validatePrePostProcessFunctions(Map.of());
165-
assertNotNull(action);
166162
action = new ConnectorAction(
167163
TEST_ACTION_TYPE,
168164
TEST_METHOD_HTTP,
@@ -173,7 +169,6 @@ public void cohereConnectorWithCorrectInBuiltPrePostProcessFunction() {
173169
COHERE_EMBEDDING
174170
);
175171
action.validatePrePostProcessFunctions(Map.of());
176-
assertNotNull(action);
177172
action = new ConnectorAction(
178173
TEST_ACTION_TYPE,
179174
TEST_METHOD_HTTP,
@@ -188,7 +183,7 @@ public void cohereConnectorWithCorrectInBuiltPrePostProcessFunction() {
188183
}
189184

190185
@Test
191-
public void cohereConnectorWithWrongInBuiltPrePostProcessFunction() {
186+
public void testValidatePrePostProcessFunctionsWithCohereConnectorWrongInBuiltPrePostProcessFunctionThrowsException() {
192187
ConnectorAction action1 = new ConnectorAction(
193188
TEST_ACTION_TYPE,
194189
TEST_METHOD_HTTP,
@@ -214,7 +209,7 @@ public void cohereConnectorWithWrongInBuiltPrePostProcessFunction() {
214209
}
215210

216211
@Test
217-
public void bedrockConnectorWithCorrectInBuiltPrePostProcessFunction() {
212+
public void testValidatePrePostProcessFunctionsWithBedrockConnectorCorrectInBuiltPrePostProcessFunctionSuccess() {
218213
ConnectorAction action = new ConnectorAction(
219214
TEST_ACTION_TYPE,
220215
TEST_METHOD_HTTP,
@@ -225,7 +220,6 @@ public void bedrockConnectorWithCorrectInBuiltPrePostProcessFunction() {
225220
BEDROCK_EMBEDDING
226221
);
227222
action.validatePrePostProcessFunctions(Map.of());
228-
assertNotNull(action);
229223
action = new ConnectorAction(
230224
TEST_ACTION_TYPE,
231225
TEST_METHOD_HTTP,
@@ -236,7 +230,6 @@ public void bedrockConnectorWithCorrectInBuiltPrePostProcessFunction() {
236230
BEDROCK_BATCH_JOB_ARN
237231
);
238232
action.validatePrePostProcessFunctions(Map.of());
239-
assertNotNull(action);
240233
action = new ConnectorAction(
241234
TEST_ACTION_TYPE,
242235
TEST_METHOD_HTTP,
@@ -247,11 +240,10 @@ public void bedrockConnectorWithCorrectInBuiltPrePostProcessFunction() {
247240
BEDROCK_RERANK
248241
);
249242
action.validatePrePostProcessFunctions(Map.of());
250-
assertNotNull(action);
251243
}
252244

253245
@Test
254-
public void bedrockConnectorWithWrongInBuiltPrePostProcessFunction() {
246+
public void testValidatePrePostProcessFunctionsWithBedrockConnectorWrongInBuiltPrePostProcessFunctionThrowsException() {
255247
ConnectorAction action1 = new ConnectorAction(
256248
TEST_ACTION_TYPE,
257249
TEST_METHOD_HTTP,
@@ -277,7 +269,7 @@ public void bedrockConnectorWithWrongInBuiltPrePostProcessFunction() {
277269
}
278270

279271
@Test
280-
public void sagemakerConnectorWithCorrectInBuiltPrePostProcessFunction() {
272+
public void testValidatePrePostProcessFunctionsWithSagemakerConnectorWithCorrectInBuiltPrePostProcessFunctionSuccess() {
281273
ConnectorAction action = new ConnectorAction(
282274
TEST_ACTION_TYPE,
283275
TEST_METHOD_HTTP,
@@ -288,7 +280,6 @@ public void sagemakerConnectorWithCorrectInBuiltPrePostProcessFunction() {
288280
DEFAULT_EMBEDDING
289281
);
290282
action.validatePrePostProcessFunctions(Map.of());
291-
assertNotNull(action);
292283
action = new ConnectorAction(
293284
TEST_ACTION_TYPE,
294285
TEST_METHOD_HTTP,
@@ -299,11 +290,10 @@ public void sagemakerConnectorWithCorrectInBuiltPrePostProcessFunction() {
299290
DEFAULT_RERANK
300291
);
301292
action.validatePrePostProcessFunctions(Map.of());
302-
assertNotNull(action);
303293
}
304294

305295
@Test
306-
public void sagemakerConnectorWithWrongInBuiltPrePostProcessFunction() {
296+
public void testValidatePrePostProcessFunctionsWithSagemakerConnectorWrongInBuiltPrePostProcessFunctionThrowsException() {
307297
ConnectorAction action1 = new ConnectorAction(
308298
TEST_ACTION_TYPE,
309299
TEST_METHOD_HTTP,

0 commit comments

Comments
 (0)