llm_cost_tracker 0.11.0 → 0.13.0

data/lib/llm_cost_tracker/generators/llm_cost_tracker/prices_generator.rb

@@ -11,11 +11,14 @@ module LlmCostTracker
     class PricesGenerator < Rails::Generators::Base
       desc "Creates a local LLM Cost Tracker price snapshot"
 
+      PRICES_PATH = "config/llm_cost_tracker_prices.yml"
+
       def create_prices_file
-        LlmCostTracker::Pricing::Sync::RegistryWriter.new.call(
-          path: File.join(destination_root, "config/llm_cost_tracker_prices.yml"),
+        payload = LlmCostTracker::Pricing::Sync::RegistryWriter.new.render(
+          path: File.join(destination_root, PRICES_PATH),
           registry: YAML.safe_load_file(LlmCostTracker::Pricing::Registry::DEFAULT_PRICES_PATH, aliases: false) || {}
         )
+        create_file(PRICES_PATH, payload)
       end
     end
   end

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

@@ -1,5 +1,4 @@
-require "llm_cost_tracker/billing/components"
-require "llm_cost_tracker/billing/cost_status"
+require "llm_cost_tracker/charges/cost_status"
 require "llm_cost_tracker/ledger/schema/adapter"
 
 class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %>
@@ -8,7 +7,7 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
       t.string  :event_id,      null: false
       t.string  :provider,      null: false
       t.string  :model,         null: false
-<% LlmCostTracker::TokenUsage.members.each do |column| -%>
+<% LlmCostTracker::Usage::TokenUsage.members.each do |column| -%>
       t.integer :<%= column %>, null: false, default: 0
 <% end -%>
       t.decimal :total_cost, precision: 20, scale: 8
@@ -21,7 +20,7 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
       t.string  :provider_workspace_id
       t.boolean :batch, null: false, default: false
       t.string  :pricing_mode
-      t.string  :cost_status, null: false, default: LlmCostTracker::Billing::CostStatus::UNKNOWN
+      t.string  :cost_status, null: false, default: LlmCostTracker::Charges::CostStatus::UNKNOWN
       if postgresql?
         t.jsonb :pricing_snapshot
       elsif mysql?
@@ -50,7 +49,7 @@ class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %
       t.decimal :rate_quantity, precision: 30, scale: 10, null: false, default: 1
       t.decimal :cost, precision: 20, scale: 8
       t.string :currency, null: false, default: "USD"
-      t.string :cost_status, null: false, default: LlmCostTracker::Billing::CostStatus::UNKNOWN
+      t.string :cost_status, null: false, default: LlmCostTracker::Charges::CostStatus::UNKNOWN
       t.string :pricing_basis
       t.string :price_key
       t.string :price_source

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

@@ -13,7 +13,7 @@ LlmCostTracker.configure do |config|
   # Tag guardrails keep accidental high-cardinality or sensitive values out of the ledger.
   # config.max_tag_count = 50
   # config.max_tag_value_bytesize = 1024
-  # config.redacted_tag_keys = %w[api_key access_token authorization credential password refresh_token secret]
+  # config.redacted_tag_keys = <%= LlmCostTracker::Configuration::DEFAULT_REDACTED_TAG_KEYS.inspect %>
 
   # Optional SDK integrations. Provider SDK gems are not installed by LLM Cost Tracker.
   # Enabled integrations are checked at boot, so enable only clients your app loads.
@@ -57,8 +57,9 @@ LlmCostTracker.configure do |config|
   # 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
+  # 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.
   # config.ingestion          = :async
   # config.ingestion_pool_size = 5

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -8,6 +8,12 @@ module LlmCostTracker
     class Batch
       BATCH_SIZE = 100
       LOCK_TIMEOUT_SECONDS = 30
+      TRANSIENT_PERSIST_ERRORS = [
+        ActiveRecord::Deadlocked,
+        ActiveRecord::LockWaitTimeout,
+        ActiveRecord::StatementTimeout,
+        ActiveRecord::ConnectionNotEstablished
+      ].freeze
 
       def initialize(identity:)
         @identity = identity
@@ -22,31 +28,61 @@ module LlmCostTracker
         rows.size
       rescue StandardError => e
         rows_to_mark = valid_rows&.any? ? valid_rows : rows
-        mark_failed(rows_to_mark, e) if rows_to_mark&.any?
+        if rows_to_mark&.any?
+          transient = valid_rows&.any? && TRANSIENT_PERSIST_ERRORS.any? { |klass| e.is_a?(klass) }
+          mark_failed_with_message(rows_to_mark, error_message_for(e), decrement_attempts: transient)
+        end
         raise
       end
 
       def pending?
-        Ingestion::InboxEntry.where("attempts < ?", Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE).exists?
+        Ingestion::InboxEntry.pending.exists?
       end
 
       def claimable?
         claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists?
       end
 
-      def mark_failed(rows, error)
-        message = "#{error.class}: #{error.message}".byteslice(0, 1_000)
+      def mark_failed_with_message(rows, message, decrement_attempts: false)
         now = Time.now.utc
-        Ingestion::InboxEntry
-          .where(id: rows.map(&:id), locked_by: identity)
-          .update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
+        scope = Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity)
+        if decrement_attempts
+          scope.update_all(
+            Ingestion::InboxEntry.sanitize_sql_array(
+              ["last_error = ?, locked_at = ?, locked_by = NULL, " \
+               "attempts = GREATEST(attempts - 1, 0), updated_at = ?", message, now, now]
+            )
+          )
+        else
+          scope.update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
+          warn_on_quarantine(rows)
+        end
       rescue StandardError => e
         LlmCostTracker::Logging.warn(
-          "Inbox mark_failed failed for #{rows.size} rows: #{e.class}: #{e.message} (original error: #{error.class})"
+          "Inbox mark_failed_with_message failed for #{rows.size} rows: #{e.class}: #{e.message} " \
+          "(attempted message: #{message.to_s.byteslice(0, 200)})"
         )
         nil
       end
 
+      def error_message_for(error)
+        "#{error.class}: #{error.message}".byteslice(0, 1_000)
+      end
+
+      def warn_on_quarantine(rows)
+        threshold = Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE
+        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})"
+        )
+      end
+
       private
 
       attr_reader :identity
@@ -69,25 +105,38 @@ module LlmCostTracker
       def decode(rows)
         valid_rows = []
         events = []
+        failures = Hash.new { |h, k| h[k] = [] }
         rows.each do |row|
           events << Ingestion::Inbox.event_from_row(row)
           valid_rows << row
         rescue StandardError => e
-          mark_failed([row], e)
+          failures[error_message_for(e)] << row
         end
+        failures.each { |message, failed_rows| mark_failed_with_message(failed_rows, message) }
         [valid_rows, events]
       end
 
-      def persist(rows, events)
+      def persist(rows, events, retry_on_conflict: true)
         LlmCostTracker::Call.transaction do
-          Ledger::Store.insert(events)
+          Ledger::Store.persist_records(events)
           Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
+        Ledger::Rollups.increment_safely!(events) if Ingestion.cache_rollups?
+      rescue ActiveRecord::RecordNotUnique
+        raise unless retry_on_conflict
+
+        already_persisted = LlmCostTracker::Call.where(event_id: events.map(&:event_id)).pluck(:event_id)
+        fresh_events = events.reject { |event| already_persisted.include?(event.event_id) }
+        LlmCostTracker::Logging.warn(
+          "Ingestion::Batch#persist: #{already_persisted.size} event_id(s) already in ledger; " \
+          "skipped duplicates and persisted #{fresh_events.size} fresh event(s)"
+        )
+        persist(rows, fresh_events, retry_on_conflict: false)
       end
 
       def claimable_scope(cutoff)
         Ingestion::InboxEntry
-          .where("attempts < ?", Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE)
+          .pending
           .where("locked_at IS NULL OR locked_at < ?", cutoff)
       end
     end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -2,14 +2,14 @@
 
 require "json"
 require "time"
+require "active_support/core_ext/hash/keys"
 
 require_relative "../event"
 require_relative "../pricing"
-require_relative "../billing/line_item"
 
 module LlmCostTracker
   module Ingestion
-    class Inbox
+    module Inbox
       PAYLOAD_SCHEMA_VERSION = 2
 
       class << self
@@ -30,8 +30,8 @@ module LlmCostTracker
         private
 
         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))
+          cost = payload[:cost] && Charges::Cost.from_h(payload[:cost])
+          token_usage = Usage::TokenUsage.build(**payload.fetch(:token_usage).slice(*Usage::TokenUsage.members))
 
           {
             event_id: payload.fetch(:event_id),
@@ -43,16 +43,15 @@ module LlmCostTracker
             tags: payload.fetch(:tags),
             latency_ms: payload[:latency_ms],
             stream: payload.fetch(:stream),
-            usage_source: payload[:usage_source]&.to_sym,
+            usage_source: payload[:usage_source],
             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: (payload[:line_items] || []).map { |attrs| Billing::LineItem.build(attrs) }
+            pricing_snapshot: payload[:pricing_snapshot]&.deep_stringify_keys,
+            line_items: (payload[:line_items] || []).map { |attrs| Charges::LineItem.build(attrs) }
           }
         end
 

data/lib/llm_cost_tracker/ingestion/pool.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/module/delegation"
+
 module LlmCostTracker
   module Ingestion
     module Pool
@@ -7,22 +9,12 @@ module LlmCostTracker
       MUTEX = Mutex.new
 
       class << self
-        def with_connection(&)
-          pool.with_connection(&)
-        end
+        delegate :with_connection, to: :pool
 
         def pool
           @pool || MUTEX.synchronize { @pool ||= connect! }
         end
 
-        def reset!
-          MUTEX.synchronize do
-            @pool&.disconnect!
-            @pool = nil
-            @handler = nil
-          end
-        end
-
         private
 
         def connect!

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -6,11 +6,10 @@ require "securerandom"
 require_relative "inbox"
 require_relative "batch"
 require_relative "lease_claim"
-require_relative "../logging"
 
 module LlmCostTracker
   module Ingestion
-    class Worker
+    module Worker
       INTERVAL_SECONDS = 0.25
       IDLE_INTERVAL_SECONDS = 1.0
       MAX_IDLE_INTERVAL_SECONDS = 5.0
@@ -66,8 +65,12 @@ module LlmCostTracker
             @generation = @generation.to_i + 1
             @thread
           end
-          wake_thread(thread)
-          thread&.join(timeout)
+          begin
+            wake_thread(thread)
+            thread&.join(timeout)
+          rescue StandardError => e
+            handle_error(e)
+          end
           drain ? flush!(timeout: timeout, require_lease: true) : true
         rescue StandardError => e
           handle_error(e)
@@ -78,19 +81,6 @@ module LlmCostTracker
           end
         end
 
-        def reset!
-          thread = MUTEX.synchronize do
-            @stop_requested = false
-            @generation = @generation.to_i + 1
-            thread = @thread
-            @thread = nil
-            @pid = nil
-            @identity = nil
-            thread
-          end
-          wake_thread(thread)
-        end
-
         def flush_timeout_seconds(timeout)
           numeric = Float(timeout, exception: false)
           return FLUSH_TIMEOUT_SECONDS unless numeric&.finite? && numeric.positive?

data/lib/llm_cost_tracker/ingestion.rb

@@ -1,17 +1,20 @@
 # frozen_string_literal: true
 
+require "active_support/notifications"
 require "securerandom"
 
 require_relative "errors"
+require_relative "check"
 require_relative "ledger"
-require_relative "ingestion/lease_claim"
-require_relative "ingestion/pool"
-require_relative "ingestion/inbox"
-require_relative "ingestion/batch"
-require_relative "ingestion/worker"
 
 module LlmCostTracker
   module Ingestion
+    autoload :LeaseClaim, "llm_cost_tracker/ingestion/lease_claim"
+    autoload :Pool, "llm_cost_tracker/ingestion/pool"
+    autoload :Inbox, "llm_cost_tracker/ingestion/inbox"
+    autoload :Batch, "llm_cost_tracker/ingestion/batch"
+    autoload :Worker, "llm_cost_tracker/ingestion/worker"
+
     VERIFY_TAG = "llm_cost_tracker_verify"
 
     class << self
@@ -19,25 +22,12 @@ module LlmCostTracker
         "llm_cost_tracker_ingestion_"
       end
 
-      CORE_SCHEMA_GUARDS = [
-        ["llm_cost_tracker_calls",           Ledger::Schema::Calls],
-        ["llm_cost_tracker_call_line_items", Ledger::Schema::CallLineItems],
-        ["llm_cost_tracker_call_tags",       Ledger::Schema::CallTags]
-      ].freeze
-
-      ROLLUPS_SCHEMA_GUARD = ["llm_cost_tracker_call_rollups", Ledger::Schema::CallRollups].freeze
-
-      ASYNC_SCHEMA_GUARDS = [
-        ["llm_cost_tracker_ingestion_inbox_entries", Ledger::Schema::IngestionInboxEntries],
-        ["llm_cost_tracker_ingestion_leases",        Ledger::Schema::IngestionLeases]
-      ].freeze
-
       def ensure_current_schema!
         unless LlmCostTracker::Call.table_exists?
           raise Error, "llm_cost_tracker_calls table is missing; run install generator and migrate"
         end
 
-        guards_for_current_config.each do |table_name, schema_module|
+        guards_for_current_config.each do |schema_module, table_name|
           errors = schema_module.current_schema_errors
           next if errors.empty?
 
@@ -55,16 +45,16 @@ module LlmCostTracker
       end
 
       def guards_for_current_config
-        guards = CORE_SCHEMA_GUARDS.dup
-        guards << ROLLUPS_SCHEMA_GUARD if cache_rollups?
-        guards += ASYNC_SCHEMA_GUARDS if async?
+        guards = Ledger::Schema::CORE_SCHEMAS.dup
+        guards << Ledger::Schema::CACHE_ROLLUPS_SCHEMA if cache_rollups?
+        guards += Ledger::Schema::ASYNC_SCHEMAS if async?
         guards
       end
 
       def verify
         unless LlmCostTracker::Call.table_exists?
           return [
-            LlmCostTracker::Doctor::Check.new(
+            LlmCostTracker::Check.new(
               :error,
               "active_record",
               "llm_cost_tracker_calls table is missing; run install generator and migrate"
@@ -74,7 +64,7 @@ module LlmCostTracker
 
         [capture_check]
       rescue StandardError => e
-        [LlmCostTracker::Doctor::Check.new(:error, "active_record", "#{e.class}: #{e.message}")]
+        [LlmCostTracker::Check.new(:error, "active_record", "#{e.class}: #{e.message}")]
       end
 
       private
@@ -88,7 +78,7 @@ module LlmCostTracker
         event = LlmCostTracker.track(
           provider: provider,
           model: model,
-          tokens: { input: 1, output: 1 },
+          tokens: { input_tokens: 1, output_tokens: 1 },
           provider_response_id: response_id,
           tags: { feature: VERIFY_TAG }
         )
@@ -97,17 +87,17 @@ module LlmCostTracker
 
         return capture_success if persisted && notifications.any?
 
-        LlmCostTracker::Doctor::Check.new(
+        LlmCostTracker::Check.new(
           :error,
           "active_record capture",
           capture_failure_message(persisted, notifications)
         )
       rescue LlmCostTracker::BudgetExceededError => e
-        LlmCostTracker::Doctor::Check.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
+        LlmCostTracker::Check.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
       rescue LlmCostTracker::Error => e
-        LlmCostTracker::Doctor::Check.new(:error, "active_record capture", e.message)
+        LlmCostTracker::Check.new(:error, "active_record capture", e.message)
       rescue StandardError => e
-        LlmCostTracker::Doctor::Check.new(:error, "active_record capture", "#{e.class}: #{e.message}")
+        LlmCostTracker::Check.new(:error, "active_record capture", "#{e.class}: #{e.message}")
       ensure
         cleanup_verification_call(response_id) if response_id
         cleanup_verification_inbox(event: event, response_id: response_id)
@@ -122,7 +112,7 @@ module LlmCostTracker
 
       def capture_success
         path = async? ? "async inbox" : "inline writer"
-        LlmCostTracker::Doctor::Check.new(
+        LlmCostTracker::Check.new(
           :ok,
           "active_record capture",
           "manual event emitted and persisted through #{path}"
@@ -138,13 +128,11 @@ module LlmCostTracker
 
       def cleanup_verification_call(response_id)
         relation = LlmCostTracker::Call.where(provider_response_id: response_id)
-        rows = relation.pluck(:id, :tracked_at, :total_cost, :pricing_snapshot, :provider)
-        return if rows.empty?
+        records = relation.select(:id, :tracked_at, :total_cost, :pricing_snapshot, :provider).to_a
+        return if records.empty?
 
         relation.delete_all
-        return unless cache_rollups?
-
-        LlmCostTracker::Ledger::Rollups.decrement!(rows)
+        LlmCostTracker::Ledger::Rollups.decrement!(records) if cache_rollups?
       end
 
       def cleanup_verification_inbox(event:, response_id:)
@@ -162,7 +150,7 @@ module LlmCostTracker
 
       def sample_priced_identity
         key = LlmCostTracker::Pricing::Registry.builtin_prices.find do |model_id, prices|
-          model_id.include?("/") && prices[:input] && prices[:output]
+          model_id.include?("/") && prices["input"] && prices["output"]
         end&.first
         provider, model = key.to_s.split("/", 2)
         [provider || "openai", model || "gpt-4o-mini"]