[WIP] Refactor AdvisedRequest and introduce AdvisedResponse #1332

tzolov · 2024-09-09T15:01:00Z

Introduce AdvisedResponse to encapsulate ChatResponse and advice context
Refactor RequestAdvisor and ResponseAdvisor interface to use AdvisedRequest/AdvisedResponse records
Update DefaultAroundAdvisorChain to work with new advisor interfaces
Modify DefaultChatClient to use new advisor framework
Adjust MessageChatMemoryAdvisor for compatibility with new advisor structure
Update integration tests to reflect changes in advisor framework
Update AdvisedRequest to include adviseContext and add toPrompt() method
Modify advisor interfaces to remove explicit context parameter
Introduce AdvisedResponse to encapsulate ChatResponse with adviseContext
Update DefaultChatClient to use new advisor interfaces and AdvisedResponse
Adjust integration tests in OpenAiPaymentTransactionIT and VertexAiGeminiPaymentTransactionIT
Refactor the RequestResponseAdvisor interface to make it a backward compatible adabror for ReqeustAdvisor and ResponseAdvisors. Aslo mark it as deprecated.
Update various advisor implementations to work with new AdvisedRequest and AdvisedResponse
Modify AroundAdvisorChain to use AdvisedResponse in both call and stream methods
Implement new adviseResponse method for streaming responses in RequestResponseAdvisor
Update MessageChatMemoryAdvisor, PromptChatMemoryAdvisor, QuestionAnswerAdvisor, SimpleLoggerAdvisor, and VectorStoreChatMemoryAdvisor to use new advisor interfaces
Add new AdvisedResponse class to encapsulate ChatResponse with advice context
Update CallAroundAdvisor, StreamAroundAdvisor, RequestAdvisor, and ResponseAdvisor interfaces to use AdvisedResponse and remove adviceContext parameter
Modify CacheAroundAdvisor and SafeGuardAroundAdvisor to work with new advisor interfaces and AdvisedResponse
Update AdvisorObservableHelper to use AdvisedResponse and remove explicit context parameter
Refactor MessageAggregator to support aggregation of AdvisedResponse objects
Rename extractRequestAdvisors and extractResponseAdvisors methods to requestAdvisors and responseAdvisors in AdvisorObservableHelper
Add null checks and return empty lists for advisors in AdvisorObservableHelper

chemicL · 2024-09-16T09:11:30Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java

-								for (ResponseAdvisor advisor : aggregateResponseAdvisors) {
-									AdvisorObservableHelper.adviseResponse(parentObservation, advisor, chatResponse,
-											advisorContext);
+							AtomicReference<Map<String, Object>> adviseContext = new AtomicReference<>();


What's happening here? Why do you want this to be stateful? The code that follows can seem confusing, as in - it seems to be imperative, where you should get the enhanced context in the following processing but that context will only get the current snapshot of alterations per each chunk. So the ceremony seems fruitless in a sense. It doesn't join the mutated context either, so it is also unnecessary for the aggregated bit.

chemicL · 2024-09-16T09:13:18Z

spring-ai-core/src/main/java/org/springframework/ai/chat/model/MessageAggregator.java


 	private static final Logger logger = LoggerFactory.getLogger(MessageAggregator.class);

+	public Flux<AdvisedResponse> aggregateAdvisedResponse(Flux<AdvisedResponse> advisedResponseFlux,


This logic looks like a duplicate from DefaultChatClient, which does not call this method, am I correct?

chemicL · 2024-09-16T09:13:47Z

spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/AdvisorsTests.java

+			.thenReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("Hello John")))))
+			.thenReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("Your name is John")))));
+
+		when(chatModel.call(promptCaptor.capture()))


Why Is there duplication here?

chemicL · 2024-09-16T09:14:15Z

spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/AdvisorsTests.java

+				.containsEntry("aroundCallAfter", "AROUND_CALL_AFTER")
+				.doesNotContainKeys("afterCall");
+
+			assertThat(advisedResponse.adviseContext()).containsEntry("key1", "value1")


This assertion is a duplicate of the above, right?

chemicL · 2024-09-16T09:14:37Z

spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/AdvisorsTests.java

+		assertThat(mockAroundAdvisor.advisedResponse.adviseContext()).containsEntry("key1", "value1")
+			.containsEntry("key2", "value2")
+			.containsEntry("before", "BEFORE")
+			.containsEntry("aroundCallBefore", "AROUND_CALL_BEFORE")


Shouldn't these be AROUND_STREAM_* ?

chemicL · 2024-09-16T09:16:33Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java

-									boolean withFinishReason = response.getResults()
+								for (AfterAdvisor advisor : onFinishElementResponseAdvisors) {
+									boolean withFinishReason = advisedResponse.response()
+										.getResults()


Is it safe to iterate over the response multiple times? I assume this bit iterates over it just to find the condition, but later processing will actually go over it to do something useful with the data. It's just a question :)

chemicL · 2024-09-16T10:35:22Z

...core/src/main/java/org/springframework/ai/chat/client/advisor/around/CacheAroundAdvisor.java

+		Flux<AdvisedResponse> advisedResponseFlux = chain.nextAroundStream(advisedRequest);

-		return new MessageAggregator().aggregate(fluxChatResponse, chatResponse -> {
+		return new MessageAggregator().aggregateAdvisedResponse(advisedResponseFlux, chatResponse -> {


It feels inefficient to apply multiple aggregators to achieve some end goal. Can this particular advisor be split into two - one being the around-advisor which can serve a cached response, while the other, the storing one, would be an aggregate-advisor. WDYT?

chemicL · 2024-09-16T10:53:05Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/AdvisedRequest.java

+		Map<String, Object> advisorParams, Map<String, Object> adviseContext) {
+
+	public AdvisedRequest augmentContext(Function<Map<String, Object>, Map<String, Object>> contextTransform) {
+		return from(this)


Consider using a compact constructor of record to assign the adviseContext to the Collections.unmodifiableMap of the passed in argument at record creation. This will encapsulate the immutability of the map. Currently, any implementor is able to access the record's map directly and modify it. Adding that safeguard will also simplify the code here.

chemicL · 2024-09-16T11:02:03Z

...ain/java/org/springframework/ai/chat/client/advisor/observation/AdvisorObservableHelper.java

 		return advisors.stream()
-			.filter(advisor -> advisor instanceof RequestAdvisor)
-			.map(a -> (RequestAdvisor) a)
+			.filter(advisor -> advisor instanceof BeforeAdvisor)


Consider pre-processing the advisors at their definition to avoid the expensive instanceof and casting at runtime, which can add up a performance hit. Perhaps specialized API methods can be added for specific advisor types to optimize the reading aspect of it?

Perhaps in ChatClient the signature ChatClientRequestSpec advisors(Consumer<AdvisorSpec> consumer) could be the one offered and the AdvisorSpec would offer type-safe, yet specific, advisor registrations for particular types. But without changing the API the DefaultChatClient could internally store particular implementations into well-defined slots upon specification by the user.

chemicL

Left a few comments. I really like the direction of having the context embedded in the Advised(Request|Response) and the reactive manipulations look well optimized. There are a few more things that can be improved I think.

chemicL · 2024-09-16T12:45:51Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/AfterAdvisor.java

 	 */
-	default Flux<ChatResponse> adviseResponse(Flux<ChatResponse> fluxResponse, Map<String, Object> adviseContext) {
-		return fluxResponse;
+	default Flux<AdvisedResponse> afterStream(Flux<AdvisedResponse> advisedResponseStream) {


This API feels that it should belong elsewhere, to a specific custom advisor type. Currently, from my understanding, this is only called for the custom case. For most regular advisors this would be confusing that this can be overridden.

chemicL · 2024-09-16T12:47:02Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/AfterAdvisor.java

 	 * @return the advised {@link ChatResponse}.
 	 */
-	ChatResponse adviseResponse(ChatResponse response, Map<String, Object> adviseContext);
+	AdvisedResponse afterCall(AdvisedResponse advisedResponse);


This name might be better aligned with BeforeAdvisor's before() signature if the Call part is skipped, or the other one also has the Call suffix.

chemicL · 2024-09-16T12:48:11Z

...i-core/src/main/java/org/springframework/ai/chat/client/advisor/api/StreamAroundAdvisor.java

 	 */
-	Flux<ChatResponse> aroundStream(AdvisedRequest advisedRequest, Map<String, Object> adviceContext,
-			AroundAdvisorChain chain);
+	Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, AroundAdvisorChain chain);


What's stopping the "Custom" implementations to be considered as the "around" advisors, making the "custom" abstraction unnecessary?

chemicL · 2024-09-16T12:50:32Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/AfterAdvisor.java

-	default StreamResponseMode getStreamResponseMode() {
-		return StreamResponseMode.ON_FINISH_ELEMENT;
+	default AfterStreamMode getAfterStreamMode() {
+		return AfterStreamMode.ON_FINISH_ELEMENT;


Would it not be more natural if a default behaviour was to apply the advisor per each chunk in the stream using the imperative implementation? Or perhaps applying it over the aggregation? If there's no easy answer, perhaps picking a default should not happen at all and stream advisors should not belong to the imperative type at all?

…psulation - Introduce AdvisedResponse to encapsulate ChatResponse and advice context - Refactor RequestAdvisor and ResponseAdvisor interface to use AdvisedRequest/AdvisedResponse records - Update DefaultAroundAdvisorChain to work with new advisor interfaces - Modify DefaultChatClient to use new advisor framework - Adjust MessageChatMemoryAdvisor for compatibility with new advisor structure - Update integration tests to reflect changes in advisor framework - Update AdvisedRequest to include adviseContext and add toPrompt() method - Modify advisor interfaces to remove explicit context parameter - Introduce AdvisedResponse to encapsulate ChatResponse with adviseContext - Update DefaultChatClient to use new advisor interfaces and AdvisedResponse - Adjust integration tests in OpenAiPaymentTransactionIT and VertexAiGeminiPaymentTransactionIT - Refactor the RequestResponseAdvisor interface to make it a backward compatible adabror for ReqeustAdvisor and ResponseAdvisors. Aslo mark it as deprecated. - Update various advisor implementations to work with new AdvisedRequest and AdvisedResponse - Modify AroundAdvisorChain to use AdvisedResponse in both call and stream methods - Implement new adviseResponse method for streaming responses in RequestResponseAdvisor - Update MessageChatMemoryAdvisor, PromptChatMemoryAdvisor, QuestionAnswerAdvisor, SimpleLoggerAdvisor, and VectorStoreChatMemoryAdvisor to use new advisor interfaces - Add new AdvisedResponse class to encapsulate ChatResponse with advice context - Update CallAroundAdvisor, StreamAroundAdvisor, RequestAdvisor, and ResponseAdvisor interfaces to use AdvisedResponse and remove adviceContext parameter - Modify CacheAroundAdvisor and SafeGuardAroundAdvisor to work with new advisor interfaces and AdvisedResponse - Update AdvisorObservableHelper to use AdvisedResponse and remove explicit context parameter - Refactor MessageAggregator to support aggregation of AdvisedResponse objects - Rename extractRequestAdvisors and extractResponseAdvisors methods to requestAdvisors and responseAdvisors in AdvisorObservableHelper - Add null checks and return empty lists for advisors in AdvisorObservableHelper

- Rename RequestAdvisor to BeforeAdvisor and ResponseAdvisor to AfterAdvisor - Change adviseRequest method to before and adviseResponse to afterCall - Update StreamResponseMode to AfterStreamMode - Adjust implementing classes and related code to use new naming conventions - Remove deprecated RequestResponseAdvisor interface

- Move AdvisedRequest from chat.client to chat.client.advisor.api - Update import statements across multiple files - Adjust related advisor classes to use the new import path - Remove unused imports and optimize imports in affected files - Add AdvisorsTests

tzolov · 2024-09-26T12:36:43Z

Replaced by #1422

tzolov added this to the 1.0.0-M3 milestone Sep 9, 2024

tzolov force-pushed the inline-advised-response-fluxed branch 2 times, most recently from a44856d to ffd645d Compare September 14, 2024 10:29

chemicL reviewed Sep 16, 2024

View reviewed changes

markpollack added the advisors label Sep 19, 2024

tzolov added 5 commits September 24, 2024 11:21

Additional tests

c218214

Address some of the review comments

65f373a

tzolov force-pushed the inline-advised-response-fluxed branch from ffd645d to 65f373a Compare September 24, 2024 17:06

tzolov closed this Sep 26, 2024


		private static final Logger logger = LoggerFactory.getLogger(MessageAggregator.class);

		public Flux<AdvisedResponse> aggregateAdvisedResponse(Flux<AdvisedResponse> advisedResponseFlux,

[WIP] Refactor AdvisedRequest and introduce AdvisedResponse #1332

[WIP] Refactor AdvisedRequest and introduce AdvisedResponse #1332

Uh oh!

Conversation

tzolov commented Sep 9, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chemicL left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tzolov commented Sep 26, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants