llm_cost_tracker 0.5.3 → 0.6.1

data/lib/llm_cost_tracker/pricing/lookup.rb

@@ -7,41 +7,74 @@ module LlmCostTracker
     module Lookup
       Match = Data.define(:source, :key, :prices, :matched_by)
       MUTEX = Monitor.new
+      CACHE_MISS = Object.new.freeze
+      NO_MATCH = Object.new.freeze
+      MAX_LOOKUP_CACHE_ENTRIES = 512
 
       class << self
         def call(provider:, model:)
           provider_name = provider.to_s
           model_name = model.to_s
+          generation = LlmCostTracker.configuration_generation
+          cache_key = [generation, provider_name, model_name]
+          cached = cached_lookup(cache_key)
+          return cached unless cached.equal?(CACHE_MISS)
+
           provider_model = provider_name.empty? ? model_name : "#{provider_name}/#{model_name}"
           normalized_model = normalize_model_name(model_name)
-          current = current_price_tables
+          current = current_price_tables(generation)
 
-          explain_table(current.fetch(:pricing_overrides), :pricing_overrides, provider_model, model_name,
-                        normalized_model) ||
+          match =
+            explain_table(current.fetch(:pricing_overrides), :pricing_overrides, provider_model, model_name,
+                          normalized_model) ||
             explain_table(current.fetch(:file_prices), :prices_file, provider_model, model_name, normalized_model) ||
             explain_table(Pricing::PRICES, :bundled, provider_model, model_name, normalized_model)
+          cache_lookup(cache_key, match)
+          match
         end
 
         private
 
-        def current_price_tables
-          file_prices = PriceRegistry.file_prices(LlmCostTracker.configuration.prices_file)
-          overrides = PriceRegistry.normalize_price_table(LlmCostTracker.configuration.pricing_overrides)
-          cache_key = [file_prices.object_id, LlmCostTracker.configuration.pricing_overrides.hash]
-
+        def current_price_tables(generation)
           cached = @prices_cache
-          return cached[:value] if cached && cached[:key] == cache_key
+          return cached[:value] if cached && cached[:generation] == generation
 
           MUTEX.synchronize do
             cached = @prices_cache
-            return cached[:value] if cached && cached[:key] == cache_key
+            return cached[:value] if cached && cached[:generation] == generation
 
+            config = LlmCostTracker.configuration
+            file_prices = PriceRegistry.file_prices(config.prices_file)
+            overrides = PriceRegistry.normalize_price_table(config.pricing_overrides)
             value = { pricing_overrides: overrides, file_prices: file_prices }.freeze
-            @prices_cache = { key: cache_key, value: value }.freeze
+            @prices_cache = { generation: generation, value: value }.freeze
             value
           end
         end
 
+        def cached_lookup(cache_key)
+          cached = @lookup_cache
+          return CACHE_MISS unless cached && cached[:generation] == cache_key.first
+          return CACHE_MISS unless cached[:values].key?(cache_key)
+
+          match = cached[:values].fetch(cache_key)
+          match.equal?(NO_MATCH) ? nil : match
+        end
+
+        def cache_lookup(cache_key, match)
+          MUTEX.synchronize do
+            cached = @lookup_cache
+            values = if cached && cached[:generation] == cache_key.first
+                       cached[:values].dup
+                     else
+                       {}
+                     end
+            values.clear if values.size >= MAX_LOOKUP_CACHE_ENTRIES
+            values[cache_key] = match || NO_MATCH
+            @lookup_cache = { generation: cache_key.first, values: values.freeze }.freeze
+          end
+        end
+
         def explain_table(table, source, provider_model, model_name, normalized_model)
           return nil if table.empty?
 

data/lib/llm_cost_tracker/railtie.rb

@@ -3,6 +3,7 @@
 module LlmCostTracker
   class Railtie < Rails::Railtie
     generators do
+      require_relative "generators/llm_cost_tracker/add_ingestion_generator"
       require_relative "generators/llm_cost_tracker/add_period_totals_generator"
       require_relative "generators/llm_cost_tracker/add_latency_ms_generator"
       require_relative "generators/llm_cost_tracker/add_provider_response_id_generator"

data/lib/llm_cost_tracker/storage/active_record_backend.rb

@@ -3,6 +3,8 @@
 require "securerandom"
 
 require_relative "registry"
+require_relative "active_record_inbox"
+require_relative "active_record_ingestor"
 require_relative "active_record_store"
 
 module LlmCostTracker
@@ -14,7 +16,11 @@ module LlmCostTracker
         def save(event)
           require_relative "../llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
 
-          ActiveRecordStore.save(event)
+          if ActiveRecordInbox.enabled?
+            ActiveRecordInbox.save(event)
+          else
+            ActiveRecordStore.save(event)
+          end
           event
         rescue LoadError => e
           raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
@@ -49,6 +55,8 @@ module LlmCostTracker
         private
 
         def active_record_capture_check
+          return active_record_inbox_capture_check if ActiveRecordInbox.enabled?
+
           provider, model = sample_priced_identity
           response_id = "lct_verify_#{SecureRandom.hex(8)}"
           notifications = []
@@ -81,17 +89,51 @@ module LlmCostTracker
           ActiveSupport::Notifications.unsubscribe(subscription) if subscription
         end
 
+        def active_record_inbox_capture_check
+          provider, model = sample_priced_identity
+          response_id = "lct_verify_#{SecureRandom.hex(8)}"
+          notifications = []
+          subscription = subscribe_to_verification(response_id, notifications)
+
+          event = LlmCostTracker.track(
+            provider: provider,
+            model: model,
+            input_tokens: 1,
+            output_tokens: 1,
+            provider_response_id: response_id,
+            feature: VERIFY_TAG
+          )
+          LlmCostTracker.flush!
+          persisted = LlmCostTracker::LlmApiCall.where(provider_response_id: response_id).exists?
+
+          if persisted && notifications.any?
+            return active_record_capture_success("manual event emitted and persisted through durable inbox")
+          end
+
+          VerificationResult.new(:error, "active_record capture", capture_failure_message(persisted, notifications))
+        rescue LlmCostTracker::BudgetExceededError => e
+          VerificationResult.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
+        rescue LlmCostTracker::Error => e
+          VerificationResult.new(:error, "active_record capture", e.message)
+        rescue StandardError => e
+          VerificationResult.new(:error, "active_record capture", "#{e.class}: #{e.message}")
+        ensure
+          cleanup_verification_call(response_id) if response_id
+          LlmCostTracker::InboxEvent.where(event_id: event.event_id).delete_all if event
+          ActiveSupport::Notifications.unsubscribe(subscription) if subscription
+        end
+
         def subscribe_to_verification(response_id, notifications)
           ActiveSupport::Notifications.subscribe(LlmCostTracker::Tracker::EVENT_NAME) do |*, payload|
             notifications << payload if payload[:provider_response_id] == response_id
           end
         end
 
-        def active_record_capture_success
+        def active_record_capture_success(message = "manual event emitted and persisted inside rollback")
           VerificationResult.new(
             :ok,
             "active_record capture",
-            "manual event emitted and persisted inside rollback"
+            message
           )
         end
 
@@ -102,6 +144,15 @@ module LlmCostTracker
           "missing #{missing.join(' and ')} for synthetic manual event"
         end
 
+        def cleanup_verification_call(response_id)
+          relation = LlmCostTracker::LlmApiCall.where(provider_response_id: response_id)
+          rows = relation.pluck(:id, :tracked_at, :total_cost)
+          return if rows.empty?
+
+          relation.delete_all
+          ActiveRecordRollups.decrement!(rows)
+        end
+
         def sample_priced_identity
           key = LlmCostTracker::PriceRegistry.builtin_prices.find do |model_id, prices|
             model_id.include?("/") && prices[:input] && prices[:output]

data/lib/llm_cost_tracker/storage/active_record_connection_cleanup.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Storage
+    module ActiveRecordConnectionCleanup
+      def self.release!
+        ActiveRecord::Base.connection_handler.clear_active_connections!
+      rescue StandardError
+        nil
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_inbox.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+require "json"
+require "time"
+
+require_relative "../active_record_adapter"
+require_relative "../cost"
+require_relative "../event"
+require_relative "../inbox_event"
+require_relative "active_record_periods"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordInbox
+      TABLE_NAME = "llm_cost_tracker_inbox_events"
+      LEASE_TABLE_NAME = "llm_cost_tracker_ingestor_leases"
+      MAX_ATTEMPTS = 5
+
+      class << self
+        def reset!
+          remove_instance_variable(:@enabled) if instance_variable_defined?(:@enabled)
+        end
+
+        def enabled?
+          return @enabled unless @enabled.nil?
+
+          model = LlmCostTracker::LlmApiCall
+          @enabled = model.columns_hash.key?("event_id") &&
+                     model.connection.data_source_exists?(TABLE_NAME) &&
+                     model.connection.data_source_exists?(LEASE_TABLE_NAME)
+        rescue StandardError
+          @enabled = false
+        end
+
+        def save(event)
+          insert_row(row_for(event))
+          ActiveRecordIngestor.ensure_started
+          event
+        end
+
+        def pending_period_totals(periods, time:)
+          return periods.to_h { |period| [period, 0.0] } unless enabled?
+
+          periods.to_h do |period|
+            [period, pending_period_total(period, time)]
+          end
+        end
+
+        def event_from_row(row)
+          payload = JSON.parse(row.payload)
+          cost = payload["cost"] && LlmCostTracker::Cost.new(**symbolize_keys(payload["cost"]))
+
+          LlmCostTracker::Event.new(
+            event_id: payload.fetch("event_id"),
+            provider: payload.fetch("provider"),
+            model: payload.fetch("model"),
+            input_tokens: payload.fetch("input_tokens"),
+            output_tokens: payload.fetch("output_tokens"),
+            total_tokens: payload.fetch("total_tokens"),
+            cache_read_input_tokens: payload.fetch("cache_read_input_tokens"),
+            cache_write_input_tokens: payload.fetch("cache_write_input_tokens"),
+            hidden_output_tokens: payload.fetch("hidden_output_tokens"),
+            pricing_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"))
+          )
+        end
+
+        private
+
+        def row_for(event)
+          now = Time.now.utc
+          {
+            event_id: event.event_id,
+            total_cost: event.cost&.total_cost,
+            tracked_at: event.tracked_at,
+            payload: JSON.generate(payload_for(event)),
+            attempts: 0,
+            created_at: now,
+            updated_at: now
+          }
+        end
+
+        def payload_for(event)
+          {
+            event_id: event.event_id,
+            provider: event.provider,
+            model: event.model,
+            input_tokens: event.input_tokens,
+            output_tokens: event.output_tokens,
+            total_tokens: event.total_tokens,
+            cache_read_input_tokens: event.cache_read_input_tokens,
+            cache_write_input_tokens: event.cache_write_input_tokens,
+            hidden_output_tokens: event.hidden_output_tokens,
+            pricing_mode: event.pricing_mode,
+            cost: event.cost&.to_h,
+            tags: event.tags || {},
+            latency_ms: event.latency_ms,
+            stream: event.stream,
+            usage_source: event.usage_source,
+            provider_response_id: event.provider_response_id,
+            tracked_at: event.tracked_at.iso8601(6)
+          }
+        end
+
+        def insert_row(row)
+          connection = LlmCostTracker::LlmApiCall.connection
+          if connection.transaction_open? && !sqlite_database?(connection)
+            insert_with_separate_connection(row)
+          else
+            execute_insert(connection, row)
+          end
+        rescue ActiveRecord::ConnectionTimeoutError => e
+          raise LlmCostTracker::Error,
+                "ActiveRecord inbox could not checkout a separate database connection: #{e.message}"
+        end
+
+        def insert_with_separate_connection(row)
+          pool = LlmCostTracker::LlmApiCall.connection_pool
+          connection = pool.checkout
+          begin
+            connection.transaction(requires_new: true) { execute_insert(connection, row) }
+          ensure
+            pool.checkin(connection)
+          end
+        end
+
+        def execute_insert(connection, row)
+          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(TABLE_NAME)
+
+          connection.execute("INSERT INTO #{table} (#{quoted_columns}) VALUES (#{quoted_values})")
+        end
+
+        def pending_period_total(period, time)
+          LlmCostTracker::InboxEvent
+            .where("attempts < ?", MAX_ATTEMPTS)
+            .where(tracked_at: period_range(period, time))
+            .sum(:total_cost)
+            .to_f
+        end
+
+        def period_range(period, time)
+          utc_time = time.to_time.utc
+          ActiveRecordPeriods.range_start(period, utc_time)..utc_time
+        end
+
+        def symbolize_keys(hash)
+          hash.transform_keys(&:to_sym)
+        end
+
+        def sqlite_database?(connection)
+          ActiveRecordAdapter.sqlite?(connection)
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_inbox_batch.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require_relative "../inbox_event"
+require_relative "active_record_inbox"
+require_relative "active_record_store"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordInboxBatch
+      BATCH_SIZE = 100
+      LOCK_TIMEOUT_SECONDS = 30
+
+      def initialize(identity:)
+        @identity = identity
+      end
+
+      def ingest
+        rows = claim
+        return 0 if rows.empty?
+
+        valid_rows, events = decode(rows)
+        persist(valid_rows, events) if events.any?
+        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?
+        raise
+      end
+
+      def pending? = model.where("attempts < ?", ActiveRecordInbox::MAX_ATTEMPTS).exists?
+
+      def claimable? = claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists?
+
+      def mark_failed(rows, error)
+        message = "#{error.class}: #{error.message}".byteslice(0, 1_000)
+        now = Time.now.utc
+        model
+          .where(id: rows.map(&:id), locked_by: identity)
+          .update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
+      rescue StandardError
+        nil
+      end
+
+      private
+
+      attr_reader :identity
+
+      def claim
+        now = Time.now.utc
+        cutoff = now - LOCK_TIMEOUT_SECONDS
+        model.transaction do
+          rows = claimable_scope(cutoff).order(:id).limit(BATCH_SIZE).lock.to_a
+          ids = rows.map(&:id)
+          next [] if ids.empty?
+
+          updates = model.sanitize_sql_array(
+            ["locked_at = ?, locked_by = ?, attempts = attempts + 1, updated_at = ?", now, identity, now]
+          )
+          model.where(id: ids).update_all(updates)
+          model.where(id: ids, locked_by: identity).order(:id).to_a
+        end
+      end
+
+      def decode(rows)
+        valid_rows = []
+        events = []
+        rows.each do |row|
+          events << ActiveRecordInbox.event_from_row(row)
+          valid_rows << row
+        rescue StandardError => e
+          mark_failed([row], e)
+        end
+        [valid_rows, events]
+      end
+
+      def persist(rows, events)
+        LlmCostTracker::LlmApiCall.transaction do
+          ActiveRecordStore.insert_many(events)
+          model.where(id: rows.map(&:id), locked_by: identity).delete_all
+        end
+      end
+
+      def claimable_scope(cutoff)
+        model
+          .where("attempts < ?", ActiveRecordInbox::MAX_ATTEMPTS)
+          .where("locked_at IS NULL OR locked_at < ?", cutoff)
+      end
+
+      def model = LlmCostTracker::InboxEvent
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_ingestor.rb

@@ -0,0 +1,174 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+require_relative "../inbox_event"
+require_relative "../logging"
+require_relative "active_record_connection_cleanup"
+require_relative "active_record_inbox"
+require_relative "active_record_inbox_batch"
+require_relative "active_record_ingestor_lease"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordIngestor
+      INTERVAL_SECONDS = 0.25
+      IDLE_INTERVAL_SECONDS = 1.0
+      MAX_IDLE_INTERVAL_SECONDS = 5.0
+      LEASE_SECONDS = 10
+      FLUSH_TIMEOUT_SECONDS = 10
+      class << self
+        def ensure_started
+          return unless ActiveRecordInbox.enabled?
+
+          thread = mutex.synchronize do
+            reset_after_fork!
+            unless @thread&.alive?
+              @stop_requested = false
+              generation = next_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=)
+            end
+            @thread
+          end
+          wake_thread(thread)
+        end
+
+        def flush!(timeout: FLUSH_TIMEOUT_SECONDS, require_lease: false)
+          return true unless ActiveRecordInbox.enabled?
+
+          deadline = Time.now.utc + timeout
+          loop do
+            return true unless pending_events?
+            return false if Time.now.utc >= deadline
+
+            processed = ingest_once(require_lease: require_lease)
+            return false if processed.zero? && !sleep_until_next_flush(deadline)
+          end
+        end
+
+        def shutdown!(timeout: FLUSH_TIMEOUT_SECONDS, drain: true)
+          ActiveRecordInbox.reset!
+          thread = mutex.synchronize do
+            @stop_requested = true
+            next_generation
+            @thread
+          end
+          wake_thread(thread)
+          thread&.join([timeout, 1].min)
+          drain ? flush!(timeout: timeout, require_lease: true) : true
+        rescue StandardError => e
+          handle_error(e)
+          false
+        ensure
+          mutex.synchronize { @thread = nil if @thread.equal?(thread) }
+        end
+
+        def reset!
+          thread = mutex.synchronize do
+            @stop_requested = true
+            next_generation
+            thread = @thread
+            @thread = nil
+            @pid = nil
+            @identity = nil
+            thread
+          end
+          wake_thread(thread)
+        end
+
+        def ingest_once(require_lease: true)
+          return 0 unless ActiveRecordInbox.enabled?
+          return 0 unless claimable_events?
+          return 0 if require_lease && !acquire_lease
+
+          inbox_batch.ingest
+        rescue StandardError => e
+          handle_error(e)
+          0
+        end
+
+        private
+
+        def mutex = @mutex ||= Mutex.new
+
+        def run(generation)
+          idle_interval = IDLE_INTERVAL_SECONDS
+          loop do
+            break if stop_requested?(generation)
+
+            processed = executor_wrap { ingest_once }
+            ActiveRecordConnectionCleanup.release!
+            if processed.zero?
+              sleep(idle_interval)
+              idle_interval = [idle_interval * 2, MAX_IDLE_INTERVAL_SECONDS].min
+            else
+              idle_interval = IDLE_INTERVAL_SECONDS
+            end
+          rescue StandardError => e
+            handle_error(e)
+            ActiveRecordConnectionCleanup.release!
+            sleep(idle_interval)
+          end
+        ensure
+          ActiveRecordConnectionCleanup.release!
+          mutex.synchronize { @thread = nil if @thread.equal?(Thread.current) }
+        end
+
+        def sleep_until_next_flush(deadline)
+          duration = [INTERVAL_SECONDS, deadline - Time.now.utc].min
+          sleep(duration) if duration.positive?
+        end
+
+        def stop_requested?(generation) = mutex.synchronize { @stop_requested || generation != @generation }
+
+        def reset_after_fork!
+          return if @pid == Process.pid
+
+          @pid = Process.pid
+          @thread = nil
+          @identity = nil
+        end
+
+        def next_generation = (@generation = @generation.to_i + 1)
+
+        def wake_thread(thread)
+          thread&.wakeup if thread&.alive?
+        rescue ThreadError
+          nil
+        end
+
+        def executor_wrap(&)
+          executor = rails_executor
+          return yield unless executor
+
+          executor.wrap(&)
+        end
+
+        def rails_executor
+          return unless defined?(Rails) && Rails.respond_to?(:application) && Rails.application.respond_to?(:executor)
+
+          Rails.application.executor
+        rescue StandardError
+          nil
+        end
+
+        def identity = @identity ||= "pid-#{Process.pid}-#{SecureRandom.hex(6)}"
+
+        def acquire_lease = ActiveRecordIngestorLease.new(identity: identity, seconds: LEASE_SECONDS).acquire
+
+        def pending_events? = inbox_batch.pending?
+
+        def claimable_events? = inbox_batch.claimable?
+
+        def inbox_batch = ActiveRecordInboxBatch.new(identity: identity)
+
+        def handle_error(error)
+          Logging.warn("ActiveRecord ingestor failed: #{error.class}: #{error.message}") unless
+            LlmCostTracker.configuration.storage_error_behavior == :ignore
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_ingestor_lease.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require_relative "../ingestor_lease"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordIngestorLease
+      LEASE_NAME = "default"
+
+      def initialize(identity:, seconds:)
+        @identity = identity
+        @seconds = seconds
+      end
+
+      def acquire
+        now = Time.now.utc
+        LlmCostTracker::IngestorLease.transaction do
+          lease = LlmCostTracker::IngestorLease.lock.find_by(name: LEASE_NAME)
+          lease ||= LlmCostTracker::IngestorLease.create!(name: LEASE_NAME)
+          next false unless available?(lease, now)
+
+          lease.update!(locked_by: identity, locked_until: now + seconds)
+          true
+        end
+      rescue ActiveRecord::RecordNotUnique
+        false
+      end
+
+      private
+
+      attr_reader :identity, :seconds
+
+      def available?(lease, now)
+        lease.locked_by.nil? || lease.locked_by == identity || lease.locked_until.nil? || lease.locked_until < now
+      end
+    end
+  end
+end