From 823766f43602e4bce732abedfca897a0e6ab5008 Mon Sep 17 00:00:00 2001 From: Ilayaperumal Gopinathan Date: Thu, 2 Jan 2025 15:23:32 +0000 Subject: [PATCH] Fix OpenAiChatModelIT#streamingWithTokenUsage - Modify the streaming vs call usage assertion to match the close to 25% value - This change is done because the OpenAI chat response for call and stream invocation results in slight change in usage. --- .../ai/openai/chat/OpenAiChatModelIT.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java index feedb1cf4..ae7432f60 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelIT.java @@ -27,6 +27,7 @@ import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import java.util.stream.IntStream; +import org.assertj.core.data.Percentage; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.junit.jupiter.params.ParameterizedTest; @@ -216,9 +217,12 @@ public class OpenAiChatModelIT extends AbstractIT { assertThat(streamingTokenUsage.getGenerationTokens()).isGreaterThan(0); assertThat(streamingTokenUsage.getTotalTokens()).isGreaterThan(0); - assertThat(streamingTokenUsage.getPromptTokens()).isEqualTo(referenceTokenUsage.getPromptTokens()); - assertThat(streamingTokenUsage.getGenerationTokens()).isEqualTo(referenceTokenUsage.getGenerationTokens()); - assertThat(streamingTokenUsage.getTotalTokens()).isEqualTo(referenceTokenUsage.getTotalTokens()); + assertThat(streamingTokenUsage.getPromptTokens()).isCloseTo(referenceTokenUsage.getPromptTokens(), + Percentage.withPercentage(25)); + assertThat(streamingTokenUsage.getGenerationTokens()).isCloseTo(referenceTokenUsage.getGenerationTokens(), + Percentage.withPercentage(25)); + assertThat(streamingTokenUsage.getTotalTokens()).isCloseTo(referenceTokenUsage.getTotalTokens(), + Percentage.withPercentage(25)); }