llm_cost_tracker 0.14.0 → 0.14.2

data/lib/llm_cost_tracker/capture/stream_tap.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative "event_window"
+require_relative "sse"
+
+module LlmCostTracker
+  module Capture
+    class StreamTap
+      MAX_PENDING_BYTES = 16 * 1_048_576
+
+      def initialize(notable: nil)
+        @window = EventWindow.new(notable: notable)
+        @reader = SSE::Reader.new { |event| @window.push(event[:data], type: event[:event]) }
+        @received = false
+        @failed = false
+      end
+
+      def <<(chunk)
+        return if @failed
+
+        @received = true
+        @reader << chunk
+        fail! if @reader.pending_bytesize > MAX_PENDING_BYTES
+      rescue StandardError
+        fail!
+      end
+
+      def received?
+        @received
+      end
+
+      def failed?
+        @failed
+      end
+
+      def overflowed?
+        @window.overflowed?
+      end
+
+      def events
+        @reader&.finish
+        @failed ? [] : @window.events
+      rescue StandardError
+        fail!
+        []
+      end
+
+      private
+
+      def fail!
+        @failed = true
+        @reader = nil
+        @window = EventWindow.new
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/capture/stream_tracker.rb

@@ -1,10 +1,7 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/object/deep_dup"
 require "active_support/core_ext/object/try"
 
-require_relative "sdk_payload"
-
 module LlmCostTracker
   module Capture
     class StreamTracker
@@ -13,8 +10,7 @@ module LlmCostTracker
         @collector = collector
         @active = active
         @finish = finish || proc { |errored| collector.finish!(errored: errored) }
-        @finished_ref = [false]
-        @attempted_ref = [false]
+        @finished = false
         @capture_failed = false
         @mutex = Mutex.new
       end
@@ -22,30 +18,21 @@ module LlmCostTracker
       def wrap
         return @stream unless @stream
 
-        iterator_wrapped = false
-        if @stream.instance_variable_defined?(:@iterator)
-          iterator = @stream.instance_variable_get(:@iterator)
-          if iterator.respond_to?(:each)
-            @stream.instance_variable_set(:@iterator,
-                                          Enumerator.new do |yielder|
-                                            each_from(iterator) { |event| yielder << event }
-                                          end)
-            iterator_wrapped = true
-          end
-        end
-        each_wrapped = false
-        if !iterator_wrapped && @stream.respond_to?(:each)
+        iterator = @stream.instance_variable_get(:@iterator)
+        if iterator.respond_to?(:each)
+          @stream.instance_variable_set(:@iterator,
+                                        Enumerator.new do |yielder|
+                                          each_from(iterator) { |event| yielder << event }
+                                        end)
+        elsif @stream.respond_to?(:each)
           wrap_each
-          each_wrapped = true
-        end
-        unless iterator_wrapped || each_wrapped
+        else
           Logging.warn(
             "stream integration found no wrappable iterator on #{@stream.class} " \
             "(missing both `@iterator` ivar and `#each`); usage will not be captured"
           )
         end
 
-        register_orphan_finalizer
         @stream
       rescue StandardError => e
         Logging.warn("stream integration failed to install wrapper: #{e.class}: #{e.message}")
@@ -77,7 +64,7 @@ module LlmCostTracker
             yield event
           end
         end
-      rescue StandardError
+      rescue Exception # rubocop:disable Lint/RescueException -- mark Interrupt/shutdown too, so it is not masked
         errored = true
         raise
       ensure
@@ -85,10 +72,7 @@ module LlmCostTracker
       end
 
       def capture(event)
-        raw_payload = event.try(:deep_to_h) || event.try(:to_h) || {}
-        payload = SdkPayload.normalize(raw_payload)
-        type = event.try(:type) || payload["type"]
-        @collector.event(payload, type: type&.to_s)
+        @collector.event(event.try(:to_h) || {}, type: event.try(:type)&.to_s)
       rescue StandardError => e
         warn_capture_failure(e)
       end
@@ -107,10 +91,9 @@ module LlmCostTracker
 
       def finish!(errored:)
         should_finish = @mutex.synchronize do
-          @attempted_ref[0] = true
-          next false if @finished_ref[0]
+          next false if @finished
 
-          @finished_ref[0] = true
+          @finished = true
           true
         end
         return unless should_finish && @active.call
@@ -118,37 +101,10 @@ module LlmCostTracker
         begin
           @finish.call(errored)
         rescue StandardError
-          @mutex.synchronize { @finished_ref[0] = false }
+          @mutex.synchronize { @finished = false }
           raise
         end
       end
-
-      def register_orphan_finalizer
-        finished_ref = @finished_ref
-        attempted_ref = @attempted_ref
-        finish_proc = @finish
-        active_proc = @active
-        mutex = @mutex
-        finalizer = lambda do |_object_id|
-          should_finish = mutex.synchronize do
-            next false if finished_ref[0] || attempted_ref[0]
-
-            finished_ref[0] = true
-            attempted_ref[0] = true
-            true
-          end
-          next unless should_finish && active_proc.call
-
-          begin
-            Rails.application.executor.wrap { finish_proc.call(false) }
-          rescue StandardError
-            nil
-          end
-        end
-        ObjectSpace.define_finalizer(@stream, finalizer)
-      rescue TypeError, ArgumentError
-        nil
-      end
     end
   end
 end

data/lib/llm_cost_tracker/capture_verifier.rb

@@ -25,7 +25,7 @@ module LlmCostTracker
       [
         enabled_check,
         *integration_checks,
-        *storage_checks
+        *LlmCostTracker::Ingestion.verify
       ].compact
     end
 
@@ -49,11 +49,5 @@ module LlmCostTracker
         check.with(name: "sdk integration #{check.name}")
       end
     end
-
-    def storage_checks
-      LlmCostTracker::Ingestion.verify
-    rescue LlmCostTracker::Error => e
-      [Check.new(:error, "storage", e.message)]
-    end
   end
 end

data/lib/llm_cost_tracker/charges/line_item.rb

@@ -58,7 +58,11 @@ module LlmCostTracker
       def self.from_token_usage(token_usage)
         return [] unless token_usage
 
-        token_usage.priced_quantities.filter_map do |key, quantity|
+        from_quantities(token_usage.priced_quantities)
+      end
+
+      def self.from_quantities(quantities)
+        quantities.filter_map do |key, quantity|
           next unless quantity.positive?
 
           dimension = Usage::Catalog.fetch(key)

data/lib/llm_cost_tracker/configuration/budgets.rb

@@ -35,7 +35,7 @@ module LlmCostTracker
       end
 
       def finalize!
-        @per_tag = deep_freeze(@per_tag || {})
+        @per_tag = deep_freeze(@per_tag)
       end
 
       private

data/lib/llm_cost_tracker/configuration/pricing.rb

@@ -29,7 +29,7 @@ module LlmCostTracker
       end
 
       def finalize!
-        @overrides = deep_freeze(@overrides || {})
+        @overrides = deep_freeze(@overrides)
       end
     end
   end

data/lib/llm_cost_tracker/configuration.rb

@@ -18,8 +18,6 @@ module LlmCostTracker
       budgets: Budgets, capture: Capture, ingestion: Ingestion, pricing: Pricing, tags: Tags
     }.freeze
 
-    SCALAR_ATTRIBUTES = %i[enabled].freeze
-
     DEPRECATED_OPTIONS = {
       monthly_budget: { to: %i[budgets monthly] },
       daily_budget: { to: %i[budgets daily] },
@@ -46,7 +44,7 @@ module LlmCostTracker
       log_level: { to: nil, note: "LlmCostTracker logs through Rails.logger, which owns the level" }
     }.freeze
 
-    attr_reader(*SCALAR_ATTRIBUTES, *SECTIONS.keys, :instrumented_integrations)
+    attr_reader(:enabled, *SECTIONS.keys, :instrumented_integrations)
 
     def initialize
       SECTIONS.each { |name, klass| instance_variable_set(:"@#{name}", klass.new(self)) }
@@ -67,11 +65,9 @@ module LlmCostTracker
       @instrumented_integrations.include?(name)
     end
 
-    SCALAR_ATTRIBUTES.each do |name|
-      define_method(:"#{name}=") do |value|
-        ensure_mutable!
-        instance_variable_set(:"@#{name}", value)
-      end
+    def enabled=(value)
+      ensure_mutable!
+      @enabled = value
     end
 
     DEPRECATED_OPTIONS.each do |old_name, spec|
@@ -101,7 +97,7 @@ module LlmCostTracker
 
     def finalize!
       SECTIONS.each_key { |name| public_send(name).finalize! }
-      @instrumented_integrations = deep_freeze(@instrumented_integrations || Set.new)
+      @instrumented_integrations = deep_freeze(@instrumented_integrations)
       @finalized = true
     end
 

data/lib/llm_cost_tracker/doctor/price_check.rb

@@ -13,14 +13,15 @@ module LlmCostTracker
       def call
         path = LlmCostTracker.configuration.pricing.file
         return bundled_check unless path
+        return Check.new(:error, "prices", "#{path} does not exist; #{REFRESH_COMMAND}") unless File.exist?(path)
 
         count = LlmCostTracker::Pricing::Registry.file_prices(path).size
         metadata = LlmCostTracker::Pricing::Registry.file_metadata(path)
         updated_at = metadata["updated_at"] || metadata[:updated_at]
         return configured_check(:warn, path, count, "metadata.updated_at missing; #{REFRESH_COMMAND}") unless updated_at
 
-        age_days = (Date.today - Date.iso8601(updated_at.to_s)).to_i
-        if age_days > STALE_AFTER_DAYS
+        file_date = Date.iso8601(updated_at.to_s)
+        if (Date.today - file_date).to_i > STALE_AFTER_DAYS
           return configured_check(
             :warn,
             path,
@@ -28,6 +29,15 @@ module LlmCostTracker
             "updated_at=#{updated_at} is older than #{STALE_AFTER_DAYS} days; #{REFRESH_COMMAND}"
           )
         end
+        bundled_at = LlmCostTracker::Pricing::Registry.metadata["updated_at"]
+        if bundled_at && file_date < Date.iso8601(bundled_at.to_s)
+          return configured_check(
+            :warn,
+            path,
+            count,
+            "updated_at=#{updated_at} is older than the bundled prices (#{bundled_at}); #{REFRESH_COMMAND}"
+          )
+        end
 
         configured_check(:ok, path, count, "updated_at=#{updated_at}")
       rescue Date::Error
@@ -48,7 +58,8 @@ module LlmCostTracker
         Check.new(
           :warn,
           "prices",
-          "using bundled prices updated_at=#{updated_at}; commit a prices_file for production releases"
+          "using bundled prices updated_at=#{updated_at}; " \
+          "commit a prices_file (config.pricing.file) for production releases"
         )
       end
 

data/lib/llm_cost_tracker/doctor/schema_check.rb

@@ -22,8 +22,7 @@ module LlmCostTracker
         Check.new(
           :error,
           @name,
-          "current schema required; #{errors.join('; ')}; " \
-          "run bin/rails generate llm_cost_tracker:install && bin/rails db:migrate"
+          "current schema required; #{errors.join('; ')}; see docs/upgrading.md"
         )
       end
     end

data/lib/llm_cost_tracker/doctor.rb

@@ -167,7 +167,7 @@ module LlmCostTracker
       return live_rollups_check unless LlmCostTracker.configuration.budgets.totals_source == :cache
 
       errors = LlmCostTracker::Ledger::Schema::CallRollups.current_schema_errors
-      return Check.new(:ok, "call rollups", "llm_cost_tracker_call_rollups exists") if errors.empty?
+      return rollups_drift_check if errors.empty?
 
       Check.new(
         :error,
@@ -176,13 +176,30 @@ module LlmCostTracker
       )
     end
 
+    def rollups_drift_check
+      month_start = LlmCostTracker::Ledger::Period.range_start(:month, Time.now)
+      today = LlmCostTracker::Ledger::Period.range_start(:day, Time.now)
+      rolled = LlmCostTracker::CallRollup.where(period: "day", period_start: month_start.to_date...today.to_date)
+                                         .sum(:total_cost)
+      recorded = LlmCostTracker::Call.where(tracked_at: month_start...today).sum(:total_cost)
+      return Check.new(:ok, "call rollups", "llm_cost_tracker_call_rollups exists") if rolled == recorded
+
+      Check.new(
+        :warn,
+        "call rollups",
+        "this month's day rollups before today total #{rolled.to_d.to_s('F')} but the calls total " \
+        "#{recorded.to_d.to_s('F')}, and monthly budgets read the rollups; " \
+        "run bin/rails llm_cost_tracker:rebuild_rollups"
+      )
+    end
+
     def live_rollups_check
       if Probe.table_exists?("llm_cost_tracker_call_rollups")
         Check.new(
           :warn,
           "call rollups",
           "budgets.totals_source=:ledger but llm_cost_tracker_call_rollups exists. " \
-          "Set config.budgets.totals_source = :cache to keep budget reads on the rollups fast path or drop the table."
+          "Set config.budgets.totals_source = :cache to use it or drop the table."
         )
       else
         Check.new(

data/lib/llm_cost_tracker/engine.rb

@@ -3,7 +3,6 @@
 require "rails"
 require_relative "../llm_cost_tracker"
 require_relative "assets"
-require "rack/files"
 
 module LlmCostTracker
   class Engine < ::Rails::Engine

data/lib/llm_cost_tracker/errors.rb

@@ -1,10 +1,21 @@
 # frozen_string_literal: true
 
+require_relative "redaction"
+
 module LlmCostTracker
   class Error < StandardError; end
 
   class InvalidFilterError < Error; end
 
+  class TransactionAbortedError < Error
+    def initialize(error)
+      super(
+        "The database rolled back the whole surrounding transaction while recording LLM usage " \
+        "(#{error.class}: #{Redaction.text(error.message)}); the caller's transaction no longer exists"
+      )
+    end
+  end
+
   class BudgetExceededError < Error
     attr_reader :total, :budget, :budget_type, :last_event, :stage, :scope
 
@@ -40,4 +51,6 @@ module LlmCostTracker
       super("No pricing configured for LLM model: #{model.inspect}")
     end
   end
+
+  CALLER_ERRORS = [BudgetExceededError, UnknownPricingError, TransactionAbortedError].freeze
 end

data/lib/llm_cost_tracker/event.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/digest/uuid"
+
 require_relative "pricing/mode"
 
 module LlmCostTracker
@@ -53,6 +55,12 @@ module LlmCostTracker
       Pricing::Mode.tokenize(pricing_mode.to_s).include?("batch")
     end
 
+    def keyed_by_response_id
+      return self unless provider_response_id
+
+      with(event_id: Digest::UUID.uuid_v5(Digest::UUID::OID_NAMESPACE, "#{provider}/#{provider_response_id}"))
+    end
+
     def self.resolve_line_items(service_items)
       Array(service_items).map do |item|
         item.is_a?(Charges::LineItem) ? item : Charges::LineItem.build(item)

data/lib/llm_cost_tracker/generators/llm_cost_tracker/async_ingestion_generator.rb

@@ -32,12 +32,6 @@ module LlmCostTracker
           inline. The doctor check warns about unused async ingestion tables.
         MSG
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/call_rollups_generator.rb

@@ -10,7 +10,7 @@ module LlmCostTracker
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Creates the optional llm_cost_tracker_call_rollups table for fast budget reads. " \
+      desc "Creates the optional llm_cost_tracker_call_rollups table of running budget totals. " \
            "Required when config.budgets.totals_source = :cache."
 
       def create_migration_file
@@ -28,16 +28,13 @@ module LlmCostTracker
               config.budgets.totals_source = :cache
             end
 
-          Without it Tracker keeps reading budget totals as live SUM aggregates over
+          Once that is deployed, run bin/rails llm_cost_tracker:rebuild_rollups. Until then
+          calls recorded before the switch are missing from monthly budget totals.
+
+          Without the setting Tracker keeps reading budget totals as live SUM aggregates over
           llm_cost_tracker_calls. The doctor check warns about an unused rollups table.
         MSG
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/install_generator.rb

@@ -53,12 +53,6 @@ module LlmCostTracker
           unauthenticated would expose spend, tags, and provider IDs to anyone.
         MSG
       end
-
-      private
-
-      def migration_version
-        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
-      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_async_ingestion.rb.erb

@@ -1,4 +1,4 @@
-class CreateLlmCostTrackerAsyncIngestion < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerAsyncIngestion < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   def change
     create_table :llm_cost_tracker_ingestion_inbox_entries do |t|
       t.string   :event_id, null: false

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_call_rollups.rb.erb

@@ -1,4 +1,4 @@
-class CreateLlmCostTrackerCallRollups < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerCallRollups < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   def change
     create_table :llm_cost_tracker_call_rollups do |t|
       t.string :period, null: false

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_cost_tracker_calls.rb.erb

@@ -1,7 +1,7 @@
 require "llm_cost_tracker/charges/cost_status"
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %>
+class CreateLlmCostTrackerCalls < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   def change
     create_table :llm_cost_tracker_calls do |t|
       t.string  :event_id,      null: false

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/initializer.rb.erb

@@ -34,8 +34,8 @@ LlmCostTracker.configure do |config|
   #   "my-custom-model" => { input: 1.00, output: 2.00 }
   # }
   # :warn (default) records token usage with nil cost when a model has no rate.
-  # Use :raise to reject models that have no rate at all; a model priced with one
-  # component rate missing lands as partial and does not trigger it.
+  # Use :raise to record and then raise for models that have no rate at all; a model
+  # priced with one component rate missing lands as partial and does not trigger it.
   config.pricing.unknown_model_behavior = :warn
 
   # Budget guardrails — cumulative monthly/daily and per-call ceilings in USD,
@@ -50,8 +50,9 @@ LlmCostTracker.configure do |config|
 
   # One budget applied to every distinct value of a tag: each tenant gets its own
   # 1000 a month, not 1000 shared between them. Declare as many tags as you need; a rule
-  # can set its own behavior and callback instead of following the global ones. Needs the
-  # cost columns from `bin/rails generate llm_cost_tracker:upgrade_per_tag_budgets`.
+  # can set its own behavior and callback instead of following the global ones. Installs created
+  # before v0.14 add the cost columns with `bin/rails generate llm_cost_tracker:upgrade_per_tag_budgets`
+  # and `bin/rails db:migrate`, then run `bin/rails llm_cost_tracker:backfill_tag_costs` so earlier calls count.
   # Budget high-cardinality tags only — a tag with a handful of values covers most of the
   # ledger, so its check cannot use an index and every call pays for a near-full scan.
   # config.budgets.per_tag = {
@@ -63,9 +64,9 @@ LlmCostTracker.configure do |config|
   # }
 
   # Where budget checks read period spend from. :ledger sums llm_cost_tracker_calls on
-  # every check; :cache keeps running totals in llm_cost_tracker_call_rollups so checks
-  # read one row instead — switch once the live SUM gets slow, and generate the table
-  # first with `bin/rails generate llm_cost_tracker:call_rollups`.
+  # every check; :cache reads earlier days from llm_cost_tracker_call_rollups and sums
+  # only today's calls. Run `bin/rails generate llm_cost_tracker:call_rollups`, deploy
+  # this setting, then run `bin/rails llm_cost_tracker:rebuild_rollups`.
   # config.budgets.totals_source = :cache
 
   # Ingestion path — :inline (default) writes events synchronously from the request
@@ -74,18 +75,19 @@ LlmCostTracker.configure do |config|
   # inbox/leases tables created by `bin/rails generate llm_cost_tracker:async_ingestion`.
   # Synchronous inbox writes use a dedicated ActiveRecord pool (defaults to 2 connections)
   # so they don't compete with request threads for the default pool when a tracked call
-  # happens inside an open caller transaction. Bump ingestion.pool_size if your Puma
-  # worker count outgrows that.
+  # happens inside an open caller transaction. The pool is per process: bump ingestion.pool_size
+  # if concurrent tracked calls in one process (e.g. Puma threads) outgrow that.
   # config.ingestion.mode = :async
   # config.ingestion.pool_size = 5
 
-  # Register OpenAI-compatible gateway hosts so their calls are attributed to a
-  # provider name instead of the generic openai_compatible bucket.
+  # Register OpenAI-compatible gateway hosts beyond the built-in OpenRouter, DeepSeek and Groq;
+  # the Faraday middleware captures only listed hosts. Its calls, and openai gem calls whose
+  # base_url is a listed host, are recorded under the provider name given here.
   # config.capture.openai_compatible_providers["llm.my-company.com"] = "internal_gateway"
 
   # Streaming calls only report token usage when the request asks for it. Leave this
   # on and the Faraday middleware adds stream_options: { include_usage: true } to
-  # OpenAI-compatible streaming request bodies that don't already set it.
+  # streaming chat-completions requests on hosts known to accept it.
   # config.capture.request_stream_usage = false
 
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_call_rollups_provider.rb.erb

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   TABLE = :llm_cost_tracker_call_rollups

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_call_tags_key_value_index.rb.erb

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   TABLE = :llm_cost_tracker_call_tags

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_image_tokens.rb.erb

@@ -1,4 +1,4 @@
-class UpgradeLlmCostTrackerImageTokens < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerImageTokens < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   TABLE = :llm_cost_tracker_calls
   COLUMNS = %i[image_input_tokens image_output_tokens].freeze
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_indexes.rb.erb

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   CALLS = :llm_cost_tracker_calls

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_per_tag_budgets.rb.erb

@@ -1,6 +1,6 @@
 require "llm_cost_tracker/ledger/schema/adapter"
 
-class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration<%= migration_version %>
+class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
   disable_ddl_transaction!
 
   TABLE = :llm_cost_tracker_call_tags

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_call_rollups_provider_generator.rb

@@ -18,12 +18,6 @@ module LlmCostTracker
           "db/migrate/upgrade_llm_cost_tracker_call_rollups_provider.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_call_tags_key_value_index_generator.rb

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