diff --git a/lib/ruby_llm/message.rb b/lib/ruby_llm/message.rb index 6e9aa938d..1ae06324f 100644 --- a/lib/ruby_llm/message.rb +++ b/lib/ruby_llm/message.rb @@ -15,7 +15,7 @@ class Message :CONTENT_FILTERED_FINISH_REASONS attr_reader :role, :model_id, :tool_calls, :tool_call_id, :raw, :thinking, :tokens, :citations, - :finish_reason + :finish_reason, :stop_sequence, :context_management attr_writer :content def initialize(options = {}) @@ -29,12 +29,16 @@ def initialize(options = {}) output: options[:output_tokens], cached: options[:cached_tokens], cache_creation: options[:cache_creation_tokens], - thinking: options[:thinking_tokens] + thinking: options[:thinking_tokens], + cache_creation_ephemeral_5m: options[:cache_creation_ephemeral_5m_tokens], + cache_creation_ephemeral_1h: options[:cache_creation_ephemeral_1h_tokens] ) @raw = options[:raw] @thinking = options[:thinking] @citations = Array(options[:citations]) @finish_reason = options[:finish_reason] + @stop_sequence = options[:stop_sequence] + @context_management = options[:context_management] ensure_valid_role end @@ -113,7 +117,9 @@ def to_h thinking: thinking&.text, thinking_signature: thinking&.signature, citations: citations.empty? ? nil : citations.map(&:to_h), - finish_reason: finish_reason + finish_reason: finish_reason, + stop_sequence: stop_sequence, + context_management: context_management }.merge(tokens ? tokens.to_h : {}).compact end diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb index f0035669e..ad460d903 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/chat.rb @@ -105,6 +105,7 @@ def parse_completion_body(data, raw:) content_blocks = data['content'] || [] usage = data['usage'] || {} + cache_creation = usage['cache_creation'] || {} Message.new( role: :assistant, @@ -115,7 +116,11 @@ def parse_completion_body(data, raw:) output_tokens: usage['output_tokens'], cached_tokens: usage['cache_read_input_tokens'], cache_creation_tokens: usage['cache_creation_input_tokens'], + cache_creation_ephemeral_5m_tokens: cache_creation['ephemeral_5m_input_tokens'], + cache_creation_ephemeral_1h_tokens: cache_creation['ephemeral_1h_input_tokens'], finish_reason: data['stop_reason'], + stop_sequence: data['stop_sequence'], + context_management: data['context_management'], model_id: data['model'], raw: raw ) diff --git a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb index 519eeaa2e..9a4d1f976 100644 --- a/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb +++ b/lib/ruby_llm/protocols/bedrock_invoke_model/streaming.rb @@ -21,7 +21,7 @@ def stream_url end def stream_response(payload, additional_headers = {}, &block) - accumulator = StreamAccumulator.new + accumulator = StreamAccumulator.new(net_cache_tokens: true) decoder = event_stream_decoder thinking_state = {} body = JSON.generate(payload) @@ -48,9 +48,19 @@ def stream_response(payload, additional_headers = {}, &block) message = accumulator.to_message(response) RubyLLM.logger.debug { "Stream completed: #{message.content}" } + log_context_management(message) message end + def log_context_management(message) + applied_edits = message.context_management && message.context_management['applied_edits'] + return unless applied_edits && !applied_edits.empty? + + RubyLLM.logger.debug do + "Bedrock InvokeModel context_management applied_edits: #{applied_edits.inspect}" + end + end + def event_stream_decoder require 'aws-eventstream' Aws::EventStream::Decoder.new @@ -161,13 +171,20 @@ def build_chunk(event, thinking_state = {}) def build_message_start_chunk(event) message = event['message'] || {} usage = message['usage'] || {} + cache_read = usage['cache_read_input_tokens'] + cache_creation = usage['cache_creation_input_tokens'] input_tok = usage['input_tokens'] + cache_creation_detail = usage['cache_creation'] || {} 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_tok&.to_i, + cached_tokens: cache_read, + cache_creation_tokens: cache_creation, + cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'], + cache_creation_ephemeral_1h_tokens: cache_creation_detail['ephemeral_1h_input_tokens'] ) end @@ -254,13 +271,20 @@ def build_content_block_stop_chunk(event, thinking_state) def build_message_delta_chunk(event) delta = event['delta'] || {} usage = event['usage'] || {} + cache_creation_detail = usage['cache_creation'] || {} Chunk.new( role: :assistant, content: nil, model_id: @model&.id, output_tokens: usage['output_tokens'], - finish_reason: delta['stop_reason'] + cached_tokens: usage['cache_read_input_tokens'], + cache_creation_tokens: usage['cache_creation_input_tokens'], + cache_creation_ephemeral_5m_tokens: cache_creation_detail['ephemeral_5m_input_tokens'], + cache_creation_ephemeral_1h_tokens: cache_creation_detail['ephemeral_1h_input_tokens'], + finish_reason: delta['stop_reason'], + stop_sequence: delta['stop_sequence'], + context_management: delta['context_management'] ) end diff --git a/lib/ruby_llm/stream_accumulator.rb b/lib/ruby_llm/stream_accumulator.rb index 60eff437d..90a803d99 100644 --- a/lib/ruby_llm/stream_accumulator.rb +++ b/lib/ruby_llm/stream_accumulator.rb @@ -8,19 +8,31 @@ module RubyLLM class StreamAccumulator attr_reader :content, :model_id, :tool_calls - def initialize + # net_cache_tokens: when true, input_tokens is netted against the final accumulated + # cached_tokens/cache_creation_tokens at to_message time instead of being taken as-is + # from whichever chunk last set it. Bedrock InvokeModel's message_start chunk carries + # input_tokens while a later message_delta chunk can still update the cache fields, so + # netting must be deferred until all chunks are in. Other providers already net + # input_tokens against cache fields within the same usage snapshot before building the + # chunk, so deferred netting stays off for them to avoid double-subtracting. + def initialize(net_cache_tokens: false) + @net_cache_tokens = net_cache_tokens @content = +'' @citations = [] @thinking_text = +'' @thinking_signature = nil @thinking_blocks = [] @tool_calls = {} - @input_tokens = nil + @raw_input_tokens = nil @output_tokens = nil @cached_tokens = nil @cache_creation_tokens = nil + @cache_creation_ephemeral_5m_tokens = nil + @cache_creation_ephemeral_1h_tokens = nil @thinking_tokens = nil @finish_reason = nil + @stop_sequence = nil + @context_management = nil @inside_think_tag = false @pending_think_tag = +'' @latest_tool_call_id = nil @@ -35,6 +47,8 @@ def add(chunk) accumulate_citations(chunk.citations) append_thinking_from_chunk(chunk) @finish_reason = chunk.finish_reason if chunk.finish_reason + @stop_sequence = chunk.stop_sequence if chunk.stop_sequence + @context_management = chunk.context_management if chunk.context_management count_tokens chunk RubyLLM.logger.debug { inspect } if RubyLLM.config.log_stream_debug end @@ -50,13 +64,17 @@ def to_message(response) blocks: @thinking_blocks.empty? ? nil : @thinking_blocks ), tokens: Tokens.build( - input: @input_tokens, + input: netted_input_tokens, output: @output_tokens, cached: @cached_tokens, cache_creation: @cache_creation_tokens, - thinking: @thinking_tokens + thinking: @thinking_tokens, + cache_creation_ephemeral_5m: @cache_creation_ephemeral_5m_tokens, + cache_creation_ephemeral_1h: @cache_creation_ephemeral_1h_tokens ), finish_reason: @finish_reason, + stop_sequence: @stop_sequence, + context_management: @context_management, model_id: model_id, tool_calls: tool_calls_from_stream, raw: response @@ -65,6 +83,12 @@ def to_message(response) private + def netted_input_tokens + return @raw_input_tokens unless @net_cache_tokens && @raw_input_tokens + + [@raw_input_tokens - @cached_tokens.to_i - @cache_creation_tokens.to_i, 0].max + end + # Providers like Perplexity repeat the full citation list on every chunk. def accumulate_citations(new_citations) new_citations.each do |citation| @@ -153,11 +177,19 @@ def find_tool_call(stream_key) end def count_tokens(chunk) - @input_tokens = chunk.input_tokens if chunk.input_tokens + @raw_input_tokens = chunk.input_tokens if chunk.input_tokens @output_tokens = chunk.output_tokens if chunk.output_tokens @cached_tokens = chunk.cached_tokens if chunk.cached_tokens @cache_creation_tokens = chunk.cache_creation_tokens if chunk.cache_creation_tokens @thinking_tokens = chunk.thinking_tokens if chunk.thinking_tokens + count_cache_creation_ttl_tokens(chunk) + end + + def count_cache_creation_ttl_tokens(chunk) + ephemeral_5m = chunk.tokens&.cache_creation_ephemeral_5m + ephemeral_1h = chunk.tokens&.cache_creation_ephemeral_1h + @cache_creation_ephemeral_5m_tokens = ephemeral_5m if ephemeral_5m + @cache_creation_ephemeral_1h_tokens = ephemeral_1h if ephemeral_1h end def handle_chunk_content(chunk) diff --git a/lib/ruby_llm/tokens.rb b/lib/ruby_llm/tokens.rb index 2ed45e7f8..903be0355 100644 --- a/lib/ruby_llm/tokens.rb +++ b/lib/ruby_llm/tokens.rb @@ -3,27 +3,39 @@ module RubyLLM # Represents token usage for a response. class Tokens - attr_reader :input, :output, :cached, :cache_creation, :thinking + attr_reader :input, :output, :cached, :cache_creation, :thinking, + :cache_creation_ephemeral_5m, :cache_creation_ephemeral_1h - def initialize(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil) + # rubocop:disable Metrics/ParameterLists + def initialize(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil, + cache_creation_ephemeral_5m: nil, cache_creation_ephemeral_1h: nil) @input = input @output = output @cached = cached @cache_creation = cache_creation @thinking = thinking + @cache_creation_ephemeral_5m = cache_creation_ephemeral_5m + @cache_creation_ephemeral_1h = cache_creation_ephemeral_1h end + # rubocop:enable Metrics/ParameterLists - def self.build(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil) - return nil if [input, output, cached, cache_creation, thinking].all?(&:nil?) + # rubocop:disable Metrics/ParameterLists + def self.build(input: nil, output: nil, cached: nil, cache_creation: nil, thinking: nil, + cache_creation_ephemeral_5m: nil, cache_creation_ephemeral_1h: nil) + return nil if [input, output, cached, cache_creation, thinking, + cache_creation_ephemeral_5m, cache_creation_ephemeral_1h].all?(&:nil?) new( input: input, output: output, cached: cached, cache_creation: cache_creation, - thinking: thinking + thinking: thinking, + cache_creation_ephemeral_5m: cache_creation_ephemeral_5m, + cache_creation_ephemeral_1h: cache_creation_ephemeral_1h ) end + # rubocop:enable Metrics/ParameterLists def to_h { @@ -31,7 +43,9 @@ def to_h output_tokens: output, cached_tokens: cached, cache_creation_tokens: cache_creation, - thinking_tokens: thinking + thinking_tokens: thinking, + cache_creation_ephemeral_5m_tokens: cache_creation_ephemeral_5m, + cache_creation_ephemeral_1h_tokens: cache_creation_ephemeral_1h }.compact end diff --git a/spec/ruby_llm/chat_loop_spec.rb b/spec/ruby_llm/chat_loop_spec.rb index 349d8ff32..038b2c549 100644 --- a/spec/ruby_llm/chat_loop_spec.rb +++ b/spec/ruby_llm/chat_loop_spec.rb @@ -141,6 +141,25 @@ def tool_call_message(name: 'echo') expect(chat.complete).to be(answer_message) expect(chat.provider).not_to have_received(:complete) end + + it 'persists distinct, non-zero usage on every intermediate tool-call round trip' do + # Each round trip (including ones that returned tool_calls) must carry its own usage, + # not the final round trip's, so a per-turn cost breakdown stays accurate. + first_round = RubyLLM::Message.new( + role: :assistant, content: '', tool_calls: tool_call_message.tool_calls, + input_tokens: 100, output_tokens: 10 + ) + final_round = RubyLLM::Message.new(role: :assistant, content: 'hello', input_tokens: 150, output_tokens: 5) + + allow(chat.provider).to receive(:complete).and_return(first_round, final_round) + chat.ask_later('Echo "hello" back to me.') + + chat.complete + + assistant_messages = chat.messages.select { |m| m.role == :assistant } + expect(assistant_messages.map(&:input_tokens)).to eq([100, 150]) + expect(assistant_messages.map(&:output_tokens)).to eq([10, 5]) + end end describe '#add_completion' do diff --git a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb index 23a1e5a67..d5264fe62 100644 --- a/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb +++ b/spec/ruby_llm/protocols/bedrock_invoke_model_spec.rb @@ -350,6 +350,55 @@ def render_payload(messages = [], **opts) expect(chat.parse_completion_body(nil, raw: nil)).to be_nil expect(chat.parse_completion_body({}, raw: nil)).to be_nil end + + it 'reports input_tokens of 0 for a fully-cached request' do + data = basic_response.merge( + 'usage' => { + 'input_tokens' => 100, + 'output_tokens' => 5, + 'cache_read_input_tokens' => 100, + 'cache_creation_input_tokens' => 0 + } + ) + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.input_tokens).to eq(0) + expect(msg.cached_tokens).to eq(100) + end + + it 'extracts cache_creation TTL breakdown when present' do + data = basic_response.merge( + 'usage' => { + 'input_tokens' => 100, + 'output_tokens' => 5, + 'cache_read_input_tokens' => 0, + 'cache_creation_input_tokens' => 30, + 'cache_creation' => { 'ephemeral_5m_input_tokens' => 20, 'ephemeral_1h_input_tokens' => 10 } + } + ) + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.tokens.cache_creation_ephemeral_5m).to eq(20) + expect(msg.tokens.cache_creation_ephemeral_1h).to eq(10) + end + + it 'extracts stop_sequence from the response' do + data = basic_response.merge('stop_reason' => 'stop_sequence', 'stop_sequence' => 'STOP') + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.stop_sequence).to eq('STOP') + end + + it 'surfaces top-level context_management applied_edits on the parsed message' do + data = basic_response.merge( + 'context_management' => { + 'applied_edits' => [ + { 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 3, 'cleared_input_tokens' => 500 } + ] + } + ) + msg = chat.parse_completion_body(data, raw: nil) + expect(msg.context_management['applied_edits']).to eq( + [{ 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 3, 'cleared_input_tokens' => 500 }] + ) + end end # --------------------------------------------------------------------------- @@ -441,6 +490,86 @@ def render_payload(messages = [], **opts) expect(chunk.finish_reason).to eq('end_turn') end + it 'extracts cache usage from message_start and passes input_tokens through un-netted' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_read_input_tokens' => 40, + 'cache_creation_input_tokens' => 10 + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.cached_tokens).to eq(40) + expect(chunk.cache_creation_tokens).to eq(10) + # Netting against cache tokens is deferred to StreamAccumulator#to_message, since a + # later message_delta chunk can still update the cache fields for this same message. + expect(chunk.input_tokens).to eq(100) + end + + it 'passes the raw input_tokens through from message_start for a fully-cached request' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_read_input_tokens' => 100, + 'cache_creation_input_tokens' => 0 + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.input_tokens).to eq(100) + expect(chunk.cached_tokens).to eq(100) + end + + it 'extracts cache_creation TTL breakdown from message_start' do + event = { + 'type' => 'message_start', + 'message' => { + 'model' => 'anthropic.claude-sonnet-4-6', + 'usage' => { + 'input_tokens' => 100, + 'cache_creation_input_tokens' => 30, + 'cache_creation' => { 'ephemeral_5m_input_tokens' => 20, 'ephemeral_1h_input_tokens' => 10 } + } + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.tokens.cache_creation_ephemeral_5m).to eq(20) + expect(chunk.tokens.cache_creation_ephemeral_1h).to eq(10) + end + + it 'extracts cumulative cache usage, stop_sequence, and context_management from message_delta' do + event = { + 'type' => 'message_delta', + 'delta' => { + 'stop_reason' => 'end_turn', + 'stop_sequence' => nil, + 'context_management' => { + 'applied_edits' => [ + { 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 2, 'cleared_input_tokens' => 300 } + ] + } + }, + 'usage' => { + 'output_tokens' => 17, + 'cache_read_input_tokens' => 40, + 'cache_creation_input_tokens' => 10 + } + } + chunk = streaming.send(:build_chunk, event) + expect(chunk.cached_tokens).to eq(40) + expect(chunk.cache_creation_tokens).to eq(10) + expect(chunk.context_management['applied_edits']).to eq( + [{ 'type' => 'clear_tool_uses_20250919', 'cleared_tool_uses' => 2, 'cleared_input_tokens' => 300 }] + ) + end + it 'returns a chunk for message_stop without error' do event = { 'type' => 'message_stop' } expect { streaming.send(:build_chunk, event) }.not_to raise_error @@ -481,6 +610,84 @@ def render_payload(messages = [], **opts) expect(message.content).to eq('Hello world') expect(message.output_tokens).to eq(2) end + + it 'plumbs cache tokens and context_management through StreamAccumulator into the final message' do + accumulator = RubyLLM::StreamAccumulator.new(net_cache_tokens: true) + + events = [ + { 'type' => 'message_start', 'message' => { + 'model' => 'test-model', + 'usage' => { 'input_tokens' => 100, 'cache_read_input_tokens' => 40, 'cache_creation_input_tokens' => 10 } + } }, + { 'type' => 'content_block_delta', 'index' => 0, + 'delta' => { 'type' => 'text_delta', 'text' => 'Hello' } }, + { 'type' => 'message_delta', + 'delta' => { + 'stop_reason' => 'end_turn', + 'context_management' => { + 'applied_edits' => [{ 'type' => 'clear_tool_uses_20250919', 'cleared_input_tokens' => 300 }] + } + }, + 'usage' => { 'output_tokens' => 2, 'cache_read_input_tokens' => 40, 'cache_creation_input_tokens' => 10 } } + ] + + events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) } + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(50) + expect(message.cached_tokens).to eq(40) + expect(message.cache_creation_tokens).to eq(10) + expect(message.context_management['applied_edits']).to eq( + [{ 'type' => 'clear_tool_uses_20250919', 'cleared_input_tokens' => 300 }] + ) + end + + it 'reports input_tokens of 0 through the full stream round trip for a fully-cached request' do + accumulator = RubyLLM::StreamAccumulator.new(net_cache_tokens: true) + + events = [ + { 'type' => 'message_start', 'message' => { + 'model' => 'test-model', + 'usage' => { 'input_tokens' => 100, 'cache_read_input_tokens' => 100, 'cache_creation_input_tokens' => 0 } + } }, + { 'type' => 'content_block_delta', 'index' => 0, + 'delta' => { 'type' => 'text_delta', 'text' => 'Hi' } }, + { 'type' => 'message_delta', 'delta' => { 'stop_reason' => 'end_turn' }, + 'usage' => { 'output_tokens' => 2 } } + ] + + events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) } + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(0) + expect(message.cached_tokens).to eq(100) + end + + it 'nets input_tokens against cache fields that are only updated by a later message_delta' do + # Regression test: message_start's usage snapshot has no cache_read_input_tokens yet + # (some Bedrock responses only report cache usage on message_delta), so a chunk-time + # netting of message_start's input_tokens would silently miss this cache usage. Since + # netting is now deferred to to_message, it correctly nets against the final values. + accumulator = RubyLLM::StreamAccumulator.new(net_cache_tokens: true) + + events = [ + { 'type' => 'message_start', 'message' => { + 'model' => 'test-model', + 'usage' => { 'input_tokens' => 100 } + } }, + { 'type' => 'content_block_delta', 'index' => 0, + 'delta' => { 'type' => 'text_delta', 'text' => 'Hi' } }, + { 'type' => 'message_delta', 'delta' => { 'stop_reason' => 'end_turn' }, + 'usage' => { 'output_tokens' => 2, 'cache_read_input_tokens' => 40, 'cache_creation_input_tokens' => 10 } } + ] + + events.each { |e| accumulator.add(streaming.send(:build_chunk, e)) } + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(50) + expect(message.cached_tokens).to eq(40) + expect(message.cache_creation_tokens).to eq(10) + end end # --------------------------------------------------------------------------- diff --git a/spec/ruby_llm/protocols/converse/chat_spec.rb b/spec/ruby_llm/protocols/converse/chat_spec.rb index 518278ac6..00a4bd98c 100644 --- a/spec/ruby_llm/protocols/converse/chat_spec.rb +++ b/spec/ruby_llm/protocols/converse/chat_spec.rb @@ -29,6 +29,51 @@ expect(message.cache_creation_tokens).to eq(10) end + it 'reports input_tokens of 0 for a fully-cached request' do + response_body = { + 'modelId' => 'anthropic.claude-sonnet-4-5-20250929-v1:0', + 'output' => { + 'message' => { + 'content' => [{ 'text' => 'Hi!' }] + } + }, + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 100, + 'cacheWriteInputTokens' => 0 + } + } + + response = instance_double(Faraday::Response, body: response_body) + message = described_class.parse_completion_response(response) + + expect(message.input_tokens).to eq(0) + expect(message.cached_tokens).to eq(100) + end + + it 'does not fabricate an ephemeral cache-creation TTL breakdown, which Converse does not return' do + response_body = { + 'output' => { + 'message' => { + 'content' => [{ 'text' => 'Hi!' }] + } + }, + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 40, + 'cacheWriteInputTokens' => 10 + } + } + + response = instance_double(Faraday::Response, body: response_body) + message = described_class.parse_completion_response(response) + + expect(message.tokens.cache_creation_ephemeral_5m).to be_nil + expect(message.tokens.cache_creation_ephemeral_1h).to be_nil + end + it 'preserves raw stopReason as finish_reason' do response_body = { 'modelId' => 'amazon.nova-lite-v1:0', diff --git a/spec/ruby_llm/protocols/converse/streaming_spec.rb b/spec/ruby_llm/protocols/converse/streaming_spec.rb index fe1cd935e..510e51c1f 100644 --- a/spec/ruby_llm/protocols/converse/streaming_spec.rb +++ b/spec/ruby_llm/protocols/converse/streaming_spec.rb @@ -70,6 +70,64 @@ expect(chunk.thinking_tokens).to eq(7) end + it 'normalizes cache read and write tokens out of input tokens from metadata usage' do + event = { + 'metadata' => { + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 40, + 'cacheWriteInputTokens' => 10 + } + } + } + + chunk = streaming.send(:build_chunk, event) + + expect(chunk.input_tokens).to eq(50) + expect(chunk.output_tokens).to eq(5) + expect(chunk.cached_tokens).to eq(40) + expect(chunk.cache_creation_tokens).to eq(10) + end + + it 'reports input_tokens of 0 through the accumulator for a fully-cached streamed request' do + accumulator = RubyLLM::StreamAccumulator.new + usage_event = { + 'metadata' => { + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 100, + 'cacheWriteInputTokens' => 0 + } + } + } + + accumulator.add(streaming.send(:build_chunk, usage_event)) + message = accumulator.to_message(nil) + + expect(message.input_tokens).to eq(0) + expect(message.cached_tokens).to eq(100) + end + + it 'does not fabricate an ephemeral cache-creation TTL breakdown, which Converse does not return' do + event = { + 'metadata' => { + 'usage' => { + 'inputTokens' => 100, + 'outputTokens' => 5, + 'cacheReadInputTokens' => 40, + 'cacheWriteInputTokens' => 10 + } + } + } + + chunk = streaming.send(:build_chunk, event) + + expect(chunk.tokens&.cache_creation_ephemeral_5m).to be_nil + expect(chunk.tokens&.cache_creation_ephemeral_1h).to be_nil + end + it 'accumulates Bedrock Converse Stream thinking deltas into the final message' do accumulator = RubyLLM::StreamAccumulator.new text_event = {