fix(mistral) Added index of tool call in the list of tool calls

Additiona fixes: API key validation and tool calling backward compatibility

- Fix API key validation in OpenAiApi builder
- Standardize API key validation using Assert.notNull
- Add backward compatibility support for FunctionCallback in tool calling
- Update integration tests to use LegacyToolCallingManager

Co-authored-by:Christian Tzolov <[email protected]>
Signed-off-by: Ricken Bazolo <[email protected]>
Signed-off-by: Christian Tzolov <[email protected]>

Author: ricken07

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

@@ -384,7 +384,7 @@ else if (message instanceof AssistantMessage assistantMessage) {
 				if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
 					toolCalls = assistantMessage.getToolCalls().stream().map(toolCall -> {
 						var function = new ChatCompletionFunction(toolCall.name(), toolCall.arguments());
-						return new ToolCall(toolCall.id(), toolCall.type(), function);
+						return new ToolCall(toolCall.id(), toolCall.type(), function, null);
 					}).toList();
 				}
 

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

@@ -833,10 +833,11 @@ public enum Role {
 		 * @param type The type of tool call the output is required for. For now, this is
 		 * always function.
 		 * @param function The function definition.
+		 * @param index The index of the tool call in the list of tool calls.
 		 */
 		@JsonInclude(Include.NON_NULL)
 		public record ToolCall(@JsonProperty("id") String id, @JsonProperty("type") String type,
-				@JsonProperty("function") ChatCompletionFunction function) {
+				@JsonProperty("function") ChatCompletionFunction function, @JsonProperty("index") Integer index) {
 
 		}
 

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

@@ -18,6 +18,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
 
@@ -74,16 +75,16 @@ private ChunkChoice merge(ChunkChoice previous, ChunkChoice current) {
 				Optional<String> id = current.delta()
 					.toolCalls()
 					.stream()
-					.filter(tool -> tool.id() != null)
-					.map(tool -> tool.id())
+					.map(ToolCall::id)
+					.filter(Objects::nonNull)
 					.findFirst();
-				if (!id.isPresent()) {
+				if (id.isEmpty()) {
 					var newId = UUID.randomUUID().toString();
 
 					var toolCallsWithID = current.delta()
 						.toolCalls()
 						.stream()
-						.map(toolCall -> new ToolCall(newId, "function", toolCall.function()))
+						.map(toolCall -> new ToolCall(newId, "function", toolCall.function(), toolCall.index()))
 						.toList();
 
 					var role = current.delta().role() != null ? current.delta().role() : Role.ASSISTANT;
@@ -151,7 +152,8 @@ private ToolCall merge(ToolCall previous, ToolCall current) {
 		String id = (current.id() != null ? current.id() : previous.id());
 		String type = (current.type() != null ? current.type() : previous.type());
 		ChatCompletionFunction function = merge(previous.function(), current.function());
-		return new ToolCall(id, type, function);
+		Integer index = (current.index() != null ? current.index() : previous.index());
+		return new ToolCall(id, type, function, index);
 	}
 
 	private ChatCompletionFunction merge(ChatCompletionFunction previous, ChatCompletionFunction current) {

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java

@@ -1675,6 +1675,12 @@ public Builder apiKey(ApiKey apiKey) {
 			return this;
 		}
 
+		public Builder apiKey(String simpleApiKey) {
+			Assert.notNull(simpleApiKey, "apiKey cannot be null");
+			this.apiKey = new SimpleApiKey(simpleApiKey);
+			return this;
+		}
+
 		public Builder headers(MultiValueMap<String, String> headers) {
 			Assert.notNull(headers, "headers cannot be null");
 			this.headers = headers;

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/MistralWithOpenAiChatModelIT.java

@@ -46,12 +46,16 @@
 import org.springframework.ai.converter.ListOutputConverter;
 import org.springframework.ai.converter.MapOutputConverter;
 import org.springframework.ai.model.Media;
+import org.springframework.ai.model.SimpleApiKey;
 import org.springframework.ai.model.function.FunctionCallback;
+import org.springframework.ai.model.tool.LegacyToolCallingManager;
 import org.springframework.ai.openai.OpenAiChatModel;
 import org.springframework.ai.openai.OpenAiChatOptions;
 import org.springframework.ai.openai.api.OpenAiApi;
 import org.springframework.ai.openai.api.tool.MockWeatherService;
 import org.springframework.ai.openai.chat.ActorsFilms;
+import org.springframework.ai.tool.function.FunctionToolCallback;
+import org.springframework.ai.tool.method.MethodToolCallback;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.SpringBootConfiguration;
@@ -61,6 +65,7 @@
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.core.io.Resource;
 import org.springframework.util.MimeTypeUtils;
+import org.springframework.util.ReflectionUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -251,8 +256,7 @@ void functionCallTest(String modelName) {
 
 		var promptOptions = OpenAiChatOptions.builder()
 			.model(modelName)
-			.functionCallbacks(List.of(FunctionCallback.builder()
-				.function("getCurrentWeather", new MockWeatherService())
+			.toolCallbacks(List.of(FunctionToolCallback.builder("getCurrentWeather", new MockWeatherService())
 				.description("Get the weather in location")
 				.inputType(MockWeatherService.Request.class)
 				.build()))
@@ -276,8 +280,7 @@ void streamFunctionCallTest(String modelName) {
 
 		var promptOptions = OpenAiChatOptions.builder()
 			.model(modelName)
-			.functionCallbacks(List.of(FunctionCallback.builder()
-				.function("getCurrentWeather", new MockWeatherService())
+			.toolCallbacks(List.of(FunctionToolCallback.builder("getCurrentWeather", new MockWeatherService())
 				.description("Get the weather in location")
 				.inputType(MockWeatherService.Request.class)
 				.build()))
@@ -388,12 +391,16 @@ static class Config {
 
 		@Bean
 		public OpenAiApi chatCompletionApi() {
-			return new OpenAiApi(MISTRAL_BASE_URL, System.getenv("MISTRAL_AI_API_KEY"));
+			return OpenAiApi.builder().baseUrl(MISTRAL_BASE_URL).apiKey(System.getenv("MISTRAL_AI_API_KEY")).build();
 		}
 
 		@Bean
 		public OpenAiChatModel openAiClient(OpenAiApi openAiApi) {
-			return new OpenAiChatModel(openAiApi, OpenAiChatOptions.builder().model(MISTRAL_DEFAULT_MODEL).build());
+			return OpenAiChatModel.builder()
+				.openAiApi(openAiApi)
+				.toolCallingManager(LegacyToolCallingManager.builder().build())
+				.defaultOptions(OpenAiChatOptions.builder().model(MISTRAL_DEFAULT_MODEL).build())
+				.build();
 		}
 
 	}

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/OllamaWithOpenAiChatModelIT.java

@@ -50,11 +50,13 @@
 import org.springframework.ai.converter.MapOutputConverter;
 import org.springframework.ai.model.Media;
 import org.springframework.ai.model.function.FunctionCallback;
+import org.springframework.ai.model.tool.LegacyToolCallingManager;
 import org.springframework.ai.openai.OpenAiChatModel;
 import org.springframework.ai.openai.OpenAiChatOptions;
 import org.springframework.ai.openai.api.OpenAiApi;
 import org.springframework.ai.openai.api.tool.MockWeatherService;
 import org.springframework.ai.openai.chat.ActorsFilms;
+import org.springframework.ai.tool.function.FunctionToolCallback;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.SpringBootConfiguration;
@@ -272,8 +274,7 @@ void functionCallTest(String modelName) {
 			// Note for Ollama you must set the tool choice to explicitly. Unlike OpenAI
 			// (which defaults to "auto") Ollama defaults to "nono"
 			.toolChoice("auto")
-			.functionCallbacks(List.of(FunctionCallback.builder()
-				.function("getCurrentWeather", new MockWeatherService())
+			.toolCallbacks(List.of(FunctionToolCallback.builder("getCurrentWeather", new MockWeatherService())
 				.description("Get the weather in location")
 				.inputType(MockWeatherService.Request.class)
 				.build()))
@@ -412,12 +413,16 @@ static class Config {
 
 		@Bean
 		public OpenAiApi chatCompletionApi() {
-			return new OpenAiApi(baseUrl, "");
+			return OpenAiApi.builder().baseUrl(baseUrl).apiKey("").build();
 		}
 
 		@Bean
 		public OpenAiChatModel openAiClient(OpenAiApi openAiApi) {
-			return new OpenAiChatModel(openAiApi, OpenAiChatOptions.builder().model(DEFAULT_OLLAMA_MODEL).build());
+			return OpenAiChatModel.builder()
+				.openAiApi(openAiApi)
+				.toolCallingManager(LegacyToolCallingManager.builder().build())
+				.defaultOptions(OpenAiChatOptions.builder().model(DEFAULT_OLLAMA_MODEL).build())
+				.build();
 		}
 
 	}

spring-ai-core/src/main/java/org/springframework/ai/model/SimpleApiKey.java

@@ -23,46 +23,28 @@
  * be refreshed or rotated.
  *
  * @author Adib Saikali
+ * @author Christian Tzolov
  * @since 1.0.0
  */
-public final class SimpleApiKey implements ApiKey {
-
-	private final String value;
+public record SimpleApiKey(String value) implements ApiKey {
 
 	/**
 	 * Create a new SimpleApiKey.
 	 * @param value the API key value, must not be null or empty
 	 * @throws IllegalArgumentException if value is null or empty
 	 */
 	public SimpleApiKey(String value) {
-		Assert.hasText(value, "API key value must not be null or empty");
+		Assert.notNull(value, "API key value must not be null or empty");
 		this.value = value;
 	}
 
 	@Override
 	public String getValue() {
-		return this.value;
+		return this.value();
 	}
 
 	@Override
 	public String toString() {
 		return "SimpleApiKey{value='***'}";
 	}
-
-	@Override
-	public boolean equals(Object o) {
-		if (this == o) {
-			return true;
-		}
-		if (!(o instanceof SimpleApiKey that)) {
-			return false;
-		}
-		return this.value.equals(that.value);
-	}
-
-	@Override
-	public int hashCode() {
-		return this.value.hashCode();
-	}
-
 }

spring-ai-core/src/main/java/org/springframework/ai/model/tool/DefaultToolCallingManager.java

@@ -214,6 +214,12 @@ else if (prompt.getOptions() instanceof FunctionCallingOptions functionOptions)
 			else if (toolCallback instanceof ToolCallback callback) {
 				returnDirect = returnDirect && callback.getToolMetadata().returnDirect();
 			}
+			else if (returnDirect == null) {
+				// This is a temporary solution to ensure backward compatibility with
+				// FunctionCallback.
+				// TODO: remove this block when FunctionCallback is removed.
+				returnDirect = false;
+			}
 
 			String toolResult;
 			try {