Add prompt caching support for AWS Bedrock Converse API

  Implements prompt caching to reduce costs on repeated content and improve
  response times. Applications with large system prompts, extensive tool
  definitions, or multi-turn conversations can see significant savings, as
  cached content costs ~90% less to process than uncached content.

  Adds five caching strategies to address different use cases:
  - `SYSTEM_ONLY`: Cache system messages (most common - stable instructions)
  - `TOOLS_ONLY`: Cache tool definitions (when tools are stable but system varies)
  - `SYSTEM_AND_TOOLS`: Cache both (when both are large and stable)
  - `CONVERSATION_HISTORY`: Cache conversation history (for chatbots and assistants)
  - `NONE`: Default, no caching

  Implementation:
  - `BedrockCacheStrategy` enum with `BedrockCacheOptions` configuration class
  - Integrated with `BedrockChatOptions` (`equals`/`hashCode`/`copy` support)
  - Cache points applied as separate blocks to satisfy AWS SDK UNION type
    constraints where each block can only contain one field type
  - Boolean flags derived from strategy to improve code readability and avoid
    repetitive conditional checks throughout request building
  - Last user message pattern for `CONVERSATION_HISTORY` enables incremental
    caching where each turn builds on the previous cached prefix
  - Cache metrics exposed via metadata `Map` to maintain provider independence
    without adding Bedrock-specific fields to shared interfaces
  - Cache hierarchy respects AWS cascade invalidation (tools → system → messages)
    to prevent stale cache combinations
  - Debug logging for troubleshooting cache point application

  Model compatibility:
  - Claude 3.x/4.x: All strategies supported
  - Amazon Nova: `SYSTEM_ONLY` and `CONVERSATION_HISTORY` only (AWS limitation on
    tool caching for Nova models)

  Testing:
  - Integration tests for all strategies using Claude 3.7 Sonnet
  - Tests handle cache TTL overlap between runs to avoid flakiness in CI

  Documentation includes usage examples, real-world use cases (legal document
  analysis, code review, customer support, multi-tenant SaaS), best practices,
  cache invalidation behavior, and cost considerations. Break-even occurs after
  one cache hit since cache reads cost ~90% less than base input tokens while
  cache writes cost ~25% more.

Signed-off-by: Soby Chacko <[email protected]>

Author: sobychacko

models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockChatOptions.java

@@ -29,6 +29,7 @@
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import org.springframework.ai.bedrock.converse.api.BedrockCacheOptions;
 import org.springframework.ai.model.tool.ToolCallingChatOptions;
 import org.springframework.ai.tool.ToolCallback;
 import org.springframework.lang.Nullable;
@@ -81,6 +82,9 @@ public class BedrockChatOptions implements ToolCallingChatOptions {
 	@JsonIgnore
 	private Boolean internalToolExecutionEnabled;
 
+	@JsonIgnore
+	private BedrockCacheOptions cacheOptions;
+
 	public static Builder builder() {
 		return new Builder();
 	}
@@ -101,6 +105,7 @@ public static BedrockChatOptions fromOptions(BedrockChatOptions fromOptions) {
 			.toolNames(new HashSet<>(fromOptions.getToolNames()))
 			.toolContext(new HashMap<>(fromOptions.getToolContext()))
 			.internalToolExecutionEnabled(fromOptions.getInternalToolExecutionEnabled())
+			.cacheOptions(fromOptions.getCacheOptions())
 			.build();
 	}
 
@@ -237,6 +242,16 @@ public void setInternalToolExecutionEnabled(@Nullable Boolean internalToolExecut
 		this.internalToolExecutionEnabled = internalToolExecutionEnabled;
 	}
 
+	@JsonIgnore
+	public BedrockCacheOptions getCacheOptions() {
+		return this.cacheOptions;
+	}
+
+	@JsonIgnore
+	public void setCacheOptions(BedrockCacheOptions cacheOptions) {
+		this.cacheOptions = cacheOptions;
+	}
+
 	@Override
 	@SuppressWarnings("unchecked")
 	public BedrockChatOptions copy() {
@@ -259,14 +274,15 @@ public boolean equals(Object o) {
 				&& Objects.equals(this.temperature, that.temperature) && Objects.equals(this.topK, that.topK)
 				&& Objects.equals(this.topP, that.topP) && Objects.equals(this.toolCallbacks, that.toolCallbacks)
 				&& Objects.equals(this.toolNames, that.toolNames) && Objects.equals(this.toolContext, that.toolContext)
-				&& Objects.equals(this.internalToolExecutionEnabled, that.internalToolExecutionEnabled);
+				&& Objects.equals(this.internalToolExecutionEnabled, that.internalToolExecutionEnabled)
+				&& Objects.equals(this.cacheOptions, that.cacheOptions);
 	}
 
 	@Override
 	public int hashCode() {
 		return Objects.hash(this.model, this.frequencyPenalty, this.maxTokens, this.presencePenalty,
 				this.requestParameters, this.stopSequences, this.temperature, this.topK, this.topP, this.toolCallbacks,
-				this.toolNames, this.toolContext, this.internalToolExecutionEnabled);
+				this.toolNames, this.toolContext, this.internalToolExecutionEnabled, this.cacheOptions);
 	}
 
 	public static final class Builder {
@@ -356,6 +372,11 @@ public Builder internalToolExecutionEnabled(@Nullable Boolean internalToolExecut
 			return this;
 		}
 
+		public Builder cacheOptions(BedrockCacheOptions cacheOptions) {
+			this.options.setCacheOptions(cacheOptions);
+			return this;
+		}
+
 		public BedrockChatOptions build() {
 			return this.options;
 		}

models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java

@@ -24,6 +24,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Base64;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -44,6 +45,7 @@
 import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
 import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient;
 import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
+import software.amazon.awssdk.services.bedrockruntime.model.CachePointBlock;
 import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
 import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
 import software.amazon.awssdk.services.bedrockruntime.model.ConverseMetrics;
@@ -70,6 +72,8 @@
 import software.amazon.awssdk.services.bedrockruntime.model.VideoFormat;
 import software.amazon.awssdk.services.bedrockruntime.model.VideoSource;
 
+import org.springframework.ai.bedrock.converse.api.BedrockCacheOptions;
+import org.springframework.ai.bedrock.converse.api.BedrockCacheStrategy;
 import org.springframework.ai.bedrock.converse.api.BedrockMediaFormat;
 import org.springframework.ai.bedrock.converse.api.ConverseApiUtils;
 import org.springframework.ai.bedrock.converse.api.ConverseChatResponseStream;
@@ -314,6 +318,8 @@ else if (prompt.getOptions() instanceof ToolCallingChatOptions toolCallingChatOp
 				.internalToolExecutionEnabled(runtimeOptions.getInternalToolExecutionEnabled() != null
 						? runtimeOptions.getInternalToolExecutionEnabled()
 						: this.defaultOptions.getInternalToolExecutionEnabled())
+				.cacheOptions(runtimeOptions.getCacheOptions() != null ? runtimeOptions.getCacheOptions()
+						: this.defaultOptions.getCacheOptions())
 				.build();
 		}
 
@@ -324,93 +330,183 @@ else if (prompt.getOptions() instanceof ToolCallingChatOptions toolCallingChatOp
 
 	ConverseRequest createRequest(Prompt prompt) {
 
-		List<Message> instructionMessages = prompt.getInstructions()
+		BedrockChatOptions updatedRuntimeOptions = prompt.getOptions().copy();
+
+		// Get cache options to determine strategy
+		BedrockCacheOptions cacheOptions = updatedRuntimeOptions.getCacheOptions();
+		boolean shouldCacheConversationHistory = cacheOptions != null
+				&& cacheOptions.getStrategy() == BedrockCacheStrategy.CONVERSATION_HISTORY;
+
+		// Get all non-system messages
+		List<org.springframework.ai.chat.messages.Message> allNonSystemMessages = prompt.getInstructions()
 			.stream()
 			.filter(message -> message.getMessageType() != MessageType.SYSTEM)
-			.map(message -> {
-				if (message.getMessageType() == MessageType.USER) {
-					List<ContentBlock> contents = new ArrayList<>();
-					if (message instanceof UserMessage userMessage) {
-						contents.add(ContentBlock.fromText(userMessage.getText()));
-
-						if (!CollectionUtils.isEmpty(userMessage.getMedia())) {
-							List<ContentBlock> mediaContent = userMessage.getMedia()
-								.stream()
-								.map(this::mapMediaToContentBlock)
-								.toList();
-							contents.addAll(mediaContent);
-						}
-					}
-					return Message.builder().content(contents).role(ConversationRole.USER).build();
+			.toList();
+
+		// Find the last user message index for CONVERSATION_HISTORY caching
+		int lastUserMessageIndex = -1;
+		if (shouldCacheConversationHistory) {
+			for (int i = allNonSystemMessages.size() - 1; i >= 0; i--) {
+				if (allNonSystemMessages.get(i).getMessageType() == MessageType.USER) {
+					lastUserMessageIndex = i;
+					break;
 				}
-				else if (message.getMessageType() == MessageType.ASSISTANT) {
-					AssistantMessage assistantMessage = (AssistantMessage) message;
-					List<ContentBlock> contentBlocks = new ArrayList<>();
-					if (StringUtils.hasText(message.getText())) {
-						contentBlocks.add(ContentBlock.fromText(message.getText()));
+			}
+			if (logger.isDebugEnabled()) {
+				logger.debug("CONVERSATION_HISTORY caching: lastUserMessageIndex={}, totalMessages={}",
+						lastUserMessageIndex, allNonSystemMessages.size());
+			}
+		}
+
+		// Build instruction messages with potential caching
+		List<Message> instructionMessages = new ArrayList<>();
+		for (int i = 0; i < allNonSystemMessages.size(); i++) {
+			org.springframework.ai.chat.messages.Message message = allNonSystemMessages.get(i);
+
+			// Determine if this message should have a cache point
+			// For CONVERSATION_HISTORY: cache point goes on the last user message
+			boolean shouldApplyCachePoint = shouldCacheConversationHistory && i == lastUserMessageIndex;
+
+			if (message.getMessageType() == MessageType.USER) {
+				List<ContentBlock> contents = new ArrayList<>();
+				if (message instanceof UserMessage) {
+					var userMessage = (UserMessage) message;
+					contents.add(ContentBlock.fromText(userMessage.getText()));
+
+					if (!CollectionUtils.isEmpty(userMessage.getMedia())) {
+						List<ContentBlock> mediaContent = userMessage.getMedia()
+							.stream()
+							.map(this::mapMediaToContentBlock)
+							.toList();
+						contents.addAll(mediaContent);
 					}
-					if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
-						for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) {
+				}
 
-							var argumentsDocument = ConverseApiUtils
-								.convertObjectToDocument(ModelOptionsUtils.jsonToMap(toolCall.arguments()));
+				// Apply cache point if this is the last user message
+				if (shouldApplyCachePoint) {
+					CachePointBlock cachePoint = CachePointBlock.builder().type("default").build();
+					contents.add(ContentBlock.fromCachePoint(cachePoint));
+					logger.debug("Applied cache point on last user message (conversation history caching)");
+				}
+
+				instructionMessages.add(Message.builder().content(contents).role(ConversationRole.USER).build());
+			}
+			else if (message.getMessageType() == MessageType.ASSISTANT) {
+				AssistantMessage assistantMessage = (AssistantMessage) message;
+				List<ContentBlock> contentBlocks = new ArrayList<>();
+				if (StringUtils.hasText(message.getText())) {
+					contentBlocks.add(ContentBlock.fromText(message.getText()));
+				}
+				if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
+					for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) {
 
-							contentBlocks.add(ContentBlock.fromToolUse(ToolUseBlock.builder()
-								.toolUseId(toolCall.id())
-								.name(toolCall.name())
-								.input(argumentsDocument)
-								.build()));
+						var argumentsDocument = ConverseApiUtils
+							.convertObjectToDocument(ModelOptionsUtils.jsonToMap(toolCall.arguments()));
+
+						contentBlocks.add(ContentBlock.fromToolUse(ToolUseBlock.builder()
+							.toolUseId(toolCall.id())
+							.name(toolCall.name())
+							.input(argumentsDocument)
+							.build()));
 
-						}
 					}
-					return Message.builder().content(contentBlocks).role(ConversationRole.ASSISTANT).build();
 				}
-				else if (message.getMessageType() == MessageType.TOOL) {
-					List<ContentBlock> contentBlocks = ((ToolResponseMessage) message).getResponses()
-						.stream()
-						.map(toolResponse -> {
+
+				instructionMessages
+					.add(Message.builder().content(contentBlocks).role(ConversationRole.ASSISTANT).build());
+			}
+			else if (message.getMessageType() == MessageType.TOOL) {
+				List<ContentBlock> contentBlocks = new ArrayList<>(
+						((ToolResponseMessage) message).getResponses().stream().map(toolResponse -> {
 							ToolResultBlock toolResultBlock = ToolResultBlock.builder()
 								.toolUseId(toolResponse.id())
 								.content(ToolResultContentBlock.builder().text(toolResponse.responseData()).build())
 								.build();
 							return ContentBlock.fromToolResult(toolResultBlock);
-						})
-						.toList();
-					return Message.builder().content(contentBlocks).role(ConversationRole.USER).build();
-				}
-				else {
-					throw new IllegalArgumentException("Unsupported message type: " + message.getMessageType());
-				}
-			})
-			.toList();
+						}).toList());
+
+				instructionMessages.add(Message.builder().content(contentBlocks).role(ConversationRole.USER).build());
+			}
+			else {
+				throw new IllegalArgumentException("Unsupported message type: " + message.getMessageType());
+			}
+		}
 
-		List<SystemContentBlock> systemMessages = prompt.getInstructions()
+		// Determine if system message caching should be applied
+		boolean shouldCacheSystem = cacheOptions != null
+				&& (cacheOptions.getStrategy() == BedrockCacheStrategy.SYSTEM_ONLY
+						|| cacheOptions.getStrategy() == BedrockCacheStrategy.SYSTEM_AND_TOOLS);
+
+		if (logger.isDebugEnabled() && cacheOptions != null) {
+			logger.debug("Cache strategy: {}, shouldCacheSystem: {}", cacheOptions.getStrategy(), shouldCacheSystem);
+		}
+
+		// Build system messages with optional caching on last message
+		List<org.springframework.ai.chat.messages.Message> systemMessageList = prompt.getInstructions()
 			.stream()
 			.filter(m -> m.getMessageType() == MessageType.SYSTEM)
-			.map(sysMessage -> SystemContentBlock.builder().text(sysMessage.getText()).build())
 			.toList();
 
-		BedrockChatOptions updatedRuntimeOptions = prompt.getOptions().copy();
+		List<SystemContentBlock> systemMessages = new ArrayList<>();
+		for (int i = 0; i < systemMessageList.size(); i++) {
+			org.springframework.ai.chat.messages.Message sysMessage = systemMessageList.get(i);
+
+			// Add the text content block
+			SystemContentBlock textBlock = SystemContentBlock.builder().text(sysMessage.getText()).build();
+			systemMessages.add(textBlock);
+
+			// Apply cache point marker after last system message if caching is enabled
+			// SystemContentBlock is a UNION type - text and cachePoint must be separate
+			// blocks
+			boolean isLastSystem = (i == systemMessageList.size() - 1);
+			if (isLastSystem && shouldCacheSystem) {
+				CachePointBlock cachePoint = CachePointBlock.builder().type("default").build();
+				SystemContentBlock cachePointBlock = SystemContentBlock.builder().cachePoint(cachePoint).build();
+				systemMessages.add(cachePointBlock);
+				logger.debug("Applied cache point after system message");
+			}
+		}
 
 		ToolConfiguration toolConfiguration = null;
 
 		// Add the tool definitions to the request's tools parameter.
 		List<ToolDefinition> toolDefinitions = this.toolCallingManager.resolveToolDefinitions(updatedRuntimeOptions);
 
+		// Determine if tool caching should be applied
+		boolean shouldCacheTools = cacheOptions != null
+				&& (cacheOptions.getStrategy() == BedrockCacheStrategy.TOOLS_ONLY
+						|| cacheOptions.getStrategy() == BedrockCacheStrategy.SYSTEM_AND_TOOLS);
+
 		if (!CollectionUtils.isEmpty(toolDefinitions)) {
-			List<Tool> bedrockTools = toolDefinitions.stream().map(toolDefinition -> {
+			List<Tool> bedrockTools = new ArrayList<>();
+
+			for (int i = 0; i < toolDefinitions.size(); i++) {
+				ToolDefinition toolDefinition = toolDefinitions.get(i);
 				var description = toolDefinition.description();
 				var name = toolDefinition.name();
 				String inputSchema = toolDefinition.inputSchema();
-				return Tool.builder()
+
+				// Create tool specification
+				Tool tool = Tool.builder()
 					.toolSpec(ToolSpecification.builder()
 						.name(name)
 						.description(description)
 						.inputSchema(ToolInputSchema.fromJson(
 								ConverseApiUtils.convertObjectToDocument(ModelOptionsUtils.jsonToMap(inputSchema))))
 						.build())
 					.build();
-			}).toList();
+				bedrockTools.add(tool);
+
+				// Apply cache point marker after last tool if caching is enabled
+				// Tool is a UNION type - toolSpec and cachePoint must be separate objects
+				boolean isLastTool = (i == toolDefinitions.size() - 1);
+				if (isLastTool && shouldCacheTools) {
+					CachePointBlock cachePoint = CachePointBlock.builder().type("default").build();
+					Tool cachePointTool = Tool.builder().cachePoint(cachePoint).build();
+					bedrockTools.add(cachePointTool);
+					logger.debug("Applied cache point after tool definitions");
+				}
+			}
 
 			toolConfiguration = ToolConfiguration.builder().tools(bedrockTools).build();
 		}
@@ -633,12 +729,23 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
 
 		ConverseMetrics metrics = response.metrics();
 
-		var chatResponseMetaData = ChatResponseMetadata.builder()
+		var metadataBuilder = ChatResponseMetadata.builder()
 			.id(response.responseMetadata() != null ? response.responseMetadata().requestId() : "Unknown")
-			.usage(usage)
-			.build();
+			.usage(usage);
+
+		// Add cache metrics if available
+		Map<String, Object> additionalMetadata = new HashMap<>();
+		if (response.usage().cacheReadInputTokens() != null) {
+			additionalMetadata.put("cacheReadInputTokens", response.usage().cacheReadInputTokens());
+		}
+		if (response.usage().cacheWriteInputTokens() != null) {
+			additionalMetadata.put("cacheWriteInputTokens", response.usage().cacheWriteInputTokens());
+		}
+		if (!additionalMetadata.isEmpty()) {
+			metadataBuilder.metadata(additionalMetadata);
+		}
 
-		return new ChatResponse(allGenerations, chatResponseMetaData);
+		return new ChatResponse(allGenerations, metadataBuilder.build());
 	}
 
 	/**