llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/pricing/sync.rb

@@ -8,7 +8,6 @@ require "rubygems"
 require_relative "registry"
 require_relative "sync/fetcher"
 require_relative "sync/registry_diff"
-require_relative "sync/registry_loader"
 require_relative "sync/registry_writer"
 
 module LlmCostTracker
@@ -39,7 +38,7 @@ module LlmCostTracker
 
         def refresh(path: DEFAULT_OUTPUT_PATH, url: DEFAULT_REMOTE_URL, preview: false, fetcher: Fetcher.new,
                     today: Date.today)
-          current = RegistryLoader.new.call(path: path, seed_path: Registry::DEFAULT_PRICES_PATH)
+          current = load_registry(path)
           response = fetcher.get(url, etag: current.dig("metadata", "source_version"))
 
           if response.not_modified
@@ -55,7 +54,10 @@ module LlmCostTracker
           end
 
           remote = normalize_remote_registry(response.body, url: url, response: response, today: today)
-          RegistryWriter.new.call(path: path, registry: remote) unless preview
+          unless preview
+            RegistryWriter.new.call(path: path, registry: remote)
+            invalidate_pricing_caches!
+          end
           refresh_result(
             path: path,
             url: url,
@@ -67,8 +69,14 @@ module LlmCostTracker
           )
         end
 
+        def invalidate_pricing_caches!
+          Pricing::Lookup.reset!
+          Pricing::Registry.reset!
+          Pricing::ServiceCharges.reset!
+        end
+
         def check(path: DEFAULT_OUTPUT_PATH, url: DEFAULT_REMOTE_URL, fetcher: Fetcher.new, today: Date.today)
-          current = RegistryLoader.new.call(path: path, seed_path: Registry::DEFAULT_PRICES_PATH)
+          current = load_registry(path)
           response = fetcher.get(url, etag: current.dig("metadata", "source_version"))
 
           if response.not_modified
@@ -82,7 +90,7 @@ module LlmCostTracker
           end
 
           remote = normalize_remote_registry(response.body, url: url, response: response, today: today)
-          changes = RegistryDiff.call(current.fetch("models", {}), remote.fetch("models", {}))
+          changes = registry_changes(current, remote)
 
           CheckResult.new(
             path: path,
@@ -118,10 +126,15 @@ module LlmCostTracker
             raise Error, "remote pricing snapshot requires llm_cost_tracker >= #{min_gem_version}"
           end
 
-          models = registry.fetch("models", {})
-          Registry.normalize_price_table(models)
+          raw_models = registry.fetch("models", {})
+          models = Registry.normalize_price_table(raw_models).each_with_object({}) do |(model, prices), normalized|
+            model_metadata = (raw_models[model] || {}).slice(*Registry::METADATA_KEYS)
+            normalized[model] = model_metadata.merge(prices.to_h { |key, value| [key.name, value] })
+          end
+          service_charges = registry["service_charges"]
+          ServiceCharges.rates_from_registry(registry, context: "remote pricing snapshot") if service_charges
 
-          registry.merge(
+          normalized = {
             "metadata" => metadata.merge(
               "schema_version" => schema_version,
               "updated_at" => metadata["updated_at"] || today.iso8601,
@@ -129,11 +142,19 @@ module LlmCostTracker
               "source_version" => response.source_version
             ),
             "models" => models
-          )
+          }
+          normalized["service_charges"] = service_charges if service_charges.present?
+          normalized
         rescue ArgumentError, TypeError => e
           raise Error, "Unable to load remote pricing snapshot: #{e.message}"
         end
 
+        def load_registry(path)
+          YAML.safe_load_file(path, aliases: false) || {}
+        rescue Errno::ENOENT, Psych::Exception, ArgumentError, TypeError => e
+          raise Error, "Unable to load pricing registry #{path.inspect}: #{e.message}"
+        end
+
         def parse_registry(body)
           registry = JSON.parse(body.to_s)
           raise Error, "remote pricing snapshot must be a JSON object" unless registry.is_a?(Hash)
@@ -148,11 +169,37 @@ module LlmCostTracker
             path: path,
             source_url: url,
             source_version: response.source_version,
-            changes: RegistryDiff.call(current.fetch("models", {}), remote.fetch("models", {})),
+            changes: registry_changes(current, remote),
             written: written,
             not_modified: not_modified
           )
         end
+
+        def registry_changes(current, remote)
+          model_changes = RegistryDiff.call(current.fetch("models", {}), remote.fetch("models", {}))
+          charge_changes = service_charges_diff(
+            current.fetch("service_charges", {}),
+            remote.fetch("service_charges", {})
+          )
+          return model_changes if charge_changes.empty?
+
+          model_changes.merge("service_charges" => charge_changes)
+        end
+
+        def service_charges_diff(current, remote)
+          (current.keys | remote.keys).sort.each_with_object({}) do |provider, changes|
+            current_rates = (current[provider] || {}).transform_keys(&:to_s)
+            remote_rates = (remote[provider] || {}).transform_keys(&:to_s)
+            (current_rates.keys | remote_rates.keys).sort.each_with_object(changes) do |component, _|
+              from = current_rates[component]
+              to = remote_rates[component]
+              next if from == to
+
+              changes[provider] ||= {}
+              changes[provider][component] = { "from" => from, "to" => to }
+            end
+          end
+        end
       end
     end
   end

data/lib/llm_cost_tracker/pricing/sync_change_printer.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Pricing
+    module SyncChangePrinter
+      class << self
+        def call(changes, output: $stdout)
+          service_changes = changes["service_charges"]
+          model_changes = changes.except("service_charges")
+
+          output.puts "  changed models: #{model_changes.size}"
+          model_changes.each do |model, fields|
+            output.puts "    - #{model}"
+            fields.each do |field, values|
+              output.puts "      #{field}: #{values['from'].inspect} -> #{values['to'].inspect}"
+            end
+          end
+
+          return if service_changes.nil? || service_changes.empty?
+
+          output.puts "  changed service charges: #{service_changes.values.sum(&:size)}"
+          service_changes.each do |provider, components|
+            components.each do |component, values|
+              output.puts "    - #{provider}.#{component}: " \
+                          "#{values['from'].inspect} -> #{values['to'].inspect}"
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/pricing.rb

@@ -1,75 +1,239 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/hash/keys"
 require "active_support/core_ext/object/blank"
+require "time"
 
-require_relative "pricing/components"
+require_relative "version"
+require_relative "token_usage"
+require_relative "billing/components"
 require_relative "pricing/registry"
 require_relative "pricing/lookup"
 require_relative "pricing/effective_prices"
 require_relative "pricing/explainer"
+require_relative "pricing/service_charges"
 
 module LlmCostTracker
   module Pricing
-    STANDARD_MODE_VALUES = %w[auto default standard standard_only].freeze
-    private_constant :STANDARD_MODE_VALUES
+    extend ServiceCharges
+
+    STANDARD_MODE_VALUES = %i[auto default standard standard_only].freeze
+    RATE_DENOMINATOR_TOKENS = 1_000_000
+    private_constant :STANDARD_MODE_VALUES, :RATE_DENOMINATOR_TOKENS
 
     class << self
       def normalize_mode(value)
-        mode = value.to_s.strip.presence
+        return nil if value.nil?
+
+        mode = value.is_a?(Symbol) ? value.downcase : normalize_string_mode(value)
         return nil unless mode
 
-        mode = mode.tr("-", "_")
         STANDARD_MODE_VALUES.include?(mode) ? nil : mode
       end
 
-      def cost_for(provider:, model:, token_usage:, pricing_mode: nil)
-        prices = lookup(provider: provider, model: model)
-        return nil unless prices
+      def cost_for(provider:, model:, tokens:, pricing_mode: nil)
+        calculation = calculation_for(
+          provider: provider,
+          model: model,
+          tokens: tokens,
+          pricing_mode: pricing_mode
+        )
+        return nil unless calculation
+
+        cost_from(calculation)
+      end
+
+      def cost_and_snapshot_for(provider:, model:, tokens:, pricing_mode: nil)
+        calculation = calculation_for(
+          provider: provider,
+          model: model,
+          tokens: tokens,
+          pricing_mode: pricing_mode
+        )
+        return [nil, nil] unless calculation
 
-        costs = calculate_costs(token_usage, prices, pricing_mode: pricing_mode)
-        return nil unless costs
+        [cost_from(calculation), snapshot_from(calculation)]
+      end
 
-        values = COMPONENTS.to_h do |component|
-          [component.cost_key, costs.fetch(component.price_key).round(8)]
+      def price_line_items(provider:, model:, line_items:, pricing_mode: nil)
+        token_usage = TokenUsage.build_from_tokens(token_attributes_from(line_items))
+        calculation = calculation_for(provider: provider, model: model, tokens: token_usage, pricing_mode: pricing_mode)
+        snapshot = calculation && snapshot_from(calculation)
+
+        priced = line_items.map do |line_item|
+          next price_token_line_item(line_item, calculation) if line_item.unit == :token
+
+          price_service_charge_line_item(line_item, provider: provider, pricing_mode: pricing_mode)
         end
 
-        values.merge(total_cost: costs.values.sum.round(8))
+        [priced, snapshot]
       end
 
-      def lookup(provider:, model:)
-        Lookup.call(provider: provider, model: model)&.prices
+      def snapshot_for(provider:, model:, tokens:, pricing_mode: nil)
+        calculation = calculation_for(
+          provider: provider,
+          model: model,
+          tokens: tokens,
+          pricing_mode: pricing_mode
+        )
+        return nil unless calculation
+
+        snapshot_from(calculation)
       end
 
-      def explain(provider:, model:, token_usage:, pricing_mode: nil)
+      def explain(provider:, model:, tokens:, pricing_mode: nil)
         Explainer.call(
           provider: provider,
           model: model,
-          token_usage: token_usage,
+          tokens: tokens,
           pricing_mode: pricing_mode
         )
       end
 
       def stored_cost_attributes(attributes)
-        attributes.to_h.symbolize_keys.slice(*COST_KEYS).compact
+        value = attributes.to_h[:total_cost]
+        value.nil? ? {} : { total_cost: value }
       end
 
       private
 
-      def calculate_costs(usage, prices, pricing_mode:)
-        effective = EffectivePrices.call(usage: usage, prices: prices, pricing_mode: pricing_mode)
-        return nil if effective.value?(nil)
+      def normalize_string_mode(value)
+        normalized = value.strip
+        return nil if normalized.empty?
+
+        normalized.downcase.tr("-", "_").to_sym
+      end
+
+      def cost_from(calculation)
+        costs = calculation[:costs]
+        values = Billing::Components::TOKEN_PRICED.each_with_object({}) do |component, result|
+          cost = costs[component.key]
+          result[component.cost_key] = cost.round(8) unless cost.nil?
+        end
+        values[:total_cost] = costs.values.compact.sum.round(8)
+        values
+      end
+
+      def snapshot_from(calculation)
+        match = calculation[:match]
+        effective = calculation[:effective]
+        token_usage = calculation[:token_usage]
+        rates = Billing::Components::TOKEN_PRICED.each_with_object({}) do |component, values|
+          quantity = token_usage.public_send(component.token_key)
+          price = effective[component.key]
+          next if quantity.zero? || price.nil?
+
+          values[component.key] = {
+            amount: price,
+            quantity: RATE_DENOMINATOR_TOKENS
+          }
+        end
+
+        {
+          schema_version: 1,
+          source: match.source,
+          source_key: match.key,
+          source_version: source_version_for(match.source),
+          matched_by: match.matched_by,
+          currency: "USD",
+          rates: rates
+        }
+      end
+
+      def calculation_for(provider:, model:, tokens:, pricing_mode:)
+        match = Lookup.call(provider: provider, model: model)
+        return nil unless match
+
+        token_usage = TokenUsage.build_from_tokens(tokens)
+        mode = normalize_mode(pricing_mode)
+        effective = EffectivePrices.call(usage: token_usage, prices: match.prices, pricing_mode: mode)
+        return nil unless any_billable_priced?(token_usage, effective)
+
+        { match: match, effective: effective, token_usage: token_usage, costs: costs_for(token_usage, effective) }
+      end
+
+      def any_billable_priced?(token_usage, effective)
+        billable = Billing::Components::TOKEN_PRICED.select { |c| token_usage.public_send(c.token_key).positive? }
+        billable.empty? || billable.any? { |c| effective[c.key] }
+      end
+
+      def costs_for(usage, effective)
+        Billing::Components::TOKEN_PRICED.to_h do |component|
+          tokens = usage.public_send(component.token_key)
+          [component.key, token_cost(tokens, effective[component.key])]
+        end
+      end
+
+      def token_attributes_from(line_items)
+        line_items.each_with_object({}) do |line_item, totals|
+          next unless line_item.unit == :token
+
+          component = component_for_line_item(line_item)
+          next unless component
+
+          totals[component.key] = (totals[component.key] || 0) + line_item.quantity.to_i
+        end
+      end
+
+      def price_token_line_item(line_item, calculation)
+        component = component_for_line_item(line_item)
+        return line_item unless component
+        return line_item.with(cost_status: Billing::CostStatus::UNKNOWN) unless calculation
+
+        effective_price = calculation[:effective][component.key]
+        return line_item.with(cost_status: Billing::CostStatus::UNKNOWN) if effective_price.nil?
+
+        cost = (line_item.quantity * BigDecimal(effective_price.to_s)) / RATE_DENOMINATOR_TOKENS
+        match = calculation[:match]
+        line_item.with(
+          rate_amount: BigDecimal(effective_price.to_s),
+          rate_quantity: BigDecimal(RATE_DENOMINATOR_TOKENS),
+          cost: cost,
+          cost_status: cost.zero? ? Billing::CostStatus::FREE : Billing::CostStatus::COMPLETE,
+          price_key: component.key,
+          price_source: match.source,
+          price_source_version: source_version_for(match.source)
+        )
+      end
+
+      def price_service_charge_line_item(line_item, provider:, pricing_mode:)
+        return line_item if line_item.priced?
+        return line_item unless line_item.billable?
+
+        rate = charge_rate(provider: provider, component: line_item.kind, pricing_mode: pricing_mode)
+        return line_item unless rate
+
+        line_item.apply_rate(rate)
+      end
+
+      def component_for_line_item(line_item)
+        Billing::Components::REGISTRY.find do |component|
+          component.kind == line_item.kind &&
+            component.direction == line_item.direction &&
+            component.modality == line_item.modality &&
+            component.cache_state == line_item.cache_state &&
+            component.unit == line_item.unit
+        end
+      end
 
-        usage.price_quantities.to_h do |key, tokens|
-          [key, token_cost(tokens, effective.fetch(key))]
+      def source_version_for(source)
+        case source
+        when :bundled
+          LlmCostTracker::VERSION
+        when :prices_file
+          path = LlmCostTracker.configuration.prices_file
+          path ? File.mtime(path).utc.iso8601 : nil
+        when :pricing_overrides
+          "configuration"
         end
+      rescue Errno::ENOENT
+        nil
       end
 
       def token_cost(tokens, per_million_price)
-        return 0.0 if tokens.to_i.zero?
+        return 0.0 if tokens.zero?
         return nil if per_million_price.nil?
 
-        (tokens.to_f / 1_000_000) * per_million_price
+        (tokens * per_million_price) / RATE_DENOMINATOR_TOKENS
       end
     end
   end

data/lib/llm_cost_tracker/railtie.rb

@@ -9,16 +9,8 @@ module LlmCostTracker
     end
 
     generators do
-      require_relative "generators/llm_cost_tracker/add_ingestion_generator"
-      require_relative "generators/llm_cost_tracker/add_period_totals_generator"
-      require_relative "generators/llm_cost_tracker/add_latency_ms_generator"
-      require_relative "generators/llm_cost_tracker/add_provider_response_id_generator"
-      require_relative "generators/llm_cost_tracker/add_streaming_generator"
-      require_relative "generators/llm_cost_tracker/add_token_usage_generator"
       require_relative "generators/llm_cost_tracker/install_generator"
       require_relative "generators/llm_cost_tracker/prices_generator"
-      require_relative "generators/llm_cost_tracker/upgrade_cost_precision_generator"
-      require_relative "generators/llm_cost_tracker/upgrade_tags_to_jsonb_generator"
     end
 
     rake_tasks do

data/lib/llm_cost_tracker/report/data.rb

@@ -32,17 +32,18 @@ module LlmCostTracker
           breakdown_limit = nil unless breakdown_limit.positive?
         end
         from = now - days.days
-        scope = Ledger::Call.where(tracked_at: from..now)
+        scope = LlmCostTracker::Call.where(tracked_at: from..now)
         tag_breakdowns ||= LlmCostTracker.configuration.report_tag_breakdowns || []
+        aggregate = totals(scope)
 
         new(
           days: days,
           from_time: from,
           to_time: now,
-          total_cost: scope.sum(:total_cost).to_f,
-          requests_count: scope.count,
-          average_latency_ms: average_latency_ms(scope),
-          unknown_pricing_count: scope.where(total_cost: nil).count,
+          total_cost: aggregate.total_cost.to_f,
+          requests_count: aggregate.requests_count.to_i,
+          average_latency_ms: aggregate.average_latency_ms&.to_f,
+          unknown_pricing_count: aggregate.unknown_pricing_count.to_i,
           cost_by_provider: scope.cost_by_provider(limit: breakdown_limit).to_a,
           cost_by_model: scope.cost_by_model(limit: breakdown_limit).to_a,
           cost_by_tags: cost_by_tags(scope, tag_breakdowns, limit: breakdown_limit),
@@ -50,8 +51,15 @@ module LlmCostTracker
         )
       end
 
-      def self.average_latency_ms(scope)
-        scope.average(:latency_ms)&.to_f
+      def self.totals(scope)
+        scope
+          .select(
+            "COALESCE(SUM(total_cost), 0) AS total_cost, " \
+            "COUNT(*) AS requests_count, " \
+            "AVG(latency_ms) AS average_latency_ms, " \
+            "COALESCE(SUM(CASE WHEN total_cost IS NULL THEN 1 ELSE 0 END), 0) AS unknown_pricing_count"
+          )
+          .take
       end
 
       def self.cost_by_tags(scope, keys, limit:)
@@ -66,7 +74,7 @@ module LlmCostTracker
           .to_a
       end
 
-      private_class_method :average_latency_ms, :cost_by_tags, :top_calls
+      private_class_method :cost_by_tags, :top_calls, :totals
     end
   end
 end

data/lib/llm_cost_tracker/report.rb

@@ -20,10 +20,6 @@ module LlmCostTracker
       rescue StandardError => e
         "Unable to build LLM cost report: #{e.class}: #{e.message}"
       end
-
-      def data(days: Data::DEFAULT_DAYS, now: Time.now.utc, tag_breakdowns: nil)
-        Data.build(days: days, now: now, tag_breakdowns: tag_breakdowns)
-      end
     end
   end
 end

data/lib/llm_cost_tracker/retention.rb

@@ -49,20 +49,20 @@ module LlmCostTracker
       end
 
       def prune_batch(cutoff, batch_size)
-        LlmCostTracker::Ledger::Call.transaction do
-          rows = LlmCostTracker::Ledger::Call
-                 .where(tracked_at: ...cutoff)
-                 .order(:id)
-                 .limit(batch_size)
-                 .lock
-                 .pluck(:id, :tracked_at, :total_cost)
+        LlmCostTracker::Call.transaction do
+          rows = pluck_prunable(cutoff, batch_size)
           next 0 if rows.empty?
 
-          deleted = LlmCostTracker::Ledger::Call.where(id: rows.map(&:first)).delete_all
+          deleted = LlmCostTracker::Call.where(id: rows.map(&:first)).delete_all
           LlmCostTracker::Ledger::Rollups.decrement!(rows) if deleted.positive?
           deleted
         end
       end
+
+      def pluck_prunable(cutoff, batch_size)
+        LlmCostTracker::Call.where(tracked_at: ...cutoff).order(:id).limit(batch_size).lock
+                            .pluck(:id, :tracked_at, :total_cost, :pricing_snapshot)
+      end
     end
   end
 end

data/lib/llm_cost_tracker/tags/context.rb

@@ -19,11 +19,9 @@ module LlmCostTracker
 
         def tags
           default_tags = LlmCostTracker.configuration.default_tags
-          default_tags = default_tags.call if default_tags.respond_to?(:call)
+          default_tags = default_tags.call.deep_dup if default_tags.respond_to?(:call)
 
-          (default_tags || {}).deep_dup.to_h.merge(
-            (ActiveSupport::IsolatedExecutionState[KEY] || []).reduce({}) { |merged, tags| merged.merge(tags) }
-          )
+          default_tags.to_h.merge(*Array(ActiveSupport::IsolatedExecutionState[KEY]))
         end
 
         def clear!

data/lib/llm_cost_tracker/tags/key.rb

@@ -4,10 +4,14 @@ module LlmCostTracker
   module Tags
     module Key
       PATTERN = /\A[\w.-]+\z/
+      MAX_BYTESIZE = 64
 
       class << self
         def validate!(key, error_class: ArgumentError)
           key = key.to_s
+          if key.bytesize > MAX_BYTESIZE
+            raise error_class, "tag key exceeds #{MAX_BYTESIZE} bytes: #{key[0, 16].inspect}..."
+          end
           return key if key.match?(PATTERN)
 
           raise error_class, "invalid tag key: #{key.inspect}"

data/lib/llm_cost_tracker/tags/sanitizer.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/string/inflections"
 require "json"
 
 module LlmCostTracker
@@ -10,39 +11,33 @@ module LlmCostTracker
       class << self
         def call(tags, config: LlmCostTracker.configuration)
           tags = (tags || {}).to_h
+          redacted = Array(config.redacted_tag_keys).map { |key| normalized_key(key) }
+          limit = [config.max_tag_value_bytesize.to_i, 0].max
           tags.first([config.max_tag_count.to_i, 0].max).each_with_object({}) do |(key, value), sanitized|
-            sanitized[key] = sanitized_value(key, value, config)
+            sanitized[key] = sanitized_value(key, value, redacted, limit)
           end
         end
 
         private
 
-        def sanitized_value(key, value, config)
-          return REDACTED_VALUE if redacted_key?(key, config)
+        def sanitized_value(key, value, redacted, limit)
+          return REDACTED_VALUE if redacted_key?(key, redacted)
 
           string = value_string(value)
-          limit = [config.max_tag_value_bytesize.to_i, 0].max
           return value if string.bytesize <= limit
 
-          string.byteslice(0, limit).to_s.encode("UTF-8", invalid: :replace, undef: :replace)
+          string.byteslice(0, limit).encode("UTF-8", invalid: :replace, undef: :replace)
         end
 
-        def redacted_key?(key, config)
+        def redacted_key?(key, redacted)
+          return false if redacted.empty?
+
           normalized = normalized_key(key)
-          Array(config.redacted_tag_keys).map { |redacted_key| normalized_key(redacted_key) }.any? do |candidate|
-            redacted_key_component?(normalized, candidate)
-          end
+          redacted.any? { |candidate| redacted_key_component?(normalized, candidate) }
         end
 
         def normalized_key(key)
-          key.to_s
-             .gsub(/([A-Z]+)([A-Z][a-z])/, '\1_\2')
-             .gsub(/([a-z\d])([A-Z])/, '\1_\2')
-             .downcase
-             .gsub(/[^a-z0-9]+/, "_")
-             .gsub(/_+/, "_")
-             .delete_prefix("_")
-             .delete_suffix("_")
+          key.to_s.underscore.gsub(/[^a-z0-9]+/, "_").delete_prefix("_").delete_suffix("_")
         end
 
         def redacted_key_component?(key, candidate)

data/lib/llm_cost_tracker/timing.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Timing
+    module_function
+
+    def now_monotonic
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
+    def elapsed_ms(started_at)
+      ((now_monotonic - started_at) * 1000).round
+    end
+  end
+end