ChatResponseMetadata does not have data of the retrieved documents from the vector store

[iAMSagar44]

Bug description
The map in ChatResponseMetadata does not contain details of the retrieved documents from the vector store. The QuestionAnswerAdvisor has details of the retrieved documents and would have set the ChatResponseMetadata (by adding the retrieved documents to this key - qa_retrieved_documents), but I am unable to retrieve this data from the ChatResponse.

Environment
Spring AI versions - tried with both the 1.0.0-M3 and the 1.0.0-SNAPSHOT versions

Steps to reproduce
Steps to reproduce the issue.

1. Create a ChatClient as follows -

this.chatClient = chatBuilder
                                .defaultSystem(
                                                """
                                                                    You are a helpful assistant. Answer the questions based on the information provided to you by the user.
                                                                    Do not provide any information that is not relevant to the user's question.
                                                                    If you are unsure about the answer, you can ask the user for more information.
                                                                    Include the source of the information in your response at the end in a new line starting with "Source: ".
                                                                """)
                                .defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory),
                                                new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults()),
                                                new SimpleLoggerAdvisor())
                                .build();

2. Check the keys in the ChatResponse object and notice that it does not contain the key qa_retrieved_documents

StreamResponseSpec streamResponseSpec = this.chatClient.prompt()
                                .user(userRequest.message())
                                .advisors(a -> a
                                                .param(PromptChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,
                                                                conversation.getId())
                                                .param(PromptChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY, 40))
                                .stream();

                streamResponseSpec.chatResponse().blockFirst().getMetadata().keySet().forEach(key -> {
                        logger.info("Key: {}", key);
                });

Expected behavior
The ChatResponseMetadata should contain details of the retrieved documents. The application might need this data for further use cases (like retrieve the Document metadata and return it along with the Assistant Response to the calling application)

[ThomasVitale]

The ChatClient can return a ChatClientResponse which gives you access to the entire context of the advisor chain execution, including the retrieved documents from the vector store. That's by using chatClientResponse() instead of chatResponse() as final method in the ChatClient.

If you use the QuestionAnswerAdvisor, you can extract the retrieved documents as follows:

var documents = chatClientResponse.context().get(QuestionAnswerAdvisor.RETRIEVED_DOCUMENTS);

[iAMSagar44]

@ThomasVitale - I got this working for a non-streaming response. But for a scenario where I want to stream the response, retrieving the documents is resulting in an error.
My use case is to stream the response, and then after the chat content is streamed, pass the file names (for citation purposes) to the client.

Something like this -

StreamResponseSpec spec = chatClient.prompt().user(userMessage).advisors(a -> {
            a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId);
            a.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 20);
        }).advisors(RetrievalAugmentationAdvisor.builder()
                .queryTransformers(
                        CompressionQueryTransformer.builder()
                                .chatClientBuilder(builder.build().mutate()).build(),
                        RewriteQueryTransformer.builder()
                                .chatClientBuilder(builder.build().mutate()).build())
                .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(vectorStore)
                        .similarityThreshold(0.5).topK(10).build())
                .queryAugmenter(ContextualQueryAugmenter.builder().allowEmptyContext(false)
                        .promptTemplate(new PromptTemplate(CONTEXT_PROMPT)).build())
                .build()).stream();


         // Get the content as a Flux<String>
        Flux<String> contentFlux = spec.content();

        // Get the first ChatClientResponse, extract documents, and emit citations
        Flux<String> citationsFlux = spec.chatClientResponse().take(1).flatMap(response -> {
            Object docsObj = response.context().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT);
            if (docsObj instanceof List<?> docsList && !docsList.isEmpty()
                    && docsList.get(0) instanceof Document) {
                @SuppressWarnings("unchecked")
                List<Document> documents = (List<Document>) docsObj;
                Set<String> fileNames = retrieveFileNames(documents);
                return Flux.just("Citations: " + fileNames);
            } else {
                logger.warn("DOCUMENT_CONTEXT is not a List<Document>: {}", docsObj);
            }
            return Flux.empty();
        }).onErrorResume(e -> {
            logger.error("Error extracting citations", e);
            return Flux.empty();
        });


        // Concatenate the content and citations
        return contentFlux.concatWith(citationsFlux);

But this is resulting in an error as follows -

java.lang.IllegalStateException: No StreamAdvisors available to execute
	at org.springframework.ai.chat.client.advisor.DefaultAroundAdvisorChain.lambda$nextStream$10(DefaultAroundAdvisorChain.java:173) ~[spring-ai-client-chat-1.0.0-M8.jar:1.0.0-M8]
	at reactor.core.publisher.FluxDeferContextual.subscribe(FluxDeferContextual.java:49) ~[reactor-core-3.7.5.jar:3.7.5]
	at reactor.core.publisher.InternalFluxOperator.subscribe(InternalFluxOperator.java:68) ~[reactor-core-3.7.5.jar:3.7.5]
	at reactor.core.publisher.FluxDeferContextual.subscribe(FluxDeferContextual.java:57) ~[reactor-core-3.7.5.jar:3.7.5]
	at reactor.core.publisher.Flux.subscribe(Flux.java:8891) ~[reactor-core-3.7.5.jar:3.7.5]

How can I manage to retrieve the documents for this use case?