refactor : refactor MessageAggregator to include toolCalls

Author: WOONBE

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

@@ -16,6 +16,7 @@
 
 package org.springframework.ai.chat.model;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -24,6 +25,7 @@
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.util.CollectionUtils;
 import reactor.core.publisher.Flux;
 
 import org.springframework.ai.chat.messages.AssistantMessage;
@@ -35,13 +37,16 @@
 import org.springframework.ai.chat.metadata.Usage;
 import org.springframework.util.StringUtils;
 
+import static org.springframework.ai.chat.messages.AssistantMessage.*;
+
 /**
  * Helper that for streaming chat responses, aggregate the chat response messages into a
  * single AssistantMessage. Job is performed in parallel to the chat response processing.
  *
  * @author Christian Tzolov
  * @author Alexandros Pappas
  * @author Thomas Vitale
+ * @author Heonwoo Kim
  * @since 1.0.0
  */
 public class MessageAggregator {
@@ -54,6 +59,7 @@ public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
 		// Assistant Message
 		AtomicReference<StringBuilder> messageTextContentRef = new AtomicReference<>(new StringBuilder());
 		AtomicReference<Map<String, Object>> messageMetadataMapRef = new AtomicReference<>();
+		AtomicReference<List<ToolCall>> toolCallsRef = new AtomicReference<>(new ArrayList<>());
 
 		// ChatGeneration Metadata
 		AtomicReference<ChatGenerationMetadata> generationMetadataRef = new AtomicReference<>(
@@ -73,6 +79,7 @@ public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
 		return fluxChatResponse.doOnSubscribe(subscription -> {
 			messageTextContentRef.set(new StringBuilder());
 			messageMetadataMapRef.set(new HashMap<>());
+			toolCallsRef.set(new ArrayList<>());
 			metadataIdRef.set("");
 			metadataModelRef.set("");
 			metadataUsagePromptTokensRef.set(0);
@@ -94,6 +101,11 @@ public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
 				if (chatResponse.getResult().getOutput().getMetadata() != null) {
 					messageMetadataMapRef.get().putAll(chatResponse.getResult().getOutput().getMetadata());
 				}
+				AssistantMessage outputMessage = chatResponse.getResult().getOutput();
+				if (!CollectionUtils.isEmpty(outputMessage.getToolCalls())) {
+					toolCallsRef.get().addAll(outputMessage.getToolCalls());
+				}
+
 			}
 			if (chatResponse.getMetadata() != null) {
 				if (chatResponse.getMetadata().getUsage() != null) {
@@ -119,6 +131,13 @@ public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
 				if (StringUtils.hasText(chatResponse.getMetadata().getModel())) {
 					metadataModelRef.set(chatResponse.getMetadata().getModel());
 				}
+				Object toolCallsFromMetadata = chatResponse.getMetadata().get("toolCalls");
+				if (toolCallsFromMetadata instanceof List) {
+					@SuppressWarnings("unchecked")
+					List<ToolCall> toolCallsList = (List<ToolCall>) toolCallsFromMetadata;
+					toolCallsRef.get().addAll(toolCallsList);
+				}
+
 			}
 		}).doOnComplete(() -> {
 
@@ -133,12 +152,25 @@ public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
 				.promptMetadata(metadataPromptMetadataRef.get())
 				.build();
 
-			onAggregationComplete.accept(new ChatResponse(List.of(new Generation(
-					new AssistantMessage(messageTextContentRef.get().toString(), messageMetadataMapRef.get()),
+			AssistantMessage finalAssistantMessage;
+			List<ToolCall> collectedToolCalls = toolCallsRef.get();
+
+			if (!CollectionUtils.isEmpty(collectedToolCalls)) {
+
+				finalAssistantMessage = new AssistantMessage(messageTextContentRef.get().toString(),
+						messageMetadataMapRef.get(), collectedToolCalls);
+			}
+			else {
+				finalAssistantMessage = new AssistantMessage(messageTextContentRef.get().toString(),
+						messageMetadataMapRef.get());
+			}
+			onAggregationComplete.accept(new ChatResponse(List.of(new Generation(finalAssistantMessage,
+
 					generationMetadataRef.get())), chatResponseMetadata));
 
 			messageTextContentRef.set(new StringBuilder());
 			messageMetadataMapRef.set(new HashMap<>());
+			toolCallsRef.set(new ArrayList<>());
 			metadataIdRef.set("");
 			metadataModelRef.set("");
 			metadataUsagePromptTokensRef.set(0);