llm_cost_tracker 0.5.3 → 0.6.0

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

data/lib/llm_cost_tracker/storage/active_record_rollups.rb

@@ -2,14 +2,13 @@
 
 require "bigdecimal"
 
+require_relative "active_record_periods"
+require_relative "active_record_rollup_batch"
+require_relative "active_record_rollup_upsert_sql"
+
 module LlmCostTracker
   module Storage
     class ActiveRecordRollups
-      PERIODS = {
-        monthly: "month",
-        daily:   "day"
-      }.freeze
-
       class << self
         def reset!
           remove_instance_variable(:@period_totals_enabled) if instance_variable_defined?(:@period_totals_enabled)
@@ -22,7 +21,21 @@ module LlmCostTracker
           model = period_total_model
           model.upsert_all(
             period_rows(event),
-            on_duplicate: total_upsert_sql(model),
+            on_duplicate: ActiveRecordRollupUpsertSql.call(model),
+            record_timestamps: true,
+            unique_by: unique_by(model, %i[period period_start])
+          )
+        end
+
+        def increment_many!(events)
+          events = Array(events).select { |event| event.cost&.total_cost }
+          return if events.empty?
+          return unless period_totals_enabled?
+
+          model = period_total_model
+          model.upsert_all(
+            ActiveRecordRollupBatch.rows(events),
+            on_duplicate: ActiveRecordRollupUpsertSql.call(model),
             record_timestamps: true,
             unique_by: unique_by(model, %i[period period_start])
           )
@@ -37,16 +50,8 @@ module LlmCostTracker
           apply_decrements(totals)
         end
 
-        def monthly_total(time: Time.now.utc)
-          period_totals(%i[monthly], time: time).fetch(:monthly)
-        end
-
-        def daily_total(time: Time.now.utc)
-          period_totals(%i[daily], time: time).fetch(:daily)
-        end
-
         def period_totals(periods, time: Time.now.utc)
-          periods = periods.map(&:to_sym).select { |period| PERIODS.key?(period) }
+          periods = ActiveRecordPeriods.valid_keys(periods)
           return {} if periods.empty?
 
           if period_totals_enabled?
@@ -59,10 +64,10 @@ module LlmCostTracker
         private
 
         def period_rows(event)
-          PERIODS.map do |period, name|
+          ActiveRecordPeriods::PERIODS.map do |period, name|
             {
               period: name,
-              period_start: bucket_for(period, event.tracked_at),
+              period_start: ActiveRecordPeriods.bucket(period, event.tracked_at),
               total_cost: event.cost.total_cost
             }
           end
@@ -73,8 +78,8 @@ module LlmCostTracker
             _id, tracked_at, total_cost = row
             next unless total_cost
 
-            PERIODS.each_key do |period|
-              totals[[period, bucket_for(period, tracked_at)]] += decimal(total_cost)
+            ActiveRecordPeriods::PERIODS.each_key do |period|
+              totals[[period, ActiveRecordPeriods.bucket(period, tracked_at)]] += BigDecimal(total_cost.to_s)
             end
           end
         end
@@ -84,28 +89,23 @@ module LlmCostTracker
           now = Time.now.utc
 
           totals.each do |(period, period_start), amount|
-            row = model.lock.find_by(period: PERIODS.fetch(period), period_start: period_start)
+            row = model.lock.find_by(period: ActiveRecordPeriods::PERIODS.fetch(period), period_start: period_start)
             next unless row
 
             row.update_columns(total_cost: decremented_total(row.total_cost, amount), updated_at: now)
           end
         end
 
-        def decremented_total(current, amount)
-          [decimal(current) - amount, BigDecimal("0")].max
-        end
-
-        def decimal(value)
-          BigDecimal(value.to_s)
-        end
+        def decremented_total(current, amount) = [BigDecimal(current.to_s) - amount, BigDecimal("0")].max
 
         def rollup_period_totals(periods, time)
-          buckets = periods.to_h { |period| [period, bucket_for(period, time)] }
-          index = buckets.to_h { |period, bucket| [[PERIODS.fetch(period), bucket], period] }
+          buckets = periods.to_h { |period| [period, ActiveRecordPeriods.bucket(period, time)] }
+          index = buckets.to_h { |period, bucket| [[ActiveRecordPeriods::PERIODS.fetch(period), bucket], period] }
           totals = periods.to_h { |period| [period, 0.0] }
 
           period_total_model
-            .where(period: periods.map { |period| PERIODS.fetch(period) }, period_start: buckets.values)
+            .where(period: periods.map { |period| ActiveRecordPeriods::PERIODS.fetch(period) },
+                   period_start: buckets.values)
             .pluck(:period, :period_start, :total_cost)
             .each do |name, start, total|
               period = index[[name, start.to_date]]
@@ -117,7 +117,7 @@ module LlmCostTracker
 
         def fallback_period_total(period, time)
           LlmCostTracker::LlmApiCall
-            .where(tracked_at: range_start_for(period, time)..time)
+            .where(tracked_at: ActiveRecordPeriods.range_start(period, time)..time)
             .sum(:total_cost)
             .to_f
         end
@@ -135,49 +135,11 @@ module LlmCostTracker
           LlmCostTracker::PeriodTotal
         end
 
-        def range_start_for(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_for(period, time)
-          utc_time = time.to_time.utc
-
-          case period
-          when :monthly then utc_time.beginning_of_month.to_date
-          when :daily   then utc_time.to_date
-          end
-        end
-
         def unique_by(model, column)
           return unless model.connection.supports_insert_conflict_target?
 
           column
         end
-
-        def total_upsert_sql(model)
-          Arel.sql(case model.connection.adapter_name
-                   when /mysql/i
-                     mysql_upsert_sql(model)
-                   else
-                     "total_cost = total_cost + excluded.total_cost, updated_at = excluded.updated_at"
-                   end)
-        end
-
-        def mysql_upsert_sql(model)
-          connection = model.connection
-          if connection.respond_to?(:supports_insert_raw_alias_syntax?, true) &&
-             connection.send(:supports_insert_raw_alias_syntax?)
-            values_reference = connection.quote_table_name("#{model.table_name}_values")
-            "total_cost = total_cost + #{values_reference}.total_cost, updated_at = #{values_reference}.updated_at"
-          else
-            "total_cost = total_cost + VALUES(total_cost), updated_at = VALUES(updated_at)"
-          end
-        end
       end
     end
   end

data/lib/llm_cost_tracker/storage/active_record_store.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "active_record_inbox"
+require_relative "active_record_period_totals"
 require_relative "active_record_rollups"
 
 module LlmCostTracker
@@ -11,8 +13,33 @@ module LlmCostTracker
         end
 
         def save(event)
-          tags = stringify_tags(event.tags || {})
           model = LlmCostTracker::LlmApiCall
+          attributes = attributes_for(event, model)
+
+          model.transaction do
+            call = model.create!(attributes)
+            ActiveRecordRollups.increment!(event)
+            call
+          end
+        end
+
+        def insert_many(events)
+          events = Array(events)
+          return [] if events.empty?
+
+          model = LlmCostTracker::LlmApiCall
+          insertable = new_events(model, events)
+
+          if insertable.any?
+            rows = insertable.map { |event| attributes_for(event, model) }
+            model.insert_all!(rows, **insert_options)
+            ActiveRecordRollups.increment_many!(insertable)
+          end
+          events
+        end
+
+        def attributes_for(event, model = LlmCostTracker::LlmApiCall)
+          tags = stringify_tags(event.tags || {})
           columns = model.columns_hash
 
           attributes = {
@@ -27,6 +54,7 @@ module LlmCostTracker
             tags:          tags_for_storage(tags, model),
             tracked_at:    event.tracked_at
           }
+          attributes[:event_id] = event.event_id if columns.key?("event_id")
           optional_attributes(event).each do |name, value|
             attributes[name] = value if columns.key?(name.to_s)
           end
@@ -35,23 +63,19 @@ module LlmCostTracker
           attributes[:usage_source] = event.usage_source if columns.key?("usage_source")
           attributes[:provider_response_id] = event.provider_response_id if columns.key?("provider_response_id")
 
-          model.transaction do
-            call = model.create!(attributes)
-            ActiveRecordRollups.increment!(event)
-            call
-          end
+          attributes
         end
 
         def monthly_total(time: Time.now.utc)
-          ActiveRecordRollups.monthly_total(time: time)
+          period_totals(%i[monthly], time: time).fetch(:monthly)
         end
 
         def daily_total(time: Time.now.utc)
-          ActiveRecordRollups.daily_total(time: time)
+          period_totals(%i[daily], time: time).fetch(:daily)
         end
 
         def period_totals(periods, time: Time.now.utc)
-          ActiveRecordRollups.period_totals(periods, time: time)
+          ActiveRecordPeriodTotals.call(periods, time: time)
         end
 
         def prune(cutoff:, batch_size:)
@@ -66,6 +90,15 @@ module LlmCostTracker
 
         private
 
+        def new_events(model, events)
+          return events unless model.columns_hash.key?("event_id")
+
+          existing_ids = model.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
+          events.reject { |event| existing_ids.include?(event.event_id) }
+        end
+
+        def insert_options = { record_timestamps: true, returning: false }
+
         def prune_batch(cutoff, batch_size)
           LlmCostTracker::LlmApiCall.transaction do
             rows = LlmCostTracker::LlmApiCall

data/lib/llm_cost_tracker/stream_collector.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "json"
 require "monitor"
 
 require_relative "stream_capture"
@@ -164,10 +163,9 @@ module LlmCostTracker
     end
 
     def capture_event(data, type:)
-      copied = ValueHelpers.deep_dup(data)
-      size = event_bytes(copied, type)
+      size = event_bytes(data, type)
       if @captured_bytes + size <= StreamCapture::LIMIT_BYTES
-        @events << { event: type, data: copied }
+        @events << { event: type, data: ValueHelpers.deep_dup(data) }
         @captured_bytes += size
       else
         @overflowed = true
@@ -176,9 +174,22 @@ module LlmCostTracker
     end
 
     def event_bytes(data, type)
-      JSON.generate(event: type, data: data).bytesize
-    rescue JSON::GeneratorError, TypeError
-      type.to_s.bytesize + data.to_s.bytesize
+      type.to_s.bytesize + estimated_bytes(data) + 32
+    end
+
+    def estimated_bytes(value)
+      case value
+      when Hash
+        value.sum { |key, nested| estimated_bytes(key) + estimated_bytes(nested) + 4 }
+      when Array
+        value.sum { |nested| estimated_bytes(nested) + 2 }
+      when String
+        value.bytesize + 2
+      when Numeric, true, false, nil
+        value.to_s.bytesize
+      else
+        value.to_s.bytesize + 2
+      end
     end
 
     def error_metadata(errored) = errored ? { stream_errored: true } : {}

data/lib/llm_cost_tracker/tag_sql.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "active_record_adapter"
 require_relative "tag_key"
 
 module LlmCostTracker
@@ -9,11 +10,10 @@ module LlmCostTracker
         key = TagKey.validate!(key)
         column = "#{table_name}.#{model.connection.quote_column_name('tags')}"
 
-        case model.connection.adapter_name
-        when /postgres/i
+        if ActiveRecordAdapter.postgresql?(model.connection)
           json_column = model.tags_jsonb_column? ? column : "(#{column})::jsonb"
           "#{json_column}->>#{model.connection.quote(key)}"
-        when /mysql/i
+        elsif ActiveRecordAdapter.mysql?(model.connection)
           "JSON_UNQUOTE(JSON_EXTRACT(#{column}, #{model.connection.quote(json_path(key))}))"
         else
           "json_extract(#{column}, #{model.connection.quote(json_path(key))})"

data/lib/llm_cost_tracker/tags_column.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "active_record_adapter"
+
 module LlmCostTracker
   module TagsColumn
     USAGE_BREAKDOWN_COLUMNS = %w[
@@ -79,7 +81,11 @@ module LlmCostTracker
     def build_lct_schema_capabilities(columns, adapter_name)
       tag_column = columns["tags"]
       tags_jsonb = tag_column && (tag_column.type == :jsonb || tag_column.sql_type.to_s.downcase == "jsonb")
-      tags_mysql_json = tag_column && !tags_jsonb && tag_column.type == :json && adapter_name.match?(/mysql/i)
+      tags_mysql_json =
+        tag_column &&
+        !tags_jsonb &&
+        tag_column.type == :json &&
+        ActiveRecordAdapter.mysql?(adapter_name)
 
       {
         tags_jsonb: tags_jsonb ? true : false,

data/lib/llm_cost_tracker/tracker.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "securerandom"
+
 require_relative "storage/dispatcher"
 
 module LlmCostTracker
@@ -74,6 +76,7 @@ module LlmCostTracker
       def build_event(provider:, model:, usage:, cost_data:, metadata:, latency_ms:, stream:, usage_source:,
                       provider_response_id:)
         Event.new(
+          event_id: SecureRandom.uuid,
           provider: provider,
           model: model,
           input_tokens: usage[:input_tokens],

data/lib/llm_cost_tracker/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LlmCostTracker
-  VERSION = "0.5.3"
+  VERSION = "0.6.0"
 end