added cache suport to assistant message

Author: Claudio-code

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

@@ -475,19 +475,26 @@ private Map<String, String> mergeHttpHeaders(Map<String, String> runtimeHttpHead
 
 	ChatCompletionRequest createRequest(Prompt prompt, boolean stream) {
 
-		Optional<Message> lastMessage = prompt.getInstructions()
+		List<Message> userMessagesList = prompt.getInstructions()
 				.stream()
-				.filter(message -> message.getMessageType() != MessageType.SYSTEM)
-				.findFirst();
+				.filter(message -> message.getMessageType() == MessageType.USER)
+				.toList();
+		Message lastUserMessage = userMessagesList.isEmpty() ? null : userMessagesList.get(userMessagesList.size() - 1);
+
+		List<Message> assistantMessageList = prompt.getInstructions()
+				.stream()
+				.filter(message -> message.getMessageType() == MessageType.ASSISTANT)
+				.toList();
+		Message lastAssistantMessage = assistantMessageList.isEmpty() ? null : assistantMessageList.get(assistantMessageList.size() - 1);
 
 		List<AnthropicMessage> userMessages = prompt.getInstructions()
 			.stream()
 			.filter(message -> message.getMessageType() != MessageType.SYSTEM)
 			.map(message -> {
+				AbstractMessage abstractMessage = (AbstractMessage) message;
 				if (message.getMessageType() == MessageType.USER) {
-					AbstractMessage abstractMessage = (AbstractMessage) message;
 					List<ContentBlock> contents;
-					boolean isLastItem = lastMessage.filter(message::equals).isPresent();
+					boolean isLastItem = message.equals(lastUserMessage);
 					if (isLastItem && abstractMessage.getCache() != null) {
 						AnthropicCacheType cacheType = AnthropicCacheType.valueOf(abstractMessage.getCache());
 						contents = new ArrayList<>(
@@ -511,8 +518,14 @@ ChatCompletionRequest createRequest(Prompt prompt, boolean stream) {
 				else if (message.getMessageType() == MessageType.ASSISTANT) {
 					AssistantMessage assistantMessage = (AssistantMessage) message;
 					List<ContentBlock> contentBlocks = new ArrayList<>();
+					boolean isLastItem = message.equals(lastAssistantMessage);
 					if (StringUtils.hasText(message.getText())) {
-						contentBlocks.add(new ContentBlock(message.getText()));
+						if (isLastItem && abstractMessage.getCache() != null) {
+							AnthropicCacheType cacheType = AnthropicCacheType.valueOf(abstractMessage.getCache());
+							contentBlocks.add(new ContentBlock(message.getText(), cacheType.cacheControl()));
+						} else {
+							contentBlocks.add(new ContentBlock(message.getText()));
+						}
 					}
 					if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
 						for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) {
@@ -551,6 +564,7 @@ else if (message.getMessageType() == MessageType.TOOL) {
 		// Add the tool definitions to the request's tools parameter.
 		List<ToolDefinition> toolDefinitions = this.toolCallingManager.resolveToolDefinitions(requestOptions);
 		if (!CollectionUtils.isEmpty(toolDefinitions)) {
+			var tool = getFunctionTools(toolDefinitions);
 			request = ModelOptionsUtils.merge(request, this.defaultOptions, ChatCompletionRequest.class);
 			request = ChatCompletionRequest.from(request).tools(getFunctionTools(toolDefinitions)).build();
 		}

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

@@ -18,6 +18,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Random;
 import java.util.stream.Collectors;
 
 import org.junit.jupiter.api.Test;
@@ -73,6 +74,39 @@ public class AnthropicApiIT {
 					}
 					""")));
 
+	@Test
+	void chatWithPromptCacheInAssistantMessage() {
+		String assistantMessageText = "It could be either a contraction of the full title Quenta Silmarillion (\"Tale of the Silmarils\") or also a plain Genitive which "
+				+ "(as in Ancient Greek) signifies reference. This genitive is translated in English with \"about\" or \"of\" "
+				+ "constructions; the titles of the chapters in The Silmarillion are examples of this genitive in poetic English "
+				+ "(Of the Sindar, Of Men, Of the Darkening of Valinor etc), where \"of\" means \"about\" or \"concerning\". "
+				+ "In the same way, Silmarillion can be taken to mean \"Of/About the Silmarils\"";
+
+		AnthropicMessage chatCompletionMessage = new AnthropicMessage(
+				List.of(new ContentBlock(assistantMessageText.repeat(20), AnthropicCacheType.EPHEMERAL.cacheControl())),
+				Role.ASSISTANT);
+
+		ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
+				.model(AnthropicApi.ChatModel.CLAUDE_3_5_HAIKU)
+				.messages(List.of(chatCompletionMessage))
+				.maxTokens(1500)
+				.temperature(0.8)
+				.stream(false)
+				.build();
+
+		AnthropicApi.Usage createdCacheToken = anthropicApi.chatCompletionEntity(chatCompletionRequest)
+				.getBody()
+				.usage();
+
+		assertThat(createdCacheToken.cacheCreationInputTokens()).isGreaterThan(0);
+		assertThat(createdCacheToken.cacheReadInputTokens()).isEqualTo(0);
+
+		AnthropicApi.Usage readCacheToken = anthropicApi.chatCompletionEntity(chatCompletionRequest).getBody().usage();
+
+		assertThat(readCacheToken.cacheCreationInputTokens()).isEqualTo(0);
+		assertThat(readCacheToken.cacheReadInputTokens()).isGreaterThan(0);
+	}
+
 	@Test
 	void chatWithPromptCache() {
 		String userMessageText = "It could be either a contraction of the full title Quenta Silmarillion (\"Tale of the Silmarils\") or also a plain Genitive which "