From 85d2fa525e44626187aabb3b9096e6a7219a435b Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 13 Jun 2024 15:37:50 +0200 Subject: [PATCH] Fix the Gemini parallel function calling behaviour. 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. --- .../gemini/VertexAiGeminiChatModel.java | 33 +++++----- ...texAiGeminiChatModelFunctionCallingIT.java | 25 ++------ .../VertexAiGeminiPaymentTransactionIT.java | 60 +++++++++++-------- .../tool/FunctionCallWithFunctionBeanIT.java | 30 ++++------ .../FunctionCallWithFunctionWrapperIT.java | 24 +++----- .../FunctionCallWithPromptFunctionIT.java | 19 +++--- .../gemini/tool/MockWeatherService.java | 1 + 7 files changed, 89 insertions(+), 103 deletions(-) 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 23269d7b7..220439de6 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 @@ -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 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()); } diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatModelFunctionCallingIT.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatModelFunctionCallingIT.java index 963c0d46d..a936fccb7 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatModelFunctionCallingIT.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiChatModelFunctionCallingIT.java @@ -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 messages = new ArrayList<>(List.of(systemMessage, userMessage)); + List 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 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 messages = new ArrayList<>(List.of(systemMessage, userMessage)); + List messages = new ArrayList<>(List.of(userMessage)); var promptOptions = VertexAiGeminiChatOptions.builder() .withModel(VertexAiGeminiChatModel.ChatModel.GEMINI_1_5_FLASH) diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiPaymentTransactionIT.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiPaymentTransactionIT.java index 348809008..1950fd40b 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiPaymentTransactionIT.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/function/VertexAiGeminiPaymentTransactionIT.java @@ -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 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) { diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java index 8b0ad0ef1..3657c8648 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java @@ -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); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionWrapperIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionWrapperIT.java index a9a2a724f..2cf31c1d4 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionWrapperIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionWrapperIT.java @@ -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"); }); } diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java index a3c88082b..28645ef5e 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java @@ -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); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/MockWeatherService.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/MockWeatherService.java index 1e25d6d54..ed34d7b0c 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/MockWeatherService.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vertexai/gemini/tool/MockWeatherService.java @@ -28,6 +28,7 @@ import com.fasterxml.jackson.annotation.JsonPropertyDescription; * * @author Christian Tzolov */ +@JsonClassDescription("Get the weather in location") public class MockWeatherService implements Function { /**