llm_cost_tracker 0.11.0 → 0.12.0

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -1,54 +1,84 @@
 # frozen_string_literal: true
 
-require "bigdecimal"
+require "bigdecimal/util"
 
 require_relative "period"
-require_relative "rollups/upsert_sql"
 
 module LlmCostTracker
   module Ledger
-    class Rollups
-      DEFAULT_CURRENCY = "USD"
-
+    module Rollups
       class << self
-        def increment!(event)
-          return unless event.total_cost
+        def increment!(events)
+          events = Array(events).select(&:total_cost)
+          return if events.empty?
 
-          upsert_call_rollups(period_rows(event))
+          LlmCostTracker::CallRollup.increment_all(period_rows_for_events(events))
         end
 
-        def increment_many!(events)
-          events = Array(events).select(&:total_cost)
-          return if events.empty?
+        ROLLUP_INCREMENT_ATTEMPTS = 3
+        ROLLUP_INCREMENT_BASE_DELAY_SECONDS = 0.05
+        REBUILD_INSERT_SLICE = 1_000
+        private_constant :ROLLUP_INCREMENT_ATTEMPTS, :ROLLUP_INCREMENT_BASE_DELAY_SECONDS, :REBUILD_INSERT_SLICE
+
+        def increment_safely!(events)
+          attempt = 0
+          begin
+            attempt += 1
+            increment!(events)
+          rescue StandardError => e
+            raise if LlmCostTracker::Call.connection.open_transactions.positive?
+
+            if attempt < ROLLUP_INCREMENT_ATTEMPTS
+              sleep(ROLLUP_INCREMENT_BASE_DELAY_SECONDS * (2**(attempt - 1)))
+              retry
+            end
 
-          upsert_call_rollups(period_rows_for_events(events))
+            LlmCostTracker::Logging.warn(
+              "Rollup increment failed for #{events.size} events after #{attempt} attempts: " \
+              "#{e.class}: #{e.message}"
+            )
+          end
         end
 
-        def decrement!(call_rows)
-          totals = period_decrement_totals(call_rows)
-          return if totals.empty?
+        def decrement!(records)
+          buckets = bucket_totals(records)
+          return if buckets.empty?
+
+          LlmCostTracker::CallRollup.decrement(buckets)
+        end
 
-          apply_decrements(totals)
+        def rebuild!
+          accumulated = accumulate_priced_calls
+          LlmCostTracker::CallRollup.transaction do
+            LlmCostTracker::CallRollup.delete_all
+            rows_from_buckets(accumulated).each_slice(REBUILD_INSERT_SLICE) do |rows|
+              LlmCostTracker::CallRollup.increment_all(rows)
+            end
+          end
+          accumulated.size
         end
 
         private
 
-        def period_rows(event)
-          currency = currency_from_snapshot(event.pricing_snapshot)
-          provider = event.provider.to_s
-          Period::PERIODS.map do |period, name|
-            {
-              period: name,
-              period_start: Period.bucket(period, event.tracked_at),
-              currency: currency,
-              provider: provider,
-              total_cost: event.total_cost
-            }
+        def accumulate_priced_calls
+          accumulated = Hash.new { |totals, key| totals[key] = BigDecimal("0") }
+          priced_calls.in_batches do |batch|
+            bucket_totals(batch).each { |key, total| accumulated[key] += total }
           end
+          accumulated
+        end
+
+        def priced_calls
+          LlmCostTracker::Call.where.not(total_cost: nil)
+                              .select(:id, :total_cost, :pricing_snapshot, :provider, :tracked_at)
         end
 
         def period_rows_for_events(events)
-          call_rollups(events).map do |(period, period_start, currency, provider), total_cost|
+          rows_from_buckets(bucket_totals(events))
+        end
+
+        def rows_from_buckets(buckets)
+          buckets.map do |(period, period_start, currency, provider), total_cost|
             {
               period: period,
               period_start: period_start,
@@ -59,84 +89,23 @@ module LlmCostTracker
           end
         end
 
-        def call_rollups(events)
-          events.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |event, totals|
-            currency = currency_from_snapshot(event.pricing_snapshot)
-            provider = event.provider.to_s
-            Period::PERIODS.each do |period, name|
-              key = [name, Period.bucket(period, event.tracked_at), currency, provider]
-              totals[key] += BigDecimal(event.total_cost.to_s)
-            end
-          end
-        end
+        def bucket_totals(entries)
+          entries.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |entry, totals|
+            next unless entry.total_cost
 
-        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, pricing_snapshot, provider = row
-            next unless total_cost
-
-            currency = currency_from_snapshot(pricing_snapshot)
-            provider_key = provider.to_s
-            Period::PERIODS.each_key do |period|
-              totals[[period, Period.bucket(period, tracked_at), currency, provider_key]] += total_cost
+            currency = currency_from_snapshot(entry.pricing_snapshot)
+            provider = entry.provider.to_s
+            Period::PERIODS.each do |period|
+              key = [period.to_s, Period.bucket(period, entry.tracked_at), currency, provider]
+              totals[key] += entry.total_cost.to_d
             end
           end
         end
 
-        def apply_decrements(totals)
-          now = Time.now.utc
-          buckets_by_period = totals.each_with_object({}) do |(key, amount), grouped|
-            period, period_start, currency, provider = key
-            grouped[[period, currency, provider]] ||= {}
-            grouped[[period, currency, provider]][period_start] = amount
-          end
-
-          conn = LlmCostTracker::CallRollup.connection
-          table = LlmCostTracker::CallRollup.quoted_table_name
-          period_col = conn.quote_column_name("period")
-          start_col = conn.quote_column_name("period_start")
-          currency_col = conn.quote_column_name("currency")
-          provider_col = conn.quote_column_name("provider")
-          total_col = conn.quote_column_name("total_cost")
-          updated_col = conn.quote_column_name("updated_at")
-
-          buckets_by_period.each do |(period, currency, provider), by_start|
-            case_clauses = by_start.map do |period_start, amount|
-              "WHEN #{start_col} = #{conn.quote(period_start)} THEN #{conn.quote(amount)}"
-            end.join(" ")
-            starts = by_start.keys.map { |period_start| conn.quote(period_start) }.join(", ")
-
-            conn.execute(
-              "UPDATE #{table} " \
-              "SET #{total_col} = GREATEST(0, #{total_col} - CASE #{case_clauses} ELSE 0 END), " \
-              "#{updated_col} = #{conn.quote(now)} " \
-              "WHERE #{period_col} = #{conn.quote(Period::PERIODS.fetch(period))} " \
-              "AND #{currency_col} = #{conn.quote(currency)} " \
-              "AND #{provider_col} = #{conn.quote(provider)} " \
-              "AND #{start_col} IN (#{starts})"
-            )
-          end
-        end
-
         def currency_from_snapshot(snapshot)
-          value = (snapshot.is_a?(Hash) && (snapshot["currency"] || snapshot[:currency])) || DEFAULT_CURRENCY
+          value = (snapshot.is_a?(Hash) && snapshot["currency"]) || LlmCostTracker::DEFAULT_CURRENCY
           value.to_s.upcase
         end
-
-        def upsert_call_rollups(rows)
-          LlmCostTracker::CallRollup.upsert_all(
-            rows,
-            on_duplicate: Ledger::Rollups::UpsertSql.call,
-            record_timestamps: true,
-            unique_by: call_rollups_unique_by
-          )
-        end
-
-        def call_rollups_unique_by
-          return unless LlmCostTracker::CallRollup.connection.supports_insert_conflict_target?
-
-          %i[period period_start currency provider]
-        end
       end
     end
   end

data/lib/llm_cost_tracker/ledger/schema/adapter.rb

@@ -32,22 +32,21 @@ module LlmCostTracker
             raise Error, "Unsupported database adapter: #{adapter_name(value)}. Use PostgreSQL or MySQL."
           end
 
-          def json_column_errors(column, adapter_value, column_name)
-            return [] unless column
+          PG_PERIOD_FORMATS = { day: "YYYY-MM-DD", month: "YYYY-MM" }.freeze
+          MYSQL_PERIOD_FORMATS = { day: "%Y-%m-%d", month: "%Y-%m" }.freeze
+          private_constant :PG_PERIOD_FORMATS, :MYSQL_PERIOD_FORMATS
 
-            expected_type = postgresql?(adapter_value) ? "jsonb" : "json"
-            return [] if json_column_type?(column, adapter_value)
-
-            ["#{column_name} column must use #{expected_type} (got #{column.sql_type})"]
-          end
-
-          def json_column_type?(column, adapter_value)
-            sql_type = column.sql_type.to_s.downcase
-            if postgresql?(adapter_value)
-              column.type == :jsonb || sql_type == "jsonb"
+          def period_bucket_sql(connection, period, column)
+            period = period.to_sym
+            if postgresql?(connection)
+              "TO_CHAR(DATE_TRUNC('#{period}', #{column}), '#{PG_PERIOD_FORMATS.fetch(period)}')"
+            elsif mysql?(connection)
+              "DATE_FORMAT(#{column}, '#{MYSQL_PERIOD_FORMATS.fetch(period)}')"
             else
-              column.type == :json || sql_type == "json" || sql_type == "longtext"
+              ensure_supported!(connection)
             end
+          rescue KeyError
+            raise ArgumentError, "invalid period: #{period.inspect}"
           end
 
           private

data/lib/llm_cost_tracker/ledger/schema/base.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/string/inflections"
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module Base
+        def model
+          @model ||= LlmCostTracker.const_get(detect_model_name)
+        end
+
+        def columns(*names)
+          @required_columns = names.map(&:to_s).freeze
+        end
+
+        def required_columns
+          @required_columns || [].freeze
+        end
+
+        def current_schema_errors
+          connection = model.connection
+          Adapter.ensure_supported!(connection)
+          table_name = model.table_name
+          return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+
+          columns_hash = model.columns_hash
+          cache = @schema_capabilities
+          return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns_hash)
+
+          errors = column_errors(columns_hash)
+          @schema_capabilities = { columns: columns_hash, errors: errors }
+          errors
+        end
+
+        private
+
+        def column_errors(columns_hash)
+          missing = required_columns - columns_hash.keys
+          missing.any? ? ["missing columns: #{missing.join(', ')}"] : []
+        end
+
+        def detect_model_name
+          name.delete_prefix("#{Schema.name}::").singularize
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/schema/call_line_items.rb

@@ -1,89 +1,34 @@
 # frozen_string_literal: true
 
-require_relative "adapter"
+require_relative "base"
 
 module LlmCostTracker
   module Ledger
     module Schema
       module CallLineItems
-        REQUIRED_COLUMNS = %w[
-          llm_cost_tracker_call_id
-          position
-          kind
-          direction
-          modality
-          cache_state
-          quantity
-          unit
-          rate_amount
-          rate_quantity
-          cost
-          currency
-          cost_status
-          pricing_basis
-          price_key
-          price_source
-          price_source_version
-          provider_field
-          provider_item_id
-          details
-          created_at
-        ].freeze
-
-        REQUIRED_INDEX_COLUMNS = [
-          %w[llm_cost_tracker_call_id position]
-        ].freeze
-
-        class << self
-          def current_schema_errors
-            connection = LlmCostTracker::Call.connection
-            Adapter.ensure_supported!(connection)
-            table_name = LlmCostTracker::CallLineItem.table_name
-            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
-
-            columns = LlmCostTracker::CallLineItem.columns_hash
-            cache = @schema_capabilities
-            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
-
-            errors = compute_errors(connection, table_name, columns)
-            @schema_capabilities = { columns: columns, errors: errors }
-            errors
-          end
-
-          private
-
-          def compute_errors(connection, table_name, columns)
-            errors = []
-            missing = REQUIRED_COLUMNS - columns.keys
-            errors << "missing columns: #{missing.join(', ')}" if missing.any?
-            errors.concat(Adapter.json_column_errors(columns["details"], connection, "details"))
-            errors.concat(missing_index_errors(connection, table_name))
-            errors << missing_fk_error(connection, table_name) if missing_fk?(connection, table_name)
-            errors.compact
-          end
-
-          def missing_index_errors(connection, table_name)
-            existing = connection.indexes(table_name).map { |index| Array(index.columns).map(&:to_s) }
-            REQUIRED_INDEX_COLUMNS.filter_map do |required|
-              next if existing.any? { |columns| columns == required }
-
-              "missing index on (#{required.join(', ')})"
-            end
-          end
-
-          def missing_fk?(connection, table_name)
-            connection.foreign_keys(table_name).none? do |fk|
-              fk.column.to_s == "llm_cost_tracker_call_id" &&
-                fk.to_table.to_s == "llm_cost_tracker_calls"
-            end
-          rescue NotImplementedError, NoMethodError
-            false
-          end
-
-          def missing_fk_error(_connection, _table_name)
-            "missing foreign key on llm_cost_tracker_call_id referencing llm_cost_tracker_calls"
-          end
-        end
+        extend Base
+
+        columns :llm_cost_tracker_call_id,
+                :position,
+                :kind,
+                :direction,
+                :modality,
+                :cache_state,
+                :quantity,
+                :unit,
+                :rate_amount,
+                :rate_quantity,
+                :cost,
+                :currency,
+                :cost_status,
+                :pricing_basis,
+                :price_key,
+                :price_source,
+                :price_source_version,
+                :provider_field,
+                :provider_item_id,
+                :details,
+                :created_at
       end
     end
   end

data/lib/llm_cost_tracker/ledger/schema/call_rollups.rb

@@ -1,46 +1,14 @@
 # frozen_string_literal: true
 
-require_relative "adapter"
+require_relative "base"
 
 module LlmCostTracker
   module Ledger
     module Schema
       module CallRollups
-        REQUIRED_COLUMNS = %w[period period_start currency provider total_cost created_at updated_at].freeze
-        UNIQUE_COLUMNS = %i[period period_start currency provider].freeze
+        extend Base
 
-        class << self
-          def current_schema_errors
-            connection = LlmCostTracker::CallRollup.connection
-            Adapter.ensure_supported!(connection)
-            table_name = LlmCostTracker::CallRollup.table_name
-            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
-
-            columns = LlmCostTracker::CallRollup.columns_hash
-            cache = @schema_capabilities
-            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
-
-            errors = compute_errors(connection, table_name, columns)
-            @schema_capabilities = { columns: columns, errors: errors }
-            errors
-          end
-
-          private
-
-          def compute_errors(connection, table_name, columns)
-            errors = []
-            missing = REQUIRED_COLUMNS - columns.keys
-            errors << "missing columns: #{missing.join(', ')}" if missing.any?
-            unless unique_period_index?(connection, table_name)
-              errors << "missing unique index: period, period_start, currency, provider"
-            end
-            errors
-          end
-
-          def unique_period_index?(connection, table_name)
-            connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
-          end
-        end
+        columns :period, :period_start, :currency, :provider, :total_cost, :created_at, :updated_at
       end
     end
   end

data/lib/llm_cost_tracker/ledger/schema/call_tags.rb

@@ -1,51 +1,14 @@
 # frozen_string_literal: true
 
+require_relative "base"
+
 module LlmCostTracker
   module Ledger
     module Schema
       module CallTags
-        REQUIRED_COLUMNS = %w[llm_cost_tracker_call_id key value].freeze
-
-        REQUIRED_INDEX_COLUMNS = [
-          %w[key value],
-          %w[llm_cost_tracker_call_id]
-        ].freeze
-
-        class << self
-          def current_schema_errors
-            connection = LlmCostTracker::Call.connection
-            Ledger::Schema::Adapter.ensure_supported!(connection)
-            table_name = LlmCostTracker::CallTag.table_name
-            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
-
-            columns = LlmCostTracker::CallTag.columns_hash
-            cache = @schema_capabilities
-            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
-
-            errors = compute_errors(connection, table_name, columns)
-            @schema_capabilities = { columns: columns, errors: errors }
-            errors
-          end
-
-          private
-
-          def compute_errors(connection, table_name, columns)
-            errors = []
-            missing = REQUIRED_COLUMNS - columns.keys
-            errors << "missing columns: #{missing.join(', ')}" if missing.any?
-            errors.concat(missing_index_errors(connection, table_name))
-            errors
-          end
-
-          def missing_index_errors(connection, table_name)
-            existing = connection.indexes(table_name).map { |index| Array(index.columns).map(&:to_s) }
-            REQUIRED_INDEX_COLUMNS.filter_map do |required|
-              next if existing.any? { |columns| (required - columns).empty? }
+        extend Base
 
-              "missing index on (#{required.join(', ')})"
-            end
-          end
-        end
+        columns :llm_cost_tracker_call_id, :key, :value
       end
     end
   end

data/lib/llm_cost_tracker/ledger/schema/calls.rb

@@ -1,109 +1,40 @@
 # frozen_string_literal: true
 
-require "llm_cost_tracker/ledger/schema/adapter"
+require_relative "base"
 
 module LlmCostTracker
   module Ledger
     module Schema
       module Calls
-        CURRENT_SCHEMA_COLUMNS = %w[
-          event_id
-          provider
-          model
-          input_tokens
-          output_tokens
-          total_tokens
-          cache_read_input_tokens
-          cache_write_input_tokens
-          cache_write_extended_input_tokens
-          audio_input_tokens
-          audio_output_tokens
-          image_input_tokens
-          image_output_tokens
-          hidden_output_tokens
-          total_cost
-          latency_ms
-          stream
-          usage_source
-          provider_response_id
-          provider_project_id
-          provider_api_key_id
-          provider_workspace_id
-          batch
-          pricing_mode
-          cost_status
-          pricing_snapshot
-          tracked_at
-        ].freeze
-
-        REQUIRED_INDEXES = [
-          { columns: :event_id, unique: true },
-          { columns: :tracked_at },
-          { columns: %i[provider tracked_at] },
-          { columns: %i[model tracked_at] },
-          { columns: :cost_status },
-          { columns: :provider_response_id }
-        ].freeze
-        private_constant :REQUIRED_INDEXES
-
-        class << self
-          def current_schema_errors
-            schema_capabilities.fetch(:current_schema_errors)
-          end
-
-          def missing_current_schema_columns
-            schema_capabilities.fetch(:missing_current_schema_columns)
-          end
-
-          private
-
-          def schema_capabilities
-            columns = LlmCostTracker::Call.columns_hash
-            adapter_name = LlmCostTracker::Call.connection.adapter_name
-            cache = @schema_capabilities
-
-            return cache.fetch(:values) if cache && cache.fetch(:columns).equal?(columns) &&
-                                           cache.fetch(:adapter_name) == adapter_name
-
-            values = build_schema_capabilities(columns, adapter_name)
-            @schema_capabilities = { columns: columns, adapter_name: adapter_name, values: values }
-            values
-          end
-
-          def build_schema_capabilities(columns, adapter_name)
-            Ledger::Schema::Adapter.ensure_supported!(adapter_name)
-
-            {
-              missing_current_schema_columns: missing_columns_for(columns),
-              current_schema_errors: schema_errors_for(columns, adapter_name)
-            }
-          end
-
-          def schema_errors_for(columns, adapter_name)
-            errors = []
-            missing = missing_columns_for(columns)
-            errors << "missing columns: #{missing.join(', ')}" if missing.any?
-            errors.concat(Adapter.json_column_errors(columns["pricing_snapshot"], adapter_name, "pricing_snapshot"))
-            errors.concat(missing_index_errors)
-            errors
-          end
-
-          def missing_index_errors
-            connection = LlmCostTracker::Call.connection
-            REQUIRED_INDEXES.filter_map do |spec|
-              next if connection.index_exists?(LlmCostTracker::Call.table_name, spec[:columns], **spec.except(:columns))
-
-              prefix = spec[:unique] ? "unique " : ""
-              "missing #{prefix}index: #{Array(spec[:columns]).join(', ')}"
-            end
-          rescue StandardError
-            []
-          end
-
-          def missing_columns_for(columns)
-            CURRENT_SCHEMA_COLUMNS - columns.keys
-          end
-        end
+        extend Base
+
+        columns :event_id,
+                :provider,
+                :model,
+                :input_tokens,
+                :output_tokens,
+                :total_tokens,
+                :cache_read_input_tokens,
+                :cache_write_input_tokens,
+                :cache_write_extended_input_tokens,
+                :audio_input_tokens,
+                :audio_output_tokens,
+                :image_input_tokens,
+                :image_output_tokens,
+                :hidden_output_tokens,
+                :total_cost,
+                :latency_ms,
+                :stream,
+                :usage_source,
+                :provider_response_id,
+                :provider_project_id,
+                :provider_api_key_id,
+                :provider_workspace_id,
+                :batch,
+                :pricing_mode,
+                :cost_status,
+                :pricing_snapshot,
+                :tracked_at
       end
     end
   end

data/lib/llm_cost_tracker/ledger/schema/ingestion/inbox_entries.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require_relative "../base"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module Ingestion
+        module InboxEntries
+          extend Base
+
+          columns :event_id,
+                  :total_cost,
+                  :tracked_at,
+                  :payload,
+                  :locked_at,
+                  :locked_by,
+                  :attempts,
+                  :last_error,
+                  :created_at,
+                  :updated_at
+        end
+      end
+    end
+  end
+end