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

Extend the CreateGeminiRequestTests for tool config tests

Browse Source
This commit is contained in:
Christian Tzolov
2024-07-16 10:00:01 +02:00
parent df7ac05235
commit 87dc3be739
1 changed files with 71 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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