llm_cost_tracker 0.12.0 → 0.14.0

data/lib/llm_cost_tracker/tags/sanitizer.rb

@@ -23,9 +23,9 @@ module LlmCostTracker
       class << self
         def call(tags, config: LlmCostTracker.configuration)
           tags = (tags || {}).to_h
-          redacted = config.normalized_redacted_tag_keys
-          limit = [config.max_tag_value_bytesize.to_i, 0].max
-          max_count = [config.max_tag_count.to_i, 0].max
+          redacted = config.tags.normalized_redacted_keys
+          limit = [config.tags.max_value_bytesize.to_i, 0].max
+          max_count = [config.tags.max_count.to_i, 0].max
           tags.to_a.last(max_count).each_with_object({}) do |(key, value), sanitized|
             next unless valid_key?(key)
 
@@ -43,7 +43,7 @@ module LlmCostTracker
 
         def cap(tags, config: LlmCostTracker.configuration)
           tags = (tags || {}).to_h
-          max_count = [config.max_tag_count.to_i, 0].max
+          max_count = [config.tags.max_count.to_i, 0].max
           return tags if tags.size <= max_count
 
           tags.to_a.last(max_count).to_h

data/lib/llm_cost_tracker/tracker.rb

@@ -13,38 +13,26 @@ module LlmCostTracker
     EVENT_NAME = "llm_request.llm_cost_tracker"
 
     class << self
-      def record(event:, latency_ms: nil, pricing_mode: nil, metadata: {}, context_tags: nil, enforce_budget: false)
+      def record(event:, latency_ms: nil, metadata: {}, context_tags: nil, enforce_budget: false)
         return unless LlmCostTracker.configuration.enabled
 
-        pricing_mode ||= event.pricing_mode
         calculation = Pricing::Calculation.for(
           provider: event.provider,
           model: event.model,
           tokens: event.token_usage,
           line_items: event.line_items,
-          pricing_mode: pricing_mode,
+          pricing_mode: event.pricing_mode,
           usage_source: event.usage_source
         )
 
-        if enforce_budget
-          Budget.enforce!(provider: event.provider,
-                          model: event.model,
-                          estimate: calculation.cost&.total,
-                          force: true)
-        end
+        tags = build_tags(context_tags: context_tags, metadata: metadata)
 
         if calculation.token_cost.nil? && event.token_usage.total_tokens.positive? &&
            calculation.priced_line_items.none?(&:priced?)
-          Pricing::Unknown.process(event.model)
+          Pricing::Unknown.process(event.model, pricing_mode: calculation.mode)
         end
 
-        event = build_event(
-          event: event,
-          calculation: calculation,
-          metadata: metadata,
-          latency_ms: latency_ms,
-          context_tags: context_tags
-        )
+        event = build_event(event: event, calculation: calculation, tags: tags, latency_ms: latency_ms)
 
         if Ingestion.async?
           Ingestion::Inbox.save(event)
@@ -55,11 +43,19 @@ module LlmCostTracker
 
         yield if block_given?
         notify_subscribers(event)
-        Budget.check!(event)
+        behavior_override = :raise if enforce_budget
+        Budget.check!(event, behavior_override: behavior_override)
+        Budget.check_persisted!([event], behavior_override: behavior_override) unless Ingestion.async?
 
         event
       end
 
+      def build_tags(context_tags:, metadata:)
+        resolved = (context_tags || LlmCostTracker::Tags::Context.tags).to_h
+        sanitized_metadata = LlmCostTracker::Tags::Sanitizer.call(metadata.to_h)
+        LlmCostTracker::Tags::Sanitizer.cap(resolved.merge(sanitized_metadata)).freeze
+      end
+
       private
 
       def notify_subscribers(event)
@@ -70,13 +66,12 @@ module LlmCostTracker
         Logging.warn("Subscriber raised on #{EVENT_NAME}: #{e.class}: #{e.message}")
       end
 
-      def build_event(event:, calculation:, metadata:, latency_ms:, context_tags:)
-        context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).to_h
+      def build_event(event:, calculation:, tags:, latency_ms:)
         event.with(
           event_id: SecureRandom.uuid,
           pricing_mode: calculation.mode,
           cost: calculation.cost,
-          tags: build_tags(context_tags: context_tags, metadata: metadata),
+          tags: tags,
           latency_ms: finite_latency_ms(latency_ms),
           tracked_at: Time.now.utc,
           cost_status: calculation.cost_status,
@@ -85,11 +80,6 @@ module LlmCostTracker
         )
       end
 
-      def build_tags(context_tags:, metadata:)
-        sanitized_metadata = LlmCostTracker::Tags::Sanitizer.call(metadata.to_h)
-        LlmCostTracker::Tags::Sanitizer.cap(context_tags.merge(sanitized_metadata)).freeze
-      end
-
       def finite_latency_ms(latency_ms)
         return nil if latency_ms.nil?
 

data/lib/llm_cost_tracker/usage/catalog.rb

@@ -16,8 +16,7 @@ module LlmCostTracker
         "request" => "per_request",
         "session" => "per_session",
         "hour" => "per_hour",
-        "minute" => "per_minute",
-        "image" => "per_image"
+        "minute" => "per_minute"
       }.freeze
 
       class << self
@@ -31,6 +30,10 @@ module LlmCostTracker
           @token_priced ||= all.select(&:token_key).freeze
         end
 
+        def find_by(kind:, direction:, modality:, cache_state:, unit:)
+          by_attributes[[kind, direction, modality, cache_state, unit]]
+        end
+
         def token_priced_for(kind:, direction:, cache_state:)
           token_priced.find do |dimension|
             dimension.kind == kind && dimension.direction == direction && dimension.cache_state == cache_state
@@ -43,6 +46,13 @@ module LlmCostTracker
           @index ||= all.to_h { |dimension| [dimension.key, dimension] }.freeze
         end
 
+        def by_attributes
+          @by_attributes ||= all.to_h do |dimension|
+            key = [dimension.kind, dimension.direction, dimension.modality, dimension.cache_state, dimension.unit]
+            [key, dimension]
+          end.freeze
+        end
+
         def load_definitions
           Psych.safe_load_file(DEFINITIONS_PATH, permitted_classes: [], symbolize_names: true)
                .map { |attributes| build(attributes) }

data/lib/llm_cost_tracker/usage/dimensions.yml

@@ -140,27 +140,3 @@
   cache_state: none
   unit: minute
   rate_basis: per_minute
-
-- key: image_generation_call
-  kind: image_generation_call
-  direction: output
-  modality: image
-  cache_state: none
-  unit: image
-  rate_basis: per_image
-
-- key: computer_call
-  kind: computer_call
-  direction: neither
-  modality: text
-  cache_state: none
-  unit: request
-  rate_basis: per_request
-
-- key: mcp_call
-  kind: mcp_call
-  direction: neither
-  modality: text
-  cache_state: none
-  unit: request
-  rate_basis: per_request

data/lib/llm_cost_tracker/usage/token_usage.rb

@@ -4,97 +4,42 @@ require_relative "catalog"
 
 module LlmCostTracker
   module Usage
-    KNOWN_TOKEN_KEYS = (
-      Catalog.token_priced.map { |dimension| dimension.token_key.to_s } + %w[total_tokens hidden_output_tokens]
-    ).freeze
-
-    TokenUsage = Data.define(
-      :input_tokens,
-      :cache_read_input_tokens,
-      :cache_write_input_tokens,
-      :cache_write_extended_input_tokens,
-      :audio_input_tokens,
-      :image_input_tokens,
-      :output_tokens,
-      :audio_output_tokens,
-      :image_output_tokens,
-      :total_tokens,
-      :hidden_output_tokens
-    ) do
+    TokenUsage = Data.define(*Catalog.token_priced.map(&:token_key), :total_tokens, :hidden_output_tokens) do
       def priced_quantities
         Catalog.token_priced.to_h { |dimension| [dimension.key, public_send(dimension.token_key)] }
       end
 
+      def self.build(**values)
+        unknown = values.keys - members
+        raise ArgumentError, "unknown token keys: #{unknown.inspect}" if unknown.any?
+
+        priced = Catalog.token_priced.to_h do |dimension|
+          [dimension.token_key, non_negative_int(values[dimension.token_key])]
+        end
+        subtotal = priced.values.sum
+        declared_total = values[:total_tokens]
+        total = declared_total ? [non_negative_int(declared_total), subtotal].max : subtotal
+        new(**priced, total_tokens: total, hidden_output_tokens: non_negative_int(values[:hidden_output_tokens]))
+      end
+
       def self.build_from_tokens(tokens)
         return tokens if tokens.is_a?(self)
         raise ArgumentError, "tokens must be a Hash, got #{tokens.class}" unless tokens.respond_to?(:to_h)
 
         values = tokens.to_h.transform_keys(&:to_s)
-        warn_on_unknown_keys(values)
-        token_attributes = Catalog.token_priced.to_h do |dimension|
-          [dimension.token_key, values.fetch(dimension.token_key.to_s, 0)]
+        known = members.map(&:to_s)
+        unknown = values.keys - known
+        if unknown.any?
+          hint = values.keys.intersect?(known) ? "" : ". Did you pass a raw provider response?"
+          raise ArgumentError, "unknown token keys: #{unknown.inspect}; expected #{known.inspect}#{hint}"
         end
 
-        build(
-          **token_attributes,
-          total_tokens: values["total_tokens"],
-          hidden_output_tokens: values.fetch("hidden_output_tokens", 0)
-        )
-      end
-
-      def self.warn_on_unknown_keys(values)
-        return if values.empty?
-        return if values.keys.intersect?(KNOWN_TOKEN_KEYS)
-
-        Logging.warn(
-          "tokens hash contains no recognized keys (#{values.keys.inspect}); " \
-          "expected one of #{KNOWN_TOKEN_KEYS.inspect}. Did you pass a raw provider response?"
-        )
+        build(**values.transform_keys(&:to_sym))
       end
 
       def self.non_negative_int(value)
         [value.to_i, 0].max
       end
-
-      def self.build(input_tokens:,
-                     output_tokens:,
-                     cache_read_input_tokens: 0,
-                     cache_write_input_tokens: 0,
-                     cache_write_extended_input_tokens: 0,
-                     audio_input_tokens: 0,
-                     audio_output_tokens: 0,
-                     image_input_tokens: 0,
-                     image_output_tokens: 0,
-                     total_tokens: nil,
-                     hidden_output_tokens: 0)
-        input = non_negative_int(input_tokens)
-        output = non_negative_int(output_tokens)
-        cache_read = non_negative_int(cache_read_input_tokens)
-        cache_write = non_negative_int(cache_write_input_tokens)
-        cache_write_extended = non_negative_int(cache_write_extended_input_tokens)
-        audio_input = non_negative_int(audio_input_tokens)
-        audio_output = non_negative_int(audio_output_tokens)
-        image_input = non_negative_int(image_input_tokens)
-        image_output = non_negative_int(image_output_tokens)
-        hidden_output = non_negative_int(hidden_output_tokens)
-        calculated_total = input + cache_read + cache_write + cache_write_extended +
-                           audio_input + image_input + output + audio_output + image_output
-        total = total_tokens ? [non_negative_int(total_tokens), calculated_total].max : calculated_total
-
-        new(
-          input_tokens: input,
-          cache_read_input_tokens: cache_read,
-          cache_write_input_tokens: cache_write,
-          cache_write_extended_input_tokens: cache_write_extended,
-          audio_input_tokens: audio_input,
-          image_input_tokens: image_input,
-          output_tokens: output,
-          audio_output_tokens: audio_output,
-          image_output_tokens: image_output,
-          total_tokens: total,
-          hidden_output_tokens: hidden_output
-        )
-      end
     end
   end
 end

data/lib/llm_cost_tracker/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LlmCostTracker
-  VERSION = "0.12.0"
+  VERSION = "0.14.0"
 end

data/lib/llm_cost_tracker.rb

@@ -30,7 +30,7 @@ require_relative "llm_cost_tracker/ingestion"
 require_relative "llm_cost_tracker/tracker"
 
 module LlmCostTracker
-  autoload :Engine, "llm_cost_tracker/engine"
+  require_relative "llm_cost_tracker/engine"
   autoload :Doctor, "llm_cost_tracker/doctor"
   autoload :CaptureVerifier, "llm_cost_tracker/capture_verifier"
   autoload :Report, "llm_cost_tracker/report"
@@ -110,7 +110,12 @@ module LlmCostTracker
                      provider_workspace_id: nil,
                      pricing_mode: nil)
       require_relative "llm_cost_tracker/capture/stream_collector"
-      Budget.enforce!(provider: provider, model: model, force: true) if enforce_budget
+      if enforce_budget
+        Budget.enforce!(provider: provider,
+                        model: model,
+                        force: true,
+                        tags: Tracker.build_tags(context_tags: nil, metadata: tags))
+      end
       collector = Capture::StreamCollector.new(
         provider: provider.to_s,
         model: model,

data/lib/tasks/llm_cost_tracker.rake

@@ -6,7 +6,15 @@ require "rails/generators"
 require_relative "../llm_cost_tracker/generators/llm_cost_tracker/install_generator"
 require_relative "../llm_cost_tracker/pricing/sync/change_printer"
 
-# rubocop:disable Metrics/BlockLength
+module LlmCostTrackerTasks
+  def self.price_refresh_output_path
+    path = LlmCostTracker::Pricing::Sync.configured_output_path
+    FileUtils.mkdir_p(File.dirname(path))
+    path
+  end
+end
+
+# rubocop:disable-next Metrics/BlockLength
 namespace :llm_cost_tracker do
   desc "Install LLM Cost Tracker with dashboard and prices, migrate, and run doctor"
   task :setup do
@@ -75,6 +83,18 @@ namespace :llm_cost_tracker do
     puts "llm_cost_tracker: rebuilt #{rows} rollup rows from the calls ledger"
   end
 
+  desc "Copy call cost and time onto existing llm_cost_tracker_call_tags rows so per-tag " \
+       "budgets count history recorded before the upgrade. Use BATCH_SIZE=N to tune."
+  task backfill_tag_costs: :environment do
+    unless LlmCostTracker::Budget::PerTag.columns?
+      abort("llm_cost_tracker: llm_cost_tracker_call_tags is missing the cost columns; " \
+            "run the upgrade_per_tag_budgets generator and migrate first")
+    end
+    batch_size = (ENV["BATCH_SIZE"] || LlmCostTracker::Budget::PerTag::DEFAULT_BACKFILL_BATCH).to_i
+    filled = LlmCostTracker::Budget::PerTag.backfill(batch_size: batch_size)
+    puts "llm_cost_tracker: filled cost and time on #{filled} tag rows"
+  end
+
   namespace :prices do
     desc(
       "Refresh the configured pricing file from the maintained LLM Cost Tracker price snapshot. " \
@@ -84,7 +104,7 @@ namespace :llm_cost_tracker do
       Rake::Task["environment"].invoke if Rake::Task.task_defined?("environment")
       require_relative "../llm_cost_tracker"
 
-      output_path = price_refresh_output_path
+      output_path = LlmCostTrackerTasks.price_refresh_output_path
       source_url = LlmCostTracker::Pricing::Sync.configured_remote_url
       preview = ENV["PREVIEW"] == "1"
       result = LlmCostTracker::Pricing::Sync.refresh(
@@ -112,7 +132,7 @@ namespace :llm_cost_tracker do
       Rake::Task["environment"].invoke if Rake::Task.task_defined?("environment")
       require_relative "../llm_cost_tracker"
 
-      output_path = price_refresh_output_path
+      output_path = LlmCostTrackerTasks.price_refresh_output_path
       source_url = LlmCostTracker::Pricing::Sync.configured_remote_url
       result = LlmCostTracker::Pricing::Sync.check(path: output_path, url: source_url)
 
@@ -125,10 +145,3 @@ namespace :llm_cost_tracker do
     end
   end
 end
-# rubocop:enable Metrics/BlockLength
-
-def price_refresh_output_path
-  path = LlmCostTracker::Pricing::Sync.configured_output_path
-  FileUtils.mkdir_p(File.dirname(path))
-  path
-end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: llm_cost_tracker
 version: !ruby/object:Gem::Version
-  version: 0.12.0
+  version: 0.14.0
 platform: ruby
 authors:
 - Sergii Khomenko
@@ -109,42 +109,42 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.42'
+        version: '1.61'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.42'
+        version: '1.61'
 - !ruby/object:Gem::Dependency
   name: nokogiri
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.16'
+        version: '1.19'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.16'
+        version: '1.19'
 - !ruby/object:Gem::Dependency
   name: openai
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.63'
+        version: '0.80'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.63'
+        version: '0.80'
 - !ruby/object:Gem::Dependency
   name: pg
   requirement: !ruby/object:Gem::Requirement
@@ -207,14 +207,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.15'
+        version: '1.16'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '1.15'
+        version: '1.16'
 - !ruby/object:Gem::Dependency
   name: simplecov
   requirement: !ruby/object:Gem::Requirement
@@ -339,6 +339,7 @@ files:
 - lib/llm_cost_tracker.rb
 - lib/llm_cost_tracker/assets.rb
 - lib/llm_cost_tracker/budget.rb
+- lib/llm_cost_tracker/budget/per_tag.rb
 - lib/llm_cost_tracker/capture/sdk_payload.rb
 - lib/llm_cost_tracker/capture/sse.rb
 - lib/llm_cost_tracker/capture/stream_collector.rb
@@ -349,7 +350,15 @@ files:
 - lib/llm_cost_tracker/charges/line_item.rb
 - lib/llm_cost_tracker/check.rb
 - lib/llm_cost_tracker/configuration.rb
+- lib/llm_cost_tracker/configuration/budgets.rb
+- lib/llm_cost_tracker/configuration/capture.rb
+- lib/llm_cost_tracker/configuration/ingestion.rb
+- lib/llm_cost_tracker/configuration/mutability.rb
+- lib/llm_cost_tracker/configuration/pricing.rb
+- lib/llm_cost_tracker/configuration/section.rb
+- lib/llm_cost_tracker/configuration/tags.rb
 - lib/llm_cost_tracker/currency.rb
+- lib/llm_cost_tracker/deprecator.rb
 - lib/llm_cost_tracker/doctor.rb
 - lib/llm_cost_tracker/doctor/ingestion_check.rb
 - lib/llm_cost_tracker/doctor/price_check.rb
@@ -369,9 +378,13 @@ files:
 - lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_call_rollups_provider.rb.erb
 - lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_call_tags_key_value_index.rb.erb
 - lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_image_tokens.rb.erb
+- lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_indexes.rb.erb
+- lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_per_tag_budgets.rb.erb
 - lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_call_rollups_provider_generator.rb
 - lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_call_tags_key_value_index_generator.rb
 - lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_image_tokens_generator.rb
+- lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_indexes_generator.rb
+- lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_per_tag_budgets_generator.rb
 - lib/llm_cost_tracker/ingestion.rb
 - lib/llm_cost_tracker/ingestion/batch.rb
 - lib/llm_cost_tracker/ingestion/inbox.rb