enhance SafeGuardAdvisor

Signed-off-by: Karson To <karsontao@hotmail.com>

Author: karsonto

spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/advisor/SafeGuardAdvisor.java

@@ -18,6 +18,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.Predicate;
 
 import reactor.core.publisher.Flux;
 
@@ -37,6 +38,7 @@
  * An advisor that blocks the call to the model provider if the user input contains any of
  * the sensitive words.
  *
+ * @author Karson To
  * @author Christian Tzolov
  * @author Ilayaperumal Gopinathan
  * @author Thomas Vitale
@@ -49,19 +51,19 @@ public class SafeGuardAdvisor implements CallAdvisor, StreamAdvisor {
 	private static final int DEFAULT_ORDER = 0;
 
 	private final String failureResponse;
-
-	private final List<String> sensitiveWords;
+	
+	private final Predicate<String> contentValidator;
 
 	private final int order;
 
-	public SafeGuardAdvisor(List<String> sensitiveWords) {
-		this(sensitiveWords, DEFAULT_FAILURE_RESPONSE, DEFAULT_ORDER);
+	public SafeGuardAdvisor(Predicate<String> contentValidator) {
+		this(contentValidator, DEFAULT_FAILURE_RESPONSE, DEFAULT_ORDER);
 	}
 
-	public SafeGuardAdvisor(List<String> sensitiveWords, String failureResponse, int order) {
-		Assert.notNull(sensitiveWords, "Sensitive words must not be null!");
+	public SafeGuardAdvisor(Predicate<String> contentValidator, String failureResponse, int order) {
+		Assert.notNull(contentValidator, "SContent validator must not be null!");
 		Assert.notNull(failureResponse, "Failure response must not be null!");
-		this.sensitiveWords = sensitiveWords;
+		this.contentValidator = contentValidator;
 		this.failureResponse = failureResponse;
 		this.order = order;
 	}
@@ -76,8 +78,7 @@ public String getName() {
 
 	@Override
 	public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
-		if (!CollectionUtils.isEmpty(this.sensitiveWords)
-				&& this.sensitiveWords.stream().anyMatch(w -> chatClientRequest.prompt().getContents().contains(w))) {
+		if (contentValidator.test(chatClientRequest.prompt().getContents())) {
 			return createFailureResponse(chatClientRequest);
 		}
 
@@ -87,8 +88,7 @@ public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAd
 	@Override
 	public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest,
 			StreamAdvisorChain streamAdvisorChain) {
-		if (!CollectionUtils.isEmpty(this.sensitiveWords)
-				&& this.sensitiveWords.stream().anyMatch(w -> chatClientRequest.prompt().getContents().contains(w))) {
+		if (contentValidator.test(chatClientRequest.prompt().getContents())) {
 			return Flux.just(createFailureResponse(chatClientRequest));
 		}
 
@@ -111,7 +111,7 @@ public int getOrder() {
 
 	public static final class Builder {
 
-		private List<String> sensitiveWords;
+		private Predicate<String> contentValidator;
 
 		private String failureResponse = DEFAULT_FAILURE_RESPONSE;
 
@@ -120,8 +120,8 @@ public static final class Builder {
 		private Builder() {
 		}
 
-		public Builder sensitiveWords(List<String> sensitiveWords) {
-			this.sensitiveWords = sensitiveWords;
+		public Builder sensitiveWords(Predicate<String> contentValidator) {
+			this.contentValidator = contentValidator;
 			return this;
 		}
 
@@ -136,7 +136,7 @@ public Builder order(int order) {
 		}
 
 		public SafeGuardAdvisor build() {
-			return new SafeGuardAdvisor(this.sensitiveWords, this.failureResponse, this.order);
+			return new SafeGuardAdvisor(this.contentValidator, this.failureResponse, this.order);
 		}
 
 	}