llm_cost_tracker 0.6.1 → 0.7.1

data/lib/llm_cost_tracker/request_url.rb

@@ -1,20 +0,0 @@
-# frozen_string_literal: true
-
-require "uri"
-
-module LlmCostTracker
-  module RequestUrl
-    class << self
-      def label(value)
-        uri = URI.parse(value.to_s)
-        uri.query = nil
-        uri.fragment = nil
-        uri.user = nil if uri.respond_to?(:user=)
-        uri.password = nil if uri.respond_to?(:password=)
-        uri.to_s
-      rescue URI::InvalidURIError
-        value.to_s.split("?", 2).first
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/storage/active_record_backend.rb

@@ -1,166 +0,0 @@
-# frozen_string_literal: true
-
-require "securerandom"
-
-require_relative "registry"
-require_relative "active_record_inbox"
-require_relative "active_record_ingestor"
-require_relative "active_record_store"
-
-module LlmCostTracker
-  module Storage
-    class ActiveRecordBackend
-      VERIFY_TAG = "llm_cost_tracker_verify"
-
-      class << self
-        def save(event)
-          require_relative "../llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
-
-          if ActiveRecordInbox.enabled?
-            ActiveRecordInbox.save(event)
-          else
-            ActiveRecordStore.save(event)
-          end
-          event
-        rescue LoadError => e
-          raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
-        end
-
-        def verify
-          require_relative "../llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
-
-          unless LlmCostTracker::LlmApiCall.table_exists?
-            return [
-              VerificationResult.new(
-                :error,
-                "active_record",
-                "llm_api_calls table is missing; run install generator and migrate"
-              )
-            ]
-          end
-
-          [active_record_capture_check]
-        rescue LoadError => e
-          [VerificationResult.new(:error, "active_record", "unavailable: #{e.message}")]
-        rescue StandardError => e
-          [VerificationResult.new(:error, "active_record", "#{e.class}: #{e.message}")]
-        end
-
-        def prune(cutoff:, batch_size:)
-          require_relative "../llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
-
-          ActiveRecordStore.prune(cutoff: cutoff, batch_size: batch_size)
-        end
-
-        private
-
-        def active_record_capture_check
-          return active_record_inbox_capture_check if ActiveRecordInbox.enabled?
-
-          provider, model = sample_priced_identity
-          response_id = "lct_verify_#{SecureRandom.hex(8)}"
-          notifications = []
-          persisted = false
-          subscription = subscribe_to_verification(response_id, notifications)
-
-          LlmCostTracker::LlmApiCall.transaction do
-            LlmCostTracker.track(
-              provider: provider,
-              model: model,
-              input_tokens: 1,
-              output_tokens: 1,
-              provider_response_id: response_id,
-              feature: VERIFY_TAG
-            )
-            persisted = LlmCostTracker::LlmApiCall.where(provider_response_id: response_id).exists?
-            raise ActiveRecord::Rollback
-          end
-
-          return active_record_capture_success if persisted && notifications.any?
-
-          VerificationResult.new(:error, "active_record capture", capture_failure_message(persisted, notifications))
-        rescue LlmCostTracker::BudgetExceededError => e
-          VerificationResult.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
-        rescue LlmCostTracker::Error => e
-          VerificationResult.new(:error, "active_record capture", e.message)
-        rescue StandardError => e
-          VerificationResult.new(:error, "active_record capture", "#{e.class}: #{e.message}")
-        ensure
-          ActiveSupport::Notifications.unsubscribe(subscription) if subscription
-        end
-
-        def active_record_inbox_capture_check
-          provider, model = sample_priced_identity
-          response_id = "lct_verify_#{SecureRandom.hex(8)}"
-          notifications = []
-          subscription = subscribe_to_verification(response_id, notifications)
-
-          event = LlmCostTracker.track(
-            provider: provider,
-            model: model,
-            input_tokens: 1,
-            output_tokens: 1,
-            provider_response_id: response_id,
-            feature: VERIFY_TAG
-          )
-          LlmCostTracker.flush!
-          persisted = LlmCostTracker::LlmApiCall.where(provider_response_id: response_id).exists?
-
-          if persisted && notifications.any?
-            return active_record_capture_success("manual event emitted and persisted through durable inbox")
-          end
-
-          VerificationResult.new(:error, "active_record capture", capture_failure_message(persisted, notifications))
-        rescue LlmCostTracker::BudgetExceededError => e
-          VerificationResult.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
-        rescue LlmCostTracker::Error => e
-          VerificationResult.new(:error, "active_record capture", e.message)
-        rescue StandardError => e
-          VerificationResult.new(:error, "active_record capture", "#{e.class}: #{e.message}")
-        ensure
-          cleanup_verification_call(response_id) if response_id
-          LlmCostTracker::InboxEvent.where(event_id: event.event_id).delete_all if event
-          ActiveSupport::Notifications.unsubscribe(subscription) if subscription
-        end
-
-        def subscribe_to_verification(response_id, notifications)
-          ActiveSupport::Notifications.subscribe(LlmCostTracker::Tracker::EVENT_NAME) do |*, payload|
-            notifications << payload if payload[:provider_response_id] == response_id
-          end
-        end
-
-        def active_record_capture_success(message = "manual event emitted and persisted inside rollback")
-          VerificationResult.new(
-            :ok,
-            "active_record capture",
-            message
-          )
-        end
-
-        def capture_failure_message(persisted, notifications)
-          missing = []
-          missing << "notification" if notifications.empty?
-          missing << "persisted row" unless persisted
-          "missing #{missing.join(' and ')} for synthetic manual event"
-        end
-
-        def cleanup_verification_call(response_id)
-          relation = LlmCostTracker::LlmApiCall.where(provider_response_id: response_id)
-          rows = relation.pluck(:id, :tracked_at, :total_cost)
-          return if rows.empty?
-
-          relation.delete_all
-          ActiveRecordRollups.decrement!(rows)
-        end
-
-        def sample_priced_identity
-          key = LlmCostTracker::PriceRegistry.builtin_prices.find do |model_id, prices|
-            model_id.include?("/") && prices[:input] && prices[:output]
-          end&.first
-          provider, model = key.to_s.split("/", 2)
-          [provider || "openai", model || "gpt-4o-mini"]
-        end
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/storage/active_record_connection_cleanup.rb

@@ -1,13 +0,0 @@
-# frozen_string_literal: true
-
-module LlmCostTracker
-  module Storage
-    module ActiveRecordConnectionCleanup
-      def self.release!
-        ActiveRecord::Base.connection_handler.clear_active_connections!
-      rescue StandardError
-        nil
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/storage/active_record_inbox.rb

@@ -1,165 +0,0 @@
-# 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_period_totals.rb

@@ -1,84 +0,0 @@
-# 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_rollup_batch.rb

@@ -1,41 +0,0 @@
-# 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

@@ -1,42 +0,0 @@
-# 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

@@ -1,146 +0,0 @@
-# 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
-      class << self
-        def reset!
-          remove_instance_variable(:@period_totals_enabled) if instance_variable_defined?(:@period_totals_enabled)
-        end
-
-        def increment!(event)
-          return unless event.cost&.total_cost
-          return unless period_totals_enabled?
-
-          model = period_total_model
-          model.upsert_all(
-            period_rows(event),
-            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])
-          )
-        end
-
-        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 = ActiveRecordPeriods.valid_keys(periods)
-          return {} if periods.empty?
-
-          if period_totals_enabled?
-            rollup_period_totals(periods, time)
-          else
-            periods.to_h { |period| [period, fallback_period_total(period, time)] }
-          end
-        end
-
-        private
-
-        def period_rows(event)
-          ActiveRecordPeriods::PERIODS.map do |period, name|
-            {
-              period: name,
-              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, 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| 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]]
-              totals[period] = total.to_f if period
-            end
-
-          totals
-        end
-
-        def fallback_period_total(period, time)
-          LlmCostTracker::LlmApiCall
-            .where(tracked_at: ActiveRecordPeriods.range_start(period, time)..time)
-            .sum(:total_cost)
-            .to_f
-        end
-
-        def period_totals_enabled?
-          return @period_totals_enabled unless @period_totals_enabled.nil?
-
-          @period_totals_enabled =
-            LlmCostTracker::LlmApiCall.connection.data_source_exists?("llm_cost_tracker_period_totals")
-        end
-
-        def period_total_model
-          require_relative "../period_total" unless defined?(LlmCostTracker::PeriodTotal)
-
-          LlmCostTracker::PeriodTotal
-        end
-
-        def unique_by(model, column)
-          return unless model.connection.supports_insert_conflict_target?
-
-          column
-        end
-      end
-    end
-  end
-end