From 87dc3be7398fbbb9ba0026ff02a8d3bed77e996f Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 16 Jul 2024 10:00:01 +0200 Subject: [PATCH] Extend the CreateGeminiRequestTests for tool config tests --- .../gemini/CreateGeminiRequestTests.java | 146 +++++++++--------- 1 file changed, 71 insertions(+), 75 deletions(-) diff --git a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/CreateGeminiRequestTests.java b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/CreateGeminiRequestTests.java index f7206818b..6c5e8d336 100644 --- a/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/CreateGeminiRequestTests.java +++ b/models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/CreateGeminiRequestTests.java @@ -29,13 +29,14 @@ import org.springframework.ai.chat.messages.Media; import org.springframework.ai.chat.messages.SystemMessage; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.function.FunctionCallbackWrapper; import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel.GeminiRequest; -import org.springframework.util.MimeType; import org.springframework.util.MimeTypeUtils; import com.google.cloud.vertexai.VertexAI; import com.google.cloud.vertexai.api.Content; import com.google.cloud.vertexai.api.Part; +import org.springframework.ai.vertexai.gemini.function.MockWeatherService; /** * @author Christian Tzolov @@ -102,97 +103,92 @@ public class CreateGeminiRequestTests { System.out.println(mediaPart); } - // @Test - // public void promptOptionsTools() { + @Test + public void promptOptionsTools() { - // final String TOOL_FUNCTION_NAME = "CurrentWeather"; + final String TOOL_FUNCTION_NAME = "CurrentWeather"; - // var client = new OpenAiChatModel(new OpenAiApi("TEST"), - // OpenAiChatOptions.builder().withModel("DEFAULT_MODEL").build()); + var client = new VertexAiGeminiChatModel(vertexAI, + VertexAiGeminiChatOptions.builder().withModel("DEFAULT_MODEL").build()); - // var request = client.createRequest(new Prompt("Test message content", - // OpenAiChatOptions.builder() - // .withModel("PROMPT_MODEL") - // .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new - // MockWeatherService()) - // .withName(TOOL_FUNCTION_NAME) - // .withDescription("Get the weather in location") - // .withResponseConverter((response) -> "" + response.temp() + response.unit()) - // .build())) - // .build()), - // false); + var request = client.createGeminiRequest(new Prompt("Test message content", + VertexAiGeminiChatOptions.builder() + .withModel("PROMPT_MODEL") + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName(TOOL_FUNCTION_NAME) + .withDescription("Get the weather in location") + .withResponseConverter((response) -> "" + response.temp() + response.unit()) + .build())) + .build())); - // assertThat(client.getFunctionCallbackRegister()).hasSize(1); - // assertThat(client.getFunctionCallbackRegister()).containsKeys(TOOL_FUNCTION_NAME); + assertThat(client.getFunctionCallbackRegister()).hasSize(1); + assertThat(client.getFunctionCallbackRegister()).containsKeys(TOOL_FUNCTION_NAME); - // assertThat(request.messages()).hasSize(1); - // assertThat(request.stream()).isFalse(); - // assertThat(request.model()).isEqualTo("PROMPT_MODEL"); + assertThat(request.contents()).hasSize(1); + assertThat(request.model().getSystemInstruction()).isNotPresent(); + assertThat(request.model().getModelName()).isEqualTo("PROMPT_MODEL"); - // assertThat(request.tools()).hasSize(1); - // assertThat(request.tools().get(0).function().name()).isEqualTo(TOOL_FUNCTION_NAME); - // } + assertThat(request.model().getTools()).hasSize(1); + assertThat(request.model().getTools().get(0).getFunctionDeclarations(0).getName()) + .isEqualTo(TOOL_FUNCTION_NAME); + } - // @Test - // public void defaultOptionsTools() { + @Test + public void defaultOptionsTools() { - // final String TOOL_FUNCTION_NAME = "CurrentWeather"; + final String TOOL_FUNCTION_NAME = "CurrentWeather"; - // var client = new OpenAiChatModel(new OpenAiApi("TEST"), - // OpenAiChatOptions.builder() - // .withModel("DEFAULT_MODEL") - // .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new - // MockWeatherService()) - // .withName(TOOL_FUNCTION_NAME) - // .withDescription("Get the weather in location") - // .withResponseConverter((response) -> "" + response.temp() + response.unit()) - // .build())) - // .build()); + var client = new VertexAiGeminiChatModel(vertexAI, + VertexAiGeminiChatOptions.builder() + .withModel("DEFAULT_MODEL") + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName(TOOL_FUNCTION_NAME) + .withDescription("Get the weather in location") + .withResponseConverter((response) -> "" + response.temp() + response.unit()) + .build())) + .build()); - // var request = client.createRequest(new Prompt("Test message content"), false); + var request = client.createGeminiRequest(new Prompt("Test message content")); - // assertThat(client.getFunctionCallbackRegister()).hasSize(1); - // assertThat(client.getFunctionCallbackRegister()).containsKeys(TOOL_FUNCTION_NAME); - // assertThat(client.getFunctionCallbackRegister().get(TOOL_FUNCTION_NAME).getDescription()) - // .isEqualTo("Get the weather in location"); + assertThat(client.getFunctionCallbackRegister()).hasSize(1); + assertThat(client.getFunctionCallbackRegister()).containsKeys(TOOL_FUNCTION_NAME); + assertThat(client.getFunctionCallbackRegister().get(TOOL_FUNCTION_NAME).getDescription()) + .isEqualTo("Get the weather in location"); - // assertThat(request.messages()).hasSize(1); - // assertThat(request.stream()).isFalse(); - // assertThat(request.model()).isEqualTo("DEFAULT_MODEL"); + assertThat(request.contents()).hasSize(1); + assertThat(request.model().getSystemInstruction()).isNotPresent(); + assertThat(request.model().getModelName()).isEqualTo("DEFAULT_MODEL"); - // assertThat(request.tools()).as("Default Options callback functions are not - // automatically enabled!") - // .isNullOrEmpty(); + assertThat(request.model().getTools()).as("Default Options callback functions are not automatically enabled!") + .isNullOrEmpty(); - // // Explicitly enable the function - // request = client.createRequest(new Prompt("Test message content", - // OpenAiChatOptions.builder().withFunction(TOOL_FUNCTION_NAME).build()), false); + // Explicitly enable the function + request = client.createGeminiRequest(new Prompt("Test message content", + VertexAiGeminiChatOptions.builder().withFunction(TOOL_FUNCTION_NAME).build())); - // assertThat(request.tools()).hasSize(1); - // assertThat(request.tools().get(0).function().name()).as("Explicitly enabled - // function") - // .isEqualTo(TOOL_FUNCTION_NAME); + assertThat(request.model().getTools()).hasSize(1); + assertThat(request.model().getTools().get(0).getFunctionDeclarations(0).getName()) + .as("Explicitly enabled function") + .isEqualTo(TOOL_FUNCTION_NAME); - // // Override the default options function with one from the prompt - // request = client.createRequest(new Prompt("Test message content", - // OpenAiChatOptions.builder() - // .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new - // MockWeatherService()) - // .withName(TOOL_FUNCTION_NAME) - // .withDescription("Overridden function description") - // .build())) - // .build()), - // false); + // Override the default options function with one from the prompt + request = client.createGeminiRequest(new Prompt("Test message content", + VertexAiGeminiChatOptions.builder() + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName(TOOL_FUNCTION_NAME) + .withDescription("Overridden function description") + .build())) + .build())); - // assertThat(request.tools()).hasSize(1); - // assertThat(request.tools().get(0).function().name()).as("Explicitly enabled - // function") - // .isEqualTo(TOOL_FUNCTION_NAME); + assertThat(request.model().getTools()).hasSize(1); + assertThat(request.model().getTools().get(0).getFunctionDeclarations(0).getName()) + .as("Explicitly enabled function") + .isEqualTo(TOOL_FUNCTION_NAME); - // assertThat(client.getFunctionCallbackRegister()).hasSize(1); - // assertThat(client.getFunctionCallbackRegister()).containsKeys(TOOL_FUNCTION_NAME); - // assertThat(client.getFunctionCallbackRegister().get(TOOL_FUNCTION_NAME).getDescription()) - // .isEqualTo("Overridden function description"); - // } + assertThat(client.getFunctionCallbackRegister()).hasSize(1); + assertThat(client.getFunctionCallbackRegister()).containsKeys(TOOL_FUNCTION_NAME); + assertThat(client.getFunctionCallbackRegister().get(TOOL_FUNCTION_NAME).getDescription()) + .isEqualTo("Overridden function description"); + } }