Flaky tests (langchain4j#1823)

katcharov · Sep 24, 2024 · c9995ab · c9995ab
1 parent 34f6712
commit c9995ab
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 40 deletions.
diff --git a/...gle-ai-gemini/src/test/java/dev/langchain4j/model/googleai/GoogleAiGeminiChatModelIT.java b/...gle-ai-gemini/src/test/java/dev/langchain4j/model/googleai/GoogleAiGeminiChatModelIT.java
@@ -2,6 +2,7 @@
 
 import com.google.gson.Gson;
 import dev.langchain4j.agent.tool.JsonSchemaProperty;
+import dev.langchain4j.agent.tool.P;
 import dev.langchain4j.agent.tool.Tool;
 import dev.langchain4j.agent.tool.ToolExecutionRequest;
 import dev.langchain4j.agent.tool.ToolSpecification;
@@ -31,6 +32,7 @@
 import org.junit.jupiter.api.Test;
 import org.junitpioneer.jupiter.RetryingTest;
 
+import java.time.Duration;
 import java.util.*;
 import java.util.stream.Collectors;
 
@@ -707,73 +709,61 @@ void should_count_tokens() {
         assertThat(listOfMsgTokenCount).isGreaterThan(0);
     }
 
-    static class Calculator {
-        @Tool
-        public int add(int a, int b) {
-            System.out.println("ADD " + a + " + " + b);
-            return a + b;
-        }
-
-        @Tool
-        public int subtract(int a, int b) {
-            System.out.println("SUBTRACT " + a + " - " + b);
-            return a - b;
-        }
-
-        @Tool
-        public int multiply(int a, int b) {
-            System.out.println("MULTIPLY " + a + " * " + b);
-            return a * b;
-        }
-
-        @Tool
-        public int divide(int a, int b) {
-            System.out.println("DIVIDE " + a + " / " + b);
-            return a / b;
+    static class Transactions {
+        @Tool("returns amount of a given transaction")
+        double getTransactionAmount(@P("ID of a transaction") String id) {
+            System.out.printf("called getTransactionAmount(%s)%n", id);
+            switch (id) {
+                case "T001":
+                    return 11.1;
+                case "T002":
+                    return 22.2;
+                default:
+                    throw new IllegalArgumentException("Unknown transaction ID: " + id);
+            }
         }
     }
 
+
     interface Assistant {
-        @dev.langchain4j.service.SystemMessage(
-            "When asked to evaluate some math expressions, " +
-            "you MUST use the `add`, `substract`, `multiply`, and `divide` tool functions."
-        )
         String chat(String userMessage);
     }
 
-    @RetryingTest(3)
+    @RetryingTest(10)
     void should_work_with_tools_with_AiServices() {
         // given
         GoogleAiGeminiChatModel gemini = GoogleAiGeminiChatModel.builder()
             .apiKey(GOOGLE_AI_GEMINI_API_KEY)
             .modelName("gemini-1.5-pro")
             .logRequestsAndResponses(true)
+            .timeout(Duration.ofMinutes(2))
+            .temperature(0.0)
+            .topP(0.0)
+            .topK(1)
             .build();
 
         // when
-        Calculator spyCalculator = spy(new Calculator());
+        Transactions spyTransactions = spy(new Transactions());
 
         MessageWindowChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(20);
         Assistant assistant = AiServices.builder(Assistant.class)
-            .tools(spyCalculator)
+            .tools(spyTransactions)
             .chatMemory(chatMemory)
             .chatLanguageModel(gemini)
             .build();
 
         // then
         String response = "";
 
-        response = assistant.chat("How much is 3 + 4?");
-        assertThat(response).containsIgnoringCase("7");
-        verify(spyCalculator).add(3, 4);
-
-        response = assistant.chat("How much is 7 * 11?");
-        assertThat(response).containsIgnoringCase("77");
-        verify(spyCalculator).multiply(7, 11);
+        response = assistant.chat("What is the amount of transaction T001?");
+        assertThat(response).containsIgnoringCase("11.1");
+        verify(spyTransactions).getTransactionAmount("T001");
 
-        verifyNoMoreInteractions(spyCalculator);
+        response = assistant.chat("What is the amount of transaction T002?");
+        assertThat(response).containsIgnoringCase("22.2");
+        verify(spyTransactions).getTransactionAmount("T002");
 
-        System.out.println("chatMemory = " + chatMemory.messages());
+        verifyNoMoreInteractions(spyTransactions);
     }
 
     @AfterEach

diff --git a/...tex-ai-gemini/src/test/java/dev/langchain4j/model/vertexai/VertexAiGeminiChatModelIT.java b/...tex-ai-gemini/src/test/java/dev/langchain4j/model/vertexai/VertexAiGeminiChatModelIT.java
@@ -645,7 +645,7 @@ void should_support_json_response_mime_type() {
         assertThat(json).isEqualToIgnoringWhitespace(expectedJson);
     }
 
-    @RetryingTest(4)
+    @RetryingTest(10)
     void should_allow_defining_safety_settings() {
         // given
         HashMap<HarmCategory, SafetyThreshold> safetySettings = new HashMap<>();
@@ -659,6 +659,10 @@ void should_allow_defining_safety_settings() {
             .location(System.getenv("GCP_LOCATION"))
             .modelName("gemini-1.5-flash-001")
             .safetySettings(safetySettings)
+            .temperature(0.0f)
+            .topP(0.0f)
+            .topK(1)
+            .seed(1234)
             .logRequests(true)
             .logResponses(true)
             .build();