llm_cost_tracker 0.5.2 → 0.6.0

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

data/lib/llm_cost_tracker/storage/active_record_period_totals.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+require_relative "active_record_inbox"
+require_relative "active_record_periods"
+require_relative "active_record_rollups"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordPeriodTotals
+      def self.call(periods, time:)
+        new(periods, time: time).totals
+      end
+
+      def initialize(periods, time:)
+        @periods = ActiveRecordPeriods.valid_keys(periods)
+        @time = time
+      end
+
+      def totals
+        return {} if periods.empty?
+        return ActiveRecordRollups.period_totals(periods, time: time) unless ActiveRecordInbox.enabled?
+
+        snapshot_totals
+      end
+
+      private
+
+      attr_reader :periods, :time
+
+      def snapshot_totals
+        values = periods.to_h { |period| [period, 0.0] }
+        connection.select_all(snapshot_sql).each do |row|
+          values[row.fetch("period_key").to_sym] = row.fetch("total_cost").to_f
+        end
+        values
+      end
+
+      def snapshot_sql
+        periods.map { |period| snapshot_select(period) }.join(" UNION ALL ")
+      end
+
+      def snapshot_select(period)
+        start = range_start_for(period)
+        "SELECT #{connection.quote(period.to_s)} AS period_key, " \
+          "(#{stored_total_sql(period, start)}) + (#{pending_total_sql(start)}) AS total_cost"
+      end
+
+      def stored_total_sql(period, start)
+        period_totals_table? ? rollup_total_sql(period) : ledger_total_sql(start)
+      end
+
+      def rollup_total_sql(period)
+        table = connection.quote_table_name("llm_cost_tracker_period_totals")
+        "COALESCE((SELECT total_cost FROM #{table} " \
+          "WHERE period = #{connection.quote(ActiveRecordPeriods::PERIODS.fetch(period))} " \
+          "AND period_start = #{connection.quote(ActiveRecordPeriods.bucket(period, time))} LIMIT 1), 0)"
+      end
+
+      def ledger_total_sql(start)
+        table = LlmCostTracker::LlmApiCall.quoted_table_name
+        total_cost = connection.quote_column_name("total_cost")
+        tracked_at = connection.quote_column_name("tracked_at")
+        "COALESCE((SELECT SUM(#{total_cost}) FROM #{table} " \
+          "WHERE #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)}), 0)"
+      end
+
+      def pending_total_sql(start)
+        table = connection.quote_table_name(ActiveRecordInbox::TABLE_NAME)
+        total_cost = connection.quote_column_name("total_cost")
+        tracked_at = connection.quote_column_name("tracked_at")
+        attempts = connection.quote_column_name("attempts")
+        "COALESCE((SELECT SUM(#{total_cost}) FROM #{table} " \
+          "WHERE #{attempts} < #{ActiveRecordInbox::MAX_ATTEMPTS} " \
+          "AND #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)}), 0)"
+      end
+
+      def period_totals_table? = connection.data_source_exists?("llm_cost_tracker_period_totals")
+
+      def range_start_for(period) = ActiveRecordPeriods.range_start(period, time)
+
+      def connection = LlmCostTracker::LlmApiCall.connection
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_periods.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Storage
+    module ActiveRecordPeriods
+      PERIODS = {
+        monthly: "month",
+        daily:   "day"
+      }.freeze
+
+      module_function
+
+      def valid_keys(periods)
+        periods.map(&:to_sym).select { |period| PERIODS.key?(period) }
+      end
+
+      def range_start(period, time)
+        utc_time = time.to_time.utc
+
+        case period
+        when :monthly then utc_time.beginning_of_month
+        when :daily then utc_time.beginning_of_day
+        end
+      end
+
+      def bucket(period, time)
+        range_start(period, time).to_date
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_rollup_batch.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "active_record_periods"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordRollupBatch
+      def self.rows(events)
+        new(events).rows
+      end
+
+      def initialize(events)
+        @events = events
+      end
+
+      def rows
+        totals.map do |(period, period_start), total_cost|
+          {
+            period: period,
+            period_start: period_start,
+            total_cost: total_cost
+          }
+        end
+      end
+
+      private
+
+      attr_reader :events
+
+      def totals
+        events.each_with_object(Hash.new { |hash, key| hash[key] = BigDecimal("0") }) do |event, rows|
+          ActiveRecordPeriods::PERIODS.each do |period, name|
+            rows[[name, ActiveRecordPeriods.bucket(period, event.tracked_at)]] += BigDecimal(event.cost.total_cost.to_s)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/active_record_rollup_upsert_sql.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require_relative "../active_record_adapter"
+
+module LlmCostTracker
+  module Storage
+    class ActiveRecordRollupUpsertSql
+      def self.call(model)
+        new(model).call
+      end
+
+      def initialize(model)
+        @model = model
+      end
+
+      def call
+        return Arel.sql(mysql_sql) if ActiveRecordAdapter.mysql?(connection)
+        return Arel.sql(postgres_sql) if ActiveRecordAdapter.postgresql?(connection)
+
+        Arel.sql("total_cost = total_cost + excluded.total_cost, updated_at = excluded.updated_at")
+      end
+
+      private
+
+      attr_reader :model
+
+      def postgres_sql
+        total_cost = connection.quote_column_name("total_cost")
+        updated_at = connection.quote_column_name("updated_at")
+
+        "#{total_cost} = #{model.quoted_table_name}.#{total_cost} + excluded.#{total_cost}, " \
+          "#{updated_at} = excluded.#{updated_at}"
+      end
+
+      def mysql_sql
+        "total_cost = total_cost + VALUES(total_cost), updated_at = VALUES(updated_at)"
+      end
+
+      def connection = model.connection
+    end
+  end
+end