2121import  java .util .Map ;
2222import  java .util .stream .Collectors ;
2323
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 ;
2629import  reactor .core .scheduler .Schedulers ;
2730
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 ;
3531import  org .springframework .ai .chat .model .ChatResponse ;
3632import  org .springframework .ai .chat .prompt .PromptTemplate ;
3733import  org .springframework .ai .document .Document ;
5349 * @author Thomas Vitale 
5450 * @since 1.0.0 
5551 */ 
56- public  class  QuestionAnswerAdvisor  implements  CallAroundAdvisor ,  StreamAroundAdvisor  {
52+ public  class  QuestionAnswerAdvisor  implements  BaseAdvisor  {
5753
5854	public  static  final  String  RETRIEVED_DOCUMENTS  = "qa_retrieved_documents" ;
5955
@@ -80,198 +76,96 @@ public class QuestionAnswerAdvisor implements CallAroundAdvisor, StreamAroundAdv
8076
8177	private  final  SearchRequest  searchRequest ;
8278
83- 	private  final  boolean   protectFromBlocking ;
79+ 	private  final  Scheduler   scheduler ;
8480
8581	private  final  int  order ;
8682
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- 	 */ 
9283	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 ,
12285				DEFAULT_ORDER );
12386	}
12487
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- 
16888	QuestionAnswerAdvisor (VectorStore  vectorStore , SearchRequest  searchRequest , @ Nullable  PromptTemplate  promptTemplate ,
169- 			boolean   protectFromBlocking , int  order ) {
89+ 			@ Nullable   Scheduler   scheduler , int  order ) {
17090		Assert .notNull (vectorStore , "vectorStore cannot be null" );
17191		Assert .notNull (searchRequest , "searchRequest cannot be null" );
17292
17393		this .vectorStore  = vectorStore ;
17494		this .searchRequest  = searchRequest ;
17595		this .promptTemplate  = promptTemplate  != null  ? promptTemplate  : DEFAULT_PROMPT_TEMPLATE ;
176- 		this .protectFromBlocking  = protectFromBlocking ;
96+ 		this .scheduler  = scheduler  !=  null  ?  scheduler  :  BaseAdvisor . DEFAULT_SCHEDULER ;
17797		this .order  = order ;
17898	}
17999
180100	public  static  Builder  builder (VectorStore  vectorStore ) {
181101		return  new  Builder (vectorStore );
182102	}
183103
184- 	@ Override 
185- 	public  String  getName () {
186- 		return  this .getClass ().getSimpleName ();
187- 	}
188- 
189104	@ Override 
190105	public  int  getOrder () {
191106		return  this .order ;
192107	}
193108
194109	@ 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 ) {
231111		// 1. Search for similar documents in the vector store. 
232112		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 () ))
235115			.build ();
236116
237117		List <Document > documents  = this .vectorStore .similaritySearch (searchRequestToUse );
238118
239119		// 2. Create the context from the documents. 
120+ 		Map <String , Object > context  = new  HashMap <>(chatClientRequest .context ());
240121		context .put (RETRIEVED_DOCUMENTS , documents );
241122
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 ()));
245125
246126		// 3. Augment the user prompt with the document context. 
247127		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 ())
249130			.variables (Map .of ("question_answer_context" , documentContext ))
250131			.build ()
251132			.render ();
252133
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 )
256138			.build ();
257- 
258- 		return  advisedRequest ;
259139	}
260140
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 ();
265155	}
266156
157+ 	@ Nullable 
267158	protected  Filter .Expression  doGetFilterExpression (Map <String , Object > context ) {
268- 
269159		if  (!context .containsKey (FILTER_EXPRESSION )
270160				|| !StringUtils .hasText (context .get (FILTER_EXPRESSION ).toString ())) {
271161			return  this .searchRequest .getFilterExpression ();
272162		}
273163		return  new  FilterExpressionTextParser ().parse (context .get (FILTER_EXPRESSION ).toString ());
164+ 	}
274165
166+ 	@ Override 
167+ 	public  Scheduler  getScheduler () {
168+ 		return  this .scheduler ;
275169	}
276170
277171	public  static  final  class  Builder  {
@@ -282,7 +176,7 @@ public static final class Builder {
282176
283177		private  PromptTemplate  promptTemplate ;
284178
285- 		private  boolean   protectFromBlocking  =  true ;
179+ 		private  Scheduler   scheduler ;
286180
287181		private  int  order  = DEFAULT_ORDER ;
288182
@@ -303,18 +197,13 @@ public Builder searchRequest(SearchRequest searchRequest) {
303197			return  this ;
304198		}
305199
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 ();
313202			return  this ;
314203		}
315204
316- 		public  Builder  protectFromBlocking ( boolean   protectFromBlocking ) {
317- 			this .protectFromBlocking  = protectFromBlocking ;
205+ 		public  Builder  scheduler ( Scheduler   scheduler ) {
206+ 			this .scheduler  = scheduler ;
318207			return  this ;
319208		}
320209
@@ -324,8 +213,8 @@ public Builder order(int order) {
324213		}
325214
326215		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 );
329218		}
330219
331220	}
0 commit comments