Spring AI VectorStoreChatMemoryAdvisor Never Persist Assistant Response in Database For Bedrock Converse Streaming

Bug description VectorStoreChatMemoryAdvisor after method is not never called for stream based advisor. This leads not having the ASSISTANT messages being persisted in in database i.e Postgres with PG Vector

Note, this issue was not present in M8 and it happened after upgrading to RC1. This probably caused by StreamAroundAdvisor removal and relying on after method instead. This behavior have changed introduced in RC1. See below for detailed investigation.

Environment This is occurs in RC1 release.

Steps to reproduce 1. Configure a VectorStoreChatMemoryAdvisor bean. Note The vector stored injected here of type `PgVectorStore'


  @Bean
    public VectorStoreChatMemoryAdvisor vectorStoreChatMemoryAdvisor(VectorStore chatMemoryStore) {
        return VectorStoreChatMemoryAdvisor.builder(chatMemoryStore)
                .systemPromptTemplate(PromptTemplate.builder()
                        .build())
                .defaultTopK(CHAT_MEMORY_RETRIEVE_SIZE)
                .build();

    }

Add this advisor to a chatClient. AWS Bedrock with Claude as LLM is used. The chatClient is initialized in RestController

        this.chatClient = chatClientBuilder.defaultAdvisors(List.of(vectorStoreChatMemoryAdvisor)).build();

Converse with chat client

    @GetMapping("/conversations/{id}/stream/")
    public Flux<String> helloWorld() {
        return chatClient.prompt()
                .system("You are a chat agent. Please response to user queries")
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.DEFAULT_CONVERSATION_ID, 10))
                .user("hi").stream().content();
    }

The database will not persist the Assistant responses. Only the User message is persisted

Further Investigation & Problem Analysis

In pre-RC1 releases such as M8, VectorStoreChatMemoryAdvisor implemented StreamAroundAdvisor where it aggeregated message reponses. Link to code

This whereobserveAfter calls this.getChatMemoryStore().write.... to write the chat Responses to the vector store.

In RC1 however, VectorStoreChatMemoryAdvisor switched to after ( StreamAroundAdvisor is deprecated & removed) where after method is defined in BaseAdvisor. After method is called in adviseStream. See here

However, for the after to be invoked, (AdvisorUtils.onFinishReason().test(response) needs to return true. So this means the condition here needs hold. Specifically, result.getMetadata().getFinishReason() needs to have a text.

finishReason is populated when ConverseApiUtils, intercepts MessageStopEvent (part of the AWS chat response) to populate that field. See here for how. However, this piece of block never get executed and I have put breaking points on it.

Though MessageStopEvent is created ( I have put breaking point there) in the in the stack but it seems that it get "lost" somewhere.

Update: Regarding to MessageStopEvent to ConverseApiUtils interception, it seems this also happens in M8 release

Expected behavior All the assistant user messages types should have been persisted in database

Comment From: ilayaperumalg

@saadswift Could you check if this is fixed via https://github.com/spring-projects/spring-ai/pull/3193 by testing against the latest 1.0.0-SNAPSHOT?

Comment From: saadswift

I can confirm with latest with SNAPSHOT that the regression fix above did the job . I can see that the MessageAggeregator is invoking adviseStream from BedrockProxyChatModel's AggeregationHandler .Thank you for bring this up @ilayaperumalg.

One thing that still a bothersome and be addressed is the MessageStopEvent needs to be addressed because eventually the metadata of resultant response still does not have the finish reason populated. This will cause issue down the line