diff --git a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java index 8bbd417bc..438482174 100644 --- a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java +++ b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java @@ -16,7 +16,6 @@ package org.springframework.ai.vertexai.gemini; -import com.google.cloud.vertexai.api.Tool.GoogleSearch; import java.util.ArrayList; import java.util.Collection; import java.util.List; @@ -24,6 +23,7 @@ import java.util.Map; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonInclude.Include; +import com.fasterxml.jackson.databind.JsonNode; import com.google.cloud.vertexai.VertexAI; import com.google.cloud.vertexai.api.Candidate; import com.google.cloud.vertexai.api.Candidate.FinishReason; @@ -33,15 +33,16 @@ import com.google.cloud.vertexai.api.FunctionDeclaration; import com.google.cloud.vertexai.api.FunctionResponse; import com.google.cloud.vertexai.api.GenerateContentResponse; import com.google.cloud.vertexai.api.GenerationConfig; -import com.google.cloud.vertexai.api.GoogleSearchRetrieval; import com.google.cloud.vertexai.api.Part; import com.google.cloud.vertexai.api.SafetySetting; import com.google.cloud.vertexai.api.Schema; import com.google.cloud.vertexai.api.Tool; +import com.google.cloud.vertexai.api.Tool.GoogleSearch; import com.google.cloud.vertexai.generativeai.GenerativeModel; import com.google.cloud.vertexai.generativeai.PartMaker; import com.google.cloud.vertexai.generativeai.ResponseStream; import com.google.protobuf.Struct; +import com.google.protobuf.Value; import com.google.protobuf.util.JsonFormat; import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; @@ -226,7 +227,8 @@ public class VertexAiGeminiChatModel implements ChatModel, DisposableBean { this.observationRegistry = observationRegistry; this.toolExecutionEligibilityPredicate = toolExecutionEligibilityPredicate; - // Wrap the provided tool calling manager in a VertexToolCallingManager to ensure + // Wrap the provided tool calling manager in a VertexToolCallingManager to + // ensure // compatibility with Vertex AI's OpenAPI schema format. if (toolCallingManager instanceof VertexToolCallingManager) { this.toolCallingManager = toolCallingManager; @@ -334,8 +336,34 @@ public class VertexAiGeminiChatModel implements ChatModel, DisposableBean { private static Struct jsonToStruct(String json) { try { - var structBuilder = Struct.newBuilder(); - JsonFormat.parser().ignoringUnknownFields().merge(json, structBuilder); + JsonNode rootNode = ModelOptionsUtils.OBJECT_MAPPER.readTree(json); + + Struct.Builder structBuilder = Struct.newBuilder(); + + if (rootNode.isArray()) { + // Handle JSON array + List values = new ArrayList<>(); + + for (JsonNode element : rootNode) { + String elementJson = element.toString(); + Struct.Builder elementBuilder = Struct.newBuilder(); + JsonFormat.parser().ignoringUnknownFields().merge(elementJson, elementBuilder); + + // Add each parsed object as a value in an array field + values.add(Value.newBuilder().setStructValue(elementBuilder.build()).build()); + } + + // Add the array to the main struct with a field name like "items" + structBuilder.putFields("items", + Value.newBuilder() + .setListValue(com.google.protobuf.ListValue.newBuilder().addAllValues(values).build()) + .build()); + } + else { + // Original behavior for single JSON object + JsonFormat.parser().ignoringUnknownFields().merge(json, structBuilder); + } + return structBuilder.build(); } catch (Exception e) { diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java index 3dec24540..89a8a85d1 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java @@ -25,11 +25,13 @@ import java.util.stream.Stream; import com.google.cloud.vertexai.Transport; import com.google.cloud.vertexai.VertexAI; +import io.micrometer.observation.ObservationRegistry; import org.jetbrains.annotations.NotNull; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.UserMessage; @@ -42,6 +44,8 @@ import org.springframework.ai.content.Media; import org.springframework.ai.converter.BeanOutputConverter; import org.springframework.ai.converter.ListOutputConverter; import org.springframework.ai.converter.MapOutputConverter; +import org.springframework.ai.model.tool.ToolCallingManager; +import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel.ChatModel; import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetySetting; import org.springframework.beans.factory.annotation.Autowired; @@ -293,6 +297,70 @@ class VertexAiGeminiChatModelIT { assertThat(response.getResult().getOutput().getText()).containsAnyOf("Spring AI", "portable API"); } + /** + * Helper method to create a VertexAI instance for tests + */ + private VertexAI vertexAiApi() { + String projectId = System.getenv("VERTEX_AI_GEMINI_PROJECT_ID"); + String location = System.getenv("VERTEX_AI_GEMINI_LOCATION"); + return new VertexAI.Builder().setProjectId(projectId) + .setLocation(location) + .setTransport(Transport.REST) + .build(); + } + + @Test + void jsonArrayToolCallingTest() { + // Test for the improved jsonToStruct method that handles JSON arrays in tool + // calling + + ToolCallingManager toolCallingManager = ToolCallingManager.builder() + .observationRegistry(ObservationRegistry.NOOP) + .build(); + + VertexAiGeminiChatModel chatModelWithTools = VertexAiGeminiChatModel.builder() + .vertexAI(vertexAiApi()) + .toolCallingManager(toolCallingManager) + .defaultOptions(VertexAiGeminiChatOptions.builder() + .model(VertexAiGeminiChatModel.ChatModel.GEMINI_2_0_FLASH) + .temperature(0.1) + .build()) + .build(); + + ChatClient chatClient = ChatClient.builder(chatModelWithTools).build(); + + // Create a prompt that will trigger the tool call with a specific request that + // should invoke the tool + String response = chatClient.prompt() + .tools(new ScientistTools()) + .user("List 3 famous scientists and their discoveries. Make sure to use the tool to get this information.") + .call() + .content(); + + assertThat(response).isNotEmpty(); + + assertThat(response).satisfiesAnyOf(content -> assertThat(content).contains("Einstein"), + content -> assertThat(content).contains("Newton"), content -> assertThat(content).contains("Curie")); + + } + + /** + * Tool class that returns a JSON array to test the jsonToStruct method's ability to + * handle JSON arrays. This specifically tests the PR changes that improve the + * jsonToStruct method to handle JSON arrays in addition to JSON objects. + */ + public static class ScientistTools { + + @Tool(description = "Get information about famous scientists and their discoveries") + public List> getScientists() { + // Return a JSON array with scientist information + return List.of(Map.of("name", "Albert Einstein", "discovery", "Theory of Relativity"), + Map.of("name", "Isaac Newton", "discovery", "Laws of Motion"), + Map.of("name", "Marie Curie", "discovery", "Radioactivity")); + } + + } + record ActorsFilmsRecord(String actor, List movies) { } diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/tool/VertexAiGeminiChatModelToolCallingIT.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/tool/VertexAiGeminiChatModelToolCallingIT.java index 9c82a819d..2f79a8b94 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/tool/VertexAiGeminiChatModelToolCallingIT.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/tool/VertexAiGeminiChatModelToolCallingIT.java @@ -126,7 +126,7 @@ public class VertexAiGeminiChatModelToolCallingIT { assertThat(chatResponse.getMetadata()).isNotNull(); assertThat(chatResponse.getMetadata().getUsage()).isNotNull(); - assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isGreaterThan(150).isLessThan(310); + assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isGreaterThan(150).isLessThan(330); ChatResponse response2 = this.chatModel .call(new Prompt("What is the payment status for transaction 696?", promptOptions)); @@ -201,7 +201,7 @@ public class VertexAiGeminiChatModelToolCallingIT { assertThat(chatResponse).isNotNull(); assertThat(chatResponse.getMetadata()).isNotNull(); assertThat(chatResponse.getMetadata().getUsage()).isNotNull(); - assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isGreaterThan(150).isLessThan(310); + assertThat(chatResponse.getMetadata().getUsage().getTotalTokens()).isGreaterThan(150).isLessThan(330); } diff --git a/spring-ai-model/src/main/java/org/springframework/ai/converter/ListOutputConverter.java b/spring-ai-model/src/main/java/org/springframework/ai/converter/ListOutputConverter.java index 3a16e2758..362f927d3 100644 --- a/spring-ai-model/src/main/java/org/springframework/ai/converter/ListOutputConverter.java +++ b/spring-ai-model/src/main/java/org/springframework/ai/converter/ListOutputConverter.java @@ -31,6 +31,10 @@ import org.springframework.lang.NonNull; */ public class ListOutputConverter extends AbstractConversionServiceOutputConverter> { + public ListOutputConverter() { + this(new DefaultConversionService()); + } + public ListOutputConverter(DefaultConversionService defaultConversionService) { super(defaultConversionService); }