Skip to content

Commit efc6ccd

Browse files
committed
added test to validate creation/read token in cache
1 parent cafb350 commit efc6ccd

File tree

3 files changed

+27
-15
lines changed

3 files changed

+27
-15
lines changed

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,6 @@ public class AnthropicApi {
7676

7777
private static final String HEADER_ANTHROPIC_BETA = "anthropic-beta";
7878

79-
public static final String BETA_PROMPT_CACHING = "prompt-caching-2024-07-31";
80-
8179
private static final Predicate<String> SSE_DONE_PREDICATE = "[DONE]"::equals;
8280

8381
private final RestClient restClient;
@@ -905,7 +903,9 @@ public record ChatCompletionResponse(
905903
public record Usage(
906904
// @formatter:off
907905
@JsonProperty("input_tokens") Integer inputTokens,
908-
@JsonProperty("output_tokens") Integer outputTokens) {
906+
@JsonProperty("output_tokens") Integer outputTokens,
907+
@JsonProperty("cache_creation_input_tokens") Integer cacheCreationInputTokens,
908+
@JsonProperty("cache_read_input_tokens") Integer cacheReadInputTokens) {
909909
// @formatter:off
910910
}
911911

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/StreamHelper.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,9 @@ else if (event.type().equals(EventType.MESSAGE_DELTA)) {
174174

175175
if (messageDeltaEvent.usage() != null) {
176176
var totalUsage = new Usage(contentBlockReference.get().usage.inputTokens(),
177-
messageDeltaEvent.usage().outputTokens());
177+
messageDeltaEvent.usage().outputTokens(),
178+
contentBlockReference.get().usage.cacheCreationInputTokens(),
179+
contentBlockReference.get().usage.cacheReadInputTokens());
178180
contentBlockReference.get().withUsage(totalUsage);
179181
}
180182
}

models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/api/AnthropicApiIT.java

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,7 @@
2727
import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionResponse;
2828
import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock;
2929
import org.springframework.ai.anthropic.api.AnthropicApi.Role;
30-
import org.springframework.ai.retry.RetryUtils;
3130
import org.springframework.http.ResponseEntity;
32-
import org.springframework.web.client.RestClient;
33-
import org.springframework.web.reactive.function.client.WebClient;
3431

3532
import static org.assertj.core.api.Assertions.assertThat;
3633
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -46,17 +43,30 @@ public class AnthropicApiIT {
4643

4744
@Test
4845
void chatWithPromptCache() {
49-
AnthropicApi anthropicApiBeta = new AnthropicApi(AnthropicApi.DEFAULT_BASE_URL,
50-
System.getenv("ANTHROPIC_API_KEY"), AnthropicApi.DEFAULT_ANTHROPIC_VERSION, RestClient.builder(),
51-
WebClient.builder(), RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER, AnthropicApi.BETA_PROMPT_CACHING);
46+
String userMessageText = "It could be either a contraction of the full title Quenta Silmarillion (\"Tale of the Silmarils\") or also a plain Genitive which "
47+
+ "(as in Ancient Greek) signifies reference. This genitive is translated in English with \"about\" or \"of\" "
48+
+ "constructions; the titles of the chapters in The Silmarillion are examples of this genitive in poetic English "
49+
+ "(Of the Sindar, Of Men, Of the Darkening of Valinor etc), where \"of\" means \"about\" or \"concerning\". "
50+
+ "In the same way, Silmarillion can be taken to mean \"Of/About the Silmarils\"";
51+
5252
AnthropicMessage chatCompletionMessage = new AnthropicMessage(
53-
List.of(new ContentBlock("Tell me a Joke?", AnthropicCacheType.EPHEMERAL.cacheControl())), Role.USER);
53+
List.of(new ContentBlock(userMessageText.repeat(20), AnthropicCacheType.EPHEMERAL.cacheControl())),
54+
Role.USER);
5455

55-
ResponseEntity<ChatCompletionResponse> response = anthropicApiBeta
56-
.chatCompletionEntity(new ChatCompletionRequest(AnthropicApi.ChatModel.CLAUDE_3_HAIKU.getValue(),
57-
List.of(chatCompletionMessage), null, 100, 0.8, false));
56+
ChatCompletionRequest chatCompletionRequest = new ChatCompletionRequest(
57+
AnthropicApi.ChatModel.CLAUDE_3_HAIKU.getValue(), List.of(chatCompletionMessage), null, 100, 0.8,
58+
false);
59+
AnthropicApi.Usage createdCacheToken = anthropicApi.chatCompletionEntity(chatCompletionRequest)
60+
.getBody()
61+
.usage();
5862

59-
assertThat(response).isNotNull();
63+
assertThat(createdCacheToken.cacheCreationInputTokens()).isGreaterThan(0);
64+
assertThat(createdCacheToken.cacheReadInputTokens()).isEqualTo(0);
65+
66+
AnthropicApi.Usage readCacheToken = anthropicApi.chatCompletionEntity(chatCompletionRequest).getBody().usage();
67+
68+
assertThat(readCacheToken.cacheCreationInputTokens()).isEqualTo(0);
69+
assertThat(readCacheToken.cacheReadInputTokens()).isGreaterThan(0);
6070
}
6171

6272
@Test

0 commit comments

Comments
 (0)