llm_cost_tracker 0.8.0 → 0.10.0

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

@@ -4,7 +4,10 @@ LlmCostTracker.configure do |config|
   # Set to false to temporarily disable tracking without removing middleware.
   config.enabled = true
 
-  # Tags are merged into every event. Use a callable for request/job-time context.
+  # LLM Cost Tracker logs warnings through Rails.logger when available.
+  config.log_level = :info
+
+  # Tags merged into every event. Use a callable for request/job-time context.
   config.default_tags = -> { { environment: Rails.env } }
 
   # Tag guardrails keep accidental high-cardinality or sensitive values out of the ledger.
@@ -18,40 +21,47 @@ LlmCostTracker.configure do |config|
   # config.instrument :anthropic
   # config.instrument :ruby_llm
 
-  # Budget behavior: :notify calls on_budget_exceeded, :raise raises after recording,
-  # :block_requests preflights monthly/daily budgets before supported requests.
-  config.budget_exceeded_behavior = :notify
-
-  # Unknown pricing records token usage with nil cost by default. Use :raise if
-  # every model must have known pricing before it can be used.
-  config.unknown_pricing_behavior = :warn
-
-  # LLM Cost Tracker logs warnings through Rails.logger when available.
-  config.log_level = :info
+  # Pricing — local file refreshed via bin/rails llm_cost_tracker:prices:refresh
+  # plus inline overrides. Rates are per 1M tokens; the snapshot's currency
+  # is read from your prices_file's `metadata.currency` (USD in the bundled
+  # snapshot — set a different code per file if you maintain non-USD prices).
 <% if options[:prices] -%>
-
-  # Local JSON/YAML pricing file generated by --prices. Keep it in source control
-  # and refresh it with bin/rails llm_cost_tracker:prices:refresh.
   config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
+<% else -%>
+  # config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
 <% end -%>
+  # config.pricing_overrides = {
+  #   "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 require known pricing for every model.
+  config.unknown_pricing_behavior = :warn
 
-  # Cumulative monthly/daily budgets and a single-call ceiling, in USD.
+  # Budget guardrails — cumulative monthly/daily and per-call ceilings in USD,
+  # plus behavior on crossing (:notify default fires on_budget_exceeded; :raise
+  # raises after recording; :block_requests preflights supported requests, also
+  # estimating the current call's input cost via chars/4 so it can block before
+  # send) and an optional callback. Cap evaluation reads from llm_cost_tracker_calls live;
+  # flip cache_rollups to true at high volume so reads hit the rollups table
+  # instead — generate the table with `bin/rails generate llm_cost_tracker:call_rollups`.
   # config.monthly_budget = 100.00
   # config.daily_budget = 10.00
   # config.per_call_budget = 1.00
-
-  # Called when :notify is selected and a monthly, daily, or per-call budget is exceeded.
+  config.budget_exceeded_behavior = :notify
   # config.on_budget_exceeded = ->(data) {
-  #   Rails.logger.warn(
-  #     "LLM #{data[:budget_type]} budget exceeded: $#{data[:total]} / $#{data[:budget]}"
-  #   )
+  #   Rails.logger.warn("LLM #{data[:budget_type]} budget exceeded: $#{data[:total]} / $#{data[:budget]}")
   # }
+  # config.cache_rollups = true
 
-  # Local pricing table and small Ruby-side overrides. Prices are USD per 1M tokens.
-  # config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.yml")
-  # config.pricing_overrides = {
-  #   "my-custom-model" => { input: 1.00, output: 2.00 }
-  # }
+  # Ingestion path — :inline (default) writes events synchronously from the request
+  # thread. Set to :async for a write-ahead inbox + background worker that batches
+  # inserts and survives caller transaction rollbacks. Requires the optional
+  # inbox/leases tables created by `bin/rails generate llm_cost_tracker:async_ingestion`.
+  # The worker uses a dedicated ActiveRecord pool (defaults to 2 connections) so it
+  # doesn't compete with request threads — bump ingestion_pool_size if your Puma
+  # worker count outgrows that.
+  # config.ingestion          = :async
+  # config.ingestion_pool_size = 5
 
   # Register OpenAI-compatible gateway hosts and choose extra tag breakdowns
   # for bin/rails llm_cost_tracker:report.

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

@@ -0,0 +1,35 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_call_rollups
+  OLD_INDEX = %i[period period_start currency].freeze
+  NEW_INDEX = %i[period period_start currency provider].freeze
+
+  def up
+    add_column TABLE, :provider, :string, null: false, default: "" unless column_exists?(TABLE, :provider)
+    add_unique_index NEW_INDEX
+    remove_index TABLE, column: OLD_INDEX, unique: true, if_exists: true
+  end
+
+  def down
+    add_unique_index OLD_INDEX
+    remove_index TABLE, column: NEW_INDEX, unique: true, if_exists: true
+    remove_column TABLE, :provider if column_exists?(TABLE, :provider)
+  end
+
+  private
+
+  def add_unique_index(columns)
+    if postgresql?
+      add_index TABLE, columns, unique: true, algorithm: :concurrently, if_not_exists: true
+    else
+      add_index TABLE, columns, unique: true, if_not_exists: true
+    end
+  end
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

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

@@ -0,0 +1,32 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerCallTagsKeyValueIndex < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_call_tags
+  INDEX_COLUMNS = %i[key value].freeze
+
+  def up
+    if postgresql?
+      add_index TABLE, INDEX_COLUMNS, algorithm: :concurrently, if_not_exists: true
+    elsif mysql?
+      add_index TABLE, INDEX_COLUMNS, length: { value: 191 }, if_not_exists: true
+    else
+      raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+    end
+  end
+
+  def down
+    remove_index TABLE, column: INDEX_COLUMNS, if_exists: true
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+
+  def mysql?
+    LlmCostTracker::Ledger::Schema::Adapter.mysql?(connection)
+  end
+end

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

@@ -0,0 +1,18 @@
+class UpgradeLlmCostTrackerImageTokens < ActiveRecord::Migration<%= migration_version %>
+  TABLE = :llm_cost_tracker_calls
+  COLUMNS = %i[image_input_tokens image_output_tokens].freeze
+
+  def up
+    COLUMNS.each do |column|
+      next if column_exists?(TABLE, column)
+
+      add_column TABLE, column, :integer, null: false, default: 0
+    end
+  end
+
+  def down
+    COLUMNS.each do |column|
+      remove_column TABLE, column if column_exists?(TABLE, column)
+    end
+  end
+end

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

@@ -0,0 +1,32 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerProviderInvoiceImportsProvider < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_provider_invoice_imports
+  OLD_INDEX = %i[source started_at].freeze
+  NEW_INDEX = %i[source provider started_at].freeze
+
+  def up
+    add_column TABLE, :provider, :string, null: false, default: "" unless column_exists?(TABLE, :provider)
+    if postgresql?
+      remove_index TABLE, column: OLD_INDEX, algorithm: :concurrently, if_exists: true
+      add_index TABLE, NEW_INDEX, algorithm: :concurrently, if_not_exists: true
+    else
+      remove_index TABLE, column: OLD_INDEX, if_exists: true
+      add_index TABLE, NEW_INDEX, if_not_exists: true
+    end
+  end
+
+  def down
+    remove_index TABLE, column: NEW_INDEX, if_exists: true
+    add_index TABLE, OLD_INDEX, if_not_exists: true
+    remove_column TABLE, :provider if column_exists?(TABLE, :provider)
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

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

@@ -0,0 +1,25 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerProviderInvoicesMetadataIndex < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_provider_invoices
+
+  def up
+    return unless postgresql?
+
+    add_index TABLE, :metadata, using: :gin, algorithm: :concurrently, if_not_exists: true
+  end
+
+  def down
+    return unless postgresql?
+
+    remove_index TABLE, column: :metadata, if_exists: true
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

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

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeCallRollupsProviderGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds the v0.9 provider column and unique index to llm_cost_tracker_call_rollups."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_call_rollups_provider.rb.erb",
+          "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

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeCallTagsKeyValueIndexGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a (key, value) composite index on llm_cost_tracker_call_tags " \
+           "so high-cardinality tag filters use an index lookup instead of a key-only scan."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_call_tags_key_value_index.rb.erb",
+          "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

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

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeImageTokensGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds image_input_tokens and image_output_tokens columns to llm_cost_tracker_calls."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_image_tokens.rb.erb",
+          "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_provider_invoice_imports_provider_generator.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeProviderInvoiceImportsProviderGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a provider column to llm_cost_tracker_provider_invoice_imports and a " \
+           "(source, provider, started_at) index so resume_cursor_for and " \
+           "last_completed_window_for can isolate per-provider state on shared sources (e.g. csv)."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_provider_invoice_imports_provider.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_provider_invoice_imports_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_provider_invoices_metadata_index_generator.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeProviderInvoicesMetadataIndexGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a GIN index on llm_cost_tracker_provider_invoices.metadata for PostgreSQL " \
+           "so Reconciliation::Diff queries that filter on metadata->>'provider' / 'row_type' / " \
+           "'match_basis' hit an index instead of a sequential scan. No-op on MySQL."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_provider_invoices_metadata_index.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_provider_invoices_metadata_index.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -40,7 +40,10 @@ module LlmCostTracker
         Ingestion::InboxEntry
           .where(id: rows.map(&:id), locked_by: identity)
           .update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
-      rescue StandardError
+      rescue StandardError => e
+        LlmCostTracker::Logging.warn(
+          "Inbox mark_failed failed for #{rows.size} rows: #{e.class}: #{e.message} (original error: #{error.class})"
+        )
         nil
       end
 
@@ -77,7 +80,7 @@ module LlmCostTracker
 
       def persist(rows, events)
         LlmCostTracker::Call.transaction do
-          Ledger::Store.insert_many(events)
+          Ledger::Store.insert(events)
           Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
       end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -15,8 +15,6 @@ module LlmCostTracker
       class << self
         def save(event)
           insert_row(row_for(event))
-          Ingestion::Worker.ensure_started
-          event
         end
 
         def event_from_row(row)
@@ -54,14 +52,10 @@ module LlmCostTracker
             tracked_at: Time.iso8601(payload.fetch(:tracked_at)),
             cost_status: payload.fetch(:cost_status),
             pricing_snapshot: payload[:pricing_snapshot],
-            line_items: line_items_from(payload)
+            line_items: (payload[:line_items] || []).map { |attrs| Billing::LineItem.build(attrs) }
           }
         end
 
-        def line_items_from(payload)
-          (payload[:line_items] || []).map { |attributes| Billing::LineItem.build(attributes) }
-        end
-
         def row_for(event)
           now = Time.now.utc
           {
@@ -86,25 +80,10 @@ module LlmCostTracker
         end
 
         def insert_row(row)
-          connection = LlmCostTracker::Call.connection
-          if connection.transaction_open?
-            insert_with_separate_connection(row)
-          else
-            execute_insert(connection, row)
-          end
+          Pool.with_connection { |connection| execute_insert(connection, row) }
         rescue ActiveRecord::ConnectionTimeoutError => e
           raise LlmCostTracker::Error,
-                "ledger inbox could not checkout a separate database connection: #{e.message}"
-        end
-
-        def insert_with_separate_connection(row)
-          pool = LlmCostTracker::Call.connection_pool
-          connection = pool.checkout
-          begin
-            connection.transaction(requires_new: true) { execute_insert(connection, row) }
-          ensure
-            pool.checkin(connection)
-          end
+                "ledger inbox could not checkout a database connection: #{e.message}"
         end
 
         def execute_insert(connection, row)
@@ -112,7 +91,6 @@ module LlmCostTracker
           quoted_columns = columns.map { |column| connection.quote_column_name(column) }.join(", ")
           quoted_values = columns.map { |column| connection.quote(row.fetch(column)) }.join(", ")
           table = connection.quote_table_name(InboxEntry.table_name)
-
           connection.execute("INSERT INTO #{table} (#{quoted_columns}) VALUES (#{quoted_values})")
         end
       end

data/lib/llm_cost_tracker/ingestion/pool.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ingestion
+    module Pool
+      DEFAULT_POOL_SIZE = 2
+      MUTEX = Mutex.new
+
+      class << self
+        def with_connection(&)
+          pool.with_connection(&)
+        end
+
+        def pool
+          @pool || MUTEX.synchronize { @pool ||= connect! }
+        end
+
+        def reset!
+          MUTEX.synchronize do
+            @pool&.disconnect!
+            @pool = nil
+            @handler = nil
+          end
+        end
+
+        private
+
+        def connect!
+          @handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+          @handler.establish_connection(connection_config)
+        end
+
+        def connection_config
+          LlmCostTracker::Call.connection_db_config.configuration_hash.merge(pool: pool_size)
+        end
+
+        def pool_size
+          configured = LlmCostTracker.configuration.ingestion_pool_size.to_i
+          configured.positive? ? configured : DEFAULT_POOL_SIZE
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -16,25 +16,31 @@ module LlmCostTracker
       MAX_IDLE_INTERVAL_SECONDS = 5.0
       LEASE_SECONDS = 10
       FLUSH_TIMEOUT_SECONDS = 10
+      MUTEX = Mutex.new
+
       class << self
         def ensure_started
-          thread = mutex.synchronize do
+          return unless Ingestion.async?
+
+          thread = MUTEX.synchronize do
             reset_after_fork!
-            unless @thread&.alive?
-              @stop_requested = false
-              @generation = @generation.to_i + 1
-              generation = @generation
-              @thread = Thread.new { run(generation) }
-              @thread.name = "llm_cost_tracker_ingestor"
-              @thread.report_on_exception = false
-            end
+            break @thread if @stop_requested || @thread&.alive?
+
+            @generation = @generation.to_i + 1
+            generation = @generation
+            @thread = Thread.new { run(generation) }
+            @thread.name = "llm_cost_tracker_ingestor"
+            @thread.report_on_exception = false
             @thread
           end
           wake_thread(thread)
         end
 
         def flush!(timeout: nil, require_lease: false)
+          return true unless Ingestion.async?
+
           Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
 
           deadline = Time.now.utc + flush_timeout_seconds(timeout)
           loop do
@@ -52,25 +58,29 @@ module LlmCostTracker
         end
 
         def shutdown!(timeout: nil, drain: true)
+          return true unless Ingestion.async?
+
           timeout ||= FLUSH_TIMEOUT_SECONDS
-          thread = mutex.synchronize do
+          thread = MUTEX.synchronize do
             @stop_requested = true
             @generation = @generation.to_i + 1
             @thread
           end
           wake_thread(thread)
-          thread&.join([timeout, 1].min)
+          thread&.join(timeout)
           drain ? flush!(timeout: timeout, require_lease: true) : true
         rescue StandardError => e
           handle_error(e)
           false
         ensure
-          mutex.synchronize { @thread = nil if @thread.equal?(thread) }
+          MUTEX.synchronize do
+            @thread = nil if @thread.equal?(thread) && !thread&.alive?
+          end
         end
 
         def reset!
-          thread = mutex.synchronize do
-            @stop_requested = true
+          thread = MUTEX.synchronize do
+            @stop_requested = false
             @generation = @generation.to_i + 1
             thread = @thread
             @thread = nil
@@ -90,6 +100,7 @@ module LlmCostTracker
 
         def ingest_once(require_lease: true)
           Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
           batch = Ingestion::Batch.new(identity: identity)
           return 0 unless batch.claimable?
           return 0 if require_lease && !Ingestion::LeaseClaim.new(identity: identity, seconds: LEASE_SECONDS).acquire
@@ -102,16 +113,12 @@ module LlmCostTracker
 
         private
 
-        def mutex
-          @mutex ||= Mutex.new
-        end
-
         def run(generation)
           idle_interval = IDLE_INTERVAL_SECONDS
           loop do
-            break if mutex.synchronize { @stop_requested || generation != @generation }
+            break if MUTEX.synchronize { @stop_requested || generation != @generation }
 
-            processed = executor_wrap { ingest_once }
+            processed = Rails.application.executor.wrap { ingest_once }
             release_connection!
             if processed.zero?
               sleep(idle_interval)
@@ -126,7 +133,7 @@ module LlmCostTracker
           end
         ensure
           release_connection!
-          mutex.synchronize { @thread = nil if @thread.equal?(Thread.current) }
+          MUTEX.synchronize { @thread = nil if @thread.equal?(Thread.current) }
         end
 
         def reset_after_fork!
@@ -143,19 +150,6 @@ module LlmCostTracker
           nil
         end
 
-        def executor_wrap(&)
-          executor = rails_executor
-          return yield unless executor
-
-          executor.wrap(&)
-        end
-
-        def rails_executor
-          Rails.application.try(:executor)
-        rescue StandardError
-          nil
-        end
-
         def identity
           @identity ||= "pid-#{Process.pid}-#{SecureRandom.hex(6)}"
         end