fix(chat): handle multiple returnDirect responses in DefaultChatClient (#4655)

Signed-off-by: Kuntal Maity <[email protected]>

Author: kuntal1461

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

@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Objects;
 import java.util.function.Consumer;
 
 import io.micrometer.observation.Observation;
@@ -47,6 +48,7 @@
 import org.springframework.ai.chat.model.ChatModel;
 import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.model.tool.ToolExecutionResult;
 import org.springframework.ai.chat.prompt.ChatOptions;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.content.Media;
@@ -521,8 +523,33 @@ private ChatClientResponse doGetObservableChatClientResponse(ChatClientRequest c
 
 		@Nullable
 		private static String getContentFromChatResponse(@Nullable ChatResponse chatResponse) {
-			return Optional.ofNullable(chatResponse)
-				.map(ChatResponse::getResult)
+			if (chatResponse == null) {
+				return null;
+			}
+			var results = chatResponse.getResults();
+			if (results == null || results.isEmpty()) {
+				return null;
+			}
+			if (results.size() == 1) {
+				return Optional.ofNullable(results.get(0))
+					.map(Generation::getOutput)
+					.map(AbstractMessage::getText)
+					.orElse(null);
+			}
+			boolean allReturnDirect = results.stream().allMatch(g -> {
+				var finish = g.getMetadata() != null ? g.getMetadata().getFinishReason() : null;
+				return finish != null && finish.equalsIgnoreCase(ToolExecutionResult.FINISH_REASON);
+			});
+			if (allReturnDirect) {
+				return results.stream()
+					.map(Generation::getOutput)
+					.map(AbstractMessage::getText)
+					.filter(Objects::nonNull)
+					.filter(StringUtils::hasText)
+					.reduce((a, b) -> a + "\n" + b)
+					.orElse(null);
+			}
+			return Optional.ofNullable(results.get(0))
 				.map(Generation::getOutput)
 				.map(AbstractMessage::getText)
 				.orElse(null);
@@ -594,10 +621,35 @@ public Flux<String> content() {
 			// @formatter:off
 			return doGetObservableFluxChatResponse(this.request)
 					.mapNotNull(ChatClientResponse::chatResponse)
-					.map(r -> Optional.ofNullable(r.getResult())
+					.map(r -> {
+						var results = r.getResults();
+						if (results == null || results.isEmpty()) {
+							return "";
+						}
+						if (results.size() == 1) {
+							return Optional.ofNullable(results.get(0))
+								.map(Generation::getOutput)
+								.map(AbstractMessage::getText)
+								.orElse("");
+						}
+						boolean allReturnDirect = results.stream().allMatch(g -> {
+							var finish = g.getMetadata() != null ? g.getMetadata().getFinishReason() : null;
+							return finish != null && finish.equalsIgnoreCase(org.springframework.ai.model.tool.ToolExecutionResult.FINISH_REASON);
+						});
+						if (allReturnDirect) {
+							return results.stream()
+								.map(Generation::getOutput)
+								.map(AbstractMessage::getText)
+								.filter(java.util.Objects::nonNull)
+								.filter(StringUtils::hasText)
+								.reduce((a, b) -> a + "\n" + b)
+								.orElse("");
+						}
+						return Optional.ofNullable(results.get(0))
 							.map(Generation::getOutput)
 							.map(AbstractMessage::getText)
-							.orElse(""))
+							.orElse("");
+					})
 					.filter(StringUtils::hasLength);
 			// @formatter:on
 		}

spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientReturnDirectAggregationTests.java

@@ -0,0 +1,63 @@
+package org.springframework.ai.chat.client;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.model.tool.ToolExecutionResult;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/*
+ * @author: Kuntal Maity
+ */
+class DefaultChatClientReturnDirectAggregationTests {
+
+	private static Generation generation(String text, String finishReason) {
+		var metadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
+		return new Generation(new AssistantMessage(text), metadata);
+	}
+
+	@Test
+	void aggregatesMultipleReturnDirectGenerationsInContent() {
+		var chatResponse = new ChatResponse(List.of(generation("DATE=2025-10-18", ToolExecutionResult.FINISH_REASON),
+				generation("TIME=12:34:56.789", ToolExecutionResult.FINISH_REASON)));
+
+		ChatModel stub = new ChatModel() {
+			@Override
+			public ChatResponse call(Prompt prompt) {
+				return chatResponse;
+			}
+		};
+
+		var client = ChatClient.builder(stub).build();
+		String content = client.prompt("now").call().content();
+
+		assertThat(content).isEqualTo("DATE=2025-10-18\nTIME=12:34:56.789");
+	}
+
+	@Test
+	void returnsFirstWhenNotAllReturnDirect() {
+		var chatResponse = new ChatResponse(
+				List.of(generation("FIRST", ToolExecutionResult.FINISH_REASON), generation("SECOND", "stop")));
+
+		ChatModel stub = new ChatModel() {
+			@Override
+			public ChatResponse call(Prompt prompt) {
+				return chatResponse;
+			}
+		};
+
+		var client = ChatClient.builder(stub).build();
+		String content = client.prompt("now").call().content();
+
+		assertThat(content).isEqualTo("FIRST");
+	}
+
+}