Fix MistralAiApi to handle thinking content blocks

models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiAssistantMessage.java

@@ -0,0 +1,181 @@
+/*
+ * Copyright 2023-2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.mistralai;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.content.Media;
+
+/**
+ * A Mistral AI specific implementation of {@link AssistantMessage} that supports
+ * additional fields returned by Magistral reasoning models.
+ *
+ * <p>
+ * Magistral models (like magistral-medium-latest and magistral-small-latest) return
+ * thinking/reasoning content alongside the regular response content. This class captures
+ * both the final response text and the intermediate reasoning process.
+ * </p>
+ *
+ * @author Kyle Kreuter
+ * @since 1.1.0
+ */
+public class MistralAiAssistantMessage extends AssistantMessage {
+
+	/**
+	 * The thinking/reasoning content from Magistral models. This contains the model's
+	 * intermediate reasoning steps before producing the final response.
+	 */
+	private String thinkingContent;
+
+	/**
+	 * Constructs a new MistralAiAssistantMessage with all fields.
+	 * @param content the main text content of the message
+	 * @param thinkingContent the thinking/reasoning content from Magistral models
+	 * @param properties additional metadata properties
+	 * @param toolCalls list of tool calls requested by the model
+	 * @param media list of media attachments
+	 */
+	protected MistralAiAssistantMessage(String content, String thinkingContent, Map<String, Object> properties,
+			List<ToolCall> toolCalls, List<Media> media) {
+		super(content, properties, toolCalls, media);
+		this.thinkingContent = thinkingContent;
+	}
+
+	/**
+	 * Returns the thinking/reasoning content from Magistral models.
+	 * @return the thinking content, or null if not available
+	 */
+	public String getThinkingContent() {
+		return this.thinkingContent;
+	}
+
+	/**
+	 * Sets the thinking/reasoning content.
+	 * @param thinkingContent the thinking content to set
+	 * @return this instance for method chaining
+	 */
+	public MistralAiAssistantMessage setThinkingContent(String thinkingContent) {
+		this.thinkingContent = thinkingContent;
+		return this;
+	}
+
+	@Override
+	public boolean equals(Object o) {
+		if (this == o) {
+			return true;
+		}
+		if (!(o instanceof MistralAiAssistantMessage that)) {
+			return false;
+		}
+		if (!super.equals(o)) {
+			return false;
+		}
+		return Objects.equals(this.thinkingContent, that.thinkingContent);
+	}
+
+	@Override
+	public int hashCode() {
+		return Objects.hash(super.hashCode(), this.thinkingContent);
+	}
+
+	@Override
+	public String toString() {
+		return "MistralAiAssistantMessage{" + "media=" + this.media + ", messageType=" + this.messageType
+				+ ", metadata=" + this.metadata + ", thinkingContent='" + this.thinkingContent + '\''
+				+ ", textContent='" + this.textContent + '\'' + '}';
+	}
+
+	/**
+	 * Builder for creating MistralAiAssistantMessage instances.
+	 */
+	public static final class Builder {
+
+		private String content;
+
+		private Map<String, Object> properties = Map.of();
+
+		private List<ToolCall> toolCalls = List.of();
+
+		private List<Media> media = List.of();
+
+		private String thinkingContent;
+
+		/**
+		 * Sets the main text content.
+		 * @param content the content to set
+		 * @return this builder
+		 */
+		public Builder content(String content) {
+			this.content = content;
+			return this;
+		}
+
+		/**
+		 * Sets the metadata properties.
+		 * @param properties the properties to set
+		 * @return this builder
+		 */
+		public Builder properties(Map<String, Object> properties) {
+			this.properties = properties;
+			return this;
+		}
+
+		/**
+		 * Sets the tool calls.
+		 * @param toolCalls the tool calls to set
+		 * @return this builder
+		 */
+		public Builder toolCalls(List<ToolCall> toolCalls) {
+			this.toolCalls = toolCalls;
+			return this;
+		}
+
+		/**
+		 * Sets the media attachments.
+		 * @param media the media to set
+		 * @return this builder
+		 */
+		public Builder media(List<Media> media) {
+			this.media = media;
+			return this;
+		}
+
+		/**
+		 * Sets the thinking/reasoning content from Magistral models.
+		 * @param thinkingContent the thinking content to set
+		 * @return this builder
+		 */
+		public Builder thinkingContent(String thinkingContent) {
+			this.thinkingContent = thinkingContent;
+			return this;
+		}
+
+		/**
+		 * Builds the MistralAiAssistantMessage instance.
+		 * @return a new MistralAiAssistantMessage
+		 */
+		public MistralAiAssistantMessage build() {
+			return new MistralAiAssistantMessage(this.content, this.thinkingContent, this.properties, this.toolCalls,
+					this.media);
+		}
+
+	}
+
+}

models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java

@@ -358,8 +358,8 @@ private Generation buildGeneration(Choice choice, Map<String, Object> metadata)
 							toolCall.function().name(), toolCall.function().arguments()))
 					.toList();
 
-		var assistantMessage = AssistantMessage.builder()
-			.content(choice.message().content())
+		var assistantMessage = new MistralAiAssistantMessage.Builder().content(choice.message().content())
+			.thinkingContent(choice.message().thinkingContent())
 			.properties(metadata)
 			.toolCalls(toolCalls)
 			.build();

models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java

@@ -63,6 +63,7 @@
  * @author Thomas Vitale
  * @author Jason Smith
  * @author Nicolas Krier
+ * @author Kyle Kreuter
  * @since 1.0.0
  */
 public class MistralAiApi {
@@ -209,6 +210,51 @@ public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chat
 			.flatMap(mono -> mono);
 	}
 
+	/**
+	 * Sealed interface for content chunks returned by Magistral reasoning models.
+	 * Magistral models can return content as an array of typed blocks instead of a simple
+	 * string.
+	 *
+	 * @since 1.0.0
+	 */
+	public sealed interface ContentChunk permits TextChunk, ThinkChunk, ReferenceChunk {
+
+	}
+
+	/**
+	 * A text content chunk containing the main response text.
+	 *
+	 * @param text the text content
+	 */
+	@JsonInclude(Include.NON_NULL)
+	@JsonIgnoreProperties(ignoreUnknown = true)
+	public record TextChunk(@JsonProperty("text") String text) implements ContentChunk {
+
+	}
+
+	/**
+	 * A thinking/reasoning content chunk from Magistral models. Contains the model's
+	 * intermediate reasoning process.
+	 *
+	 * @param thinking the thinking/reasoning content
+	 */
+	@JsonInclude(Include.NON_NULL)
+	@JsonIgnoreProperties(ignoreUnknown = true)
+	public record ThinkChunk(@JsonProperty("thinking") String thinking) implements ContentChunk {
+
+	}
+
+	/**
+	 * A reference content chunk containing citation reference IDs.
+	 *
+	 * @param referenceIds list of reference IDs for citations
+	 */
+	@JsonInclude(Include.NON_NULL)
+	@JsonIgnoreProperties(ignoreUnknown = true)
+	public record ReferenceChunk(@JsonProperty("reference_ids") List<Integer> referenceIds) implements ContentChunk {
+
+	}
+
 	/**
 	 * The reason the model stopped generating tokens.
 	 */
@@ -806,7 +852,9 @@ public ChatCompletionMessage(Object content, Role role) {
 		}
 
 		/**
-		 * Get message content as String.
+		 * Returns the text content of the message. For reasoning models (Magistral),
+		 * extracts the text block from the content array.
+		 * @return the text content or null if not available
 		 */
 		public String content() {
 			if (this.rawContent == null) {
@@ -815,7 +863,132 @@ public String content() {
 			if (this.rawContent instanceof String text) {
 				return text;
 			}
-			throw new IllegalStateException("The content is not a string!");
+			if (this.rawContent instanceof List<?> blocks) {
+				StringBuilder textBuilder = new StringBuilder();
+				for (Object block : blocks) {
+					if (block instanceof Map<?, ?> map && "text".equals(map.get("type"))) {
+						Object text = map.get("text");
+						if (text instanceof String s) {
+							if (!textBuilder.isEmpty()) {
+								textBuilder.append("\n");
+							}
+							textBuilder.append(s);
+						}
+					}
+				}
+				return textBuilder.isEmpty() ? null : textBuilder.toString();
+			}
+			throw new IllegalStateException("Unexpected content type: " + rawContent.getClass());
+		}
+
+		/**
+		 * Returns the thinking/reasoning content from Magistral models. For non-Magistral
+		 * models or when no thinking content is present, returns null.
+		 * @return the thinking content or null if not available
+		 */
+		public String thinkingContent() {
+			if (this.rawContent == null) {
+				return null;
+			}
+			if (this.rawContent instanceof String) {
+				return null;
+			}
+			if (this.rawContent instanceof List<?> blocks) {
+				StringBuilder thinkingBuilder = new StringBuilder();
+				for (Object block : blocks) {
+					if (block instanceof Map<?, ?> map && "thinking".equals(map.get("type"))) {
+						Object thinking = map.get("thinking");
+						if (thinking instanceof List<?> thinkingBlocks) {
+							for (Object thinkingBlock : thinkingBlocks) {
+								if (thinkingBlock instanceof Map<?, ?> thinkingMap
+										&& "text".equals(thinkingMap.get("type"))) {
+									Object text = thinkingMap.get("text");
+									if (text instanceof String s) {
+										if (!thinkingBuilder.isEmpty()) {
+											thinkingBuilder.append("\n");
+										}
+										thinkingBuilder.append(s);
+									}
+								}
+							}
+						}
+						else if (thinking instanceof String s) {
+							if (!thinkingBuilder.isEmpty()) {
+								thinkingBuilder.append("\n");
+							}
+							thinkingBuilder.append(s);
+						}
+					}
+				}
+				return thinkingBuilder.isEmpty() ? null : thinkingBuilder.toString();
+			}
+			return null;
+		}
+
+		/**
+		 * Parses the raw content into a list of typed ContentChunk objects. For string
+		 * content, returns a single TextChunk. For array content from Magistral models,
+		 * parses each block into its appropriate type.
+		 * @return list of ContentChunk objects, or empty list if content is null
+		 */
+		@SuppressWarnings("unchecked")
+		public List<ContentChunk> contentChunks() {
+			if (this.rawContent == null) {
+				return List.of();
+			}
+			if (this.rawContent instanceof String text) {
+				return List.of(new TextChunk(text));
+			}
+			if (this.rawContent instanceof List<?> blocks) {
+				List<ContentChunk> chunks = new java.util.ArrayList<>();
+				for (Object block : blocks) {
+					if (block instanceof Map<?, ?> map) {
+						String type = (String) map.get("type");
+						if ("text".equals(type)) {
+							String text = (String) map.get("text");
+							if (text != null) {
+								chunks.add(new TextChunk(text));
+							}
+						}
+						else if ("thinking".equals(type)) {
+							Object thinking = map.get("thinking");
+							if (thinking instanceof List<?> thinkingBlocks) {
+								StringBuilder thinkingBuilder = new StringBuilder();
+								for (Object thinkingBlock : thinkingBlocks) {
+									if (thinkingBlock instanceof Map<?, ?> thinkingMap
+											&& "text".equals(thinkingMap.get("type"))) {
+										Object text = thinkingMap.get("text");
+										if (text instanceof String s) {
+											if (!thinkingBuilder.isEmpty()) {
+												thinkingBuilder.append("\n");
+											}
+											thinkingBuilder.append(s);
+										}
+									}
+								}
+								if (!thinkingBuilder.isEmpty()) {
+									chunks.add(new ThinkChunk(thinkingBuilder.toString()));
+								}
+							}
+							else if (thinking instanceof String s) {
+								chunks.add(new ThinkChunk(s));
+							}
+						}
+						else if ("reference".equals(type)) {
+							Object refIds = map.get("reference_ids");
+							if (refIds instanceof List<?> ids) {
+								List<Integer> referenceIds = ((List<Object>) ids).stream()
+									.filter(id -> id instanceof Number)
+									.map(id -> ((Number) id).intValue())
+									.toList();
+								chunks.add(new ReferenceChunk(referenceIds));
+							}
+						}
+					}
+				}
+				return chunks;
+			}
+			return List.of();
 		}
 
 		/**