Skip to content

Commit 45421b1

Browse files
ricken07ilayaperumalg
authored andcommitted
improved ability to use openAiApi to support DeepSeek
Signed-off-by: Ricken Bazolo <[email protected]>
1 parent 35101e7 commit 45421b1

File tree

2 files changed

+47
-13
lines changed

2 files changed

+47
-13
lines changed

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

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1412,17 +1412,28 @@ public record TopLogProbs(// @formatter:off
14121412
* completion).
14131413
* @param promptTokensDetails Breakdown of tokens used in the prompt.
14141414
* @param completionTokenDetails Breakdown of tokens used in a completion.
1415+
* @param promptCacheHitTokens Number of tokens in the prompt that were served from
1416+
* (util for
1417+
* <a href="https://api-docs.deepseek.com/api/create-chat-completion">DeepSeek</a>
1418+
* support).
1419+
* @param promptCacheMissTokens Number of tokens in the prompt that were not served
1420+
* (util for
1421+
* <a href="https://api-docs.deepseek.com/api/create-chat-completion">DeepSeek</a>
1422+
* support).
14151423
*/
14161424
@JsonInclude(Include.NON_NULL)
1425+
@JsonIgnoreProperties(ignoreUnknown = true)
14171426
public record Usage(// @formatter:off
14181427
@JsonProperty("completion_tokens") Integer completionTokens,
14191428
@JsonProperty("prompt_tokens") Integer promptTokens,
14201429
@JsonProperty("total_tokens") Integer totalTokens,
14211430
@JsonProperty("prompt_tokens_details") PromptTokensDetails promptTokensDetails,
1422-
@JsonProperty("completion_tokens_details") CompletionTokenDetails completionTokenDetails) { // @formatter:on
1431+
@JsonProperty("completion_tokens_details") CompletionTokenDetails completionTokenDetails,
1432+
@JsonProperty("prompt_cache_hit_tokens") Integer promptCacheHitTokens,
1433+
@JsonProperty("prompt_cache_miss_tokens") Integer promptCacheMissTokens) { // @formatter:on
14231434

14241435
public Usage(Integer completionTokens, Integer promptTokens, Integer totalTokens) {
1425-
this(completionTokens, promptTokens, totalTokens, null, null);
1436+
this(completionTokens, promptTokens, totalTokens, null, null, null, null);
14261437
}
14271438

14281439
/**

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

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import org.junit.jupiter.api.Test;
2020

2121
import org.springframework.ai.chat.metadata.DefaultUsage;
22-
import org.springframework.ai.chat.metadata.Usage;
2322
import org.springframework.ai.openai.api.OpenAiApi;
2423

2524
import static org.assertj.core.api.Assertions.assertThat;
@@ -81,7 +80,7 @@ void whenTotalTokensIsNull() {
8180

8281
@Test
8382
void whenPromptAndCompletionTokensDetailsIsNull() {
84-
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null, null);
83+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null, null, null, null);
8584
DefaultUsage usage = getDefaultUsage(openAiUsage);
8685
assertThat(usage.getTotalTokens()).isEqualTo(300);
8786
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
@@ -91,7 +90,7 @@ void whenPromptAndCompletionTokensDetailsIsNull() {
9190

9291
@Test
9392
void whenCompletionTokenDetailsIsNull() {
94-
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null, null);
93+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null, null, null, null);
9594
DefaultUsage usage = getDefaultUsage(openAiUsage);
9695
assertThat(usage.getTotalTokens()).isEqualTo(300);
9796
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
@@ -101,7 +100,7 @@ void whenCompletionTokenDetailsIsNull() {
101100
@Test
102101
void whenReasoningTokensIsNull() {
103102
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
104-
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
103+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null), null, null);
105104
DefaultUsage usage = getDefaultUsage(openAiUsage);
106105
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
107106
assertThat(nativeUsage.completionTokenDetails().reasoningTokens()).isEqualTo(null);
@@ -110,7 +109,7 @@ void whenReasoningTokensIsNull() {
110109
@Test
111110
void whenCompletionTokenDetailsIsPresent() {
112111
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
113-
new OpenAiApi.Usage.CompletionTokenDetails(50, null, null, null));
112+
new OpenAiApi.Usage.CompletionTokenDetails(50, null, null, null), null, null);
114113
DefaultUsage usage = getDefaultUsage(openAiUsage);
115114
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
116115
assertThat(nativeUsage.completionTokenDetails().reasoningTokens()).isEqualTo(50);
@@ -122,7 +121,7 @@ void whenCompletionTokenDetailsIsPresent() {
122121
@Test
123122
void whenAcceptedPredictionTokensIsPresent() {
124123
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
125-
new OpenAiApi.Usage.CompletionTokenDetails(null, 75, null, null));
124+
new OpenAiApi.Usage.CompletionTokenDetails(null, 75, null, null), null, null);
126125
DefaultUsage usage = getDefaultUsage(openAiUsage);
127126
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
128127
assertThat(nativeUsage.completionTokenDetails().reasoningTokens()).isEqualTo(null);
@@ -134,7 +133,7 @@ void whenAcceptedPredictionTokensIsPresent() {
134133
@Test
135134
void whenAudioTokensIsPresent() {
136135
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
137-
new OpenAiApi.Usage.CompletionTokenDetails(null, null, 125, null));
136+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, 125, null), null, null);
138137
DefaultUsage usage = getDefaultUsage(openAiUsage);
139138
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
140139
assertThat(nativeUsage.completionTokenDetails().reasoningTokens()).isEqualTo(null);
@@ -146,7 +145,7 @@ void whenAudioTokensIsPresent() {
146145
@Test
147146
void whenRejectedPredictionTokensIsNull() {
148147
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
149-
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null));
148+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, null), null, null);
150149
DefaultUsage usage = getDefaultUsage(openAiUsage);
151150
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
152151
assertThat(nativeUsage.completionTokenDetails().reasoningTokens()).isEqualTo(null);
@@ -160,7 +159,7 @@ void whenRejectedPredictionTokensIsNull() {
160159
@Test
161160
void whenRejectedPredictionTokensIsPresent() {
162161
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300, null,
163-
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, 25));
162+
new OpenAiApi.Usage.CompletionTokenDetails(null, null, null, 25), null, null);
164163
DefaultUsage usage = getDefaultUsage(openAiUsage);
165164
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
166165
assertThat(nativeUsage.completionTokenDetails().reasoningTokens()).isEqualTo(null);
@@ -172,7 +171,7 @@ void whenRejectedPredictionTokensIsPresent() {
172171
@Test
173172
void whenCacheTokensIsNull() {
174173
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300,
175-
new OpenAiApi.Usage.PromptTokensDetails(null, null), null);
174+
new OpenAiApi.Usage.PromptTokensDetails(null, null), null, null, null);
176175
DefaultUsage usage = getDefaultUsage(openAiUsage);
177176
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
178177
assertThat(nativeUsage.promptTokensDetails().audioTokens()).isEqualTo(null);
@@ -182,11 +181,35 @@ void whenCacheTokensIsNull() {
182181
@Test
183182
void whenCacheTokensIsPresent() {
184183
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300,
185-
new OpenAiApi.Usage.PromptTokensDetails(99, 15), null);
184+
new OpenAiApi.Usage.PromptTokensDetails(99, 15), null, null, null);
186185
DefaultUsage usage = getDefaultUsage(openAiUsage);
187186
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
188187
assertThat(nativeUsage.promptTokensDetails().audioTokens()).isEqualTo(99);
189188
assertThat(nativeUsage.promptTokensDetails().cachedTokens()).isEqualTo(15);
190189
}
191190

191+
@Test
192+
void whenPromptCacheHitTokensIsPresent() {
193+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300,
194+
new OpenAiApi.Usage.PromptTokensDetails(99, 15), null, 150, null);
195+
DefaultUsage usage = getDefaultUsage(openAiUsage);
196+
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
197+
assertThat(nativeUsage.promptTokensDetails().audioTokens()).isEqualTo(99);
198+
assertThat(nativeUsage.promptTokensDetails().cachedTokens()).isEqualTo(15);
199+
assertThat(nativeUsage.promptCacheHitTokens()).isEqualTo(150);
200+
assertThat(nativeUsage.promptCacheMissTokens()).isNull();
201+
}
202+
203+
@Test
204+
void whenPromptCacheMissTokensIsPresent() {
205+
OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300,
206+
new OpenAiApi.Usage.PromptTokensDetails(99, 15), null, null, 80);
207+
DefaultUsage usage = getDefaultUsage(openAiUsage);
208+
OpenAiApi.Usage nativeUsage = (OpenAiApi.Usage) usage.getNativeUsage();
209+
assertThat(nativeUsage.promptTokensDetails().audioTokens()).isEqualTo(99);
210+
assertThat(nativeUsage.promptTokensDetails().cachedTokens()).isEqualTo(15);
211+
assertThat(nativeUsage.promptCacheMissTokens()).isEqualTo(80);
212+
assertThat(nativeUsage.promptCacheHitTokens()).isNull();
213+
}
214+
192215
}

0 commit comments

Comments
 (0)