Fix Geminie GenerativeModel handling between calls

Resolves #560
spring-projects · Apr 7, 2024 · e268975 · e268975
1 parent 2b421a4
commit e268975
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 5 deletions.
diff --git a/...gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClient.java b/...gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatClient.java
@@ -78,8 +78,6 @@ public class VertexAiGeminiChatClient
 
 	private final GenerationConfig generationConfig;
 
-	private GenerativeModel generativeModel;
-
 	public enum GeminiMessageType {
 
 		USER("user"),
@@ -140,7 +138,6 @@ public VertexAiGeminiChatClient(VertexAI vertexAI, VertexAiGeminiChatOptions opt
 		this.vertexAI = vertexAI;
 		this.defaultOptions = options;
 		this.generationConfig = toGenerationConfig(options);
-		this.generativeModel = new GenerativeModel(options.getModel(), vertexAI);
 	}
 
 	// https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini
@@ -204,7 +201,8 @@ private GeminiRequest createGeminiRequest(Prompt prompt) {
 		Set<String> functionsForThisRequest = new HashSet<>();
 
 		GenerationConfig generationConfig = this.generationConfig;
-		GenerativeModel generativeModel = this.generativeModel;
+
+		GenerativeModel generativeModel = new GenerativeModel(this.defaultOptions.getModel(), this.vertexAI);
 
 		VertexAiGeminiChatOptions updatedRuntimeOptions = null;
 

diff --git a/...pringframework/ai/vertexai/gemini/function/VertexAiGeminiChatClientFunctionCallingIT.java b/...pringframework/ai/vertexai/gemini/function/VertexAiGeminiChatClientFunctionCallingIT.java
@@ -146,7 +146,7 @@ public void functionCallTestInferredOpenApiSchema() {
 	public void functionCallTestInferredOpenApiSchemaStream() {
 
 		UserMessage userMessage = new UserMessage(
-				"What's the weather like in San Francisco, in Paris and in Tokyo? Use Multi-turn function calling.");
+				"What's the weather like in San Francisco, in Paris and in Tokyo, Japan? Use Multi-turn function calling. Provide answer for all requested locations.");
 
 		List<Message> messages = new ArrayList<>(List.of(userMessage));
 

diff --git a/...springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java b/...springframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithFunctionBeanIT.java
@@ -83,6 +83,13 @@ void functionCallTest() {
 
 				assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
 
+				response = chatClient
+					.call(new Prompt(List.of(systemMessage, userMessage), VertexAiGeminiChatOptions.builder().build()));
+
+				logger.info("Response: {}", response);
+
+				assertThat(response.getResult().getOutput().getContent()).doesNotContain("30", "10", "15");
+
 			});
 	}
 

diff --git a/...ringframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java b/...ringframework/ai/autoconfigure/vertexai/gemini/tool/FunctionCallWithPromptFunctionIT.java
@@ -77,6 +77,15 @@ void functionCallTest() {
 				logger.info("Response: {}", response);
 
 				assertThat(response.getResult().getOutput().getContent()).contains("30", "10", "15");
+
+				// Verify that no function call is made.
+				response = chatClient
+					.call(new Prompt(List.of(systemMessage, userMessage), VertexAiGeminiChatOptions.builder().build()));
+
+				logger.info("Response: {}", response);
+
+				assertThat(response.getResult().getOutput().getContent()).doesNotContain("30", "10", "15");
+
 			});
 	}