1919
2020import org .springframework .ai .chat .client .advisor .api .AdvisedRequest ;
2121import org .springframework .ai .chat .client .advisor .api .AdvisedResponse ;
22- import org .springframework .ai .chat .client .advisor .api .CallAroundAdvisorChain ;
2322import org .springframework .ai .chat .client .advisor .api .CallAroundAdvisor ;
23+ import org .springframework .ai .chat .client .advisor .api .CallAroundAdvisorChain ;
2424import org .springframework .ai .chat .client .advisor .api .StreamAroundAdvisor ;
2525import org .springframework .ai .chat .client .advisor .api .StreamAroundAdvisorChain ;
26+ import org .springframework .ai .chat .messages .AssistantMessage ;
2627import org .springframework .ai .chat .model .ChatResponse ;
28+ import org .springframework .ai .chat .model .Generation ;
29+ import org .springframework .util .Assert ;
2730import org .springframework .util .CollectionUtils ;
2831
2932import reactor .core .publisher .Flux ;
3740 */
3841public class SafeGuardAdvisor implements CallAroundAdvisor , StreamAroundAdvisor {
3942
43+ private final static String DEFAULT_FAILURE_RESPONSE = "I'm unable to respond to that due to sensitive content. Could we rephrase or discuss something else?" ;
44+
45+ private final static int DEFAULT_ORDER = 0 ;
46+
47+ private final String failureResponse ;
48+
4049 private final List <String > sensitiveWords ;
4150
51+ private final int order ;
52+
4253 public SafeGuardAdvisor (List <String > sensitiveWords ) {
54+ this (sensitiveWords , DEFAULT_FAILURE_RESPONSE , DEFAULT_ORDER );
55+ }
56+
57+ public SafeGuardAdvisor (List <String > sensitiveWords , String failureResponse , int order ) {
58+ Assert .notNull (sensitiveWords , "Sensitive words must not be null!" );
59+ Assert .notNull (failureResponse , "Failure response must not be null!" );
4360 this .sensitiveWords = sensitiveWords ;
61+ this .failureResponse = failureResponse ;
62+ this .order = order ;
4463 }
4564
4665 public String getName () {
@@ -51,9 +70,9 @@ public String getName() {
5170 public AdvisedResponse aroundCall (AdvisedRequest advisedRequest , CallAroundAdvisorChain chain ) {
5271
5372 if (!CollectionUtils .isEmpty (this .sensitiveWords )
54- && sensitiveWords .stream ().anyMatch (w -> advisedRequest .userText ().contains (w ))) {
55- return new AdvisedResponse ( ChatResponse . builder (). withGenerations ( List . of ()). build (),
56- advisedRequest . adviseContext () );
73+ && this . sensitiveWords .stream ().anyMatch (w -> advisedRequest .userText ().contains (w ))) {
74+
75+ return createFailureResponse ( advisedRequest );
5776 }
5877
5978 return chain .nextAroundCall (advisedRequest );
@@ -64,16 +83,57 @@ public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamA
6483
6584 if (!CollectionUtils .isEmpty (this .sensitiveWords )
6685 && sensitiveWords .stream ().anyMatch (w -> advisedRequest .userText ().contains (w ))) {
67- return Flux .empty ( );
86+ return Flux .just ( createFailureResponse ( advisedRequest ) );
6887 }
6988
7089 return chain .nextAroundStream (advisedRequest );
90+ }
7191
92+ private AdvisedResponse createFailureResponse (AdvisedRequest advisedRequest ) {
93+ return new AdvisedResponse (ChatResponse .builder ()
94+ .withGenerations (List .of (new Generation (new AssistantMessage (this .failureResponse ))))
95+ .build (), advisedRequest .adviseContext ());
7296 }
7397
7498 @ Override
7599 public int getOrder () {
76- return 0 ;
100+ return this .order ;
101+ }
102+
103+ public static Builder builder () {
104+ return new Builder ();
105+ }
106+
107+ public static class Builder {
108+
109+ private List <String > sensitiveWords ;
110+
111+ private String failureResponse = DEFAULT_FAILURE_RESPONSE ;
112+
113+ private int order = DEFAULT_ORDER ;
114+
115+ private Builder () {
116+ }
117+
118+ public Builder withSensitiveWords (List <String > sensitiveWords ) {
119+ this .sensitiveWords = sensitiveWords ;
120+ return this ;
121+ }
122+
123+ public Builder withFailureResponse (String failureResponse ) {
124+ this .failureResponse = failureResponse ;
125+ return this ;
126+ }
127+
128+ public Builder withOrder (int order ) {
129+ this .order = order ;
130+ return this ;
131+ }
132+
133+ public SafeGuardAdvisor build () {
134+ return new SafeGuardAdvisor (this .sensitiveWords , this .failureResponse , this .order );
135+ }
136+
77137 }
78138
79139}
0 commit comments