guardrails bug fixes and IT for creating guardrails (#2269) (#2275)

Signed-off-by: Jing Zhang <[email protected]>
(cherry picked from commit be56bcf)

Co-authored-by: Jing Zhang <[email protected]>

Author: opensearch-trigger-bot[bot]

common/src/main/java/org/opensearch/ml/common/model/Guardrails.java

@@ -22,26 +22,22 @@
 @Getter
 public class Guardrails implements ToXContentObject {
     public static final String TYPE_FIELD = "type";
-    public static final String ENGLISH_DETECTION_ENABLED_FIELD = "english_detection_enabled";
     public static final String INPUT_GUARDRAIL_FIELD = "input_guardrail";
     public static final String OUTPUT_GUARDRAIL_FIELD = "output_guardrail";
 
     private String type;
-    private Boolean engDetectionEnabled;
     private Guardrail inputGuardrail;
     private Guardrail outputGuardrail;
 
     @Builder(toBuilder = true)
-    public Guardrails(String type, Boolean engDetectionEnabled, Guardrail inputGuardrail, Guardrail outputGuardrail) {
+    public Guardrails(String type, Guardrail inputGuardrail, Guardrail outputGuardrail) {
         this.type = type;
-        this.engDetectionEnabled = engDetectionEnabled;
         this.inputGuardrail = inputGuardrail;
         this.outputGuardrail = outputGuardrail;
     }
 
     public Guardrails(StreamInput input) throws IOException {
         type = input.readString();
-        engDetectionEnabled = input.readBoolean();
         if (input.readBoolean()) {
             inputGuardrail = new Guardrail(input);
         }
@@ -52,7 +48,6 @@ public Guardrails(StreamInput input) throws IOException {
 
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(type);
-        out.writeBoolean(engDetectionEnabled);
         if (inputGuardrail != null) {
             out.writeBoolean(true);
             inputGuardrail.writeTo(out);
@@ -73,9 +68,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (type != null) {
             builder.field(TYPE_FIELD, type);
         }
-        if (engDetectionEnabled != null) {
-            builder.field(ENGLISH_DETECTION_ENABLED_FIELD, engDetectionEnabled);
-        }
         if (inputGuardrail != null) {
             builder.field(INPUT_GUARDRAIL_FIELD, inputGuardrail);
         }
@@ -88,7 +80,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
     public static Guardrails parse(XContentParser parser) throws IOException {
         String type = null;
-        Boolean engDetectionEnabled = null;
         Guardrail inputGuardrail = null;
         Guardrail outputGuardrail = null;
 
@@ -101,9 +92,6 @@ public static Guardrails parse(XContentParser parser) throws IOException {
                 case TYPE_FIELD:
                     type = parser.text();
                     break;
-                case ENGLISH_DETECTION_ENABLED_FIELD:
-                    engDetectionEnabled = parser.booleanValue();
-                    break;
                 case INPUT_GUARDRAIL_FIELD:
                     inputGuardrail = Guardrail.parse(parser);
                     break;
@@ -117,7 +105,6 @@ public static Guardrails parse(XContentParser parser) throws IOException {
         }
         return Guardrails.builder()
                 .type(type)
-                .engDetectionEnabled(engDetectionEnabled)
                 .inputGuardrail(inputGuardrail)
                 .outputGuardrail(outputGuardrail)
                 .build();

common/src/main/java/org/opensearch/ml/common/model/MLGuard.java

@@ -83,11 +83,11 @@ private void fillStopWordsToMap(@NonNull Guardrail guardrail, Map<String, List<S
         }
     }
 
-    public Boolean validate(String input, int type) {
+    public Boolean validate(String input, Type type) {
         switch (type) {
-            case 0: // validate input
+            case INPUT: // validate input
                 return validateRegexList(input, inputRegexPattern) && validateStopWords(input, stopWordsIndicesInput);
-            case 1: // validate output
+            case OUTPUT: // validate output
                 return validateRegexList(input, outputRegexPattern) && validateStopWords(input, stopWordsIndicesOutput);
             default:
                 throw new IllegalArgumentException("Unsupported type to validate for guardrails.");
@@ -159,4 +159,9 @@ public Boolean validateStopWordsSingleIndex(String input, String indexName, List
         }
         return hitStopWords.get();
     }
+
+    public enum Type {
+        INPUT,
+        OUTPUT
+    }
 }

common/src/test/java/org/opensearch/ml/common/model/GuardrailsTests.java

@@ -40,26 +40,24 @@ public void setUp() {
 
     @Test
     public void writeTo() throws IOException {
-        Guardrails guardrails = new Guardrails("test_type", false, inputGuardrail, outputGuardrail);
+        Guardrails guardrails = new Guardrails("test_type", inputGuardrail, outputGuardrail);
         BytesStreamOutput output = new BytesStreamOutput();
         guardrails.writeTo(output);
         Guardrails guardrails1 = new Guardrails(output.bytes().streamInput());
 
         Assert.assertEquals(guardrails.getType(), guardrails1.getType());
-        Assert.assertEquals(guardrails.getEngDetectionEnabled(), guardrails1.getEngDetectionEnabled());
         Assert.assertEquals(guardrails.getInputGuardrail(), guardrails1.getInputGuardrail());
         Assert.assertEquals(guardrails.getOutputGuardrail(), guardrails1.getOutputGuardrail());
     }
 
     @Test
     public void toXContent() throws IOException {
-        Guardrails guardrails = new Guardrails("test_type", false, inputGuardrail, outputGuardrail);
+        Guardrails guardrails = new Guardrails("test_type", inputGuardrail, outputGuardrail);
         XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
         guardrails.toXContent(builder, ToXContent.EMPTY_PARAMS);
         String content = TestHelper.xContentBuilderToString(builder);
 
         Assert.assertEquals("{\"type\":\"test_type\"," +
-                "\"english_detection_enabled\":false," +
                 "\"input_guardrail\":{\"stop_words\":[{\"index_name\":\"test_index\",\"source_fields\":[\"test_field\"]}],\"regex\":[\"regex1\"]}," +
                 "\"output_guardrail\":{\"stop_words\":[{\"index_name\":\"test_index\",\"source_fields\":[\"test_field\"]}],\"regex\":[\"regex1\"]}}",
                 content);
@@ -68,7 +66,6 @@ public void toXContent() throws IOException {
     @Test
     public void parse() throws IOException {
         String jsonStr = "{\"type\":\"test_type\"," +
-                "\"english_detection_enabled\":false," +
                 "\"input_guardrail\":{\"stop_words\":[{\"index_name\":\"test_index\",\"source_fields\":[\"test_field\"]}],\"regex\":[\"regex1\"]}," +
                 "\"output_guardrail\":{\"stop_words\":[{\"index_name\":\"test_index\",\"source_fields\":[\"test_field\"]}],\"regex\":[\"regex1\"]}}";
         XContentParser parser = XContentType.JSON.xContent().createParser(new NamedXContentRegistry(new SearchModule(Settings.EMPTY,
@@ -77,7 +74,6 @@ public void parse() throws IOException {
         Guardrails guardrails = Guardrails.parse(parser);
 
         Assert.assertEquals(guardrails.getType(), "test_type");
-        Assert.assertEquals(guardrails.getEngDetectionEnabled(), false);
         Assert.assertEquals(guardrails.getInputGuardrail(), inputGuardrail);
         Assert.assertEquals(guardrails.getOutputGuardrail(), outputGuardrail);
     }

common/src/test/java/org/opensearch/ml/common/model/MLGuardTests.java

@@ -75,22 +75,22 @@ public void setUp() {
         regexPatterns = List.of(Pattern.compile("(.|\n)*stop words(.|\n)*"));
         inputGuardrail = new Guardrail(List.of(stopWords), regex);
         outputGuardrail = new Guardrail(List.of(stopWords), regex);
-        guardrails = new Guardrails("test_type", false, inputGuardrail, outputGuardrail);
+        guardrails = new Guardrails("test_type", inputGuardrail, outputGuardrail);
         mlGuard = new MLGuard(guardrails, xContentRegistry, client);
     }
 
     @Test
     public void validateInput() {
         String input = "\n\nHuman:hello stop words.\n\nAssistant:";
-        Boolean res = mlGuard.validate(input, 0);
+        Boolean res = mlGuard.validate(input, MLGuard.Type.INPUT);
 
         Assert.assertFalse(res);
     }
 
     @Test
     public void validateOutput() {
         String input = "\n\nHuman:hello stop words.\n\nAssistant:";
-        Boolean res = mlGuard.validate(input, 1);
+        Boolean res = mlGuard.validate(input, MLGuard.Type.OUTPUT);
 
         Assert.assertFalse(res);
     }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java

@@ -150,7 +150,7 @@ public void invokeRemoteModel(MLInput mlInput, Map<String, String> parameters, S
                 throw new OpenSearchStatusException("No response from model", RestStatus.BAD_REQUEST);
             }
             String modelResponse = responseBuilder.toString();
-            if (getMlGuard() != null && !getMlGuard().validate(modelResponse, 1)) {
+            if (getMlGuard() != null && !getMlGuard().validate(modelResponse, MLGuard.Type.OUTPUT)) {
                 throw new IllegalArgumentException("guardrails triggered for LLM output");
             }
             if (statusCode < 200 || statusCode >= 300) {

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java

@@ -139,7 +139,7 @@ public void invokeRemoteModel(MLInput mlInput, Map<String, String> parameters, S
                 return null;
             });
             String modelResponse = responseRef.get();
-            if (getMlGuard() != null && !getMlGuard().validate(modelResponse, 1)) {
+            if (getMlGuard() != null && !getMlGuard().validate(modelResponse, MLGuard.Type.OUTPUT)) {
                 throw new IllegalArgumentException("guardrails triggered for LLM output");
             }
             Integer statusCode = statusCodeRef.get();

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java

@@ -142,7 +142,7 @@ && getUserRateLimiterMap().get(user.getName()) != null
                 RestStatus.TOO_MANY_REQUESTS
             );
         } else {
-            if (getMlGuard() != null && !getMlGuard().validate(payload, 0)) {
+            if (getMlGuard() != null && !getMlGuard().validate(payload, MLGuard.Type.INPUT)) {
                 throw new IllegalArgumentException("guardrails triggered for user input");
             }
             invokeRemoteModel(mlInput, parameters, payload, tensorOutputs);