llm_cost_tracker 0.13.0 → 0.14.0

data/lib/llm_cost_tracker/budget/per_tag.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+require_relative "../ledger/schema/adapter"
+require_relative "../ledger/tags/encoding"
+
+module LlmCostTracker
+  module Budget
+    module PerTag
+      COST_COLUMN = "total_cost"
+      TIME_COLUMN = "tracked_at"
+      WINDOW_STARTS = { daily: :beginning_of_day, weekly: :beginning_of_week, monthly: :beginning_of_month }.freeze
+      WINDOW_NEXTS = { daily: :next_day, weekly: :next_week, monthly: :next_month }.freeze
+      SLOW_READ_SECONDS = 0.1
+      DEFAULT_BACKFILL_BATCH = 5_000
+
+      Rule = Data.define(:key, :value, :windows, :behavior, :on_exceeded)
+
+      class << self
+        def configured
+          LlmCostTracker.configuration.budgets.per_tag
+        end
+
+        def active?
+          return false if configured.empty?
+          return true if columns?
+
+          warn_missing_columns
+          false
+        end
+
+        def blocking?
+          active? && configured.each_value.any? { |entry| behavior_for(entry) == :block_requests }
+        end
+
+        def rules_for(tags, blocking_only: false)
+          return [] unless active?
+
+          normalized = (tags || {}).to_h.transform_keys(&:to_s)
+          configured.filter_map do |key, entry|
+            raw = normalized[key]
+            next if raw.nil?
+            next if blocking_only && behavior_for(entry) != :block_requests
+
+            Rule.new(
+              key: key,
+              value: Ledger::Tags::Encoding.encode(raw),
+              windows: entry[:windows],
+              behavior: behavior_for(entry),
+              on_exceeded: on_exceeded_for(entry)
+            )
+          end
+        end
+
+        def rules_for_events(events)
+          events.each_with_object({}) do |event, grouped|
+            rules_for(event.tags).each { |rule| (grouped[rule] ||= []) << event }
+          end
+        end
+
+        def spend(key, value, window, time:)
+          spend_by_value(key, [value], window, window_start(window, time)).fetch(value, 0).to_d
+        end
+
+        def spend_by_value(key, values, window, bucket)
+          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          totals = LlmCostTracker::CallTag
+                   .where(key: key, value: values, TIME_COLUMN => window_range(window, bucket))
+                   .group(:value).sum(COST_COLUMN)
+          warn_slow_read(key, window, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at)
+          totals
+        end
+
+        def window_start(window, time)
+          time.to_time.utc.public_send(WINDOW_STARTS.fetch(window))
+        end
+
+        def backfill(batch_size: DEFAULT_BACKFILL_BATCH)
+          return 0 unless columns?
+
+          filled = 0
+          loop do
+            copied = copy_next_batch(batch_size)
+            break if copied.zero?
+
+            filled += copied
+          end
+          filled
+        end
+
+        def columns?
+          (LlmCostTracker::CallTag.column_names & [COST_COLUMN, TIME_COLUMN]).size == 2
+        end
+
+        private
+
+        def window_range(window, bucket)
+          bucket...bucket.public_send(WINDOW_NEXTS.fetch(window))
+        end
+
+        def copy_next_batch(batch_size)
+          ids = LlmCostTracker::CallTag.where(TIME_COLUMN => nil).limit(batch_size).pluck(:id)
+          return 0 if ids.empty?
+
+          LlmCostTracker::CallTag.connection.update(copy_sql(ids))
+        end
+
+        def copy_sql(ids)
+          tags = LlmCostTracker::CallTag.quoted_table_name
+          calls = LlmCostTracker::Call.quoted_table_name
+          list = ids.join(",")
+          return <<~SQL.squish unless Ledger::Schema::Adapter.postgresql?(LlmCostTracker::CallTag.connection)
+            UPDATE #{tags} t JOIN #{calls} c ON c.id = t.llm_cost_tracker_call_id
+               SET t.#{COST_COLUMN} = c.#{COST_COLUMN}, t.#{TIME_COLUMN} = c.#{TIME_COLUMN}
+             WHERE t.id IN (#{list})
+          SQL
+
+          <<~SQL.squish
+            UPDATE #{tags} AS t
+               SET #{COST_COLUMN} = c.#{COST_COLUMN}, #{TIME_COLUMN} = c.#{TIME_COLUMN}
+              FROM #{calls} AS c
+             WHERE c.id = t.llm_cost_tracker_call_id AND t.id IN (#{list})
+          SQL
+        end
+
+        def behavior_for(entry)
+          entry[:behavior] || LlmCostTracker.configuration.budgets.exceeded_behavior
+        end
+
+        def on_exceeded_for(entry)
+          entry[:on_exceeded] || LlmCostTracker.configuration.budgets.on_exceeded
+        end
+
+        def warn_slow_read(key, window, seconds)
+          return if seconds < SLOW_READ_SECONDS
+
+          @slow_keys ||= Set.new
+          return unless @slow_keys.add?(key)
+
+          Logging.warn(
+            "config.budgets.per_tag[#{key.inspect}] #{window} read took #{(seconds * 1000).round} ms. " \
+            "A tag with few distinct values covers most of the ledger, so its budget check cannot use " \
+            "an index effectively. Budget high-cardinality tags such as a tenant or user id."
+          )
+        end
+
+        def warn_missing_columns
+          return if @missing_columns_warned
+
+          @missing_columns_warned = true
+          Logging.warn(
+            "config.budgets.per_tag is set but llm_cost_tracker_call_tags is missing " \
+            "#{COST_COLUMN} / #{TIME_COLUMN}; per-tag budgets are not enforced. Run the " \
+            "upgrade_per_tag_budgets generator and migrate, or clear config.budgets.per_tag."
+          )
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/budget.rb

@@ -4,42 +4,94 @@ require "bigdecimal"
 
 require_relative "ledger"
 require_relative "pricing/estimator"
+require_relative "budget/per_tag"
+require_relative "tags/context"
 
 module LlmCostTracker
   module Budget
     BUDGET_TYPE_TO_PERIOD = { monthly: :month, daily: :day }.freeze
 
     class << self
-      def enforce!(provider: nil, model: nil, request: nil, estimate: nil, force: false)
+      def enforce!(provider: nil, model: nil, request: nil, estimate: nil, tags: nil, force: false)
         config = LlmCostTracker.configuration
         return unless config.enabled
-        return unless force || config.budget_exceeded_behavior == :block_requests
+
+        globally = force || config.budgets.exceeded_behavior == :block_requests
+        per_tag = force || PerTag.blocking?
+        return unless globally || per_tag
 
         estimate ||= estimate_cost(provider: provider, model: model, request: request)
-        raise_per_call_pre_send(estimate, config.per_call_budget) if config.per_call_budget && estimate.positive?
-
-        check_windowed({ monthly: config.monthly_budget, daily: config.daily_budget }.compact,
-                       time: Time.now.utc,
-estimate: estimate) do |budget_type, total, budget|
-          raise BudgetExceededError.new(**budget_payload(
-            budget_type: budget_type, total: total, budget: budget, last_event: nil, stage: :pre_send
-          ))
+        now = Time.now.utc
+        enforce_globally(config, estimate: estimate, time: now) if globally
+        return unless per_tag
+
+        check_per_tag(tags || Tags::Context.tags,
+                      time: now,
+                      estimate: estimate,
+                      blocking_only: !force) do |rule, window, total, limit|
+          raise_pre_send(budget_type: window, total: total, budget: limit, scope: scope_for(rule))
         end
       end
 
-      def check!(event)
+      def check!(event, behavior_override: nil)
         config = LlmCostTracker.configuration
         return unless event.total_cost
 
-        check_per_call_budget(event, config)
-        check_windowed({ daily: config.daily_budget, monthly: config.monthly_budget }.compact,
+        check_per_call_budget(event, config, behavior_override)
+        check_windowed({ daily: config.budgets.daily, monthly: config.budgets.monthly }.compact,
                        time: event.tracked_at) do |budget_type, total, budget|
-          handle_exceeded(budget_type: budget_type, total: total, budget: budget, last_event: event)
+          handle_exceeded(budget_type: budget_type,
+                          total: total,
+                          budget: budget,
+                          previous_total: total - event.total_cost,
+                          last_event: event,
+                          behavior: behavior_override)
+        end
+      end
+
+      def check_persisted!(events, behavior_override: nil)
+        by_rule = PerTag.rules_for_events(events.select(&:total_cost))
+        by_rule = by_rule.reject { |rule, _| rule.on_exceeded.nil? } if behavior_override == :notify
+        window_buckets(by_rule).each do |(key, window, bucket), scored|
+          totals = PerTag.spend_by_value(key, scored.keys.map(&:value), window, bucket)
+          scored.each do |rule, recorded|
+            total = totals.fetch(rule.value, 0).to_d
+            limit = rule.windows.fetch(window)
+            next if total < limit
+
+            handle_exceeded(
+              budget_type: window,
+              total: total,
+              budget: limit,
+              previous_total: total - recorded.sum(&:total_cost),
+              last_event: recorded.last,
+              scope: scope_for(rule),
+              behavior: behavior_override || rule.behavior,
+              on_exceeded: rule.on_exceeded
+            )
+          end
         end
       end
 
+      def notify_persisted_safely!(events)
+        check_persisted!(events, behavior_override: :notify)
+      rescue StandardError => e
+        Logging.warn("Per-tag budget check failed after ingest: #{e.class}: #{e.message}")
+      end
+
       private
 
+      def window_buckets(by_rule)
+        by_rule.each_with_object({}) do |(rule, events), grouped|
+          rule.windows.each_key do |window|
+            events.group_by { |event| PerTag.window_start(window, event.tracked_at) }
+                  .each do |bucket, bucket_events|
+              (grouped[[rule.key, window, bucket]] ||= {})[rule] = bucket_events
+            end
+          end
+        end
+      end
+
       def estimate_cost(provider:, model:, request:)
         return BigDecimal("0") unless provider && model && request
 
@@ -54,14 +106,48 @@ estimate: estimate) do |budget_type, total, budget|
         ))
       end
 
-      def check_per_call_budget(event, config)
-        budget = config.per_call_budget
+      def check_per_call_budget(event, config, behavior_override)
+        budget = config.budgets.per_call
         return unless budget
 
         total = event.total_cost
         return unless total >= budget
 
-        handle_exceeded(budget_type: :per_call, total: total, budget: budget, last_event: event)
+        handle_exceeded(budget_type: :per_call,
+                        total: total,
+                        budget: budget,
+                        previous_total: nil,
+                        last_event: event,
+                        behavior: behavior_override)
+      end
+
+      def enforce_globally(config, estimate:, time:)
+        raise_per_call_pre_send(estimate, config.budgets.per_call) if config.budgets.per_call && estimate.positive?
+
+        check_windowed({ monthly: config.budgets.monthly, daily: config.budgets.daily }.compact,
+                       time: time,
+                       estimate: estimate) do |budget_type, total, budget|
+          raise_pre_send(budget_type: budget_type, total: total, budget: budget)
+        end
+      end
+
+      def raise_pre_send(budget_type:, total:, budget:, scope: nil)
+        raise BudgetExceededError.new(**budget_payload(
+          budget_type: budget_type, total: total, budget: budget, last_event: nil, stage: :pre_send, scope: scope
+        ))
+      end
+
+      def check_per_tag(tags, time:, estimate: BigDecimal("0"), blocking_only: false)
+        PerTag.rules_for(tags, blocking_only: blocking_only).each do |rule|
+          rule.windows.each do |window, limit|
+            total = PerTag.spend(rule.key, rule.value, window, time: time) + estimate
+            yield(rule, window, total, limit) if total >= limit
+          end
+        end
+      end
+
+      def scope_for(rule)
+        { key: rule.key, value: rule.value }
       end
 
       def check_windowed(budgets, time:, estimate: BigDecimal("0"))
@@ -82,38 +168,40 @@ estimate: estimate) do |budget_type, total, budget|
         period_for.transform_values { |period| period_totals.fetch(period) }
       end
 
-      def handle_exceeded(budget_type:, total:, budget:, last_event: nil)
+      def handle_exceeded(budget_type:,
+                          total:,
+                          budget:,
+                          previous_total:,
+                          last_event: nil,
+                          scope: nil,
+                          behavior: nil,
+                          on_exceeded: nil)
         config = LlmCostTracker.configuration
+        behavior ||= config.budgets.exceeded_behavior
+        on_exceeded ||= config.budgets.on_exceeded
         payload = budget_payload(
           budget_type: budget_type,
           total: total,
           budget: budget,
           last_event: last_event,
-          stage: :post_spend
+          stage: :post_spend,
+          scope: scope
         )
 
-        if notify_exceeded?(config, budget_type: budget_type, total: total, budget: budget, last_event: last_event)
-          config.on_budget_exceeded&.call(payload)
-        end
-        raise BudgetExceededError.new(**payload) if %i[raise block_requests].include?(config.budget_exceeded_behavior)
+        on_exceeded.call(payload) if on_exceeded && (previous_total.nil? || previous_total < budget)
+        raise BudgetExceededError.new(**payload) if %i[raise block_requests].include?(behavior)
       end
 
-      def budget_payload(budget_type:, total:, budget:, last_event:, stage:)
+      def budget_payload(budget_type:, total:, budget:, last_event:, stage:, scope: nil)
         {
           budget_type: budget_type,
           total: total,
           budget: budget,
           last_event: last_event,
-          stage: stage
+          stage: stage,
+          scope: scope
         }
       end
-
-      def notify_exceeded?(config, budget_type:, total:, budget:, last_event:)
-        return false unless config.on_budget_exceeded
-        return true if !last_event&.total_cost || budget_type == :per_call
-
-        total - last_event.total_cost < budget
-      end
     end
   end
 end

data/lib/llm_cost_tracker/charges/cost_status.rb

@@ -16,7 +16,7 @@ module LlmCostTracker
         "#{total_cost} IS NULL OR #{cost_status} IN (#{statuses})"
       end
 
-      # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+      # rubocop:disable-next Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
       def self.call(token_usage:,
                     usage_source:,
                     token_cost:,
@@ -43,9 +43,10 @@ module LlmCostTracker
         return UNKNOWN if unpriced && !priced
         return PARTIAL if unpriced
 
-        total_cost.nil? || total_cost.zero? ? FREE : COMPLETE
+        return UNKNOWN if total_cost.nil?
+
+        total_cost.zero? ? FREE : COMPLETE
       end
-      # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
     end
   end
 end

data/lib/llm_cost_tracker/configuration/budgets.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "section"
+
+module LlmCostTracker
+  class Configuration
+    class Budgets < Section
+      EXCEEDED_BEHAVIORS = %i[notify raise block_requests].freeze
+      TOTALS_SOURCES = %i[ledger cache].freeze
+      PER_TAG_WINDOWS = %i[daily weekly monthly].freeze
+      PER_TAG_OPTIONS = %i[behavior on_exceeded].freeze
+
+      attributes :monthly, :daily, :per_call, :on_exceeded
+      enum_attribute :exceeded_behavior, allowed: EXCEEDED_BEHAVIORS, default: :notify
+      enum_attribute :totals_source, allowed: TOTALS_SOURCES, default: :ledger
+
+      attr_reader :per_tag
+
+      def initialize(owner)
+        super
+        @monthly = nil
+        @daily = nil
+        @per_call = nil
+        @on_exceeded = nil
+        @per_tag = {}
+        self.exceeded_behavior = :notify
+        self.totals_source = :ledger
+      end
+
+      def per_tag=(value)
+        ensure_mutable!
+        @per_tag = (value || {}).to_h.to_h { |key, entry| [validated_key(key), validated_entry(key, entry)] }
+      end
+
+      def finalize!
+        @per_tag = deep_freeze(@per_tag || {})
+      end
+
+      private
+
+      def validated_key(key)
+        LlmCostTracker::Tags::Key.validate!(key, error_class: Error)
+      end
+
+      def validated_entry(key, entry)
+        raise Error, "budgets.per_tag[#{key.inspect}] must be a hash" unless entry.is_a?(Hash)
+
+        normalized = entry.to_h.transform_keys(&:to_sym)
+        windows, options = normalized.partition { |name, _| !PER_TAG_OPTIONS.include?(name) }.map(&:to_h)
+        if windows.empty?
+          raise Error,
+                "budgets.per_tag[#{key.inspect}] needs at least one of: #{PER_TAG_WINDOWS.join(', ')}"
+        end
+
+        {
+          windows: windows.to_h do |window, limit|
+                     [validated_window(key, window), validated_limit(key, window, limit)]
+                   end,
+          behavior: validated_behavior(key, options[:behavior]),
+          on_exceeded: options[:on_exceeded]
+        }
+      end
+
+      def validated_behavior(key, behavior)
+        return nil if behavior.nil?
+        return behavior.to_sym if EXCEEDED_BEHAVIORS.include?(behavior.to_sym)
+
+        raise Error,
+              "Unknown budgets.per_tag[#{key.inspect}] behavior: #{behavior.inspect}. " \
+              "Use one of: #{EXCEEDED_BEHAVIORS.join(', ')}"
+      end
+
+      def validated_window(key, window)
+        return window.to_sym if PER_TAG_WINDOWS.include?(window.to_sym)
+
+        raise Error,
+              "Unknown budgets.per_tag[#{key.inspect}] window: #{window.inspect}. " \
+              "Use one of: #{PER_TAG_WINDOWS.join(', ')}"
+      end
+
+      def validated_limit(key, window, limit)
+        numeric = Float(limit, exception: false)
+        return BigDecimal(limit.to_s) if numeric&.positive?
+
+        raise Error,
+              "budgets.per_tag[#{key.inspect}][#{window.inspect}] must be a positive number, " \
+              "got #{limit.inspect}"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/configuration/capture.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require_relative "section"
+
+module LlmCostTracker
+  class Configuration
+    class Capture < Section
+      OPENAI_COMPATIBLE_PROVIDERS = {
+        "openrouter.ai" => "openrouter",
+        "api.deepseek.com" => "deepseek",
+        "api.groq.com" => "groq"
+      }.freeze
+
+      attributes :request_stream_usage
+
+      attr_reader :openai_compatible_providers
+
+      def initialize(owner)
+        super
+        @request_stream_usage = true
+        self.openai_compatible_providers = OPENAI_COMPATIBLE_PROVIDERS
+      end
+
+      def openai_compatible_providers=(providers)
+        ensure_mutable!
+        @openai_compatible_providers = normalize_providers(providers)
+      end
+
+      def finalize!
+        @openai_compatible_providers = deep_freeze(normalize_providers(@openai_compatible_providers))
+      end
+
+      private
+
+      def normalize_providers(providers)
+        (providers || {}).each_with_object({}) do |(host, provider), normalized|
+          normalized[host.to_s.downcase] = provider.to_s
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/configuration/ingestion.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require_relative "section"
+
+module LlmCostTracker
+  class Configuration
+    class Ingestion < Section
+      MODES = %i[inline async].freeze
+
+      attributes :pool_size
+      enum_attribute :mode, allowed: MODES, default: :inline
+
+      def initialize(owner)
+        super
+        @pool_size = nil
+        self.mode = :inline
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/configuration/mutability.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  class Configuration
+    module Mutability
+      private
+
+      def ensure_mutable!
+        return unless finalized?
+
+        raise FrozenError, "can't modify frozen LlmCostTracker::Configuration"
+      end
+
+      def deep_freeze(value)
+        case value
+        when Hash
+          value.each do |key, nested_value|
+            deep_freeze(key)
+            deep_freeze(nested_value)
+          end
+          value.frozen? ? value : value.freeze
+        when Array, Set
+          value.each { |nested_value| deep_freeze(nested_value) }
+          value.frozen? ? value : value.freeze
+        when String
+          value.frozen? ? value : value.freeze
+        else
+          value
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/configuration/pricing.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative "section"
+
+module LlmCostTracker
+  class Configuration
+    class Pricing < Section
+      UNKNOWN_BEHAVIORS = %i[ignore warn raise].freeze
+
+      attributes :file
+      enum_attribute :unknown_model_behavior, allowed: UNKNOWN_BEHAVIORS, default: :warn
+
+      attr_reader :overrides
+
+      def initialize(owner)
+        super
+        @file = nil
+        self.overrides = {}
+        self.unknown_model_behavior = :warn
+      end
+
+      def overrides=(value)
+        ensure_mutable!
+        @overrides = LlmCostTracker::Pricing::Registry.normalize_price_entries(
+          value || {}, context: "pricing_overrides"
+        )
+      rescue ArgumentError, TypeError => e
+        raise Error, "invalid pricing.overrides: #{e.message}"
+      end
+
+      def finalize!
+        @overrides = deep_freeze(@overrides || {})
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/configuration/section.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/module/delegation"
+require "active_support/core_ext/string/inflections"
+
+require_relative "mutability"
+
+module LlmCostTracker
+  class Configuration
+    class Section
+      include Mutability
+
+      class << self
+        def attributes(*names)
+          names.each do |name|
+            attr_reader name
+
+            define_method(:"#{name}=") do |value|
+              ensure_mutable!
+              instance_variable_set(:"@#{name}", value)
+            end
+          end
+        end
+
+        def enum_attribute(name, allowed:, default:)
+          attr_reader name
+
+          define_method(:"#{name}=") do |value|
+            ensure_mutable!
+            instance_variable_set(:"@#{name}", normalize_enum(name, value, allowed, default))
+          end
+        end
+      end
+
+      private attr_reader :owner
+
+      delegate :finalized?, to: :owner
+
+      def initialize(owner)
+        @owner = owner
+      end
+
+      def finalize!; end
+
+      private
+
+      def section_name
+        self.class.name.split("::").last.underscore
+      end
+
+      def normalize_enum(name, value, allowed, default)
+        value = default if value.nil?
+        return value if allowed.include?(value)
+
+        raise Error, "Unknown #{section_name}.#{name}: #{value.inspect}. Use one of: #{allowed.join(', ')}"
+      end
+    end
+  end
+end