fix stopwords npe (opensearch-project#2311)

Signed-off-by: Jing Zhang <[email protected]>
TrungBui59 · Apr 12, 2024 · acfb4a9 · acfb4a9
1 parent abba0ec
commit acfb4a9
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/common/src/main/java/org/opensearch/ml/common/model/MLGuard.java b/common/src/main/java/org/opensearch/ml/common/model/MLGuard.java
@@ -77,7 +77,9 @@ private void fillStopWordsToMap(@NonNull Guardrail guardrail, Map<String, List<S
             return;
         }
         for (StopWords e : stopWords) {
-            map.put(e.getIndex(), Arrays.asList(e.getSourceFields()));
+            if (e.getIndex() != null && e.getSourceFields() != null) {
+                map.put(e.getIndex(), Arrays.asList(e.getSourceFields()));
+            }
         }
     }
 

diff --git a/common/src/test/java/org/opensearch/ml/common/model/MLGuardTests.java b/common/src/test/java/org/opensearch/ml/common/model/MLGuardTests.java
@@ -87,6 +87,21 @@ public void validateInput() {
         Assert.assertFalse(res);
     }
 
+    @Test
+    public void validateInitializedStopWordsEmpty() {
+        stopWords = new StopWords(null, null);
+        regex = List.of("(.|\n)*stop words(.|\n)*").toArray(new String[0]);
+        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", inputGuardrail, outputGuardrail);
+        mlGuard = new MLGuard(guardrails, xContentRegistry, client);
+
+        String input = "\n\nHuman:hello good words.\n\nAssistant:";
+        Boolean res = mlGuard.validate(input, MLGuard.Type.INPUT);
+        Assert.assertTrue(res);
+    }
+
     @Test
     public void validateOutput() {
         String input = "\n\nHuman:hello stop words.\n\nAssistant:";