llm_cost_tracker 0.5.2 → 0.6.0

data/lib/llm_cost_tracker/storage/active_record_rollups.rb

@@ -1,13 +1,14 @@
 # frozen_string_literal: true
 
+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)
@@ -20,22 +21,37 @@ 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 monthly_total(time: Time.now.utc)
-          period_totals(%i[monthly], time: time).fetch(:monthly)
+        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])
+          )
         end
 
-        def daily_total(time: Time.now.utc)
-          period_totals(%i[daily], time: time).fetch(:daily)
+        def decrement!(call_rows)
+          return unless period_totals_enabled?
+
+          totals = period_decrement_totals(call_rows)
+          return if totals.empty?
+
+          apply_decrements(totals)
         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?
@@ -48,22 +64,48 @@ 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
         end
 
+        def period_decrement_totals(call_rows)
+          call_rows.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |row, totals|
+            _id, tracked_at, total_cost = row
+            next unless total_cost
+
+            ActiveRecordPeriods::PERIODS.each_key do |period|
+              totals[[period, ActiveRecordPeriods.bucket(period, tracked_at)]] += BigDecimal(total_cost.to_s)
+            end
+          end
+        end
+
+        def apply_decrements(totals)
+          model = period_total_model
+          now = Time.now.utc
+
+          totals.each do |(period, period_start), amount|
+            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) = [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]]
@@ -75,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
@@ -93,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,27 +63,58 @@ 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:)
+          deleted = 0
+          loop do
+            batch = prune_batch(cutoff, batch_size)
+            deleted += batch
+            break if batch < batch_size
+          end
+          deleted
         end
 
         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
+                   .where(tracked_at: ...cutoff)
+                   .order(:id)
+                   .limit(batch_size)
+                   .lock
+                   .pluck(:id, :tracked_at, :total_cost)
+            next 0 if rows.empty?
+
+            deleted = LlmCostTracker::LlmApiCall.where(id: rows.map(&:first)).delete_all
+            ActiveRecordRollups.decrement!(rows) if deleted.positive?
+            deleted
+          end
+        end
+
         def stringify_tags(tags)
           tags.transform_keys(&:to_s).transform_values { |value| stringify_tag_value(value) }
         end

data/lib/llm_cost_tracker/storage/custom_backend.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+require_relative "registry"
+
+module LlmCostTracker
+  module Storage
+    class CustomBackend
+      class << self
+        def save(event)
+          result = LlmCostTracker.configuration.custom_storage&.call(event)
+          result == false ? false : event
+        end
+
+        def verify
+          if LlmCostTracker.configuration.custom_storage.respond_to?(:call)
+            return [
+              VerificationResult.new(
+                :ok,
+                "storage",
+                "custom storage callable configured; external sink was not invoked"
+              )
+            ]
+          end
+
+          [
+            VerificationResult.new(:error, "storage", "custom storage backend requires config.custom_storage")
+          ]
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/dispatcher.rb

@@ -1,18 +1,17 @@
 # frozen_string_literal: true
 
 require_relative "../logging"
+require_relative "registry"
+require_relative "active_record_backend"
+require_relative "custom_backend"
+require_relative "log_backend"
 
 module LlmCostTracker
   module Storage
     class Dispatcher
       class << self
         def save(event)
-          config = LlmCostTracker.configuration
-          case config.storage_backend
-          when :log           then log_event(event, config)
-          when :active_record then active_record_save(event)
-          when :custom        then custom_save(event, config)
-          end
+          backend.save(event)
         rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
           raise
         rescue StandardError => e
@@ -22,34 +21,8 @@ module LlmCostTracker
 
         private
 
-        def log_event(event, config)
-          message = "#{event.provider}/#{event.model} " \
-                    "tokens=#{event.total_tokens} " \
-                    "cost=#{log_cost_label(event)}"
-          message += " latency=#{event.latency_ms}ms" if event.latency_ms
-          message += " stream=#{event.stream}" if event.stream
-          message += " source=#{event.usage_source}" if event.usage_source
-          message += " tags=#{event.tags}" unless event.tags.empty?
-
-          Logging.log(config.log_level, message)
-          event
-        end
-
-        def log_cost_label(event) = event.cost ? "$#{format('%.6f', event.cost.total_cost)}" : "unknown"
-
-        def active_record_save(event)
-          require_relative "../llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
-          require_relative "active_record_store" unless defined?(LlmCostTracker::Storage::ActiveRecordStore)
-
-          ActiveRecordStore.save(event)
-          event
-        rescue LoadError => e
-          raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
-        end
-
-        def custom_save(event, config)
-          result = config.custom_storage&.call(event)
-          result == false ? false : event
+        def backend
+          Registry.fetch(LlmCostTracker.configuration.storage_backend)
         end
 
         def handle_error(error)
@@ -64,5 +37,9 @@ module LlmCostTracker
         end
       end
     end
+
+    Registry.register(:log, LogBackend)
+    Registry.register(:active_record, ActiveRecordBackend)
+    Registry.register(:custom, CustomBackend)
   end
 end

data/lib/llm_cost_tracker/storage/log_backend.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require_relative "../logging"
+require_relative "registry"
+
+module LlmCostTracker
+  module Storage
+    class LogBackend
+      class << self
+        def save(event)
+          config = LlmCostTracker.configuration
+          message = "#{event.provider}/#{event.model} " \
+                    "tokens=#{event.total_tokens} " \
+                    "cost=#{cost_label(event)}"
+          message += " latency=#{event.latency_ms}ms" if event.latency_ms
+          message += " stream=#{event.stream}" if event.stream
+          message += " source=#{event.usage_source}" if event.usage_source
+          message += " tags=#{event.tags}" unless event.tags.empty?
+
+          Logging.log(config.log_level, message)
+          event
+        end
+
+        def verify
+          [
+            VerificationResult.new(:ok, "storage", "log backend configured; capture writes to logs only")
+          ]
+        end
+
+        private
+
+        def cost_label(event)
+          event.cost ? "$#{format('%.6f', event.cost.total_cost)}" : "unknown"
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/storage/registry.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+require_relative "../errors"
+
+module LlmCostTracker
+  module Storage
+    VerificationResult = Data.define(:status, :name, :message)
+
+    module Registry
+      MUTEX = Monitor.new
+
+      class << self
+        def register(name, backend)
+          name = normalize_name(name)
+          validate_backend!(backend)
+          MUTEX.synchronize { @backends = backends.merge(name => backend).freeze }
+          backend
+        end
+
+        def fetch(name)
+          key = normalize_name(name)
+          backends.fetch(key) do
+            raise Error, "Unknown storage_backend: #{key.inspect}. Use one of: #{names.join(', ')}"
+          end
+        end
+
+        def registered?(name)
+          backends.key?(normalize_name(name))
+        end
+
+        def names
+          backends.keys
+        end
+
+        private
+
+        def backends
+          @backends || MUTEX.synchronize { @backends ||= {}.freeze }
+        end
+
+        def normalize_name(name)
+          name.to_sym
+        end
+
+        def validate_backend!(backend)
+          return if backend.respond_to?(:save)
+
+          raise ArgumentError, "storage backend must respond to save"
+        end
+      end
+    end
+
+    def self.register(name, backend)
+      Registry.register(name, backend)
+    end
+
+    def self.backends
+      Registry.names
+    end
+  end
+end

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

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require_relative "active_record_adapter"
+require_relative "tag_key"
+
+module LlmCostTracker
+  module TagSql
+    class << self
+      def value_expression(model, key, table_name:)
+        key = TagKey.validate!(key)
+        column = "#{table_name}.#{model.connection.quote_column_name('tags')}"
+
+        if ActiveRecordAdapter.postgresql?(model.connection)
+          json_column = model.tags_jsonb_column? ? column : "(#{column})::jsonb"
+          "#{json_column}->>#{model.connection.quote(key)}"
+        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))})"
+        end
+      end
+
+      def value_label(value)
+        value.nil? || value == "" ? "(untagged)" : value.to_s
+      end
+
+      private
+
+      def json_path(key)
+        "$.\"#{key}\""
+      end
+    end
+  end
+end

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.2"
+  VERSION = "0.6.0"
 end