diff --git a/lib/ruby_llm/configuration.rb b/lib/ruby_llm/configuration.rb index 52c1b65b4..f61e4fb28 100644 --- a/lib/ruby_llm/configuration.rb +++ b/lib/ruby_llm/configuration.rb @@ -70,6 +70,11 @@ def defaults = @defaults ||= {} option :log_stream_debug, -> { ENV['RUBYLLM_STREAM_DEBUG'] == 'true' } option :log_regexp_timeout, -> { Regexp.respond_to?(:timeout) ? (Regexp.timeout || 1.0) : nil } + # Auto-inject Bedrock InvokeModel prompt-cache breakpoints (system + tail). See + # RubyLLM::Protocols::BedrockInvokeModel::Chat#render_payload. Defaults on because the + # only production consumer of the InvokeModel path expects native caching to take over. + option :bedrock_invoke_model_prompt_caching, true + def initialize self.class.send(:defaults).each do |key, default| value = default.respond_to?(:call) ? instance_exec(&default) : default diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb index f0035669e..7886d1b56 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb @@ -9,6 +9,14 @@ class BedrockInvokeModel module Chat BEDROCK_INLINE_DOCUMENT_LIMIT = 4_500_000 + # Bedrock allows at most 4 cache_control breakpoints per request, counted across + # system, messages, and tools combined. + MAX_CACHE_BREAKPOINTS = 4 + + # Block types Bedrock will attach a cache breakpoint to. Thinking/redacted_thinking + # are intentionally excluded — Anthropic does not support caching on them. + CACHEABLE_BLOCK_TYPES = %w[text image tool_use tool_result document].freeze + module_function def completion_url @@ -70,11 +78,76 @@ def render_payload(messages, tools:, temperature:, model:, stream: false, add_tool_fields(payload, tools, tool_prefs) add_thinking_fields(payload, thinking) add_beta_fields(payload) + inject_cache_breakpoints(payload) payload end # rubocop:enable Metrics/ParameterLists,Lint/UnusedMethodArgument + # Injects the two auto cache breakpoints (system tail + final-message tail) when + # bedrock_invoke_model_prompt_caching is enabled. Never removes a breakpoint that + # arrived via a translated cachePoint — only ever adds up to the 4-breakpoint budget. + # Deliberately does not estimate token counts against the model's minimum cacheable + # prefix (e.g. 4,096 tokens): a breakpoint below the minimum silently doesn't cache + # and costs nothing, so that estimation is left to the caller. + def inject_cache_breakpoints(payload) + return unless @config.bedrock_invoke_model_prompt_caching + + remaining = MAX_CACHE_BREAKPOINTS - count_cache_breakpoints(payload) + return if remaining <= 0 + + # Tail takes priority over system when only one breakpoint slot remains: the tail + # breakpoint is what makes the in-flight tool loop cheap round-trip to round-trip, + # while the system breakpoint only protects against a single oversized turn. + inject_system_cache_breakpoint?(payload[:system]) if remaining >= 2 + inject_tail_cache_breakpoint(payload[:messages]) + end + + def count_cache_breakpoints(payload) + count = 0 + count += count_blocks_with_cache_control(payload[:system]) + count += (payload[:tools] || []).count { |t| block_cache_control(t) } + (payload[:messages] || []).each { |m| count += count_blocks_with_cache_control(m[:content]) } + count + end + + def count_blocks_with_cache_control(blocks) + (blocks || []).count { |b| b.is_a?(Hash) && block_cache_control(b) } + end + + def block_cache_control(block) + block[:cache_control] || block['cache_control'] + end + + def block_type(block) + block[:type] || block['type'] + end + + def inject_system_cache_breakpoint?(system_blocks) + return false if system_blocks.nil? || system_blocks.empty? + return false if system_blocks.any? { |b| b.is_a?(Hash) && block_cache_control(b) } + + system_blocks.last[:cache_control] = { type: 'ephemeral' } + true + end + + def inject_tail_cache_breakpoint(messages) + return unless messages + + block = messages.reverse_each.lazy.filter_map { |m| last_cacheable_block(m[:content]) }.first + return unless block + return if block_cache_control(block) + + block[:cache_control] = { type: 'ephemeral' } + end + + def last_cacheable_block(blocks) + (blocks || []).reverse_each do |block| + return block if block.is_a?(Hash) && CACHEABLE_BLOCK_TYPES.include?(block_type(block)) + end + nil + end + def add_tool_fields(payload, tools, tool_prefs) return unless tools.any? @@ -161,7 +234,7 @@ def format_non_tool_message(msg) def format_message_content(msg) if msg.content.is_a?(RubyLLM::Content::Raw) raw = msg.content.value - return raw.is_a?(Array) ? raw : [raw] + return translate_raw_blocks(raw.is_a?(Array) ? raw : [raw]) end blocks = [] @@ -182,8 +255,56 @@ def format_message_content(msg) blocks end + # Translates Content::Raw values into Anthropic Messages blocks at every point they + # enter this protocol. The consumer app persists system messages as Converse-format + # block arrays (e.g. [{ text: "..." }, { cachePoint: { type: 'default' } }]); those + # arrive here wrapped in Content::Raw with string keys after a DB round-trip or + # symbol keys in-memory. Anthropic-format blocks (already carrying a `type` key, + # including any cache_control) pass through unchanged. + def translate_raw_blocks(blocks) + result = [] + + blocks.each do |block| + translate_raw_block(block, result) + end + + result + end + + def translate_raw_block(block, result) + return result << block unless block.is_a?(Hash) + return result << block.dup if block.key?(:type) || block.key?('type') + + text = block[:text] || block['text'] + return result << { type: 'text', text: text } if text + + cache_point = block[:cachePoint] || block['cachePoint'] + return attach_cache_point(result, cache_point) if cache_point + + result << block + end + + # Replaces the previously emitted block (result.last) with a duped copy carrying + # cache_control, so the mutation never touches the caller's original hash object. + def attach_cache_point(result, cache_point) + block = result.last + return unless block + + ttl = cache_point[:ttl] || cache_point['ttl'] + cache_control = { type: 'ephemeral' } + cache_control[:ttl] = ttl if ttl + + result[-1] = block.dup.merge(cache_control: cache_control) + end + def format_text_and_media(content) # rubocop:disable Metrics/PerceivedComplexity return [] if content.nil? || (content.respond_to?(:empty?) && content.empty?) + + if content.is_a?(RubyLLM::Content::Raw) + raw = content.value + return translate_raw_blocks(raw.is_a?(Array) ? raw : [raw]) + end + return [{ type: 'text', text: content.to_json }] if content.is_a?(Hash) || content.is_a?(Array) return [{ type: 'text', text: content }] unless content.is_a?(RubyLLM::Content) @@ -241,7 +362,10 @@ def format_tool_result_block(msg) end def format_tool_result_content(content) - return content.value if content.is_a?(RubyLLM::Content::Raw) + if content.is_a?(RubyLLM::Content::Raw) + raw = content.value + return translate_raw_blocks(raw.is_a?(Array) ? raw : [raw]) + end return [{ type: 'text', text: content.to_json }] if content.is_a?(Hash) || content.is_a?(Array) return content_to_blocks_or_fallback(content) if content.is_a?(RubyLLM::Content) diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb index 519eeaa2e..62b72bd8b 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb @@ -161,13 +161,14 @@ def build_chunk(event, thinking_state = {}) def build_message_start_chunk(event) message = event['message'] || {} usage = message['usage'] || {} - input_tok = usage['input_tokens'] Chunk.new( role: :assistant, content: nil, model_id: message['model'] || @model&.id, - input_tokens: input_tok ? [input_tok.to_i, 0].max : nil + input_tokens: input_tokens(usage), + cached_tokens: usage['cache_read_input_tokens'], + cache_creation_tokens: usage['cache_creation_input_tokens'] ) end diff --git a/lib/ruby_llm/providers/bedrock.rb b/lib/ruby_llm/providers/bedrock.rb index 51bc46c2b..48f3b3204 100644 --- a/lib/ruby_llm/providers/bedrock.rb +++ b/lib/ruby_llm/providers/bedrock.rb @@ -59,6 +59,7 @@ def configuration_options bedrock_use_invoke_model anthropic_beta anthropic_context_management + bedrock_invoke_model_prompt_caching ] end diff --git a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb index 23a1e5a67..0178d1c35 100644 --- a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb +++ b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb @@ -377,6 +377,24 @@ def render_payload(messages = [], **opts) expect(chunk.input_tokens).to eq(42) end + it 'extracts cache usage fields from message_start, netting them out of input_tokens' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_read_input_tokens' => 60, + 'cache_creation_input_tokens' => 30 + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.cached_tokens).to eq(60) + expect(chunk.cache_creation_tokens).to eq(30) + expect(chunk.input_tokens).to eq(10) + end + it 'extracts model_id from message_start' do event = { 'type' => 'message_start', @@ -906,6 +924,260 @@ def model_double(id, metadata: {}) # URL image source rejection # --------------------------------------------------------------------------- + # --------------------------------------------------------------------------- + # Prompt caching — auto-injected breakpoints + # --------------------------------------------------------------------------- + + describe 'Chat#render_payload prompt caching' do + subject(:chat) { described_class::Chat } + + it 'injects cache_control on the last system block and last cacheable block of the final ' \ + 'message when caching is enabled' do + sys = RubyLLM::Message.new(role: :system, content: 'You are helpful') + msg = RubyLLM::Message.new(role: :user, content: 'Hello') + result = render_payload([sys, msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(result[:system].last[:cache_control]).to eq({ type: 'ephemeral' }) + expect(result[:messages].last[:content].last[:cache_control]).to eq({ type: 'ephemeral' }) + end + + it 'injects no cache_control anywhere when caching is disabled' do + sys = RubyLLM::Message.new(role: :system, content: 'You are helpful') + msg = RubyLLM::Message.new(role: :user, content: 'Hello') + result = render_payload([sys, msg], config_overrides: { bedrock_invoke_model_prompt_caching: false }) + + expect(result[:system].any? { |b| b[:cache_control] }).to be(false) + expect(result[:messages].flat_map { |m| m[:content] }.any? { |b| b[:cache_control] }).to be(false) + end + + it 'defaults to enabled' do + sys = RubyLLM::Message.new(role: :system, content: 'You are helpful') + result = render_payload([sys]) + expect(result[:system].last[:cache_control]).to eq({ type: 'ephemeral' }) + end + + it 'translates a symbol-keyed Converse raw system message and attaches cache_control from cachePoint' do + raw = RubyLLM::Content::Raw.new([{ text: 'PROMPT' }, { cachePoint: { type: 'default' } }]) + sys = RubyLLM::Message.new(role: :system, content: raw) + result = render_payload([sys], config_overrides: { bedrock_invoke_model_prompt_caching: false }) + + expect(result[:system]).to eq([{ type: 'text', text: 'PROMPT', cache_control: { type: 'ephemeral' } }]) + end + + it 'translates a string-keyed Converse raw system message and attaches cache_control from cachePoint' do + raw = RubyLLM::Content::Raw.new([{ 'text' => 'PROMPT' }, { 'cachePoint' => { 'type' => 'default' } }]) + sys = RubyLLM::Message.new(role: :system, content: raw) + result = render_payload([sys], config_overrides: { bedrock_invoke_model_prompt_caching: false }) + + expect(result[:system]).to eq([{ type: 'text', text: 'PROMPT', cache_control: { type: 'ephemeral' } }]) + end + + it 'passes through a ttl carried on the cachePoint' do + raw = RubyLLM::Content::Raw.new([{ text: 'PROMPT' }, { cachePoint: { type: 'default', ttl: '1h' } }]) + sys = RubyLLM::Message.new(role: :system, content: raw) + result = render_payload([sys], config_overrides: { bedrock_invoke_model_prompt_caching: false }) + + expect(result[:system]).to eq( + [{ type: 'text', text: 'PROMPT', cache_control: { type: 'ephemeral', ttl: '1h' } }] + ) + end + + it 'does not inject a breakpoint when 4 translated breakpoints already exist (budget exhausted)' do + raw = RubyLLM::Content::Raw.new( + [{ type: 'text', text: 'a', cache_control: { type: 'ephemeral' } }] + ) + sys = RubyLLM::Message.new(role: :system, content: raw) + + user_raw = RubyLLM::Content::Raw.new( + [ + { type: 'text', text: 'b', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'c', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'd', cache_control: { type: 'ephemeral' } } + ] + ) + msg = RubyLLM::Message.new(role: :user, content: user_raw) + + result = render_payload([sys, msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + total = count_cache_controls(result) + expect(total).to eq(4) + end + + it 'injects only the tail breakpoint (not system) when exactly one budget slot remains' do + raw = RubyLLM::Content::Raw.new( + [{ type: 'text', text: 'a', cache_control: { type: 'ephemeral' } }] + ) + sys_translated = RubyLLM::Message.new(role: :system, content: raw) + + user_raw = RubyLLM::Content::Raw.new( + [ + { type: 'text', text: 'b', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'c', cache_control: { type: 'ephemeral' } } + ] + ) + other_sys = RubyLLM::Message.new(role: :system, content: 'plain system text') + msg = RubyLLM::Message.new(role: :user, content: user_raw) + tail_msg = RubyLLM::Message.new(role: :user, content: 'final turn text') + + result = render_payload([sys_translated, other_sys, msg, tail_msg], + config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(result[:system].last[:cache_control]).to be_nil + expect(result[:messages].last[:content].last[:cache_control]).to eq({ type: 'ephemeral' }) + expect(count_cache_controls(result)).to eq(4) + end + + it 'skips thinking blocks and places the tail breakpoint on the preceding tool_result/text block' do + thinking = RubyLLM::Thinking.build( + blocks: [{ 'type' => 'redacted_thinking', 'data' => 'opaque' }] + ) + msg = RubyLLM::Message.new(role: :assistant, content: 'reply text', thinking: thinking) + + inst = make_instance(config_overrides: { bedrock_invoke_model_prompt_caching: true }) + model = inst.instance_variable_get(:@model) + result = inst.send(:render_payload, [msg], tools: {}, temperature: nil, model: model) + + content = result[:messages].last[:content] + thinking_block = content.find { |b| b['type'] == 'redacted_thinking' } + text_block = content.find { |b| b[:type] == 'text' } + + expect(thinking_block[:cache_control]).to be_nil + expect(text_block[:cache_control]).to eq({ type: 'ephemeral' }) + end + + it 'passes Anthropic-format raw blocks through byte-identical, preserving existing cache_control' do + original = [{ type: 'text', text: 'already anthropic', cache_control: { type: 'ephemeral', ttl: '1h' } }] + raw = RubyLLM::Content::Raw.new(original) + sys = RubyLLM::Message.new(role: :system, content: raw) + + result = render_payload([sys], config_overrides: { bedrock_invoke_model_prompt_caching: false }) + + expect(result[:system]).to eq(original) + end + + it 'preserves a pre-existing cache_control (including ttl) on the tail block instead of overwriting it' do + raw = RubyLLM::Content::Raw.new( + [{ type: 'text', text: 'a', cache_control: { type: 'ephemeral', ttl: '1h' } }] + ) + msg = RubyLLM::Message.new(role: :user, content: raw) + + result = render_payload([msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(result[:messages].last[:content].last[:cache_control]).to eq({ type: 'ephemeral', ttl: '1h' }) + end + + it 'does not mutate the caller-owned Content::Raw hash objects passed in via render_payload' do + original_system = [{ text: 'PROMPT' }] + original_message = [{ text: 'Hello' }] + sys = RubyLLM::Message.new(role: :system, content: RubyLLM::Content::Raw.new(original_system)) + msg = RubyLLM::Message.new(role: :user, content: RubyLLM::Content::Raw.new(original_message)) + + render_payload([sys, msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(original_system.first).not_to have_key(:cache_control) + expect(original_message.first).not_to have_key(:cache_control) + end + + it 'does not mutate a caller-owned Anthropic-format hash that already has a type key' do + original = { type: 'text', text: 'already anthropic' } + raw = RubyLLM::Content::Raw.new([original]) + msg = RubyLLM::Message.new(role: :user, content: raw) + + render_payload([msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(original).not_to have_key(:cache_control) + end + + it 'counts a string-keyed cache_control block toward the 4-breakpoint budget' do + raw = RubyLLM::Content::Raw.new( + [{ 'type' => 'text', 'text' => 'a', 'cache_control' => { 'type' => 'ephemeral' } }] + ) + sys = RubyLLM::Message.new(role: :system, content: raw) + + user_raw = RubyLLM::Content::Raw.new( + [ + { type: 'text', text: 'b', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'c', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'd', cache_control: { type: 'ephemeral' } } + ] + ) + msg = RubyLLM::Message.new(role: :user, content: user_raw) + + result = render_payload([sys, msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(count_cache_controls(result)).to eq(4) + end + + it 'dedupes against a string-keyed cache_control when deciding whether to inject the system breakpoint' do + raw = RubyLLM::Content::Raw.new( + [{ 'type' => 'text', 'text' => 'PROMPT', 'cache_control' => { 'type' => 'ephemeral' } }] + ) + sys = RubyLLM::Message.new(role: :system, content: raw) + msg = RubyLLM::Message.new(role: :user, content: 'Hello') + + result = render_payload([sys, msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + expect(result[:system].count { |b| b['cache_control'] || b[:cache_control] }).to eq(1) + end + + it 'finds a string-keyed type block as the tail cacheable block' do + raw = RubyLLM::Content::Raw.new([{ 'type' => 'text', 'text' => 'final turn' }]) + msg = RubyLLM::Message.new(role: :user, content: raw) + + result = render_payload([msg], config_overrides: { bedrock_invoke_model_prompt_caching: true }) + + block = result[:messages].last[:content].last + expect(block['cache_control'] || block[:cache_control]).to eq({ type: 'ephemeral' }) + end + + def count_cache_controls(payload) + message_blocks = (payload[:messages] || []).flat_map { |m| m[:content] } + all_blocks = (payload[:system] || []) + (payload[:tools] || []) + message_blocks + all_blocks.count { |b| b.is_a?(Hash) && (b[:cache_control] || b['cache_control']) } + end + end + + # --------------------------------------------------------------------------- + # Raw translation — Converse-format block translation + # --------------------------------------------------------------------------- + + describe 'Chat.translate_raw_blocks' do + subject(:chat) { described_class::Chat } + + it 'converts a symbol-keyed Converse text block to Anthropic text block' do + result = chat.translate_raw_blocks([{ text: 'hi' }]) + expect(result).to eq([{ type: 'text', text: 'hi' }]) + end + + it 'converts a string-keyed Converse text block to Anthropic text block' do + result = chat.translate_raw_blocks([{ 'text' => 'hi' }]) + expect(result).to eq([{ type: 'text', text: 'hi' }]) + end + + it 'drops a cachePoint with no preceding block' do + result = chat.translate_raw_blocks([{ cachePoint: { type: 'default' } }]) + expect(result).to eq([]) + end + + it 'leaves an unrecognized block untouched' do + block = { foo: 'bar' } + result = chat.translate_raw_blocks([block]) + expect(result).to eq([block]) + end + end + + describe 'Chat#format_tool_result_content prompt caching translation' do + subject(:chat) { described_class::Chat } + + it 'translates Converse-format raw tool_result content' do + raw = RubyLLM::Content::Raw.new([{ text: 'result' }, { cachePoint: { type: 'default' } }]) + msg = RubyLLM::Message.new(role: :tool, content: raw, tool_call_id: 'call_1') + result = chat.format_messages([msg]) + block = result.first[:content].first + expect(block[:content]).to eq([{ type: 'text', text: 'result', cache_control: { type: 'ephemeral' } }]) + end + end + describe 'Chat#format_image_attachment' do subject(:chat) { described_class::Chat }