llm_cost_tracker 0.7.3 → 0.9.0

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

@@ -0,0 +1,55 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class CreateLlmCostTrackerReconciliation < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :llm_cost_tracker_provider_invoices, if_not_exists: true do |t|
+      t.string :source, null: false
+      t.date :period_start, null: false
+      t.date :period_end, null: false
+      t.string :external_id, null: false
+      t.decimal :billed_amount, precision: 20, scale: 8
+      t.string :currency, null: false, default: "USD"
+      if postgresql?
+        t.jsonb :metadata, null: false, default: {}
+      elsif mysql?
+        t.json :metadata, null: false
+      else
+        raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+      end
+      t.datetime :imported_at, null: false
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_provider_invoice_imports, if_not_exists: true do |t|
+      t.string :source, null: false
+      t.string :cursor
+      t.date :window_start
+      t.date :window_end
+      t.string :state, null: false
+      t.text :last_error
+      t.integer :rows_imported, null: false, default: 0
+      t.datetime :started_at, null: false
+      t.datetime :finished_at
+
+      t.timestamps
+    end
+
+    add_index :llm_cost_tracker_provider_invoices, :external_id, unique: true,
+              if_not_exists: true
+    add_index :llm_cost_tracker_provider_invoices, %i[source currency period_start],
+              if_not_exists: true
+    add_index :llm_cost_tracker_provider_invoice_imports, %i[source started_at],
+              if_not_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/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,40 @@ 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. Prices are USD per 1M tokens.
 <% 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) 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 — false (default) writes events synchronously from the request
+  # thread. Flip to true 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:durable_ingestion`.
+  # config.durable_ingestion = true
 
   # 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,20 @@
+class UpgradeLlmCostTrackerCallRollupsProvider < ActiveRecord::Migration<%= migration_version %>
+  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
+    unless column_exists?(TABLE, :provider)
+      execute "DELETE FROM #{TABLE}"
+      add_column TABLE, :provider, :string, null: false, default: ""
+    end
+    remove_index TABLE, column: OLD_INDEX, unique: true if index_exists?(TABLE, OLD_INDEX, unique: true)
+    add_index TABLE, NEW_INDEX, unique: true unless index_exists?(TABLE, NEW_INDEX, unique: true)
+  end
+
+  def down
+    remove_index TABLE, column: NEW_INDEX, unique: true if index_exists?(TABLE, NEW_INDEX, unique: true)
+    add_index TABLE, OLD_INDEX, unique: true unless index_exists?(TABLE, OLD_INDEX, unique: true)
+    remove_column TABLE, :provider if column_exists?(TABLE, :provider)
+  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 %>
+  TABLE = :llm_cost_tracker_call_tags
+  INDEX_COLUMNS = %i[key value].freeze
+
+  def up
+    return if index_exists?(TABLE, INDEX_COLUMNS)
+
+    if postgresql?
+      add_index TABLE, INDEX_COLUMNS
+    elsif mysql?
+      add_index TABLE, INDEX_COLUMNS, length: { value: 191 }
+    else
+      raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+    end
+  end
+
+  def down
+    remove_index TABLE, column: INDEX_COLUMNS if index_exists?(TABLE, INDEX_COLUMNS)
+  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/upgrade_call_rollups_provider_generator.rb

@@ -0,0 +1,38 @@
+# 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
+
+      def warn_about_rollups_truncation
+        say(<<~MSG, :yellow)
+          The migration clears existing llm_cost_tracker_call_rollups rows before adding the
+          provider column. Budget reads fall back to live aggregation from
+          llm_cost_tracker_calls until new events repopulate the rollups under their provider
+          keys. See docs/upgrading.md for details.
+        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/{add_provider_response_id_generator.rb → upgrade_call_tags_key_value_index_generator.rb}

@@ -5,17 +5,18 @@ require "rails/generators/active_record"
 
 module LlmCostTracker
   module Generators
-    class AddProviderResponseIdGenerator < Rails::Generators::Base
+    class UpgradeCallTagsKeyValueIndexGenerator < Rails::Generators::Base
       include ActiveRecord::Generators::Migration
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Creates a migration to add llm_api_calls.provider_response_id"
+      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(
-          "add_provider_response_id_to_llm_api_calls.rb.erb",
-          "db/migrate/add_provider_response_id_to_llm_api_calls.rb"
+          "upgrade_call_tags_key_value_index.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_call_tags_key_value_index.rb"
         )
       end
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/{add_streaming_generator.rb → upgrade_image_tokens_generator.rb}

@@ -5,17 +5,17 @@ require "rails/generators/active_record"
 
 module LlmCostTracker
   module Generators
-    class AddStreamingGenerator < Rails::Generators::Base
+    class UpgradeImageTokensGenerator < Rails::Generators::Base
       include ActiveRecord::Generators::Migration
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Creates a migration to add llm_api_calls.stream and llm_api_calls.usage_source"
+      desc "Adds image_input_tokens and image_output_tokens columns to llm_cost_tracker_calls."
 
       def create_migration_file
         migration_template(
-          "add_streaming_to_llm_api_calls.rb.erb",
-          "db/migrate/add_streaming_to_llm_api_calls.rb"
+          "upgrade_image_tokens.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_image_tokens.rb"
         )
       end
 

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -27,7 +27,7 @@ module LlmCostTracker
       end
 
       def pending?
-        Ingestion::Event.where("attempts < ?", Ingestion::Event::MAX_ATTEMPTS).exists?
+        Ingestion::InboxEntry.where("attempts < ?", Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE).exists?
       end
 
       def claimable?
@@ -37,7 +37,7 @@ module LlmCostTracker
       def mark_failed(rows, error)
         message = "#{error.class}: #{error.message}".byteslice(0, 1_000)
         now = Time.now.utc
-        Ingestion::Event
+        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
@@ -51,16 +51,15 @@ module LlmCostTracker
       def claim
         now = Time.now.utc
         cutoff = now - LOCK_TIMEOUT_SECONDS
-        Ingestion::Event.transaction do
+        Ingestion::InboxEntry.transaction do
           rows = claimable_scope(cutoff).order(:id).limit(BATCH_SIZE).lock.to_a
-          ids = rows.map(&:id)
-          next [] if ids.empty?
+          next [] if rows.empty?
 
-          updates = Ingestion::Event.sanitize_sql_array(
+          updates = Ingestion::InboxEntry.sanitize_sql_array(
             ["locked_at = ?, locked_by = ?, attempts = attempts + 1, updated_at = ?", now, identity, now]
           )
-          Ingestion::Event.where(id: ids).update_all(updates)
-          Ingestion::Event.where(id: ids, locked_by: identity).order(:id).to_a
+          Ingestion::InboxEntry.where(id: rows.map(&:id)).update_all(updates)
+          rows
         end
       end
 
@@ -77,15 +76,15 @@ module LlmCostTracker
       end
 
       def persist(rows, events)
-        LlmCostTracker::Ledger::Call.transaction do
+        LlmCostTracker::Call.transaction do
           Ledger::Store.insert_many(events)
-          Ingestion::Event.where(id: rows.map(&:id), locked_by: identity).delete_all
+          Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
       end
 
       def claimable_scope(cutoff)
-        Ingestion::Event
-          .where("attempts < ?", Ingestion::Event::MAX_ATTEMPTS)
+        Ingestion::InboxEntry
+          .where("attempts < ?", Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE)
           .where("locked_at IS NULL OR locked_at < ?", cutoff)
       end
     end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -5,11 +5,12 @@ require "time"
 
 require_relative "../event"
 require_relative "../pricing"
+require_relative "../billing/line_item"
 
 module LlmCostTracker
   module Ingestion
     class Inbox
-      PAYLOAD_SCHEMA_VERSION = 1
+      PAYLOAD_SCHEMA_VERSION = 2
 
       class << self
         def save(event)
@@ -19,32 +20,47 @@ module LlmCostTracker
         end
 
         def event_from_row(row)
-          payload = JSON.parse(row.payload)
-          schema_version = payload.fetch("schema_version", 0)
-          unless [0, PAYLOAD_SCHEMA_VERSION].include?(schema_version)
+          payload = JSON.parse(row.payload, symbolize_names: true)
+          schema_version = payload[:schema_version]
+          unless schema_version == PAYLOAD_SCHEMA_VERSION
             raise LlmCostTracker::Error, "unsupported ledger inbox payload schema version #{schema_version.inspect}"
           end
 
-          cost = payload["cost"] && Pricing.stored_cost_attributes(payload["cost"])
-          token_usage = payload["token_usage"] || payload
+          LlmCostTracker::Event.new(**event_attributes_from(payload))
+        end
+
+        private
 
-          LlmCostTracker::Event.new(
-            event_id: payload.fetch("event_id"),
-            provider: payload.fetch("provider"),
-            model: payload.fetch("model"),
-            token_usage: TokenUsage.from_hash(token_usage),
-            pricing_mode: payload["pricing_mode"],
+        def event_attributes_from(payload)
+          cost = payload[:cost] && Pricing.stored_cost_attributes(payload[:cost])
+          token_usage = TokenUsage.build(**payload.fetch(:token_usage).slice(*TokenUsage.members))
+
+          {
+            event_id: payload.fetch(:event_id),
+            provider: payload.fetch(:provider),
+            model: payload.fetch(:model),
+            token_usage: token_usage,
+            pricing_mode: Pricing.normalize_mode(payload[:pricing_mode]),
             cost: cost,
-            tags: payload.fetch("tags"),
-            latency_ms: payload["latency_ms"],
-            stream: payload.fetch("stream"),
-            usage_source: payload["usage_source"],
-            provider_response_id: payload["provider_response_id"],
-            tracked_at: Time.iso8601(payload.fetch("tracked_at"))
-          )
+            tags: payload.fetch(:tags),
+            latency_ms: payload[:latency_ms],
+            stream: payload.fetch(:stream),
+            usage_source: payload[:usage_source]&.to_sym,
+            provider_response_id: payload[:provider_response_id],
+            provider_project_id: payload[:provider_project_id],
+            provider_api_key_id: payload[:provider_api_key_id],
+            provider_workspace_id: payload[:provider_workspace_id],
+            batch: payload.fetch(:batch),
+            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)
+          }
         end
 
-        private
+        def line_items_from(payload)
+          (payload[:line_items] || []).map { |attributes| Billing::LineItem.build(attributes) }
+        end
 
         def row_for(event)
           now = Time.now.utc
@@ -70,7 +86,7 @@ module LlmCostTracker
         end
 
         def insert_row(row)
-          connection = LlmCostTracker::Ledger::Call.connection
+          connection = LlmCostTracker::Call.connection
           if connection.transaction_open?
             insert_with_separate_connection(row)
           else
@@ -82,7 +98,7 @@ module LlmCostTracker
         end
 
         def insert_with_separate_connection(row)
-          pool = LlmCostTracker::Ledger::Call.connection_pool
+          pool = LlmCostTracker::Call.connection_pool
           connection = pool.checkout
           begin
             connection.transaction(requires_new: true) { execute_insert(connection, row) }
@@ -95,8 +111,7 @@ module LlmCostTracker
           columns = row.keys
           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(Event.table_name)
-
+          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/inline.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative "../ledger/store"
+
+module LlmCostTracker
+  module Ingestion
+    module Inline
+      class << self
+        def save(event)
+          persist(event)
+          event
+        end
+
+        private
+
+        def persist(event)
+          Ledger::Store.insert_many([event])
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -18,6 +18,8 @@ module LlmCostTracker
       FLUSH_TIMEOUT_SECONDS = 10
       class << self
         def ensure_started
+          return unless Ingestion.durable?
+
           thread = mutex.synchronize do
             reset_after_fork!
             unless @thread&.alive?
@@ -25,18 +27,20 @@ module LlmCostTracker
               @generation = @generation.to_i + 1
               generation = @generation
               @thread = Thread.new { run(generation) }
-              @thread.name = "llm_cost_tracker_ingestor" if @thread.respond_to?(:name=)
-              @thread.report_on_exception = false if @thread.respond_to?(:report_on_exception=)
+              @thread.name = "llm_cost_tracker_ingestor"
+              @thread.report_on_exception = false
             end
             @thread
           end
           wake_thread(thread)
         end
 
-        def flush!(timeout: FLUSH_TIMEOUT_SECONDS, require_lease: false)
+        def flush!(timeout: nil, require_lease: false)
+          return true unless Ingestion.durable?
+
           Ingestion.ensure_current_schema!
 
-          deadline = Time.now.utc + timeout
+          deadline = Time.now.utc + flush_timeout_seconds(timeout)
           loop do
             return true unless Ingestion::Batch.new(identity: identity).pending?
             return false if Time.now.utc >= deadline
@@ -51,20 +55,25 @@ module LlmCostTracker
           end
         end
 
-        def shutdown!(timeout: FLUSH_TIMEOUT_SECONDS, drain: true)
+        def shutdown!(timeout: nil, drain: true)
+          return true unless Ingestion.durable?
+
+          timeout ||= FLUSH_TIMEOUT_SECONDS
           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!
@@ -80,7 +89,15 @@ module LlmCostTracker
           wake_thread(thread)
         end
 
+        def flush_timeout_seconds(timeout)
+          numeric = Float(timeout, exception: false)
+          return FLUSH_TIMEOUT_SECONDS unless numeric&.finite? && numeric.positive?
+
+          numeric
+        end
+
         def ingest_once(require_lease: true)
+          Ingestion.ensure_current_schema!
           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