Return IDs of persisting ChatMemory messages

[philip-ft]

Hello,

I have implemented a custom ChatMemory with my own repository to persist messages in a table format that suits our project. I then set this MessageChatMemoryAdvisor in the default advisor in the ChatClient builder which built my ChatClient instance to use. I was wondering if anyone done it before if there is an easy way to return the persistent IDs of the user and response from LLM messages so it can be returned to the frontend chatbot.

I know there is a chat context to store them and retrieve them in the Advisors but this means that I need to implement my own advisor to make this happen. Does anyone have any ideas?

[192902649]

Currently, when using Spring AI, the execution flow is as follows: the responseId is first generated in the Advisor.before() phase and stored in the context. Then, in the Advisor.after() phase, new records are inserted. This approach ensures that the response ID for the current AI conversation is returned to the frontend in advance, making the whole process more controllable.

---

public ChatClientRequest before(ChatClientRequest chatClientRequest, AdvisorChain advisorChain) {
Map<String, Object> context = chatClientRequest.context();
Long userId = (Long) context.get(YanSenChatMemory.USER_ID);
String conversationId = getConversationId(chatClientRequest.context(), this.defaultConversationId);

	// 1. Retrieve the chat memory for the current conversation.
	// 获取聊天记忆
	List<Message> memoryMessages = this.chatMemory.get(conversationId,userId);

	// 2. Advise the request messages list.
	List<Message> processedMessages = new ArrayList<>(memoryMessages);
	processedMessages.addAll(chatClientRequest.prompt().getInstructions());

	// 3. Create a new request with the advised messages.
	ChatClientRequest processedChatClientRequest = chatClientRequest.mutate()
		.prompt(chatClientRequest.prompt().mutate().messages(processedMessages).build())
		.build();
	// 4. Add the new user message to the conversation memory.
	UserMessage userMessage = processedChatClientRequest.prompt().getUserMessage();
	//this.chatMemory.add(conversationId,userId, userMessage);
	processedChatClientRequest.context().put(MessageType.USER.getValue(), userMessage);
	// save response message id
	processedChatClientRequest.context().put(RESPONSE_ID, System.currentTimeMillis()+Long.valueOf(conversationId)+   userId);
	return processedChatClientRequest;

}

---

public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {
List assistantMessages = new ArrayList<>();
if (chatClientResponse.chatResponse() != null) {
assistantMessages = chatClientResponse.chatResponse()
.getResults()
.stream()
.map(g -> (Message) g.getOutput())
.toList();
}
Map<String, Object> context = chatClientResponse.context();
//
UserMessage userMessage = (UserMessage) context.get(MessageType.USER.getValue());
List newAssistantMessages = new ArrayList<>();
newAssistantMessages.add(userMessage);
newAssistantMessages.addAll(assistantMessages);
this.chatMemory.add(context,newAssistantMessages);
return chatClientResponse;
}

---

public void save(Map<String, Object> context, List messages) {
Long userId = (Long) context.get(YanSenChatMemory.USER_ID);
String conversationId = (String) context.get(YanSenChatMemory.CONVERSATION_ID);
Long responseId = (Long) context.get( YanSenMessageChatMemoryAdvisor.RESPONSE_ID);

    Message message = messages.get(0);
    AiChatMemory aiChatMemory = new AiChatMemory();
    aiChatMemory.setType(message.getMessageType().name());
    aiChatMemory.setContent(message.getText());
    aiChatMemory.setCreateTime(LocalDateTime.now());
    aiChatMemory.setConversationId(conversationId);

    aiChatMemory.setUserId(userId);
    aiChatMemoryMapper.insert(aiChatMemory); // 插入聊天
    message = messages.get(1);
    AiChatMemory aiChatResp = new AiChatMemory();
    aiChatResp.setId(responseId);
    aiChatResp.setType(message.getMessageType().name());
    aiChatResp.setContent(message.getText());
    aiChatResp.setCreateTime(LocalDateTime.now());
    aiChatResp.setConversationId(conversationId);
    aiChatResp.setUserId(userId);
    aiChatMemoryMapper.insert(aiChatResp);
    List<Document> documents = (List<Document>) context.get(DOCUMENT_CONTEXT);
    if (!CollectionUtils.isEmpty(documents)) {

        StringBuilder ids = new StringBuilder();
        StringBuilder scores = new StringBuilder();

        HashSet<String> fileUrlsSet = new HashSet();
        Document document = documents.get(0);
        String fileUrl = (String) document.getMetadata().get("fileUrl");
        fileUrlsSet.add(fileUrl);
        String id = document.getId();
        Double score = document.getScore();

        ids.append(id);
        scores.append(score.toString());
        for (int i = 1; i < documents.size(); i++) {
            document = documents.get(i);
             id = document.getId();
             score = document.getScore();
            ids.append(","+id);
            scores.append(","+score.toString());
             fileUrl = (String) document.getMetadata().get("fileUrl");
            fileUrlsSet.add(fileUrl);
        }
        ChatMemoryRagKnowledgeHit chatMemoryRagKnowledgeHit = new ChatMemoryRagKnowledgeHit();
        String files = fileUrlsSet.stream().collect(Collectors.joining(","));
        chatMemoryRagKnowledgeHit.setFileUrls(files);
        chatMemoryRagKnowledgeHit.setMessageId(responseId);
        chatMemoryRagKnowledgeHit.setKnowledgeId(ids.toString());
        chatMemoryRagKnowledgeHit.setSimilarityScore(scores.toString());
        chatMemoryRagKnowledgeHit.setRetrievalTime(LocalDateTime.now());
        chatMemoryRagKnowledgeHitMapper.insert(chatMemoryRagKnowledgeHit);
    }

---