llm_cost_tracker 0.7.0 → 0.7.2

data/lib/llm_cost_tracker/ledger/period/totals.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require_relative "../period"
+
+module LlmCostTracker
+  module Ledger
+    module Period
+      class Totals
+        def self.call(periods, time:)
+          new(periods, time: time).totals
+        end
+
+        def initialize(periods, time:)
+          @periods = Period.valid_keys(periods)
+          @time = time
+        end
+
+        def totals
+          return {} if periods.empty?
+
+          snapshot_totals
+        end
+
+        private
+
+        attr_reader :periods, :time
+
+        def snapshot_totals
+          values = periods.to_h { |period| [period, 0.0] }
+          sql = periods.map { |period| snapshot_select(period) }.join(" UNION ALL ")
+          LlmCostTracker::Ledger::Call.find_by_sql(sql).each do |row|
+            values[row.period_key.to_sym] = row.total_cost.to_f
+          end
+          values
+        end
+
+        def snapshot_select(period)
+          start = Period.range_start(period, time)
+          "SELECT #{connection.quote(period.to_s)} AS period_key, " \
+            "(#{rollup_total_sql(period)}) + (#{pending_total_sql(start)}) AS total_cost"
+        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(Period::PERIODS.fetch(period))} " \
+            "AND period_start = #{connection.quote(Period.bucket(period, time))} LIMIT 1), 0)"
+        end
+
+        def pending_total_sql(start)
+          table = connection.quote_table_name(Ingestion::Event.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} < #{Ingestion::Event::MAX_ATTEMPTS} " \
+            "AND #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)}), 0)"
+        end
+
+        def connection
+          LlmCostTracker::Ledger::Call.connection
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/{storage/active_record_periods.rb → ledger/period.rb}

@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
 module LlmCostTracker
-  module Storage
-    module ActiveRecordPeriods
+  module Ledger
+    module Period
       PERIODS = {
         monthly: "month",
         daily:   "day"

data/lib/llm_cost_tracker/ledger/rollups/batch.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "../period"
+
+module LlmCostTracker
+  module Ledger
+    class Rollups
+      class Batch
+        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|
+            Period::PERIODS.each do |period, name|
+              rows[[name, Period.bucket(period, event.tracked_at)]] += BigDecimal(event.total_cost.to_s)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/rollups/upsert_sql.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative "../schema/adapter"
+
+module LlmCostTracker
+  module Ledger
+    class Rollups
+      class UpsertSql
+        def self.call(model)
+          new(model).call
+        end
+
+        def initialize(model)
+          @model = model
+        end
+
+        def call
+          return Arel.sql(mysql_sql) if Ledger::Schema::Adapter.mysql?(connection)
+          return Arel.sql(postgres_sql) if Ledger::Schema::Adapter.postgresql?(connection)
+
+          Ledger::Schema::Adapter.ensure_supported!(connection)
+        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
+  end
+end

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "period"
+require_relative "rollups/batch"
+require_relative "rollups/upsert_sql"
+
+module LlmCostTracker
+  module Ledger
+    class Rollups
+      class << self
+        def increment!(event)
+          return unless event.total_cost
+
+          Period::Total.upsert_all(
+            period_rows(event),
+            on_duplicate: Ledger::Rollups::UpsertSql.call(Period::Total),
+            record_timestamps: true,
+            unique_by: unique_by(Period::Total, %i[period period_start])
+          )
+        end
+
+        def increment_many!(events)
+          events = Array(events).select(&:total_cost)
+          return if events.empty?
+
+          Period::Total.upsert_all(
+            Ledger::Rollups::Batch.rows(events),
+            on_duplicate: Ledger::Rollups::UpsertSql.call(Period::Total),
+            record_timestamps: true,
+            unique_by: unique_by(Period::Total, %i[period period_start])
+          )
+        end
+
+        def decrement!(call_rows)
+          totals = period_decrement_totals(call_rows)
+          return if totals.empty?
+
+          apply_decrements(totals)
+        end
+
+        private
+
+        def period_rows(event)
+          Period::PERIODS.map do |period, name|
+            {
+              period: name,
+              period_start: Period.bucket(period, event.tracked_at),
+              total_cost: event.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
+
+            Period::PERIODS.each_key do |period|
+              totals[[period, Period.bucket(period, tracked_at)]] += BigDecimal(total_cost.to_s)
+            end
+          end
+        end
+
+        def apply_decrements(totals)
+          now = Time.now.utc
+
+          totals.each do |(period, period_start), amount|
+            row = Period::Total.lock.find_by(period: Period::PERIODS.fetch(period),
+                                             period_start: period_start)
+            next unless row
+
+            row.update_columns(total_cost: [BigDecimal(row.total_cost.to_s) - amount, BigDecimal("0")].max,
+                               updated_at: now)
+          end
+        end
+
+        def unique_by(model, column)
+          return unless model.connection.supports_insert_conflict_target?
+
+          column
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require_relative "../../errors"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module Adapter
+        MYSQL_ADAPTERS = %w[
+          ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter
+          ActiveRecord::ConnectionAdapters::Mysql2Adapter
+          ActiveRecord::ConnectionAdapters::TrilogyAdapter
+        ].freeze
+        POSTGRESQL_ADAPTERS = %w[
+          ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+        ].freeze
+        MYSQL_PATTERN = /mysql|trilogy|mariadb/i
+        POSTGRESQL_PATTERN = /postgres/i
+
+        class << self
+          def mysql?(value)
+            adapter_instance?(value, MYSQL_ADAPTERS) || adapter_name(value).match?(MYSQL_PATTERN)
+          end
+
+          def postgresql?(value)
+            adapter_instance?(value, POSTGRESQL_ADAPTERS) || adapter_name(value).match?(POSTGRESQL_PATTERN)
+          end
+
+          def ensure_supported!(value)
+            return if mysql?(value) || postgresql?(value)
+
+            raise Error, "Unsupported database adapter: #{adapter_name(value)}. Use PostgreSQL or MySQL."
+          end
+
+          private
+
+          def adapter_instance?(value, class_names)
+            class_names.any? do |class_name|
+              adapter_class = class_name.safe_constantize
+              adapter_class && value.is_a?(adapter_class)
+            end
+          end
+
+          def adapter_name(value)
+            value.try(:adapter_name).presence || value.to_s
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+require "llm_cost_tracker/ledger/schema/adapter"
+
+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_1h_input_tokens
+          hidden_output_tokens
+          input_cost
+          output_cost
+          total_cost
+          cache_read_input_cost
+          cache_write_input_cost
+          cache_write_1h_input_cost
+          latency_ms
+          stream
+          usage_source
+          provider_response_id
+          pricing_mode
+          tags
+          tracked_at
+        ].freeze
+
+        class << self
+          def current_schema?
+            current_schema_errors.empty?
+          end
+
+          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 = Ledger::Call.columns_hash
+            adapter_name = Ledger::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?
+
+            tag_column = columns["tags"]
+            if tag_column
+              postgresql = Ledger::Schema::Adapter.postgresql?(adapter_name)
+              expected_type = postgresql ? "jsonb" : "json"
+              valid_type =
+                if postgresql
+                  tag_column.type == :jsonb || tag_column.sql_type.to_s.downcase == "jsonb"
+                else
+                  tag_column.type == :json
+                end
+
+              errors << "tags column must use #{expected_type}" unless valid_type
+            end
+
+            errors
+          end
+
+          def missing_columns_for(columns)
+            CURRENT_SCHEMA_COLUMNS - columns.keys
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module PeriodTotals
+        REQUIRED_COLUMNS = %w[period period_start total_cost].freeze
+        UNIQUE_COLUMNS = %i[period period_start].freeze
+
+        class << self
+          def current_schema_errors
+            connection = Ledger::Call.connection
+            table_name = Ledger::Period::Total.table_name
+            return ["llm_cost_tracker_period_totals table is missing"] unless connection.data_source_exists?(table_name)
+
+            errors = []
+            missing = REQUIRED_COLUMNS - Ledger::Period::Total.columns_hash.keys
+            errors << "missing columns: #{missing.join(', ')}" if missing.any?
+            errors << "missing unique index: period, period_start" unless unique_period_index?(connection, table_name)
+            errors
+          end
+
+          private
+
+          def unique_period_index?(connection, table_name)
+            connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require_relative "../pricing"
+require_relative "rollups"
+
+module LlmCostTracker
+  module Ledger
+    class Store
+      class << self
+        def insert_many(events)
+          events = Array(events)
+          return [] if events.empty?
+
+          model = LlmCostTracker::Ledger::Call
+          insertable = new_events(model, events)
+
+          if insertable.any?
+            rows = insertable.map { |event| attributes_for(event) }
+            model.insert_all!(rows, record_timestamps: true, returning: false)
+            Ledger::Rollups.increment_many!(insertable)
+          end
+          events
+        end
+
+        private
+
+        def attributes_for(event)
+          tags = (event.tags || {}).transform_keys(&:to_s).transform_values { |value| stringify_tag_value(value) }
+          usage = event.token_usage.stored_attributes
+
+          attributes = {
+            event_id: event.event_id,
+            provider: event.provider,
+            model: event.model,
+            tags: tags,
+            tracked_at: event.tracked_at,
+            pricing_mode: event.pricing_mode,
+            latency_ms: event.latency_ms,
+            stream: event.stream,
+            usage_source: event.usage_source,
+            provider_response_id: event.provider_response_id
+          }
+
+          attributes.merge(usage).merge(Pricing.stored_cost_attributes(event.cost || {}))
+        end
+
+        def new_events(model, events)
+          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 stringify_tag_value(value)
+          return value.transform_values { |nested| stringify_tag_value(nested) } if value.is_a?(Hash)
+
+          value.to_s
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/tags/query.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "json"
+
+require_relative "../schema/adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Tags
+      module Query
+        class << self
+          def apply(model, tags)
+            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
+            return model.all if normalized_tags.empty?
+
+            connection = model.connection
+            json = normalized_tags.to_json
+
+            if Schema::Adapter.postgresql?(connection)
+              model.where("tags @> ?::jsonb", json)
+            else
+              model.where("JSON_CONTAINS(tags, ?)", json)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/tags/sql.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require_relative "../schema/adapter"
+require_relative "../../tags/key"
+
+module LlmCostTracker
+  module Ledger
+    module Tags
+      module Sql
+        class << self
+          def value_expression(model, key, table_name:)
+            key = LlmCostTracker::Tags::Key.validate!(key)
+            column = "#{table_name}.#{model.connection.quote_column_name('tags')}"
+
+            if Ledger::Schema::Adapter.postgresql?(model.connection)
+              "#{column}->>#{model.connection.quote(key)}"
+            elsif Ledger::Schema::Adapter.mysql?(model.connection)
+              "JSON_UNQUOTE(JSON_EXTRACT(#{column}, #{model.connection.quote(json_path(key))}))"
+            else
+              Ledger::Schema::Adapter.ensure_supported!(model.connection)
+            end
+          end
+
+          private
+
+          def json_path(key)
+            "$.\"#{key}\""
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative "ledger/schema/adapter"
+require_relative "ledger/schema/calls"
+require_relative "ledger/schema/period_totals"
+require_relative "ledger/tags/query"
+require_relative "ledger/tags/sql"
+require_relative "ledger/period"
+require_relative "ledger/rollups/batch"
+require_relative "ledger/rollups/upsert_sql"
+require_relative "ledger/rollups"
+require_relative "ledger/store"
+require_relative "ledger/period/totals"

data/lib/llm_cost_tracker/logging.rb

@@ -19,9 +19,10 @@ module LlmCostTracker
 
       def log(level, message)
         message = prefixed(message)
+        logger = Rails.logger
 
-        if rails_logger
-          rails_logger.public_send(level, message)
+        if logger
+          logger.try(level, message)
         else
           Kernel.warn(message)
         end
@@ -35,10 +36,6 @@ module LlmCostTracker
 
         "#{PREFIX} #{message}"
       end
-
-      def rails_logger
-        Rails.logger if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
-      end
     end
   end
 end