llm_cost_tracker 0.3.3 → 0.4.1

data/lib/llm_cost_tracker/budget.rb

@@ -7,62 +7,113 @@ module LlmCostTracker
     class << self
       def enforce!
         config = LlmCostTracker.configuration
-        return unless config.monthly_budget
         return unless config.budget_exceeded_behavior == :block_requests
         return unless config.active_record?
 
-        monthly_total = active_record_monthly_total
-        return unless monthly_total >= config.monthly_budget
+        budgets = enforce_period_budgets(config)
+        return if budgets.empty?
 
-        handle_exceeded(monthly_total: monthly_total)
+        totals = active_record_totals(budgets.keys, time: Time.now.utc)
+
+        budgets.each do |period, budget|
+          total = totals.fetch(period)
+
+          handle_exceeded(budget_type: period, total: total, budget: budget) if total >= budget
+        end
       end
 
       def check!(event)
         config = LlmCostTracker.configuration
-        return unless config.monthly_budget
         return unless event.cost
 
-        monthly_total = if config.active_record?
-                          active_record_monthly_total(time: event.tracked_at)
-                        else
-                          event.cost.total_cost
-                        end
-        return unless monthly_total >= config.monthly_budget
+        check_per_call_budget(event, config)
+        budgets = check_period_budgets(config)
+        totals = totals_for_check(event, config, budgets)
 
-        handle_exceeded(monthly_total: monthly_total, last_event: event)
+        budgets.each do |period, budget|
+          total = totals.fetch(period)
+
+          handle_exceeded(budget_type: period, total: total, budget: budget, last_event: event) if total >= budget
+        end
       end
 
       private
 
-      def active_record_monthly_total(time: Time.now.utc)
+      def check_per_call_budget(event, config)
+        budget = config.per_call_budget
+        return unless budget
+
+        call_cost = event.cost.total_cost
+        return unless call_cost >= budget
+
+        handle_exceeded(budget_type: :per_call, total: call_cost, budget: budget, last_event: event)
+      end
+
+      def enforce_period_budgets(config)
+        {
+          monthly: config.monthly_budget,
+          daily: config.daily_budget
+        }.compact
+      end
+
+      def check_period_budgets(config)
+        {
+          daily: config.daily_budget,
+          monthly: config.monthly_budget
+        }.compact
+      end
+
+      def totals_for_check(event, config, budgets)
+        return {} if budgets.empty?
+        return active_record_totals(budgets.keys, time: event.tracked_at) if config.active_record?
+
+        budgets.to_h { |period, _budget| [period, event.cost.total_cost] }
+      end
+
+      def active_record_totals(periods, time:)
         require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
         require_relative "storage/active_record_store" unless defined?(LlmCostTracker::Storage::ActiveRecordStore)
 
-        LlmCostTracker::Storage::ActiveRecordStore.monthly_total(time: time)
+        LlmCostTracker::Storage::ActiveRecordStore.period_totals(periods, time: time)
       rescue LoadError => e
         raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
       end
 
-      def handle_exceeded(monthly_total:, last_event: nil)
+      def handle_exceeded(budget_type:, total:, budget:, last_event: nil)
         config = LlmCostTracker.configuration
-        payload = {
-          monthly_total: monthly_total,
-          budget: config.monthly_budget,
+        payload = budget_payload(
+          budget_type: budget_type,
+          total: total,
+          budget: budget,
           last_event: last_event
-        }
+        )
 
-        if notify_exceeded?(config, monthly_total: monthly_total, last_event: last_event)
+        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 raise_on_exceeded?(config)
       end
 
-      def notify_exceeded?(config, monthly_total:, last_event:)
+      def budget_payload(budget_type:, total:, budget:, last_event:)
+        payload = {
+          budget_type: budget_type,
+          total: total,
+          budget: budget,
+          last_event: last_event
+        }
+        payload[:monthly_total] = total if budget_type == :monthly
+        payload[:daily_total] = total if budget_type == :daily
+        payload[:call_cost] = total if budget_type == :per_call
+        payload
+      end
+
+      def notify_exceeded?(config, budget_type:, total:, budget:, last_event:)
         return false unless config.on_budget_exceeded
         return true unless config.budget_exceeded_behavior == :notify
         return true unless last_event&.cost
+        return true if budget_type == :per_call
 
-        monthly_total - last_event.cost.total_cost < config.monthly_budget
+        total - last_event.cost.total_cost < budget
       end
 
       def raise_on_exceeded?(config)

data/lib/llm_cost_tracker/configuration.rb

@@ -19,6 +19,8 @@ module LlmCostTracker
       custom_storage
       on_budget_exceeded
       monthly_budget
+      daily_budget
+      per_call_budget
       log_level
       prices_file
     ].freeze
@@ -48,6 +50,8 @@ module LlmCostTracker
       @default_tags       = {}
       @on_budget_exceeded = nil
       @monthly_budget     = nil
+      @daily_budget       = nil
+      @per_call_budget    = nil
       self.budget_exceeded_behavior = :notify
       self.storage_error_behavior = :warn
       self.unknown_pricing_behavior = :warn

data/lib/llm_cost_tracker/cost.rb

@@ -3,9 +3,8 @@
 module LlmCostTracker
   Cost = Data.define(
     :input_cost,
-    :cached_input_cost,
     :cache_read_input_cost,
-    :cache_creation_input_cost,
+    :cache_write_input_cost,
     :output_cost,
     :total_cost,
     :currency

data/lib/llm_cost_tracker/errors.rb

@@ -6,14 +6,33 @@ module LlmCostTracker
   class InvalidFilterError < Error; end
 
   class BudgetExceededError < Error
-    attr_reader :monthly_total, :budget, :last_event
+    attr_reader :monthly_total, :daily_total, :call_cost, :total, :budget, :budget_type, :last_event
 
-    def initialize(monthly_total:, budget:, last_event: nil)
+    def initialize(budget:, last_event: nil, budget_type: nil, total: nil, monthly_total: nil, daily_total: nil,
+                   call_cost: nil)
       @monthly_total = monthly_total
+      @daily_total = daily_total
+      @call_cost = call_cost
+      @total = total || monthly_total || daily_total || call_cost
       @budget = budget
+      @budget_type = budget_type || inferred_budget_type
       @last_event = last_event
 
-      super("LLM monthly budget exceeded: $#{format('%.6f', monthly_total)} / $#{format('%.6f', budget)}")
+      super("LLM #{budget_label} budget exceeded: $#{format('%.6f', @total)} / $#{format('%.6f', budget)}")
+    end
+
+    private
+
+    def inferred_budget_type
+      return :monthly if monthly_total
+      return :daily if daily_total
+      return :per_call if call_cost
+
+      :unknown
+    end
+
+    def budget_label
+      budget_type.to_s.tr("_", "-")
     end
   end
 

data/lib/llm_cost_tracker/event.rb

@@ -7,6 +7,10 @@ module LlmCostTracker
     :input_tokens,
     :output_tokens,
     :total_tokens,
+    :cache_read_input_tokens,
+    :cache_write_input_tokens,
+    :hidden_output_tokens,
+    :pricing_mode,
     :cost,
     :tags,
     :latency_ms,

data/lib/llm_cost_tracker/event_metadata.rb

@@ -3,32 +3,31 @@
 module LlmCostTracker
   module EventMetadata
     INTERNAL_TAG_KEYS = %w[
-      cache_creation_input_tokens
-      cache_creation_tokens
       cache_read_input_tokens
-      cache_read_tokens
-      cached_input_tokens
+      cache_write_input_tokens
+      hidden_output_tokens
       input_tokens
       output_tokens
+      pricing_mode
       provider_response_id
-      reasoning_tokens
       total_tokens
     ].freeze
 
     class << self
       def usage_data(input_tokens, output_tokens, metadata)
         metadata = metadata.to_h.symbolize_keys
-        cache_read = first_integer(metadata, :cache_read_input_tokens, :cache_read_tokens)
-        cache_creation = first_integer(metadata, :cache_creation_input_tokens, :cache_creation_tokens)
-
-        {
-          input_tokens: input_tokens.to_i,
-          output_tokens: output_tokens.to_i,
-          cached_input_tokens: metadata[:cached_input_tokens].to_i,
+        cache_read = first_integer(metadata, :cache_read_input_tokens)
+        cache_write = first_integer(metadata, :cache_write_input_tokens)
+        hidden_output = first_integer(metadata, :hidden_output_tokens)
+        breakdown = UsageBreakdown.build(
+          input_tokens: input_tokens,
+          output_tokens: output_tokens,
           cache_read_input_tokens: cache_read,
-          cache_creation_input_tokens: cache_creation,
-          total_tokens: input_tokens.to_i + output_tokens.to_i + cache_read + cache_creation
-        }
+          cache_write_input_tokens: cache_write,
+          hidden_output_tokens: hidden_output
+        )
+
+        breakdown.to_h.merge(pricing_mode: normalized_pricing_mode(metadata[:pricing_mode])).compact
       end
 
       def tags(metadata)
@@ -41,6 +40,13 @@ module LlmCostTracker
         keys.each { |key| return metadata[key].to_i unless metadata[key].nil? }
         0
       end
+
+      def normalized_pricing_mode(value)
+        return nil if value.nil?
+
+        mode = value.to_s.strip
+        mode.empty? || mode == "standard" ? nil : mode
+      end
     end
   end
 end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/{add_monthly_totals_generator.rb → add_period_totals_generator.rb}

@@ -5,17 +5,17 @@ require "rails/generators/active_record"
 
 module LlmCostTracker
   module Generators
-    class AddMonthlyTotalsGenerator < Rails::Generators::Base
+    class AddPeriodTotalsGenerator < Rails::Generators::Base
       include ActiveRecord::Generators::Migration
 
       source_root File.expand_path("templates", __dir__)
 
-      desc "Creates a migration to add llm_cost_tracker_monthly_totals"
+      desc "Creates a migration to add llm_cost_tracker_period_totals"
 
       def create_migration_file
         migration_template(
-          "add_monthly_totals_to_llm_cost_tracker.rb.erb",
-          "db/migrate/add_monthly_totals_to_llm_cost_tracker.rb"
+          "add_period_totals_to_llm_cost_tracker.rb.erb",
+          "db/migrate/add_period_totals_to_llm_cost_tracker.rb"
         )
       end
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/add_usage_breakdown_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class AddUsageBreakdownGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates a migration to add usage and cost breakdown columns to llm_api_calls"
+
+      def create_migration_file
+        migration_template(
+          "add_usage_breakdown_to_llm_api_calls.rb.erb",
+          "db/migrate/add_usage_breakdown_to_llm_api_calls.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/add_period_totals_to_llm_cost_tracker.rb.erb

@@ -0,0 +1,96 @@
+class AddPeriodTotalsToLlmCostTracker < ActiveRecord::Migration<%= migration_version %>
+  def up
+    create_table :llm_cost_tracker_period_totals do |t|
+      t.string :period, null: false
+      t.date :period_start, null: false
+      t.decimal :total_cost, precision: 20, scale: 8, null: false, default: 0
+
+      t.timestamps
+    end unless table_exists?(:llm_cost_tracker_period_totals)
+
+    backfill_period_totals
+
+    add_index :llm_cost_tracker_period_totals, [:period, :period_start],
+              unique: true unless index_exists?(:llm_cost_tracker_period_totals, [:period, :period_start])
+  end
+
+  def down
+    remove_index :llm_cost_tracker_period_totals, [:period, :period_start] if index_exists?(:llm_cost_tracker_period_totals, [:period, :period_start])
+    drop_table :llm_cost_tracker_period_totals if table_exists?(:llm_cost_tracker_period_totals)
+  end
+
+  private
+
+  def backfill_period_totals
+    backfill_legacy_monthly_totals if table_exists?(:llm_cost_tracker_monthly_totals)
+    return unless table_exists?(:llm_api_calls)
+
+    backfill_period_total("day", day_bucket_sql)
+    backfill_period_total("month", month_bucket_sql)
+  end
+
+  def backfill_legacy_monthly_totals
+    execute <<~SQL
+      INSERT INTO llm_cost_tracker_period_totals (period, period_start, total_cost, created_at, updated_at)
+      SELECT #{connection.quote("month")} AS period,
+             month AS period_start,
+             total_cost,
+             CURRENT_TIMESTAMP,
+             CURRENT_TIMESTAMP
+      FROM llm_cost_tracker_monthly_totals legacy
+      WHERE NOT EXISTS (
+        SELECT 1
+        FROM llm_cost_tracker_period_totals existing
+        WHERE existing.period = #{connection.quote("month")}
+          AND existing.period_start = legacy.month
+      )
+    SQL
+  end
+
+  def backfill_period_total(period, bucket_sql)
+    execute <<~SQL
+      INSERT INTO llm_cost_tracker_period_totals (period, period_start, total_cost, created_at, updated_at)
+      SELECT aggregated.period,
+             aggregated.period_start,
+             aggregated.total_cost,
+             CURRENT_TIMESTAMP,
+             CURRENT_TIMESTAMP
+      FROM (
+        SELECT #{connection.quote(period)} AS period,
+               #{bucket_sql} AS period_start,
+               SUM(total_cost) AS total_cost
+        FROM llm_api_calls
+        WHERE total_cost IS NOT NULL
+        GROUP BY #{bucket_sql}
+      ) aggregated
+      WHERE NOT EXISTS (
+        SELECT 1
+        FROM llm_cost_tracker_period_totals existing
+        WHERE existing.period = aggregated.period
+          AND existing.period_start = aggregated.period_start
+      )
+    SQL
+  end
+
+  def day_bucket_sql
+    case connection.adapter_name
+    when /postgres/i
+      "DATE_TRUNC('day', tracked_at)::date"
+    when /mysql/i
+      "DATE(tracked_at)"
+    else
+      "date(tracked_at)"
+    end
+  end
+
+  def month_bucket_sql
+    case connection.adapter_name
+    when /postgres/i
+      "DATE_TRUNC('month', tracked_at)::date"
+    when /mysql/i
+      "DATE_FORMAT(tracked_at, '%Y-%m-01')"
+    else
+      "strftime('%Y-%m-01', tracked_at)"
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/add_usage_breakdown_to_llm_api_calls.rb.erb

@@ -0,0 +1,29 @@
+class AddUsageBreakdownToLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
+  def up
+    unless column_exists?(:llm_api_calls, :cache_read_input_tokens)
+      add_column :llm_api_calls, :cache_read_input_tokens, :integer, null: false, default: 0
+    end
+    unless column_exists?(:llm_api_calls, :cache_write_input_tokens)
+      add_column :llm_api_calls, :cache_write_input_tokens, :integer, null: false, default: 0
+    end
+    unless column_exists?(:llm_api_calls, :hidden_output_tokens)
+      add_column :llm_api_calls, :hidden_output_tokens, :integer, null: false, default: 0
+    end
+    unless column_exists?(:llm_api_calls, :cache_read_input_cost)
+      add_column :llm_api_calls, :cache_read_input_cost, :decimal, precision: 20, scale: 8
+    end
+    unless column_exists?(:llm_api_calls, :cache_write_input_cost)
+      add_column :llm_api_calls, :cache_write_input_cost, :decimal, precision: 20, scale: 8
+    end
+    add_column :llm_api_calls, :pricing_mode, :string unless column_exists?(:llm_api_calls, :pricing_mode)
+  end
+
+  def down
+    remove_column :llm_api_calls, :pricing_mode if column_exists?(:llm_api_calls, :pricing_mode)
+    remove_column :llm_api_calls, :cache_write_input_cost if column_exists?(:llm_api_calls, :cache_write_input_cost)
+    remove_column :llm_api_calls, :cache_read_input_cost if column_exists?(:llm_api_calls, :cache_read_input_cost)
+    remove_column :llm_api_calls, :hidden_output_tokens if column_exists?(:llm_api_calls, :hidden_output_tokens)
+    remove_column :llm_api_calls, :cache_write_input_tokens if column_exists?(:llm_api_calls, :cache_write_input_tokens)
+    remove_column :llm_api_calls, :cache_read_input_tokens if column_exists?(:llm_api_calls, :cache_read_input_tokens)
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_api_calls.rb.erb

@@ -6,13 +6,19 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
       t.integer :input_tokens,  null: false, default: 0
       t.integer :output_tokens, null: false, default: 0
       t.integer :total_tokens,  null: false, default: 0
+      t.integer :cache_read_input_tokens,  null: false, default: 0
+      t.integer :cache_write_input_tokens, null: false, default: 0
+      t.integer :hidden_output_tokens,     null: false, default: 0
       t.decimal :input_cost,    precision: 20, scale: 8
+      t.decimal :cache_read_input_cost,  precision: 20, scale: 8
+      t.decimal :cache_write_input_cost, precision: 20, scale: 8
       t.decimal :output_cost,   precision: 20, scale: 8
       t.decimal :total_cost,    precision: 20, scale: 8
       t.integer :latency_ms
       t.boolean :stream,        null: false, default: false
       t.string  :usage_source
       t.string  :provider_response_id
+      t.string  :pricing_mode
       if postgresql?
         t.jsonb :tags, null: false, default: {}
       else
@@ -23,22 +29,22 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
       t.timestamps
     end
 
-    create_table :llm_cost_tracker_monthly_totals do |t|
-      t.date :month_start, null: false
+    create_table :llm_cost_tracker_period_totals do |t|
+      t.string :period, null: false
+      t.date :period_start, null: false
       t.decimal :total_cost, precision: 20, scale: 8, null: false, default: 0
 
       t.timestamps
     end
 
-    add_index :llm_api_calls, :provider
-    add_index :llm_api_calls, :model
     add_index :llm_api_calls, :tracked_at
     add_index :llm_api_calls, [:provider, :tracked_at]
+    add_index :llm_api_calls, [:model, :tracked_at]
     add_index :llm_api_calls, :stream
     add_index :llm_api_calls, :usage_source
     add_index :llm_api_calls, :provider_response_id
     add_index :llm_api_calls, :tags, using: :gin if postgresql?
-    add_index :llm_cost_tracker_monthly_totals, :month_start, unique: true
+    add_index :llm_cost_tracker_period_totals, [:period, :period_start], unique: true
   end
 
   private

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/initializer.rb.erb

@@ -12,6 +12,8 @@ LlmCostTracker.configure do |config|
 
   # Monthly budget in USD. Set to nil to disable budget alerts.
   # config.monthly_budget = 100.00
+  # config.daily_budget = 10.00
+  # config.per_call_budget = 1.00
   # config.budget_exceeded_behavior = :notify # :notify, :raise, or :block_requests
 
   # What to do when storage fails.
@@ -23,7 +25,7 @@ LlmCostTracker.configure do |config|
   # Callback when monthly budget is exceeded.
   # config.on_budget_exceeded = ->(data) {
   #   Rails.logger.warn "[LlmCostTracker] Budget exceeded! " \
-  #     "Monthly total: $#{data[:monthly_total]}, Budget: $#{data[:budget]}"
+  #     "#{data[:budget_type]} total: $#{data[:total]}, Budget: $#{data[:budget]}"
   #   # Or send a Slack notification, email, etc.
   # }
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/llm_cost_tracker_prices.yml.erb

@@ -8,9 +8,9 @@
 # Supported price keys:
 # - input
 # - output
-# - cached_input
 # - cache_read_input
-# - cache_creation_input
+# - cache_write_input
+# - mode_input / mode_output / mode_cache_read_input / mode_cache_write_input
 #
 # Optional metadata keys, ignored by cost calculation:
 # - _source
@@ -24,10 +24,18 @@
 # models:
 #   "ft:gpt-4o-mini:my-org":
 #     input: 0.30
-#     cached_input: 0.15
+#     cache_read_input: 0.15
 #     output: 1.20
 #     _notes: "Internal fine-tune rate"
 #
+# Example: alternate pricing mode
+# models:
+#   "batchable-model":
+#     input: 1.00
+#     output: 2.00
+#     batch_input: 0.50
+#     batch_output: 1.00
+#
 # Example: negotiated provider discount
 # models:
 #   "gpt-4o":

data/lib/llm_cost_tracker/parsed_usage.rb

@@ -7,10 +7,9 @@ module LlmCostTracker
     :input_tokens,
     :output_tokens,
     :total_tokens,
-    :cached_input_tokens,
     :cache_read_input_tokens,
-    :cache_creation_input_tokens,
-    :reasoning_tokens,
+    :cache_write_input_tokens,
+    :hidden_output_tokens,
     :stream,
     :usage_source,
     :provider_response_id
@@ -34,11 +33,10 @@ module LlmCostTracker
         model: attributes.fetch(:model),
         input_tokens: attributes.fetch(:input_tokens).to_i,
         output_tokens: attributes.fetch(:output_tokens).to_i,
-        total_tokens: attributes.fetch(:total_tokens, 0).to_i,
-        cached_input_tokens: attributes[:cached_input_tokens],
+        total_tokens: attributes.fetch(:total_tokens, usage_breakdown(attributes).total_tokens).to_i,
         cache_read_input_tokens: attributes[:cache_read_input_tokens],
-        cache_creation_input_tokens: attributes[:cache_creation_input_tokens],
-        reasoning_tokens: attributes[:reasoning_tokens],
+        cache_write_input_tokens: attributes[:cache_write_input_tokens],
+        hidden_output_tokens: attributes[:hidden_output_tokens],
         stream: attributes[:stream] || false,
         usage_source: attributes[:usage_source],
         provider_response_id: attributes[:provider_response_id]
@@ -52,5 +50,16 @@ module LlmCostTracker
     def to_h
       super.compact
     end
+
+    def self.usage_breakdown(attributes)
+      UsageBreakdown.build(
+        input_tokens: attributes.fetch(:input_tokens),
+        output_tokens: attributes.fetch(:output_tokens),
+        cache_read_input_tokens: attributes[:cache_read_input_tokens],
+        cache_write_input_tokens: attributes[:cache_write_input_tokens],
+        hidden_output_tokens: attributes[:hidden_output_tokens]
+      )
+    end
+    private_class_method :usage_breakdown
   end
 end