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

Fix the Gemini parallel function calling behaviour.

Browse Source 
Enable handling multiple function calls at once.
  Change the Gemini model names from preview to gemini-1.5-pro-001 and gemini-1.5-flash-001.
  Simplify the Gemini function calling ITs. drop the multi-turn instructions.
This commit is contained in:
Christian Tzolov
2024-06-13 15:37:50 +02:00
parent dbccfda7c5
commit 85d2fa525e
7 changed files with 89 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -19,6 +19,7 @@ import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.google.cloud.vertexai.VertexAI;
import com.google.cloud.vertexai.api.Content;
import com.google.cloud.vertexai.api.Content.Builder;
import com.google.cloud.vertexai.api.FunctionCall;
import com.google.cloud.vertexai.api.FunctionDeclaration;
import com.google.cloud.vertexai.api.FunctionResponse;
@@ -102,9 +103,9 @@ public class VertexAiGeminiChatModel
GEMINI_PRO("gemini-pro"),
GEMINI_1_5_PRO("gemini-1.5-pro-preview-0514"),
GEMINI_1_5_PRO("gemini-1.5-pro-001"),
GEMINI_1_5_FLASH("gemini-1.5-flash-preview-0514");
GEMINI_1_5_FLASH("gemini-1.5-flash-001");
ChatModel(String value) {
this.value = value;
@@ -416,27 +417,31 @@ public class VertexAiGeminiChatModel
protected GeminiRequest doCreateToolResponseRequest(GeminiRequest previousRequest, Content responseMessage,
List<Content> conversationHistory) {
FunctionCall functionCall = responseMessage.getPartsList().iterator().next().getFunctionCall();
var iterator = responseMessage.getPartsList().iterator();
var functionName = functionCall.getName();
String functionArguments = structToJson(functionCall.getArgs());
Builder builder = Content.newBuilder();
while (iterator.hasNext()) {
if (!this.functionCallbackRegister.containsKey(functionName)) {
throw new IllegalStateException("No function callback found for function name: " + functionName);
}
FunctionCall functionCall = iterator.next().getFunctionCall();
String functionResponse = this.functionCallbackRegister.get(functionName).call(functionArguments);
var functionName = functionCall.getName();
String functionArguments = structToJson(functionCall.getArgs());
Content contentFnResp = Content.newBuilder()
.addParts(Part.newBuilder()
if (!this.functionCallbackRegister.containsKey(functionName)) {
throw new IllegalStateException("No function callback found for function name: " + functionName);
}
String functionResponse = this.functionCallbackRegister.get(functionName).call(functionArguments);
builder.addParts(Part.newBuilder()
.setFunctionResponse(FunctionResponse.newBuilder()
.setName(functionCall.getName())
.setResponse(jsonToStruct(functionResponse))
.build())
.build())
.build();
.build());
conversationHistory.add(contentFnResp);
}
conversationHistory.add(builder.build());
return new GeminiRequest(conversationHistory, previousRequest.model());
}

25
models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatModelFunctionCallingIT.java
View File

@@ -60,16 +60,10 @@ public class VertexAiGeminiChatModelFunctionCallingIT {
@Test
public void functionCallExplicitOpenApiSchema() {
var systemMessage = new SystemMessage("""
Use Multi-turn function calling.
Answer for all listed locations.
If the information was not fetched call the function again. Repeat at most 3 times.
""");
UserMessage userMessage = new UserMessage(
"What's the weather like in San Francisco, Paris and in Tokyo? Perform multiple funciton execution if necessary. Return the temperature in Celsius.");
"What's the weather like in San Francisco, Paris and in Tokyo? Return the temperature in Celsius.");
List<Message> messages = new ArrayList<>(List.of(systemMessage, userMessage));
List<Message> messages = new ArrayList<>(List.of(userMessage));
String openApiSchema = """
{
@@ -90,8 +84,7 @@ public class VertexAiGeminiChatModelFunctionCallingIT {
""";
var promptOptions = VertexAiGeminiChatOptions.builder()
.withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH)
// .withModel(VertexAiGeminiModelCall.ChatModel.GEMINI_PRO_1_5_PRO)
// .withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH)
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
.withName("get_current_weather")
.withDescription("Get the current weather in a given location")
@@ -114,8 +107,7 @@ public class VertexAiGeminiChatModelFunctionCallingIT {
List<Message> messages = new ArrayList<>(List.of(userMessage));
var promptOptions = VertexAiGeminiChatOptions.builder()
.withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_PRO)
// .withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_PRO)
.withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH)
.withFunctionCallbacks(List.of(
FunctionCallbackWrapper.builder(new MockWeatherService())
.withSchemaType(SchemaType.OPEN_API_SCHEMA)
@@ -148,15 +140,10 @@ public class VertexAiGeminiChatModelFunctionCallingIT {
@Test
public void functionCallTestInferredOpenApiSchemaStream() {
var systemMessage = new SystemMessage("""
Use Multi-turn function calling.
Answer for all listed locations.
If the information was not fetched call the function again. Repeat at most 3 times.
""");
UserMessage userMessage = new UserMessage(
"What's the weather like in San Francisco, Paris and in Tokyo? Perform multiple funciton execution if necessary. Return the temperature in Celsius.");
"What's the weather like in San Francisco, Paris and in Tokyo? Return the temperature in Celsius.");
List<Message> messages = new ArrayList<>(List.of(systemMessage, userMessage));
List<Message> messages = new ArrayList<>(List.of(userMessage));
var promptOptions = VertexAiGeminiChatOptions.builder()
.withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH)

60
models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiPaymentTransactionIT.java
View File

@@ -16,6 +16,8 @@
package org.springframework.ai.vertexai.gemini.function;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
@@ -89,39 +91,46 @@ public class VertexAiGeminiPaymentTransactionIT {
@Test
public void paymentStatuses() {
String content = this.chatClient.prompt().advisors(new LoggingAdvisor()).functions("paymentStatus").user("""
// @formatter:off
String content = this.chatClient.prompt()
.advisors(new LoggingAdvisor())
.functions("paymentStatus")
.user("""
What is the status of my payment transactions 001, 002 and 003?
To answer this question invoke the 'paymentStatus' function per transaction.
If requred invoke the function per transaction.
""").call().content();
logger.info("" + content);
assertThat(content).contains("001", "002", "003");
assertThat(content).contains("pending", "approved", "rejected");
}
@RepeatedTest(5)
public void streamingPaymentStatuses() {
Flux<String> streamContent = this.chatClient.prompt()
.advisors(new LoggingAdvisor())
.functions("paymentStatus")
// .functions("paymentStatuses")
.user("""
What is the status of my payment transactions 001, 002 and 003?
To answer this question invoke the paymentStatus function per transaction.
Return the transaction id and the transaction status for each transaction.
""")
.stream()
.content();
.advisors(new LoggingAdvisor())
.functions("paymentStatus")
// .functions("paymentStatuses")
.user("""
What is the status of my payment transactions 001, 002 and 003?
If requred invoke the function per transaction.
""")
.stream()
.content();
String content = streamContent.collectList().block().stream().collect(Collectors.joining());
logger.info(content);
assertThat(content).contains("001", "002", "003");
assertThat(content).contains("pending", "approved", "rejected");
// Quota rate
try {
Thread.sleep(2000);
}
catch (InterruptedException e) {
Thread.sleep(1000);
} catch (InterruptedException e) {
}
}
@@ -173,10 +182,10 @@ public class VertexAiGeminiPaymentTransactionIT {
String location = System.getenv("VERTEX_AI_GEMINI_LOCATION");
return new VertexAI.Builder().setLocation(location)
.setProjectId(projectId)
.setTransport(Transport.REST)
// .setTransport(Transport.GRPC)
.build();
.setProjectId(projectId)
.setTransport(Transport.REST)
// .setTransport(Transport.GRPC)
.build();
}
@Bean
@@ -186,16 +195,15 @@ public class VertexAiGeminiPaymentTransactionIT {
return new VertexAiGeminiChatModel(vertexAi,
VertexAiGeminiChatOptions.builder()
.withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH)
// .withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_PRO_1_5_PRO)
.withTemperature(0.1f)
// .withResponseMimeType(ResponseMimeType.JSON)
.build(),
.withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH)
.withTemperature(0.1f)
.build(),
functionCallbackContext);
}
/**
* Because of the OPEN_API_SCHEMA type, the FunctionCallbackContext instance must
* Because of the OPEN_API_SCHEMA type, the FunctionCallbackContext instance
* must
* different from the other JSON schema types.
*/
private FunctionCallbackContext springAiFunctionManager(ApplicationContext context) {

30
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java
View File

@@ -15,6 +15,8 @@
*/
package org.springframework.ai.autoconfigure.vertexai.gemini.tool;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.function.Function;
@@ -22,11 +24,9 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.autoconfigure.vertexai.gemini.VertexAiGeminiAutoConfiguration;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel;
import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatOptions;
@@ -36,8 +36,6 @@ import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;
import static org.assertj.core.api.Assertions.assertThat;
@EnabledIfEnvironmentVariable(named = "VERTEX_AI_GEMINI_PROJECT_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "VERTEX_AI_GEMINI_LOCATION", matches = ".*")
class FunctionCallWithFunctionBeanIT {
@@ -55,34 +53,26 @@ class FunctionCallWithFunctionBeanIT {
void functionCallTest() {
contextRunner.withPropertyValues("spring.ai.vertex.ai.gemini.chat.options.model="
// + VertexAiGeminiChatModel.ChatModel.GEMINI_PRO.getValue())
// + VertexAiGeminiChatModel.ChatModel.GEMINI_PRO_1_5_PRO.getValue())
+ VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH.getValue())
.run(context -> {
VertexAiGeminiChatModel chatModel = context.getBean(VertexAiGeminiChatModel.class);
var systemMessage = new SystemMessage("""
Use Multi-turn function calling.
Answer for all listed locations.
If the information was not fetched call the function again. Repeat at most 3 times.
var userMessage = new UserMessage("""
What's the weather like in San Francisco, Paris and in Tokyo?
Return the temperature in Celsius.
Perform multiple funciton execution if necessary.
""");
var userMessage = new UserMessage(
// "What's the weather like in San Francisco, Paris and in Tokyo?
// Please let me know how many function calls you've preformed.");
"What's the weather like in San Francisco, Paris and in Tokyo? Perform multiple funciton execution if necessary. Return the temperature in Celsius.");
ChatResponse response = chatModel.call(new Prompt(List.of(systemMessage, userMessage),
ChatResponse response = chatModel.call(new Prompt(List.of(userMessage),
VertexAiGeminiChatOptions.builder().withFunction("weatherFunction").build()));
// ChatResponse response = chatModel.call(new
// Prompt(List.of(userMessage),
// VertexAiGeminiChatOptions.builder().withFunction("weatherFunction").build()));
logger.info("Response: {}", response);
assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
response = chatModel.call(new Prompt(List.of(systemMessage, userMessage),
response = chatModel.call(new Prompt(List.of(userMessage),
VertexAiGeminiChatOptions.builder().withFunction("weatherFunction3").build()));
logger.info("Response: {}", response);
@@ -90,7 +80,7 @@ class FunctionCallWithFunctionBeanIT {
assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
response = chatModel
.call(new Prompt(List.of(systemMessage, userMessage), VertexAiGeminiChatOptions.builder().build()));
.call(new Prompt(List.of(userMessage), VertexAiGeminiChatOptions.builder().build()));
logger.info("Response: {}", response);

24
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionWrapperIT.java
View File

@@ -15,17 +15,17 @@
*/
package org.springframework.ai.autoconfigure.vertexai.gemini.tool;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.autoconfigure.vertexai.gemini.VertexAiGeminiAutoConfiguration;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
@@ -37,8 +37,6 @@ import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import static org.assertj.core.api.Assertions.assertThat;
@EnabledIfEnvironmentVariable(named = "VERTEX_AI_GEMINI_PROJECT_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "VERTEX_AI_GEMINI_LOCATION", matches = ".*")
public class FunctionCallWithFunctionWrapperIT {
@@ -60,23 +58,17 @@ public class FunctionCallWithFunctionWrapperIT {
VertexAiGeminiChatModel chatModel = context.getBean(VertexAiGeminiChatModel.class);
var systemMessage = new SystemMessage("""
Use Multi-turn function calling.
Answer for all listed locations.
If the information was not fetched call the function again. Repeat at most 3 times.
var userMessage = new UserMessage("""
What's the weather like in San Francisco, Paris and in Tokyo?
Return the temperature in Celsius.
""");
var userMessage = new UserMessage(
"What's the weather like in San Francisco, Paris and in Tokyo? Perform multiple funciton execution if necessary. Return the temperature in Celsius.");
ChatResponse response = chatModel.call(new Prompt(List.of(systemMessage, userMessage),
ChatResponse response = chatModel.call(new Prompt(List.of(userMessage),
VertexAiGeminiChatOptions.builder().withFunction("WeatherInfo").build()));
logger.info("Response: {}", response);
assertThat(response.getResult().getOutput().getContent()).containsAnyOf("30.0", "30");
assertThat(response.getResult().getOutput().getContent()).containsAnyOf("10.0", "10");
assertThat(response.getResult().getOutput().getContent()).containsAnyOf("15", "15.0");
assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
});
}

19
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java
View File

@@ -56,13 +56,16 @@ public class FunctionCallWithPromptFunctionIT {
VertexAiGeminiChatModel chatModel = context.getBean(VertexAiGeminiChatModel.class);
var systemMessage = new SystemMessage("""
Use Multi-turn function calling.
Answer for all listed locations.
If the information was not fetched call the function again. Repeat at most 3 times.
// var systemMessage = new SystemMessage("""
// Use Multi-turn function calling.
// Answer for all listed locations.
// If the information was not fetched call the function again. Repeat at
// most 3 times.
// """);
var userMessage = new UserMessage("""
What's the weather like in San Francisco, Paris and in Tokyo?
Return the temperature in Celsius.
""");
UserMessage userMessage = new UserMessage(
"What's the weather like in San Francisco, in Paris and in Tokyo? Perform multiple funciton execution if necessary. Return the temperature in Celsius.");
var promptOptions = VertexAiGeminiChatOptions.builder()
.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
@@ -72,7 +75,7 @@ public class FunctionCallWithPromptFunctionIT {
.build()))
.build();
ChatResponse response = chatModel.call(new Prompt(List.of(systemMessage, userMessage), promptOptions));
ChatResponse response = chatModel.call(new Prompt(List.of(userMessage), promptOptions));
logger.info("Response: {}", response);
@@ -80,7 +83,7 @@ public class FunctionCallWithPromptFunctionIT {
// Verify that no function call is made.
response = chatModel
.call(new Prompt(List.of(systemMessage, userMessage), VertexAiGeminiChatOptions.builder().build()));
.call(new Prompt(List.of(userMessage), VertexAiGeminiChatOptions.builder().build()));
logger.info("Response: {}", response);

1
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/MockWeatherService.java
View File

@@ -28,6 +28,7 @@ import com.fasterxml.jackson.annotation.JsonPropertyDescription;
*
* @author Christian Tzolov
*/
@JsonClassDescription("Get the weather in location")
public class MockWeatherService implements Function<MockWeatherService.Request, MockWeatherService.Response> {
/**