Skip to content

Commit 2c17577

Browse files
dafriztzolov
authored andcommitted
Add support for prompt token details in OpenAI usage stats
- Add PromptTokensDetails record to track cached tokens in prompt - Update Usage record to include promptTokensDetails field - Add getCachedTokens() method to OpenAiUsage - Add test cases for cached tokens handling Resolves #1506
1 parent 8b1882b commit 2c17577

File tree

3 files changed

+43
-6
lines changed

3 files changed

+43
-6
lines changed

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
* @author Michael Lavelle
5757
* @author Mariusz Bernacki
5858
* @author Thomas Vitale
59+
* @author David Frizelle
5960
*/
6061
public class OpenAiApi {
6162

@@ -938,17 +939,29 @@ public record TopLogProbs(// @formatter:off
938939
* @param promptTokens Number of tokens in the prompt.
939940
* @param totalTokens Total number of tokens used in the request (prompt +
940941
* completion).
941-
* @param completionTokenDetails Breakdown of tokens used in a completion
942+
* @param promptTokensDetails Breakdown of tokens used in the prompt.
943+
* @param completionTokenDetails Breakdown of tokens used in a completion.
942944
*/
943945
@JsonInclude(Include.NON_NULL)
944946
public record Usage(// @formatter:off
945947
@JsonProperty("completion_tokens") Integer completionTokens,
946948
@JsonProperty("prompt_tokens") Integer promptTokens,
947949
@JsonProperty("total_tokens") Integer totalTokens,
950+
@JsonProperty("prompt_tokens_details") PromptTokensDetails promptTokensDetails,
948951
@JsonProperty("completion_tokens_details") CompletionTokenDetails completionTokenDetails) {// @formatter:on
949952

950953
public Usage(Integer completionTokens, Integer promptTokens, Integer totalTokens) {
951-
this(completionTokens, promptTokens, totalTokens, null);
954+
this(completionTokens, promptTokens, totalTokens, null, null);
955+
}
956+
957+
/**
958+
* Breakdown of tokens used in the prompt
959+
*
960+
* @param cachedTokens Cached tokens present in the prompt.
961+
*/
962+
@JsonInclude(Include.NON_NULL)
963+
public record PromptTokensDetails(// @formatter:off
964+
@JsonProperty("cached_tokens") Integer cachedTokens) {// @formatter:on
952965
}
953966

954967
/**

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
*
2525
* @author John Blum
2626
* @author Thomas Vitale
27+
* @author David Frizelle
2728
* @since 0.7.0
2829
* @see <a href=
2930
* "https://platform.openai.com/docs/api-reference/completions/object">Completion
@@ -58,6 +59,12 @@ public Long getGenerationTokens() {
5859
return generationTokens != null ? generationTokens.longValue() : 0;
5960
}
6061

62+
public Long getCachedTokens() {
63+
OpenAiApi.Usage.PromptTokensDetails promptTokenDetails = getUsage().promptTokensDetails();
64+
Integer cachedTokens = promptTokenDetails != null ? promptTokenDetails.cachedTokens() : null;
65+
return cachedTokens != null ? cachedTokens.longValue() : 0;
66+
}
67+
6168
public Long getReasoningTokens() {
6269
OpenAiApi.Usage.CompletionTokenDetails completionTokenDetails = getUsage().completionTokenDetails();
6370
Integer reasoningTokens = completionTokenDetails != null ? completionTokenDetails.reasoningTokens() : null;

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

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,27 +55,44 @@ void whenTotalTokensIsNull() {
5555
}
5656

5757
@Test
58-
void whenCompletionTokenDetailsIsNull() {
59-
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null);
58+
void whenPromptAndCompletionTokensDetailsIsNull() {
59+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null, null);
6060
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
6161
assertThat(usage.getTotalTokens()).isEqualTo(300);
62+
assertThat(usage.getCachedTokens()).isEqualTo(0);
6263
assertThat(usage.getReasoningTokens()).isEqualTo(0);
6364
}
6465

6566
@Test
6667
void whenReasoningTokensIsNull() {
67-
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300,
68+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
6869
new OpenAiApi.Usage.CompletionTokenDetails(null));
6970
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
7071
assertThat(usage.getReasoningTokens()).isEqualTo(0);
7172
}
7273

7374
@Test
7475
void whenCompletionTokenDetailsIsPresent() {
75-
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300,
76+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
7677
new OpenAiApi.Usage.CompletionTokenDetails(50));
7778
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
7879
assertThat(usage.getReasoningTokens()).isEqualTo(50);
7980
}
8081

82+
@Test
83+
void whenCacheTokensIsNull() {
84+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, new OpenAiApi.Usage.PromptTokensDetails(null),
85+
null);
86+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
87+
assertThat(usage.getCachedTokens()).isEqualTo(0);
88+
}
89+
90+
@Test
91+
void whenCacheTokensIsPresent() {
92+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, new OpenAiApi.Usage.PromptTokensDetails(15),
93+
null);
94+
OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
95+
assertThat(usage.getCachedTokens()).isEqualTo(15);
96+
}
97+
8198
}

0 commit comments

Comments
 (0)