Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(vertex-ai-gemini): enhance jsonToStruct to support JSON arrays

Browse Source 
- Improve the jsonToStruct method in VertexAiGeminiChatModel to handle JSON arrays
  in addition to JSON objects. When a JSON array is detected, it's now properly
  converted to a Protobuf Struct with an items field containing the array elements.
- Added test

Resolves #2647 , #2849

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov
2025-05-03 11:02:57 +03:00
committed by Mark Pollack
parent 2cbfb22db6
commit 78d90cd134
4 changed files with 107 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java
View File

@@ -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<Value> 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) {

68
models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java
View File

@@ -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<Map<String, String>> 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<String> movies) {
}

4
models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/tool/VertexAiGeminiChatModelToolCallingIT.java
View File

@@ -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);
}

4
spring-ai-model/src/main/java/org/springframework/ai/converter/ListOutputConverter.java
View File

@@ -31,6 +31,10 @@ import org.springframework.lang.NonNull;
*/
public class ListOutputConverter extends AbstractConversionServiceOutputConverter<List<String>> {
public ListOutputConverter() {
this(new DefaultConversionService());
}
public ListOutputConverter(DefaultConversionService defaultConversionService) {
super(defaultConversionService);
}