From 14ec6b2dc00017d3c5e3e03b18052ef14878fb65 Mon Sep 17 00:00:00 2001 From: Ilayaperumal Gopinathan Date: Wed, 29 Jan 2025 10:58:57 +0000 Subject: [PATCH] Use Usage#getCompletionTokens instead of deprecated Usage#getGenerationTokens --- .../ai/openai/chat/proxy/DeepSeekWithOpenAiChatModelIT.java | 6 +++--- .../springframework/ai/chat/metadata/DefaultUsageTests.java | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/DeepSeekWithOpenAiChatModelIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/DeepSeekWithOpenAiChatModelIT.java index 88f153c01..84c00209a 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/DeepSeekWithOpenAiChatModelIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/DeepSeekWithOpenAiChatModelIT.java @@ -128,11 +128,11 @@ class DeepSeekWithOpenAiChatModelIT { var referenceTokenUsage = this.chatModel.call(prompt).getMetadata().getUsage(); assertThat(streamingTokenUsage.getPromptTokens()).isGreaterThan(0); - assertThat(streamingTokenUsage.getGenerationTokens()).isGreaterThan(0); + assertThat(streamingTokenUsage.getCompletionTokens()).isGreaterThan(0); assertThat(streamingTokenUsage.getTotalTokens()).isGreaterThan(0); assertThat(streamingTokenUsage.getPromptTokens()).isEqualTo(referenceTokenUsage.getPromptTokens()); - assertThat(streamingTokenUsage.getGenerationTokens()).isEqualTo(referenceTokenUsage.getGenerationTokens()); + assertThat(streamingTokenUsage.getCompletionTokens()).isEqualTo(referenceTokenUsage.getCompletionTokens()); assertThat(streamingTokenUsage.getTotalTokens()).isEqualTo(referenceTokenUsage.getTotalTokens()); } @@ -316,7 +316,7 @@ class DeepSeekWithOpenAiChatModelIT { assertThat(response.getMetadata().getId()).isNotEmpty(); assertThat(response.getMetadata().getModel()).containsIgnoringCase(model); assertThat(response.getMetadata().getUsage().getPromptTokens()).isPositive(); - assertThat(response.getMetadata().getUsage().getGenerationTokens()).isPositive(); + assertThat(response.getMetadata().getUsage().getCompletionTokens()).isPositive(); assertThat(response.getMetadata().getUsage().getTotalTokens()).isPositive(); } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java index 6861881fb..3a9172117 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/chat/metadata/DefaultUsageTests.java @@ -184,7 +184,7 @@ public class DefaultUsageTests { assertThat(deserialized.getPromptTokens()).isEqualTo(usage.getPromptTokens()); assertThat(deserialized.getCompletionTokens()).isEqualTo(usage.getCompletionTokens()); assertThat(deserialized.getTotalTokens()).isEqualTo(usage.getTotalTokens()); - assertThat(deserialized.getGenerationTokens()).isEqualTo(usage.getGenerationTokens()); + assertThat(deserialized.getCompletionTokens()).isEqualTo(usage.getCompletionTokens()); @SuppressWarnings("unchecked") Map deserializedMap = (Map) deserialized.getNativeUsage();