Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Improve usage field to include new properties

Browse Source 
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.
This commit is contained in:
VictorZalevski
2024-11-06 09:16:55 +03:00
committed by Mark Pollack
parent 2a6f55a8c5
commit cd886437eb
5 changed files with 80 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java
View File

@@ -23,6 +23,7 @@ import java.util.function.Consumer;
import java.util.function.Predicate;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
@@ -1253,8 +1254,12 @@ public class OpenAiApi {
* @param reasoningTokens Number of tokens generated by the model for reasoning.
*/
@JsonInclude(Include.NON_NULL)
@JsonIgnoreProperties(ignoreUnknown = true)
public record CompletionTokenDetails(// @formatter:off
@JsonProperty("reasoning_tokens") Integer reasoningTokens) { // @formatter:on
@JsonProperty("reasoning_tokens") Integer reasoningTokens,
@JsonProperty("accepted_prediction_tokens") Integer acceptedPredictionTokens,
@JsonProperty("audio_tokens") Integer audioTokens,
@JsonProperty("rejected_prediction_tokens") Integer rejectedPredictionTokens) { // @formatter:on
}
}

20
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java
View File

@@ -72,6 +72,26 @@ public class OpenAiUsage implements Usage {
return reasoningTokens != null ? reasoningTokens.longValue() : 0;
}
public Long getAcceptedPredictionTokens() {
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
Integer acceptedPredictionTokens = completionTokenDetails != null
? completionTokenDetails.acceptedPredictionTokens() : null;
return acceptedPredictionTokens != null ? acceptedPredictionTokens.longValue() : 0;
}
public Long getAudioTokens() {
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
Integer audioTokens = completionTokenDetails != null ? completionTokenDetails.audioTokens() : null;
return audioTokens != null ? audioTokens.longValue() : 0;
}
public Long getRejectedPredictionTokens() {
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
Integer rejectedPredictionTokens = completionTokenDetails != null
? completionTokenDetails.rejectedPredictionTokens() : null;
return rejectedPredictionTokens != null ? rejectedPredictionTokens.longValue() : 0;
}
@Override
public Long getTotalTokens() {
Integer totalTokens = getUsage().totalTokens();

52
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/metadata/OpenAiUsageTests.java
View File

@@ -83,7 +83,7 @@ class OpenAiUsageTests {
@Test
void whenReasoningTokensIsNull() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null));
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getReasoningTokens()).isEqualTo(0);
}
@@ -91,11 +91,59 @@ class OpenAiUsageTests {
@Test
void whenCompletionTokenDetailsIsPresent() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(50));
new OpenAiApi.Usage.CompletionTokenDetails(50, null, null, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getReasoningTokens()).isEqualTo(50);
}
@Test
void whenAcceptedPredictionTokensIsNull() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getAcceptedPredictionTokens()).isEqualTo(0);
}
@Test
void whenAcceptedPredictionTokensIsPresent() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null, 75, null, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getAcceptedPredictionTokens()).isEqualTo(75);
}
@Test
void whenAudioTokensIsNull() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getAudioTokens()).isEqualTo(0);
}
@Test
void whenAudioTokensIsPresent() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null, null, 125, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getAudioTokens()).isEqualTo(125);
}
@Test
void whenRejectedPredictionTokensIsNull() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getRejectedPredictionTokens()).isEqualTo(0);
}
@Test
void whenRejectedPredictionTokensIsPresent() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, 25));
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
assertThat(usage.getRejectedPredictionTokens()).isEqualTo(25);
}
@Test
void whenCacheTokensIsNull() {
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, new OpenAiApi.Usage.PromptTokensDetails(null),

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

@@ -73,7 +73,7 @@ public class VertexAiTextEmbeddingRetryTests {
@BeforeEach
public void setUp() {
this.retryTemplate = RetryUtils.DEFAULT_RETRY_TEMPLATE;
this.retryTemplate = RetryUtils.SHORT_RETRY_TEMPLATE;
this.retryListener = new TestRetryListener();
this.retryTemplate.registerListener(this.retryListener);

5
spring-ai-retry/src/main/java/org/springframework/ai/retry/RetryUtils.java
View File

@@ -84,7 +84,8 @@ public abstract class RetryUtils {
.build();
/**
* Useful in testing scenarios where you don't want to wait long for retry.
* Useful in testing scenarios where you don't want to wait long for retry and now
* show stack trace
*/
public static final RetryTemplate SHORT_RETRY_TEMPLATE = RetryTemplate.builder()
.maxAttempts(10)
@@ -95,7 +96,7 @@ public abstract class RetryUtils {
@Override
public <T extends Object, E extends Throwable> void onError(RetryContext context,
RetryCallback<T, E> callback, Throwable throwable) {
logger.warn("Retry error. Retry count:" + context.getRetryCount(), throwable);
logger.warn("Retry error. Retry count:" + context.getRetryCount());
}
})
.build();