Skip to content

Commit cd88643

Browse files
VictorZZZZmarkpollack
VictorZZZZ
authored and
markpollack
committed
Improve usage field to include new properties
OpenAI's API returns additional token usage metrics that provide deeper insight into API consumption. This adds support for: - acceptedPredictionTokens: Tokens from accepted model predictions - audioTokens: Tokens used for audio processing - rejectedPredictionTokens: Tokens from rejected model predictions These fields help track resource utilization and costs more accurately by breaking down token usage by type. Added @JsonIgnoreProperties to maintain compatibility with future OpenAI API additions. Fixes warning logging in RetryUtils.SHORT_RETRY_TEMPLATE to reduce noise in test output.
1 parent 2a6f55a commit cd88643

File tree

5 files changed

+80
-6
lines changed

5 files changed

+80
-6
lines changed

‎models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.function.Predicate;
2424

2525
import com.fasterxml.jackson.annotation.JsonIgnore;
26+
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
2627
import com.fasterxml.jackson.annotation.JsonInclude;
2728
import com.fasterxml.jackson.annotation.JsonInclude.Include;
2829
import com.fasterxml.jackson.annotation.JsonProperty;
@@ -1253,8 +1254,12 @@ public record PromptTokensDetails(// @formatter:off
12531254
* @param reasoningTokens Number of tokens generated by the model for reasoning.
12541255
*/
12551256
@JsonInclude(Include.NON_NULL)
1257+
@JsonIgnoreProperties(ignoreUnknown = true)
12561258
public record CompletionTokenDetails(// @formatter:off
1257-
@JsonProperty("reasoning_tokens") Integer reasoningTokens) { // @formatter:on
1259+
@JsonProperty("reasoning_tokens") Integer reasoningTokens,
1260+
@JsonProperty("accepted_prediction_tokens") Integer acceptedPredictionTokens,
1261+
@JsonProperty("audio_tokens") Integer audioTokens,
1262+
@JsonProperty("rejected_prediction_tokens") Integer rejectedPredictionTokens) { // @formatter:on
12581263
}
12591264

12601265
}

‎models/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,26 @@ public Long getReasoningTokens() {
7272
return reasoningTokens != null ? reasoningTokens.longValue() : 0;
7373
}
7474

75+
public Long getAcceptedPredictionTokens() {
76+
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
77+
Integer acceptedPredictionTokens = completionTokenDetails != null
78+
? completionTokenDetails.acceptedPredictionTokens() : null;
79+
return acceptedPredictionTokens != null ? acceptedPredictionTokens.longValue() : 0;
80+
}
81+
82+
public Long getAudioTokens() {
83+
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
84+
Integer audioTokens = completionTokenDetails != null ? completionTokenDetails.audioTokens() : null;
85+
return audioTokens != null ? audioTokens.longValue() : 0;
86+
}
87+
88+
public Long getRejectedPredictionTokens() {
89+
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
90+
Integer rejectedPredictionTokens = completionTokenDetails != null
91+
? completionTokenDetails.rejectedPredictionTokens() : null;
92+
return rejectedPredictionTokens != null ? rejectedPredictionTokens.longValue() : 0;
93+
}
94+
7595
@Override
7696
public Long getTotalTokens() {
7797
Integer totalTokens = getUsage().totalTokens();

‎models/spring-ai-openai/src/test/java/org/springframework/ai/openai/metadata/OpenAiUsageTests.java‎

Lines changed: 50 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,19 +83,67 @@ void whenPromptAndCompletionTokensDetailsIsNull() {
8383
@Test
8484
void whenReasoningTokensIsNull() {
8585
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
86-
new OpenAiApi.Usage.CompletionTokenDetails(null));
86+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
8787
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
8888
assertThat(usage.getReasoningTokens()).isEqualTo(0);
8989
}
9090

9191
@Test
9292
void whenCompletionTokenDetailsIsPresent() {
9393
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
94-
new OpenAiApi.Usage.CompletionTokenDetails(50));
94+
new OpenAiApi.Usage.CompletionTokenDetails(50, null, null, null));
9595
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
9696
assertThat(usage.getReasoningTokens()).isEqualTo(50);
9797
}
9898

99+
@Test
100+
void whenAcceptedPredictionTokensIsNull() {
101+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
102+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
103+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
104+
assertThat(usage.getAcceptedPredictionTokens()).isEqualTo(0);
105+
}
106+
107+
@Test
108+
void whenAcceptedPredictionTokensIsPresent() {
109+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
110+
new OpenAiApi.Usage.CompletionTokenDetails(null, 75, null, null));
111+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
112+
assertThat(usage.getAcceptedPredictionTokens()).isEqualTo(75);
113+
}
114+
115+
@Test
116+
void whenAudioTokensIsNull() {
117+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
118+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
119+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
120+
assertThat(usage.getAudioTokens()).isEqualTo(0);
121+
}
122+
123+
@Test
124+
void whenAudioTokensIsPresent() {
125+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
126+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, 125, null));
127+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
128+
assertThat(usage.getAudioTokens()).isEqualTo(125);
129+
}
130+
131+
@Test
132+
void whenRejectedPredictionTokensIsNull() {
133+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
134+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
135+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
136+
assertThat(usage.getRejectedPredictionTokens()).isEqualTo(0);
137+
}
138+
139+
@Test
140+
void whenRejectedPredictionTokensIsPresent() {
141+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
142+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, 25));
143+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
144+
assertThat(usage.getRejectedPredictionTokens()).isEqualTo(25);
145+
}
146+
99147
@Test
100148
void whenCacheTokensIsNull() {
101149
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, new OpenAiApi.Usage.PromptTokensDetails(null),

‎models/spring-ai-vertex-ai-embedding/src/test/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingRetryTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ public class VertexAiTextEmbeddingRetryTests {
7373

7474
@BeforeEach
7575
public void setUp() {
76-
this.retryTemplate = RetryUtils.DEFAULT_RETRY_TEMPLATE;
76+
this.retryTemplate = RetryUtils.SHORT_RETRY_TEMPLATE;
7777
this.retryListener = new TestRetryListener();
7878
this.retryTemplate.registerListener(this.retryListener);
7979

‎spring-ai-retry/src/main/java/org/springframework/ai/retry/RetryUtils.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,8 @@ public <T extends Object, E extends Throwable> void onError(RetryContext context
8484
.build();
8585

8686
/**
87-
* Useful in testing scenarios where you don't want to wait long for retry.
87+
* Useful in testing scenarios where you don't want to wait long for retry and now
88+
* show stack trace
8889
*/
8990
public static final RetryTemplate SHORT_RETRY_TEMPLATE = RetryTemplate.builder()
9091
.maxAttempts(10)
@@ -95,7 +96,7 @@ public <T extends Object, E extends Throwable> void onError(RetryContext context
9596
@Override
9697
public <T extends Object, E extends Throwable> void onError(RetryContext context,
9798
RetryCallback<T, E> callback, Throwable throwable) {
98-
logger.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
99+
logger.warn("Retry error. Retry count:" + context.getRetryCount());
99100
}
100101
})
101102
.build();

0 commit comments

Comments
 (0)