From c1e487c4fea8b794d9afa301dbadcdac7c65970f Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Fri, 13 Nov 2026 13:10:59 -0400 Subject: [PATCH] fix: include ruby_llm cache tokens in total_tokens ruby_llm reports cache reads and writes apart from input. The total was input plus output, so those tokens were missing. Stamp the total after a stream merges its chunks, so a later output-only chunk does not drop cache already counted. --- .../providers/ruby_llm_provider.rb | 28 ++++++++++++++++--- test/providers/ruby_llm/response_test.rb | 6 ++++ test/providers/ruby_llm/streaming_test.rb | 4 +++ 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/lib/active_agent/providers/ruby_llm_provider.rb b/lib/active_agent/providers/ruby_llm_provider.rb index e5e6908ac..23d390682 100644 --- a/lib/active_agent/providers/ruby_llm_provider.rb +++ b/lib/active_agent/providers/ruby_llm_provider.rb @@ -185,13 +185,17 @@ def process_stream_chunk(chunk) # Counts arrive in separate chunks (Anthropic sends input at the start # and output at the end). They are cumulative for this turn, so merge - # reported fields rather than adding each chunk's counts. + # reported fields rather than adding each chunk's counts. The total is + # stamped after the merge: a later chunk often carries only output, and + # its own total would drop the cache already counted. if (usage = ruby_llm_usage(chunk)) - (@stream_response[:usage] ||= {}).merge!(usage) do |field, previous, current| + merged = (@stream_response[:usage] ||= {}) + merged.merge!(usage) do |field, previous, current| # Chat Completions synthesizes zero when cache writes are omitted. # That cannot erase writes already counted in this turn. field == :cache_creation_tokens && current == 0 ? previous : current end + merged[:total_tokens] = ruby_llm_total_tokens(merged) end if chunk.respond_to?(:finish_reason) && chunk.finish_reason @stream_response[:stop_reason] = ruby_llm_stop_reason(chunk, tool_calls: message[:tool_calls]) @@ -517,7 +521,10 @@ def normalize_ruby_llm_response(response, model_id) end usage = ruby_llm_usage(response) - hash[:usage] = usage if usage + if usage + usage[:total_tokens] = ruby_llm_total_tokens(usage) + hash[:usage] = usage + end hash[:model] = model_id if model_id @@ -592,7 +599,8 @@ def ruby_llm_stop_reason(response, tool_calls: response.tool_calls) # none: the message has no Tokens then in ruby_llm 1.x. RubyLLM counts # cached tokens apart from the input, so input_tokens leaves them out. # A count the provider left out is left out here too, and Usage reads it - # as zero. + # as zero. Total is not set here: a stream merges these hashes field by + # field and stamps the total once the merge is done. # # @param response [RubyLLM::Message] # @return [Hash, nil] Common::Usage attributes @@ -619,6 +627,18 @@ def ruby_llm_usage(response) }.compact end end + + # Every token the response processed. ruby_llm's input excludes cache + # reads and writes: OpenAI and Gemini subtract them from the prompt + # count, and Anthropic already reports them apart from input_tokens. + # Output already includes thinking when that is billed as output, so + # reasoning_tokens is not added again. + # + # @param usage [Hash] fields from {#ruby_llm_usage}, after a stream merge + # @return [Integer] + def ruby_llm_total_tokens(usage) + %i[input_tokens output_tokens cached_tokens cache_creation_tokens].sum { |key| usage[key] || 0 } + end end end end diff --git a/test/providers/ruby_llm/response_test.rb b/test/providers/ruby_llm/response_test.rb index 13aa05828..4a2c07de9 100644 --- a/test/providers/ruby_llm/response_test.rb +++ b/test/providers/ruby_llm/response_test.rb @@ -48,6 +48,8 @@ class RubyLLMResponseTest < ActiveSupport::TestCase assert_equal 30, usage.cached_tokens assert_equal 4, usage.cache_creation_tokens assert_equal 2, usage.reasoning_tokens + # input excludes cache, and thinking is already inside output when billed. + assert_equal 49, usage.total_tokens end # ruby_llm 1.16 has no token counts at all on such a message (its tokens is @@ -83,6 +85,8 @@ class RubyLLMResponseTest < ActiveSupport::TestCase assert_equal 3, usage.public_send(field) assert_equal 0, usage.input_tokens assert_equal 0, usage.output_tokens + # A lone thinking count is a breakdown of output, not an extra bucket. + assert_equal(count == :thinking ? 0 : 3, usage.total_tokens) end end @@ -142,6 +146,7 @@ class RubyLLMResponseTest < ActiveSupport::TestCase assert_equal 7, usage.input_tokens assert_equal 4, usage.output_tokens assert_equal 5, usage.cached_tokens + assert_equal 16, usage.total_tokens end test "reports the tokens Anthropic counted, cache reads and writes included" do @@ -155,6 +160,7 @@ class RubyLLMResponseTest < ActiveSupport::TestCase assert_equal 4, usage.output_tokens assert_equal 5, usage.cached_tokens assert_equal 2, usage.cache_creation_tokens + assert_equal 23, usage.total_tokens end # --- Stop reason --- diff --git a/test/providers/ruby_llm/streaming_test.rb b/test/providers/ruby_llm/streaming_test.rb index f237ab2de..400f6ed7b 100644 --- a/test/providers/ruby_llm/streaming_test.rb +++ b/test/providers/ruby_llm/streaming_test.rb @@ -239,6 +239,7 @@ def message_delta(stop_reason) assert_not_nil response.usage assert_equal(stream.name == :chat_completions ? 7 : 12, response.usage.input_tokens) assert_equal 4, response.usage.output_tokens + assert_equal 16, response.usage.total_tokens assert_equal "end_turn", response.finish_reason assert_equal "Done.", response.messages.last.content assert_equal 1, response.messages.count { |message| message.role == "assistant" } @@ -414,6 +415,7 @@ def message_delta(stop_reason) assert_equal 5, response.usage.output_tokens assert_equal 5, response.usage.cached_tokens assert_equal 1, response.usage.reasoning_tokens + assert_equal 27, response.usage.total_tokens assert_equal "end_turn", response.finish_reason end end @@ -435,6 +437,7 @@ def message_delta(stop_reason) assert_equal 7, usage.input_tokens assert_equal 4, usage.output_tokens assert_equal 5, usage.cache_creation_tokens + assert_equal 16, usage.total_tokens end test "updates positive cumulative cache writes and starts a fresh count on each tool turn" do @@ -453,6 +456,7 @@ def message_delta(stop_reason) assert_equal 17, response.usage.input_tokens assert_equal 5, response.usage.output_tokens assert_equal 9, response.usage.cache_creation_tokens + assert_equal 31, response.usage.total_tokens assert_equal [ 9, 0 ], response.usages.map(&:cache_creation_tokens) end end