llm_cost_tracker 0.13.0 → 0.14.1

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -128,10 +128,16 @@ module LlmCostTracker
         def each(&block)
           return enum_for(:each) unless block
 
+          deferred = nil
           @raw_stream.each do |response|
-            LlmCostTracker::Integrations::Anthropic.record_batch_result(response)
+            begin
+              LlmCostTracker::Integrations::Anthropic.record_batch_result(response)
+            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
+              deferred ||= e
+            end
             block.call(response)
           end
+          raise deferred if deferred
         end
 
         def respond_to_missing?(name, include_private = false)

data/lib/llm_cost_tracker/integrations/base.rb

@@ -27,6 +27,7 @@ module LlmCostTracker
 
       def install
         validate_contract!
+        Logging.warn(untested_version_message) if untested_version?
         patch_targets.each do |target|
           target_class = target.fetch(:constant_name).to_s.safe_constantize
           install_patch(target_class, target.fetch(:patch)) if target_class
@@ -39,6 +40,7 @@ module LlmCostTracker
         if problems.any?
           return Check.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
         end
+        return Check.new(:warn, name, untested_version_message) if untested_version?
 
         installed = patch_targets.reject { |target| target.fetch(:optional) }.all? do |target|
           target.fetch(:constant_name).to_s.safe_constantize&.ancestors&.include?(target.fetch(:patch))
@@ -60,7 +62,7 @@ module LlmCostTracker
 
       def record_safely
         yield
-      rescue LlmCostTracker::Error
+      rescue *LlmCostTracker::CALLER_ERRORS
         raise
       rescue StandardError => e
         Logging.warn("#{integration_name} integration failed to record usage: #{e.class}: #{e.message}")
@@ -124,6 +126,11 @@ module LlmCostTracker
         @minimum_version
       end
 
+      def maximum_version(value = nil)
+        @maximum_version = value if value
+        @maximum_version
+      end
+
       def gem_version
         Gem.loaded_specs[integration_name.to_s]&.version
       end
@@ -160,6 +167,15 @@ module LlmCostTracker
         ["#{name} >= #{minimum_version} is required, detected #{version}"]
       end
 
+      def untested_version?
+        maximum_version && gem_version && gem_version >= Gem::Version.new(maximum_version)
+      end
+
+      def untested_version_message
+        "#{integration_name} #{gem_version} is newer than the tested range (< #{maximum_version}); " \
+          "its calls may not be recorded"
+      end
+
       def target_problems
         patch_targets.flat_map do |target|
           constant_name = target.fetch(:constant_name)

data/lib/llm_cost_tracker/integrations/openai/batch_capture.rb

@@ -13,36 +13,36 @@ module LlmCostTracker
         class << self
           def maybe_capture(batch, resource:)
             return unless Openai.active?
-            return unless batch.respond_to?(:status) && batch.status.to_s == "completed"
-
-            output_file_id = batch.respond_to?(:output_file_id) ? batch.output_file_id : nil
-            return unless output_file_id
-
-            batch_id = batch.respond_to?(:id) ? batch.id : nil
-            return unless batch_id && claim(batch_id)
+            return unless batch.status.to_s == "completed"
+            return unless batch.output_file_id && batch.id
+            return if captured?(batch.id)
 
             client = resource.instance_variable_get(:@client)
             host = Openai.client_host_for(resource)
             Openai.record_safely do
-              io = client.files.content(output_file_id)
-              capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+              io = client.files.content(batch.output_file_id)
+              deferred = capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+              mark_captured(batch.id)
+              raise deferred if deferred
             end
           end
 
           private
 
-          def claim(batch_id)
+          def captured?(batch_id)
+            MUTEX.synchronize { @dedup&.include?(batch_id) || false }
+          end
+
+          def mark_captured(batch_id)
             MUTEX.synchronize do
               @dedup ||= Set.new
-              next false if @dedup.include?(batch_id)
-
               @dedup.clear if @dedup.size >= DEDUP_LIMIT
               @dedup.add(batch_id)
-              true
             end
           end
 
           def capture_jsonl(jsonl, host:)
+            deferred = nil
             jsonl.each_line do |line|
               line = line.strip
               next if line.empty?
@@ -54,7 +54,10 @@ module LlmCostTracker
               next unless response.is_a?(Hash) && response["usage"]
 
               record_result(response, host: host)
+            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
+              deferred ||= e
             end
+            deferred
           end
 
           def parse_line(line)

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
@@ -45,7 +46,7 @@ module LlmCostTracker
 
         def recorded_sql(period, start)
           calls = "COALESCE(#{sum_sql(LlmCostTracker::Call.between(start, time))}, 0)"
-          return calls unless LlmCostTracker.configuration.cache_rollups
+          return calls unless Rollups.cache_active?
 
           rollup = "COALESCE(#{sum_sql(rollup_scope(period))}, 0)"
           "GREATEST(#{rollup}, #{calls})"

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -2,13 +2,24 @@
 
 require "bigdecimal/util"
 
+require_relative "isolation"
 require_relative "period"
 
 module LlmCostTracker
   module Ledger
     module Rollups
       class << self
+        def cache_active?
+          return false unless LlmCostTracker.configuration.budgets.totals_source == :cache
+          return true if LlmCostTracker::CallRollup.table_exists?
+
+          warn_missing_table
+          false
+        end
+
         def increment!(events)
+          return unless cache_active?
+
           events = Array(events).select(&:total_cost)
           return if events.empty?
 
@@ -18,29 +29,45 @@ module LlmCostTracker
         ROLLUP_INCREMENT_ATTEMPTS = 3
         ROLLUP_INCREMENT_BASE_DELAY_SECONDS = 0.05
         REBUILD_INSERT_SLICE = 1_000
-        private_constant :ROLLUP_INCREMENT_ATTEMPTS, :ROLLUP_INCREMENT_BASE_DELAY_SECONDS, :REBUILD_INSERT_SLICE
+        private_constant :ROLLUP_INCREMENT_ATTEMPTS,
+                         :ROLLUP_INCREMENT_BASE_DELAY_SECONDS,
+                         :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
+            if retryable && attempt < ROLLUP_INCREMENT_ATTEMPTS && rolled_back?(e)
               sleep(ROLLUP_INCREMENT_BASE_DELAY_SECONDS * (2**(attempt - 1)))
               retry
             end
 
             LlmCostTracker::Logging.warn(
-              "Rollup increment failed for #{events.size} events after #{attempt} attempts: " \
-              "#{e.class}: #{e.message}"
+              "Rollup increment failed for #{events.size} events after #{attempt} attempt(s): " \
+              "#{e.class}: #{e.message}. Budget reads fall back to the calls ledger; " \
+              "run bin/rails llm_cost_tracker:rebuild_rollups to resync the cache."
             )
           end
         end
 
+        def rolled_back?(error)
+          error.is_a?(ActiveRecord::Deadlocked) || error.is_a?(ActiveRecord::LockWaitTimeout)
+        end
+
         def decrement!(records)
+          return unless cache_active?
+
           buckets = bucket_totals(records)
           return if buckets.empty?
 
@@ -60,6 +87,17 @@ module LlmCostTracker
 
         private
 
+        def warn_missing_table
+          return if @missing_table_warned
+
+          @missing_table_warned = true
+          Logging.warn(
+            "config.budgets.totals_source = :cache is set but llm_cost_tracker_call_rollups is missing; " \
+            "budget and dashboard totals fall back to aggregating llm_cost_tracker_calls. " \
+            "Run the call_rollups generator and migrate, or set config.budgets.totals_source = :ledger."
+          )
+        end
+
         def accumulate_priced_calls
           accumulated = Hash.new { |totals, key| totals[key] = BigDecimal("0") }
           priced_calls.in_batches do |batch|

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

@@ -4,7 +4,9 @@ require "json"
 
 require_relative "../pricing"
 require_relative "rollups"
+require_relative "../budget/per_tag"
 require_relative "tags/encoding"
+require_relative "storable"
 
 module LlmCostTracker
   module Ledger
@@ -15,18 +17,20 @@ module LlmCostTracker
           return if events.empty?
 
           persist_records(events)
-          Ledger::Rollups.increment_safely!(events) if LlmCostTracker.configuration.cache_rollups
+          Ledger::Rollups.increment_safely!(events)
         end
 
         def persist_records(events)
           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
 
@@ -61,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)
@@ -89,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,
@@ -111,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)
@@ -120,8 +122,8 @@ 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
           return if rows.empty?
@@ -129,6 +131,12 @@ module LlmCostTracker
           LlmCostTracker::CallTag.insert_all!(rows, record_timestamps: false, returning: false)
         end
 
+        def budget_columns_for(event)
+          return {} unless LlmCostTracker::Budget::PerTag.columns?
+
+          { total_cost: event.cost&.total, tracked_at: event.tracked_at }
+        end
+
         def stored_details(details)
           (details || {}).transform_keys(&:to_s).transform_values { |value| Tags::Encoding.normalize_value(value) }
         end

data/lib/llm_cost_tracker/ledger/tags/encoding.rb

@@ -6,12 +6,22 @@ module LlmCostTracker
   module Ledger
     module Tags
       module Encoding
+        INDEXABLE_BYTES = 2_048
+
         def self.encode(value)
-          case value
-          when Hash then JSON.generate(normalize_hash(value))
-          when Array then JSON.generate(normalize_array(value))
-          else value.to_s
-          end
+          encoded = case value
+                    when Hash then JSON.generate(normalize_hash(value))
+                    when Array then JSON.generate(normalize_array(value))
+                    else value.to_s
+                    end
+          truncate(encoded)
+        end
+
+        def self.truncate(string)
+          limit = [LlmCostTracker.configuration.tags.max_value_bytesize.to_i, INDEXABLE_BYTES].min
+          return string if string.bytesize <= limit
+
+          string.byteslice(0, limit).encode("UTF-8", invalid: :replace, undef: :replace)
         end
 
         def self.normalize_hash(hash)

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,17 +1,17 @@
 # frozen_string_literal: true
 
+require_relative "redaction"
+
 module LlmCostTracker
   module Logging
     class << self
-      def debug(message) = tagged.debug(message)
+      def debug(message) = Rails.logger.debug(prefixed(message))
 
-      def warn(message) = tagged.warn(message)
+      def warn(message) = Rails.logger.warn(prefixed(message))
 
       private
 
-      def tagged
-        Rails.logger.tagged(LlmCostTracker.name)
-      end
+      def prefixed(message) = "[#{LlmCostTracker.name}] #{Redaction.text(message)}"
     end
   end
 end