llm_cost_tracker 0.14.0 → 0.14.2

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_image_tokens_generator.rb

@@ -18,12 +18,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_image_tokens.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_indexes_generator.rb

@@ -10,8 +10,9 @@ module LlmCostTracker
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Adds a partial index over unpriced calls and drops the unused ingestion inbox lock index" \
-           "scanning the whole ledger."
+      desc "Adds the unpriced-calls index (partial on PostgreSQL, so backfill_unknown_pricing stops scanning " \
+           "the whole ledger; a plain id index on MySQL, which has no partial indexes) and drops the unused " \
+           "ingestion inbox lock index."
 
       def create_migration_file
         migration_template(
@@ -19,12 +20,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_indexes.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_per_tag_budgets_generator.rb

@@ -10,8 +10,9 @@ module LlmCostTracker
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Carries each call's cost and time onto its tag rows so per-tag budgets read " \
-           "llm_cost_tracker_call_tags without a join. Required when config.budgets.per_tag is set."
+      desc "Adds cost and time columns to llm_cost_tracker_call_tags so per-tag budgets read it " \
+           "without a join; run llm_cost_tracker:backfill_tag_costs afterwards so earlier calls count. " \
+           "Required for config.budgets.per_tag on installs created before v0.14."
 
       def create_migration_file
         migration_template(
@@ -19,12 +20,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_per_tag_budgets.rb"
         )
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -13,7 +13,8 @@ module LlmCostTracker
         ActiveRecord::Deadlocked,
         ActiveRecord::LockWaitTimeout,
         ActiveRecord::StatementTimeout,
-        ActiveRecord::ConnectionNotEstablished
+        ActiveRecord::ConnectionNotEstablished,
+        LlmCostTracker::TransactionAbortedError
       ].freeze
 
       def initialize(identity:)
@@ -25,7 +26,7 @@ module LlmCostTracker
         return 0 if rows.empty?
 
         valid_rows, events = decode(rows)
-        persist(valid_rows, events) if events.any?
+        persist_batch(valid_rows, events) if events.any?
         rows.size
       rescue StandardError => e
         rows_to_mark = valid_rows&.any? ? valid_rows : rows
@@ -67,7 +68,7 @@ module LlmCostTracker
       end
 
       def error_message_for(error)
-        "#{error.class}: #{error.message}".byteslice(0, 1_000)
+        "#{error.class}: #{Redaction.text(error.message)}".byteslice(0, 1_000).scrub("")
       end
 
       def warn_on_quarantine(rows)
@@ -75,12 +76,10 @@ module LlmCostTracker
         quarantined = rows.select { |row| row.attempts.to_i + 1 >= threshold }
         return if quarantined.empty?
 
-        sample = quarantined.first(10).map(&:id).join(", ")
-        sample += "..." if quarantined.size > 10
         LlmCostTracker::Logging.warn(
           "Ingestion::Batch: #{quarantined.size} inbox row(s) reached " \
           "MAX_ATTEMPTS_BEFORE_QUARANTINE=#{threshold} and will be skipped " \
-          "on the next claim cycle (ids: #{sample})"
+          "on the next claim cycle (ids: #{id_sample(quarantined)})"
         )
       end
 
@@ -88,6 +87,11 @@ module LlmCostTracker
 
       attr_reader :identity
 
+      def id_sample(rows)
+        sample = rows.first(10).map(&:id).join(", ")
+        rows.size > 10 ? "#{sample}..." : sample
+      end
+
       def claim
         now = Time.now.utc
         cutoff = now - LOCK_TIMEOUT_SECONDS
@@ -117,13 +121,38 @@ module LlmCostTracker
         [valid_rows, events]
       end
 
+      def persist_batch(rows, events)
+        landed = []
+        failed = Hash.new { |hash, message| hash[message] = [] }
+        landed.concat(persist(rows, events))
+      rescue *TRANSIENT_PERSIST_ERRORS
+        raise
+      rescue StandardError
+        rows.zip(events) do |row, event|
+          landed.concat(persist([row], [event]))
+        rescue *TRANSIENT_PERSIST_ERRORS
+          raise
+        rescue StandardError => e
+          failed[error_message_for(e)] << row
+        end
+      ensure
+        failed.each do |message, failed_rows|
+          LlmCostTracker::Logging.warn(
+            "Ingestion::Batch: #{failed_rows.size} inbox row(s) could not be stored " \
+            "(ids: #{id_sample(failed_rows)}): #{message}"
+          )
+          mark_failed_with_message(failed_rows, message)
+        end
+        Ledger::Rollups.increment_safely!(landed)
+        Budget.notify_persisted_safely!(landed)
+      end
+
       def persist(rows, events, retry_on_conflict: true)
         LlmCostTracker::Call.transaction do
           Ledger::Store.persist_records(events)
           Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
-        Ledger::Rollups.increment_safely!(events)
-        Budget.notify_persisted_safely!(events)
+        events
       rescue ActiveRecord::RecordNotUnique
         raise unless retry_on_conflict
 

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -18,7 +18,7 @@ module LlmCostTracker
         end
 
         def event_from_row(row)
-          payload = JSON.parse(row.payload, symbolize_names: true)
+          payload = JSON.parse(row.payload, symbolize_names: true, allow_duplicate_key: true)
           schema_version = payload[:schema_version]
           unless schema_version == PAYLOAD_SCHEMA_VERSION
             raise LlmCostTracker::Error, "unsupported ledger inbox payload schema version #{schema_version.inspect}"
@@ -69,13 +69,7 @@ module LlmCostTracker
         end
 
         def payload_for(event)
-          event.to_h.merge(
-            schema_version: PAYLOAD_SCHEMA_VERSION,
-            event_id: event.event_id,
-            provider: event.provider,
-            model: event.model,
-            tracked_at: event.tracked_at.iso8601(6)
-          )
+          event.to_h.merge(schema_version: PAYLOAD_SCHEMA_VERSION, tracked_at: event.tracked_at.iso8601(6))
         end
 
         def insert_row(row)

data/lib/llm_cost_tracker/ingestion.rb

@@ -28,6 +28,8 @@ module LlmCostTracker
         end
 
         guards_for_current_config.each do |schema_module, table_name|
+          next if schema_module == Ledger::Schema::CallRollups && !Ledger::Rollups.cache_active?
+
           errors = schema_module.current_schema_errors
           next if errors.empty?
 
@@ -124,7 +126,7 @@ module LlmCostTracker
 
       def cleanup_verification_call(response_id)
         relation = LlmCostTracker::Call.where(provider_response_id: response_id)
-        records = relation.select(:id, :tracked_at, :total_cost, :pricing_snapshot, :provider).to_a
+        records = relation.select(:id, *LlmCostTracker::Ledger::Rollups::SOURCE_COLUMNS).to_a
         return if records.empty?
 
         relation.delete_all

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../providers/anthropic/usage_extractor"
 require_relative "../providers/anthropic/response_parser"
 
 module LlmCostTracker
@@ -17,7 +16,8 @@ module LlmCostTracker
             patch_target("Anthropic::Resources::Messages", with: MessagesPatch),
             patch_target("Anthropic::Resources::Beta::Messages", with: MessagesPatch, optional: true),
             patch_target("Anthropic::Resources::Messages::Batches", with: BatchesPatch, optional: true),
-            patch_target("Anthropic::Resources::Beta::Messages::Batches", with: BatchesPatch, optional: true)
+            patch_target("Anthropic::Resources::Beta::Messages::Batches", with: BatchesPatch, optional: true),
+            patch_target("Anthropic::BetaRefusalFallbackMiddleware", with: FallbackMiddlewarePatch, optional: true)
           ]
         end
 
@@ -37,7 +37,8 @@ module LlmCostTracker
                 model: message.model || request[:model],
                 provider_response_id: message.id,
                 usage_source: Usage::Source::SDK_RESPONSE,
-                request: request
+                request: request,
+                **response_fields(message)
               ),
               latency_ms: latency_ms
             )
@@ -61,21 +62,52 @@ module LlmCostTracker
             next unless usage
             next if usage.input_tokens.nil? && usage.output_tokens.nil?
 
-            usage_hash = usage.deep_to_h
-            LlmCostTracker::Tracker.record(
-              event: Providers::Anthropic::ResponseParser.event_from_usage(
-                usage: usage_hash,
+            record_once(
+              Providers::Anthropic::ResponseParser.event_from_usage(
+                usage: usage.deep_to_h.merge(service_tier: "batch"),
                 model: message.model,
                 provider_response_id: message.id,
                 usage_source: Usage::Source::SDK_BATCH_RESULT,
-                pricing_mode: "batch"
+                **response_fields(message)
               )
             )
           end
         end
 
-        def stream_pricing_mode(request)
-          Providers::Anthropic::UsageExtractor.pricing_mode(request: request || {}, usage: nil)
+        def record_refused_hop(request, response)
+          record_safely { record_message(response.parse, request: request.body, latency_ms: nil) }
+        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
+          nil
+        end
+
+        def record_refused_stream_hop(hop)
+          return unless active?
+
+          usage = hop.dig(:refused, :usage).to_h.deep_symbolize_keys
+          output = usage[:output_tokens].to_i.positive?
+          usage = usage.slice(:server_tool_use) if output
+          return if output && Providers::Anthropic::UsageExtractor.service_line_items(usage).empty?
+
+          record_safely do
+            LlmCostTracker::Tracker.record(
+              event: Providers::Anthropic::ResponseParser.event_from_usage(
+                usage: usage,
+                model: hop[:model],
+                provider_response_id: nil,
+                usage_source: Usage::Source::STREAM_FINAL,
+                stream: true,
+                stop_reason: ("refusal" unless output),
+                refusal_category: hop.dig(:refused, :stop_details, "category")
+              )
+            )
+          end
+        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
+          nil
+        end
+
+        def response_fields(message)
+          { stop_reason: message.stop_reason, refusal_category: message.stop_details&.category,
+            content: message.deep_to_h[:content] }
         end
       end
 
@@ -109,6 +141,27 @@ module LlmCostTracker
         end
       end
 
+      module FallbackMiddlewarePatch
+        def call(req, nxt)
+          return super if req.streaming? || !LlmCostTracker::Integrations::Anthropic.active?
+
+          hops = []
+          response = super(req, ->(hop_req) { nxt.call(hop_req).tap { |hop| hops << [hop_req, hop] } })
+          answered = hops.select { |_hop_req, hop| hop.status < 300 }
+          answered[...-1].each { |hop| LlmCostTracker::Integrations::Anthropic.record_refused_hop(*hop) }
+          response
+        end
+
+        private
+
+        def consume_hop(*args, **kwargs)
+          refused = Thread.current[:llm_cost_tracker_refused_hop]
+          Thread.current[:llm_cost_tracker_refused_hop] = nil
+          LlmCostTracker::Integrations::Anthropic.record_refused_stream_hop(refused) if refused && kwargs[:splice]
+          super.tap { |hop| Thread.current[:llm_cost_tracker_refused_hop] = hop if hop[:refused] }
+        end
+      end
+
       module BatchesPatch
         def results_streaming(*args, **kwargs)
           raw = super
@@ -128,10 +181,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

@@ -16,20 +16,19 @@ module LlmCostTracker
         @integration_name ||= name.demodulize.underscore.to_sym
       end
 
-      def provider(value = nil)
-        @provider = value.to_s if value
-        @provider ||= integration_name.to_s
-      end
+      def provider = integration_name.to_s
 
       def active?
-        LlmCostTracker.configuration.instrumented?(integration_name)
+        LlmCostTracker.configuration.enabled && LlmCostTracker.configuration.instrumented?(integration_name)
       end
 
       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
+          patch = target.fetch(:patch)
+          target_class.prepend(patch) if target_class && !target_class.ancestors.include?(patch)
         end
       end
 
@@ -39,6 +38,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,12 +60,46 @@ 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}")
       end
 
+      def record_passthrough(provider:,
+                             model:,
+                             response:,
+                             latency_ms:,
+                             service_line_items: [],
+                             usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE,
+                             pricing_mode: nil,
+                             **token_attributes)
+        return unless active?
+
+        record_safely do
+          LlmCostTracker::Tracker.record(
+            event: LlmCostTracker::Event.build(
+              provider: provider,
+              model: model,
+              token_usage: LlmCostTracker::Usage::TokenUsage.build(**token_attributes),
+              usage_source: usage_source,
+              pricing_mode: pricing_mode,
+              provider_response_id: provider_response_id_for(response),
+              service_line_items: service_line_items
+            ),
+            latency_ms: latency_ms
+          )
+        end
+      end
+
+      def record_once(event)
+        LlmCostTracker::Tracker.record(event: event.keyed_by_response_id)
+      rescue ActiveRecord::RecordNotUnique
+        nil
+      end
+
+      def provider_response_id_for(response) = response&.try(:id)
+
       def request_params(args, kwargs)
         params =
           case args.first
@@ -106,24 +140,24 @@ module LlmCostTracker
         ).wrap
       end
 
-      def stream_collector(request, provider: self.provider)
+      def stream_collector(request)
         LlmCostTracker::Capture::StreamCollector.new(
           provider: provider,
           model: request[:model],
-          pricing_mode: stream_pricing_mode(request),
           request: request
         )
       end
 
-      def stream_pricing_mode(_request)
-        nil
-      end
-
       def minimum_version(value = nil)
         @minimum_version = value if value
         @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 +194,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)
@@ -180,12 +223,6 @@ module LlmCostTracker
           "#{target.fetch(:constant_name)}##{method_name} is not available"
         end
       end
-
-      def install_patch(target, patch)
-        return if target&.ancestors&.include?(patch)
-
-        target.prepend(patch)
-      end
     end
   end
 end

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

@@ -13,16 +13,17 @@ module LlmCostTracker
         class << self
           def maybe_capture(batch, resource:)
             return unless Openai.active?
-            return unless batch.status.to_s == "completed"
+            return unless %w[completed expired cancelled].include?(batch.status.to_s)
             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(batch.output_file_id)
-              capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+              jsonl = client.files.content(batch.output_file_id).read
+              deferred = capture_jsonl(jsonl, host: host, model: batch.model)
               mark_captured(batch.id)
+              raise deferred if deferred
             end
           end
 
@@ -40,19 +41,20 @@ module LlmCostTracker
             end
           end
 
-          def capture_jsonl(jsonl, host:)
+          def capture_jsonl(jsonl, host:, model:)
+            deferred = nil
             jsonl.each_line do |line|
-              line = line.strip
-              next if line.empty?
-
               entry = parse_line(line)
               next unless entry
 
               response = entry.dig("response", "body")
               next unless response.is_a?(Hash) && response["usage"]
 
-              record_result(response, host: host)
+              record_result({ "id" => entry["id"] }.merge(response), host: host, model: model)
+            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
+              deferred ||= e
             end
+            deferred
           end
 
           def parse_line(line)
@@ -61,19 +63,25 @@ module LlmCostTracker
             nil
           end
 
-          def record_result(response, host:)
+          def record_result(response, host:, model:)
             provider = Openai.provider_for_host(host)
             return if LlmCostTracker::Call.already_recorded?(provider: provider, provider_response_id: response["id"])
 
-            event = LlmCostTracker::Providers::Openai::ResponseParser.event_from_response(
+            parser = LlmCostTracker::Providers::Openai::ResponseParser
+            event = parser.event_from_response(
               response: response,
-              request: {},
+              request: { "model" => model },
               provider: provider,
               host: host,
               usage_source: LlmCostTracker::Usage::Source::SDK_BATCH_RESULT,
-              pricing_mode: "batch"
+              # /v1/batches runs regional processing only on the us and eu hosts.
+              pricing_mode: parser.combined_pricing_mode(
+                host: (host if host.to_s.match?(/\A(?:us|eu)\./i)),
+                model: response["model"] || model,
+                service_tier: "batch"
+              )
             )
-            LlmCostTracker::Tracker.record(event: event) if event
+            Openai.record_once(event)
           end
         end
       end

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

@@ -45,6 +45,14 @@ module LlmCostTracker
         include PatchBuilder.build(record_method: :record_response, methods: %i[create])
         include PatchBuilder.build_stream(methods: %i[stream stream_raw])
 
+        def retrieve(response_id, *args, **kwargs)
+          response = super
+          LlmCostTracker::Integrations::Openai.record_retrieved_response(
+            response, host: LlmCostTracker::Integrations::Openai.client_host_for(self)
+          )
+          response
+        end
+
         def retrieve_streaming(response_id, *args, **kwargs)
           LlmCostTracker::Integrations::Openai.wrap_stream(
             args, kwargs, **LlmCostTracker::Integrations::Openai.stream_seam(self)