21
21
import java .util .Map ;
22
22
import java .util .stream .Collectors ;
23
23
24
- import reactor .core .publisher .Flux ;
25
- import reactor .core .publisher .Mono ;
24
+ import org .springframework .ai .chat .client .ChatClientRequest ;
25
+ import org .springframework .ai .chat .client .ChatClientResponse ;
26
+ import org .springframework .ai .chat .client .advisor .api .AdvisorChain ;
27
+ import org .springframework .ai .chat .client .advisor .api .BaseAdvisor ;
28
+ import reactor .core .scheduler .Scheduler ;
26
29
import reactor .core .scheduler .Schedulers ;
27
30
28
- import org .springframework .ai .chat .client .advisor .api .AdvisedRequest ;
29
- import org .springframework .ai .chat .client .advisor .api .AdvisedResponse ;
30
- import org .springframework .ai .chat .client .advisor .api .AdvisedResponseStreamUtils ;
31
- import org .springframework .ai .chat .client .advisor .api .CallAroundAdvisor ;
32
- import org .springframework .ai .chat .client .advisor .api .CallAroundAdvisorChain ;
33
- import org .springframework .ai .chat .client .advisor .api .StreamAroundAdvisor ;
34
- import org .springframework .ai .chat .client .advisor .api .StreamAroundAdvisorChain ;
35
31
import org .springframework .ai .chat .model .ChatResponse ;
36
32
import org .springframework .ai .chat .prompt .PromptTemplate ;
37
33
import org .springframework .ai .document .Document ;
53
49
* @author Thomas Vitale
54
50
* @since 1.0.0
55
51
*/
56
- public class QuestionAnswerAdvisor implements CallAroundAdvisor , StreamAroundAdvisor {
52
+ public class QuestionAnswerAdvisor implements BaseAdvisor {
57
53
58
54
public static final String RETRIEVED_DOCUMENTS = "qa_retrieved_documents" ;
59
55
@@ -80,198 +76,96 @@ public class QuestionAnswerAdvisor implements CallAroundAdvisor, StreamAroundAdv
80
76
81
77
private final SearchRequest searchRequest ;
82
78
83
- private final boolean protectFromBlocking ;
79
+ private final Scheduler scheduler ;
84
80
85
81
private final int order ;
86
82
87
- /**
88
- * The QuestionAnswerAdvisor retrieves context information from a Vector Store and
89
- * combines it with the user's text.
90
- * @param vectorStore The vector store to use
91
- */
92
83
public QuestionAnswerAdvisor (VectorStore vectorStore ) {
93
- this (vectorStore , SearchRequest .builder ().build (), DEFAULT_PROMPT_TEMPLATE , true , DEFAULT_ORDER );
94
- }
95
-
96
- /**
97
- * The QuestionAnswerAdvisor retrieves context information from a Vector Store and
98
- * combines it with the user's text.
99
- * @param vectorStore The vector store to use
100
- * @param searchRequest The search request defined using the portable filter
101
- * expression syntax
102
- * @deprecated in favor of the builder: {@link #builder(VectorStore)}
103
- */
104
- @ Deprecated
105
- public QuestionAnswerAdvisor (VectorStore vectorStore , SearchRequest searchRequest ) {
106
- this (vectorStore , searchRequest , DEFAULT_PROMPT_TEMPLATE , true , DEFAULT_ORDER );
107
- }
108
-
109
- /**
110
- * The QuestionAnswerAdvisor retrieves context information from a Vector Store and
111
- * combines it with the user's text.
112
- * @param vectorStore The vector store to use
113
- * @param searchRequest The search request defined using the portable filter
114
- * expression syntax
115
- * @param userTextAdvise The user text to append to the existing user prompt. The text
116
- * should contain a placeholder named "question_answer_context".
117
- * @deprecated in favor of the builder: {@link #builder(VectorStore)}
118
- */
119
- @ Deprecated
120
- public QuestionAnswerAdvisor (VectorStore vectorStore , SearchRequest searchRequest , String userTextAdvise ) {
121
- this (vectorStore , searchRequest , PromptTemplate .builder ().template (userTextAdvise ).build (), true ,
84
+ this (vectorStore , SearchRequest .builder ().build (), DEFAULT_PROMPT_TEMPLATE , BaseAdvisor .DEFAULT_SCHEDULER ,
122
85
DEFAULT_ORDER );
123
86
}
124
87
125
- /**
126
- * The QuestionAnswerAdvisor retrieves context information from a Vector Store and
127
- * combines it with the user's text.
128
- * @param vectorStore The vector store to use
129
- * @param searchRequest The search request defined using the portable filter
130
- * expression syntax
131
- * @param userTextAdvise The user text to append to the existing user prompt. The text
132
- * should contain a placeholder named "question_answer_context".
133
- * @param protectFromBlocking If true the advisor will protect the execution from
134
- * blocking threads. If false the advisor will not protect the execution from blocking
135
- * threads. This is useful when the advisor is used in a non-blocking environment. It
136
- * is true by default.
137
- * @deprecated in favor of the builder: {@link #builder(VectorStore)}
138
- */
139
- @ Deprecated
140
- public QuestionAnswerAdvisor (VectorStore vectorStore , SearchRequest searchRequest , String userTextAdvise ,
141
- boolean protectFromBlocking ) {
142
- this (vectorStore , searchRequest , PromptTemplate .builder ().template (userTextAdvise ).build (), protectFromBlocking ,
143
- DEFAULT_ORDER );
144
- }
145
-
146
- /**
147
- * The QuestionAnswerAdvisor retrieves context information from a Vector Store and
148
- * combines it with the user's text.
149
- * @param vectorStore The vector store to use
150
- * @param searchRequest The search request defined using the portable filter
151
- * expression syntax
152
- * @param userTextAdvise The user text to append to the existing user prompt. The text
153
- * should contain a placeholder named "question_answer_context".
154
- * @param protectFromBlocking If true the advisor will protect the execution from
155
- * blocking threads. If false the advisor will not protect the execution from blocking
156
- * threads. This is useful when the advisor is used in a non-blocking environment. It
157
- * is true by default.
158
- * @param order The order of the advisor.
159
- * @deprecated in favor of the builder: {@link #builder(VectorStore)}
160
- */
161
- @ Deprecated
162
- public QuestionAnswerAdvisor (VectorStore vectorStore , SearchRequest searchRequest , String userTextAdvise ,
163
- boolean protectFromBlocking , int order ) {
164
- this (vectorStore , searchRequest , PromptTemplate .builder ().template (userTextAdvise ).build (), protectFromBlocking ,
165
- order );
166
- }
167
-
168
88
QuestionAnswerAdvisor (VectorStore vectorStore , SearchRequest searchRequest , @ Nullable PromptTemplate promptTemplate ,
169
- boolean protectFromBlocking , int order ) {
89
+ @ Nullable Scheduler scheduler , int order ) {
170
90
Assert .notNull (vectorStore , "vectorStore cannot be null" );
171
91
Assert .notNull (searchRequest , "searchRequest cannot be null" );
172
92
173
93
this .vectorStore = vectorStore ;
174
94
this .searchRequest = searchRequest ;
175
95
this .promptTemplate = promptTemplate != null ? promptTemplate : DEFAULT_PROMPT_TEMPLATE ;
176
- this .protectFromBlocking = protectFromBlocking ;
96
+ this .scheduler = scheduler != null ? scheduler : BaseAdvisor . DEFAULT_SCHEDULER ;
177
97
this .order = order ;
178
98
}
179
99
180
100
public static Builder builder (VectorStore vectorStore ) {
181
101
return new Builder (vectorStore );
182
102
}
183
103
184
- @ Override
185
- public String getName () {
186
- return this .getClass ().getSimpleName ();
187
- }
188
-
189
104
@ Override
190
105
public int getOrder () {
191
106
return this .order ;
192
107
}
193
108
194
109
@ Override
195
- public AdvisedResponse aroundCall (AdvisedRequest advisedRequest , CallAroundAdvisorChain chain ) {
196
-
197
- AdvisedRequest advisedRequest2 = before (advisedRequest );
198
-
199
- AdvisedResponse advisedResponse = chain .nextAroundCall (advisedRequest2 );
200
-
201
- return after (advisedResponse );
202
- }
203
-
204
- @ Override
205
- public Flux <AdvisedResponse > aroundStream (AdvisedRequest advisedRequest , StreamAroundAdvisorChain chain ) {
206
-
207
- // This can be executed by both blocking and non-blocking Threads
208
- // E.g. a command line or Tomcat blocking Thread implementation
209
- // or by a WebFlux dispatch in a non-blocking manner.
210
- Flux <AdvisedResponse > advisedResponses = (this .protectFromBlocking ) ?
211
- // @formatter:off
212
- Mono .just (advisedRequest )
213
- .publishOn (Schedulers .boundedElastic ())
214
- .map (this ::before )
215
- .flatMapMany (request -> chain .nextAroundStream (request ))
216
- : chain .nextAroundStream (before (advisedRequest ));
217
- // @formatter:on
218
-
219
- return advisedResponses .map (ar -> {
220
- if (AdvisedResponseStreamUtils .onFinishReason ().test (ar )) {
221
- ar = after (ar );
222
- }
223
- return ar ;
224
- });
225
- }
226
-
227
- private AdvisedRequest before (AdvisedRequest request ) {
228
-
229
- var context = new HashMap <>(request .adviseContext ());
230
-
110
+ public ChatClientRequest before (ChatClientRequest chatClientRequest , AdvisorChain advisorChain ) {
231
111
// 1. Search for similar documents in the vector store.
232
112
var searchRequestToUse = SearchRequest .from (this .searchRequest )
233
- .query (request . userText ())
234
- .filterExpression (doGetFilterExpression (context ))
113
+ .query (chatClientRequest . prompt (). getUserMessage (). getText ())
114
+ .filterExpression (doGetFilterExpression (chatClientRequest . context () ))
235
115
.build ();
236
116
237
117
List <Document > documents = this .vectorStore .similaritySearch (searchRequestToUse );
238
118
239
119
// 2. Create the context from the documents.
120
+ Map <String , Object > context = new HashMap <>(chatClientRequest .context ());
240
121
context .put (RETRIEVED_DOCUMENTS , documents );
241
122
242
- String documentContext = documents .stream ()
243
- .map (Document ::getText )
244
- .collect (Collectors .joining (System .lineSeparator ()));
123
+ String documentContext = documents == null ? ""
124
+ : documents .stream ().map (Document ::getText ).collect (Collectors .joining (System .lineSeparator ()));
245
125
246
126
// 3. Augment the user prompt with the document context.
247
127
String augmentedUserText = this .promptTemplate .mutate ()
248
- .template (request .userText () + System .lineSeparator () + this .promptTemplate .getTemplate ())
128
+ .template (chatClientRequest .prompt ().getUserMessage ().getText () + System .lineSeparator ()
129
+ + this .promptTemplate .getTemplate ())
249
130
.variables (Map .of ("question_answer_context" , documentContext ))
250
131
.build ()
251
132
.render ();
252
133
253
- AdvisedRequest advisedRequest = AdvisedRequest .from (request )
254
- .userText (augmentedUserText )
255
- .adviseContext (context )
134
+ // 4. Update ChatClientRequest with augmented prompt.
135
+ return chatClientRequest .mutate ()
136
+ .prompt (chatClientRequest .prompt ().augmentUserMessage (augmentedUserText ))
137
+ .context (context )
256
138
.build ();
257
-
258
- return advisedRequest ;
259
139
}
260
140
261
- private AdvisedResponse after (AdvisedResponse advisedResponse ) {
262
- ChatResponse .Builder chatResponseBuilder = ChatResponse .builder ().from (advisedResponse .response ());
263
- chatResponseBuilder .metadata (RETRIEVED_DOCUMENTS , advisedResponse .adviseContext ().get (RETRIEVED_DOCUMENTS ));
264
- return new AdvisedResponse (chatResponseBuilder .build (), advisedResponse .adviseContext ());
141
+ @ Override
142
+ public ChatClientResponse after (ChatClientResponse chatClientResponse , AdvisorChain advisorChain ) {
143
+ ChatResponse .Builder chatResponseBuilder ;
144
+ if (chatClientResponse .chatResponse () == null ) {
145
+ chatResponseBuilder = ChatResponse .builder ();
146
+ }
147
+ else {
148
+ chatResponseBuilder = ChatResponse .builder ().from (chatClientResponse .chatResponse ());
149
+ }
150
+ chatResponseBuilder .metadata (RETRIEVED_DOCUMENTS , chatClientResponse .context ().get (RETRIEVED_DOCUMENTS ));
151
+ return ChatClientResponse .builder ()
152
+ .chatResponse (chatResponseBuilder .build ())
153
+ .context (chatClientResponse .context ())
154
+ .build ();
265
155
}
266
156
157
+ @ Nullable
267
158
protected Filter .Expression doGetFilterExpression (Map <String , Object > context ) {
268
-
269
159
if (!context .containsKey (FILTER_EXPRESSION )
270
160
|| !StringUtils .hasText (context .get (FILTER_EXPRESSION ).toString ())) {
271
161
return this .searchRequest .getFilterExpression ();
272
162
}
273
163
return new FilterExpressionTextParser ().parse (context .get (FILTER_EXPRESSION ).toString ());
164
+ }
274
165
166
+ @ Override
167
+ public Scheduler getScheduler () {
168
+ return this .scheduler ;
275
169
}
276
170
277
171
public static final class Builder {
@@ -282,7 +176,7 @@ public static final class Builder {
282
176
283
177
private PromptTemplate promptTemplate ;
284
178
285
- private boolean protectFromBlocking = true ;
179
+ private Scheduler scheduler ;
286
180
287
181
private int order = DEFAULT_ORDER ;
288
182
@@ -303,18 +197,13 @@ public Builder searchRequest(SearchRequest searchRequest) {
303
197
return this ;
304
198
}
305
199
306
- /**
307
- * @deprecated in favour of {@link #promptTemplate(PromptTemplate)}
308
- */
309
- @ Deprecated
310
- public Builder userTextAdvise (String userTextAdvise ) {
311
- Assert .hasText (userTextAdvise , "The userTextAdvise must not be empty!" );
312
- this .promptTemplate = PromptTemplate .builder ().template (userTextAdvise ).build ();
200
+ public Builder protectFromBlocking (boolean protectFromBlocking ) {
201
+ this .scheduler = protectFromBlocking ? BaseAdvisor .DEFAULT_SCHEDULER : Schedulers .immediate ();
313
202
return this ;
314
203
}
315
204
316
- public Builder protectFromBlocking ( boolean protectFromBlocking ) {
317
- this .protectFromBlocking = protectFromBlocking ;
205
+ public Builder scheduler ( Scheduler scheduler ) {
206
+ this .scheduler = scheduler ;
318
207
return this ;
319
208
}
320
209
@@ -324,8 +213,8 @@ public Builder order(int order) {
324
213
}
325
214
326
215
public QuestionAnswerAdvisor build () {
327
- return new QuestionAnswerAdvisor (this .vectorStore , this .searchRequest , this .promptTemplate ,
328
- this .protectFromBlocking , this . order );
216
+ return new QuestionAnswerAdvisor (this .vectorStore , this .searchRequest , this .promptTemplate , this . scheduler ,
217
+ this .order );
329
218
}
330
219
331
220
}
0 commit comments