llm_cost_tracker 0.14.0 → 0.14.1

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -8,10 +8,14 @@ module LlmCostTracker
       extend Base
 
       minimum_version "1.15.0"
+      maximum_version "3.0.0"
 
       class << self
         def patch_targets
-          [patch_target("RubyLLM::Provider", with: ProviderPatch)]
+          [
+            patch_target("RubyLLM::Provider", with: ProviderPatch),
+            patch_target("RubyLLM::Providers::Gemini::Transcription", with: GeminiTranscriptionPatch, optional: true)
+          ]
         end
 
         def record_completion(provider, response, request:, latency_ms:, has_block:)
@@ -36,25 +40,29 @@ module LlmCostTracker
         end
 
         def record_transcription(provider, response, request:, latency_ms:)
+          model = response_model_id(response) || model_id_from_request(request[:model])
+          match = LlmCostTracker::Pricing::Matcher.lookup(provider: provider.slug.to_s, model: model)
           record_usage(
             provider: provider.slug.to_s,
-            model: response_model_id(response) || model_id_from_request(request[:model]),
+            model: model,
             response: response,
             latency_ms: latency_ms,
-            stream: false
+            stream: false,
+            audio_input: match&.prices&.key?("audio_input")
           )
         end
 
         def record_image(provider, response, request:, latency_ms:)
-          usage = response.usage.with_indifferent_access
+          image = response.is_a?(Array) ? response.first : response
+          usage = image_usage(image)
           raw_input = usage[:input_tokens].to_i
           raw_output = usage[:output_tokens].to_i
           image_input = image_token_detail(usage, :input)
           image_output = image_token_detail(usage, :output)
           record_passthrough(
             provider: provider.slug.to_s,
-            model: response_model_id(response) || model_id_from_request(request[:model]),
-            response: response,
+            model: response_model_id(image) || model_id_from_request(request[:model]),
+            response: image,
             latency_ms: latency_ms,
             input_tokens: [raw_input - image_input, 0].max,
             image_input_tokens: image_input,
@@ -74,6 +82,12 @@ module LlmCostTracker
           )
         end
 
+        def image_usage(image)
+          usage = image.try(:usage)
+          usage = image.send(:raw_usage) if !usage.is_a?(Hash) && image.respond_to?(:raw_usage, true)
+          (usage.is_a?(Hash) ? usage : {}).with_indifferent_access
+        end
+
         def image_token_detail(usage, direction)
           container_key = direction == :input ? :input_tokens_details : :output_tokens_details
           details = usage[container_key]
@@ -111,27 +125,29 @@ module LlmCostTracker
           end
         end
 
-        def record_usage(provider:, model:, response:, latency_ms:, stream:, output_tokens: nil)
+        def record_usage(provider:, model:, response:, latency_ms:, stream:, output_tokens: nil, audio_input: false)
           return unless active?
 
           record_safely do
-            input_tokens = response.input_tokens
-            output_tokens = response.output_tokens if output_tokens.nil?
+            counts = token_counts(response)
+            input_tokens = counts[:input]
+            output_tokens = counts[:output] if output_tokens.nil?
             next if input_tokens.nil? && output_tokens.nil?
 
-            cache_write_5m, cache_write_1h = cache_creation_split(provider, response)
+            cache_write_5m, cache_write_1h = cache_write_split(provider, response, counts[:cache_write])
             LlmCostTracker::Tracker.record(
               event: Event.build(
                 provider: provider,
                 model: model,
                 pricing_mode: pricing_mode_for(provider: provider, response: response),
                 token_usage: Usage::TokenUsage.build(
-                  input_tokens: input_tokens.to_i,
+                  input_tokens: audio_input ? 0 : input_tokens.to_i,
+                  audio_input_tokens: audio_input ? input_tokens.to_i : 0,
                   output_tokens: output_tokens.to_i,
-                  cache_read_input_tokens: response.try(:cached_tokens).to_i,
+                  cache_read_input_tokens: counts[:cache_read].to_i,
                   cache_write_input_tokens: cache_write_5m,
                   cache_write_extended_input_tokens: cache_write_1h,
-                  hidden_output_tokens: response.try(:thinking_tokens).to_i
+                  hidden_output_tokens: counts[:thinking].to_i
                 ),
                 stream: stream,
                 usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE,
@@ -142,11 +158,22 @@ module LlmCostTracker
           end
         end
 
-        def cache_creation_split(provider, response)
-          return [response.try(:cache_creation_tokens).to_i, 0] unless provider == "anthropic"
+        def token_counts(response)
+          tokens = response.try(:tokens)
+          return { input: response.try(:input_tokens), output: response.try(:output_tokens) } unless tokens
+
+          {
+            input: tokens.input,
+            output: tokens.output,
+            cache_read: tokens.cache_read,
+            cache_write: tokens.cache_write,
+            thinking: tokens.thinking
+          }
+        end
 
-          cache = raw_body(response).dig("usage", "cache_creation")
-          return [response.try(:cache_creation_tokens).to_i, 0] unless cache.is_a?(Hash)
+        def cache_write_split(provider, response, cache_write)
+          cache = raw_body(response).dig("usage", "cache_creation") if provider == "anthropic"
+          return [cache_write.to_i, 0] unless cache.is_a?(Hash)
 
           [cache["ephemeral_5m_input_tokens"].to_i, cache["ephemeral_1h_input_tokens"].to_i]
         end
@@ -164,7 +191,8 @@ module LlmCostTracker
         end
 
         def raw_body(response)
-          body = response.try(:raw)&.body
+          raw = response.try(:raw)
+          body = raw.respond_to?(:body) ? raw.body : raw
           body.is_a?(Hash) ? body : {}
         end
 
@@ -217,6 +245,13 @@ module LlmCostTracker
           LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
         end
       end
+
+      module GeminiTranscriptionPatch
+        def transcribe(*args, **kwargs)
+          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_transcription)
+          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
+        end
+      end
     end
   end
 end

data/lib/llm_cost_tracker/ledger/isolation.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require_relative "../errors"
+
+module LlmCostTracker
+  module Ledger
+    module Isolation
+      class << self
+        def guard(model = LlmCostTracker::Call, &)
+          connection = model.connection
+          nested = connection.transaction_open?
+          return yield unless nested
+
+          model.transaction(requires_new: true, &)
+        rescue ActiveRecord::TransactionRollbackError => e
+          raise unless nested && connection.savepoint_errors_invalidate_transactions?
+
+          raise TransactionAbortedError, e
+        end
+
+        def after_commit(&)
+          current = LlmCostTracker::Call.connection.current_transaction
+          return yield unless current.open? && current.joinable? && current.respond_to?(:after_commit)
+
+          current.after_commit(&)
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/period/totals.rb

@@ -2,6 +2,7 @@
 
 require "bigdecimal/util"
 
+require_relative "../isolation"
 require_relative "../period"
 
 module LlmCostTracker
@@ -22,7 +23,7 @@ module LlmCostTracker
 
           values = periods.to_h { |period| [period, BigDecimal("0")] }
           period_by_name = periods.to_h { |period| [period.to_s, period] }
-          LlmCostTracker::Call.find_by_sql(union_sql).each do |row|
+          Isolation.guard { LlmCostTracker::Call.find_by_sql(union_sql) }.each do |row|
             values[period_by_name.fetch(row.period_key)] = row.total_cost.to_d
           end
           values

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -2,6 +2,7 @@
 
 require "bigdecimal/util"
 
+require_relative "isolation"
 require_relative "period"
 
 module LlmCostTracker
@@ -33,14 +34,21 @@ module LlmCostTracker
                          :REBUILD_INSERT_SLICE
 
         def increment_safely!(events)
+          return unless LlmCostTracker.configuration.budgets.totals_source == :cache
+
+          Isolation.after_commit { increment_with_retries(events) }
+        end
+
+        def increment_with_retries(events)
+          retryable = !LlmCostTracker::Call.connection.transaction_open?
           attempt = 0
           begin
             attempt += 1
-            increment!(events)
+            Isolation.guard { increment!(events) }
+          rescue LlmCostTracker::Error
+            raise
           rescue StandardError => e
-            raise if LlmCostTracker::Call.connection.open_transactions.positive?
-
-            if attempt < ROLLUP_INCREMENT_ATTEMPTS && rolled_back?(e)
+            if retryable && attempt < ROLLUP_INCREMENT_ATTEMPTS && rolled_back?(e)
               sleep(ROLLUP_INCREMENT_BASE_DELAY_SECONDS * (2**(attempt - 1)))
               retry
             end

data/lib/llm_cost_tracker/ledger/storable.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ledger
+    module Storable
+      def self.clean(value)
+        case value
+        when Hash then value.to_h { |key, nested| [clean(key), clean(nested)] }
+        when Array then value.map { |nested| clean(nested) }
+        when String then value.dup.force_encoding(Encoding::UTF_8).scrub.delete("\u0000")
+        else value
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -6,6 +6,7 @@ require_relative "../pricing"
 require_relative "rollups"
 require_relative "../budget/per_tag"
 require_relative "tags/encoding"
+require_relative "storable"
 
 module LlmCostTracker
   module Ledger
@@ -23,11 +24,13 @@ module LlmCostTracker
           events = Array(events)
           return if events.empty?
 
-          LlmCostTracker::Call.transaction do
-            rows = events.map { |event| attributes_for(event) }
-            call_ids = insert_calls_returning_ids(rows, events)
-            insert_line_items(events, call_ids)
-            insert_call_tags(events, call_ids)
+          Isolation.guard do
+            LlmCostTracker::Call.transaction do
+              rows = events.map { |event| attributes_for(event) }
+              call_ids = insert_calls_returning_ids(rows, events)
+              insert_line_items(events, call_ids)
+              insert_call_tags(events, call_ids)
+            end
           end
         end
 
@@ -62,9 +65,7 @@ module LlmCostTracker
             pricing_snapshot: event.pricing_snapshot
           }
 
-          attributes
-            .merge(event.token_usage.to_h)
-            .merge(total_cost: event.cost&.total)
+          Storable.clean(attributes.merge(event.token_usage.to_h).merge(total_cost: event.cost&.total))
         end
 
         def call_ids_for(events)
@@ -90,7 +91,7 @@ module LlmCostTracker
         end
 
         def line_item_attributes(call_id:, line_item:, position:)
-          {
+          Storable.clean(
             llm_cost_tracker_call_id: call_id,
             position: position,
             kind: line_item.kind,
@@ -112,7 +113,7 @@ module LlmCostTracker
             provider_item_id: line_item.provider_item_id,
             details: stored_details(line_item.details),
             created_at: Time.now.utc
-          }
+          )
         end
 
         def insert_call_tags(events, call_ids)
@@ -121,7 +122,7 @@ module LlmCostTracker
               {
                 llm_cost_tracker_call_id: call_ids.fetch(event.event_id),
                 key: key.to_s,
-                value: Tags::Encoding.encode(value)
+                value: Tags::Encoding.encode(Storable.clean(value))
               }.merge(budget_columns_for(event))
             end
           end

data/lib/llm_cost_tracker/ledger.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "ledger/isolation"
 require_relative "ledger/schema"
 require_relative "ledger/period"
 require_relative "ledger/rollups"

data/lib/llm_cost_tracker/logging.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "redaction"
+
 module LlmCostTracker
   module Logging
     class << self
@@ -9,7 +11,7 @@ module LlmCostTracker
 
       private
 
-      def prefixed(message) = "[#{LlmCostTracker.name}] #{message}"
+      def prefixed(message) = "[#{LlmCostTracker.name}] #{Redaction.text(message)}"
     end
   end
 end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,10 +2,10 @@
 
 require "faraday"
 require "json"
-require "stringio"
 require "uri"
 
 require_relative "../capture/sse"
+require_relative "../capture/stream_tap"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -27,16 +27,17 @@ module LlmCostTracker
         if streaming
           request_body = inject_stream_usage_flag(request_env, parser, request_url, request_parsed) || request_body
         end
-        stream_buffer = install_stream_tap(request_env) if streaming
+        stream_buffer = install_stream_tap(request_env, parser) if streaming
 
         if parser
+          context_tags, metadata = tag_snapshot(request_env)
           Budget.enforce!(
             provider: parser.provider_for(request_url),
             model: parser.model_for(request_url, request_parsed),
-            request: request_parsed
+            request: request_parsed,
+            tags: Tracker.build_tags(context_tags: context_tags, metadata: metadata)
           )
         end
-        context_tags, metadata = tag_snapshot(request_env) if parser
         started_at = LlmCostTracker::Timing.now_monotonic
 
         invoke_app_with_capture(
@@ -95,7 +96,7 @@ module LlmCostTracker
 
       def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
         return nil unless LlmCostTracker.configuration.capture.request_stream_usage
-        return nil unless parser&.auto_enable_stream_usage?(request_url)
+        return nil unless parser.auto_enable_stream_usage?(request_url, request_parsed)
 
         stream_options = request_parsed["stream_options"]
         return nil if stream_options.is_a?(Hash) && stream_options.key?("include_usage")
@@ -116,21 +117,22 @@ module LlmCostTracker
         request = parser.safe_json_parse(request_body)
         event = Event.build(
           provider: parser.provider_for(request_url),
-          model: request["model"] || Event::UNKNOWN_MODEL,
+          model: parser.model_for(request_url, request) || Event::UNKNOWN_MODEL,
           token_usage: Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: Usage::Source::UNKNOWN
         )
-        merged_metadata = (metadata || {}).merge(
-          stream_interrupted: true,
-          stream_interrupted_error: "#{error.class}: #{error.message}"
-        )
+        merged_metadata = (metadata || {}).merge(stream_interrupted: true, stream_interrupted_error: error.class.name)
+        status = error.try(:response_status)
+        merged_metadata[:stream_interrupted_status] = status if status
         Tracker.record(
           event: event,
           latency_ms: latency_ms,
           metadata: merged_metadata,
           context_tags: context_tags
         )
+      rescue LlmCostTracker::TransactionAbortedError
+        raise
       rescue StandardError => e
         Logging.warn("Error recording interrupted stream: #{e.class}: #{e.message}")
       end
@@ -171,7 +173,7 @@ module LlmCostTracker
           metadata: metadata,
           context_tags: context_tags
         )
-      rescue LlmCostTracker::Error
+      rescue *LlmCostTracker::CALLER_ERRORS
         raise
       rescue StandardError => e
         Logging.warn("Error processing response: #{e.class}: #{e.message}")
@@ -198,32 +200,28 @@ module LlmCostTracker
       end
 
       def parse_stream(parser:, request_url:, request_body:, response_env:, stream_buffer:)
-        overflowed = stream_buffer&.dig(:overflowed) == true
-        Logging.warn(capture_warning(request_url, stream_buffer)) if overflowed
-
-        body = stream_buffer&.dig(:buffer)&.string
-        body = read_body(response_env.body) if body.blank?
-
-        if body.blank?
-          Logging.warn(capture_warning(request_url, stream_buffer)) unless overflowed
-          return parser.parse_stream(
-            request_url: request_url,
-            request_body: request_body,
-            response_status: response_env.status,
-            response_headers: response_env.response_headers
-          )
-        end
-
-        events = overflowed ? [] : Capture::SSE.parse(body)
         parser.parse_stream(
           request_url: request_url,
           request_body: request_body,
           response_status: response_env.status,
-          events: events,
+          events: stream_events(request_url, response_env, stream_buffer),
           response_headers: response_env.response_headers
         )
       end
 
+      def stream_events(request_url, response_env, stream_buffer)
+        if stream_buffer&.received?
+          events = stream_buffer.events
+          return events unless stream_buffer.overflowed? || stream_buffer.failed?
+        else
+          body = read_body(response_env.body)
+          return Capture::SSE.parse(body) if body.present?
+        end
+
+        Logging.warn(capture_warning(request_url, stream_buffer))
+        []
+      end
+
       def forward_on_data_chunk(callable, chunk, size, env)
         arity = callable.arity
         return callable.call(chunk, size, env) if arity.negative?
@@ -235,28 +233,19 @@ module LlmCostTracker
         end
       end
 
-      def install_stream_tap(request_env)
+      def install_stream_tap(request_env, parser)
         request = request_env.request
         return nil unless request
 
         original = request.on_data
         return nil unless original
 
-        state = { buffer: StringIO.new, bytes: 0, overflowed: false }
+        tap = Capture::StreamTap.new(notable: parser.method(:retain_stream_event?))
         request.on_data = proc do |chunk, size, env|
-          chunk = chunk.to_s
-          remaining = Capture::SSE::LIMIT_BYTES - state[:bytes]
-          if chunk.bytesize <= remaining
-            state[:buffer] << chunk
-            state[:bytes] += chunk.bytesize
-          else
-            state[:buffer] << chunk.byteslice(0, remaining) if remaining.positive?
-            state[:bytes] += [remaining, 0].max
-            state[:overflowed] = true
-          end
+          tap << chunk
           forward_on_data_chunk(original, chunk, size, env)
         end
-        state
+        tap
       rescue StandardError => e
         Logging.warn("Unable to install streaming tap: #{e.class}: #{e.message}")
         nil
@@ -295,7 +284,7 @@ module LlmCostTracker
         suffix = "recording usage_source=#{Usage::Source::UNKNOWN}. " \
                  "Use LlmCostTracker.track_stream for manual capture."
         label = request_url_label(request_url)
-        return "Unable to capture streaming response for #{label}; #{suffix}" unless stream_buffer&.dig(:overflowed)
+        return "Unable to capture streaming response for #{label}; #{suffix}" unless stream_buffer&.overflowed?
 
         "Streaming response for #{label} exceeded #{Capture::SSE::LIMIT_BYTES} bytes; #{suffix}"
       end

data/lib/llm_cost_tracker/parsers.rb

@@ -104,7 +104,11 @@ module LlmCostTracker
         nil
       end
 
-      def auto_enable_stream_usage?(_request_url)
+      def retain_stream_event?(_data)
+        false
+      end
+
+      def auto_enable_stream_usage?(_request_url, _request_parsed)
         false
       end