llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/pricing/lookup.rb

@@ -7,42 +7,18 @@ module LlmCostTracker
       MUTEX = Mutex.new
       CACHE_MISS = Object.new.freeze
       NO_MATCH = Object.new.freeze
-      MAX_LOOKUP_CACHE_ENTRIES = 512
 
       class << self
         def call(provider:, model:)
           provider_name = provider.to_s.presence
           model_name = model.to_s
+          return nil if model_name.empty?
+
           cache_key = [provider_name, model_name]
           cached = cached_lookup(cache_key)
           return cached unless cached.equal?(CACHE_MISS)
 
-          provider_model = provider_name ? "#{provider_name}/#{model_name}" : model_name
-          normalized_model = normalize_model_name(model_name)
-          current = current_price_tables
-
-          match =
-            explain_table(
-              table: current.fetch(:pricing_overrides),
-              source: :pricing_overrides,
-              provider_model: provider_model,
-              model_name: model_name,
-              normalized_model: normalized_model
-            ) ||
-            explain_table(
-              table: current.fetch(:file_prices),
-              source: :prices_file,
-              provider_model: provider_model,
-              model_name: model_name,
-              normalized_model: normalized_model
-            ) ||
-            explain_table(
-              table: Registry.builtin_prices,
-              source: :bundled,
-              provider_model: provider_model,
-              model_name: model_name,
-              normalized_model: normalized_model
-            )
+          match = lookup_match(provider_name: provider_name, model_name: model_name)
           cache_lookup(cache_key, match)
           match
         end
@@ -57,6 +33,32 @@ module LlmCostTracker
 
         private
 
+        def lookup_match(provider_name:, model_name:)
+          provider_model = provider_name ? "#{provider_name}/#{model_name}" : model_name
+          normalized_model = normalize_model_name(model_name)
+          current = current_price_tables
+
+          ordered_table_lookups(current).each do |source, table|
+            match = explain_table(
+              table: table,
+              source: source,
+              provider_model: provider_model,
+              model_name: model_name,
+              normalized_model: normalized_model
+            )
+            return match if match
+          end
+          nil
+        end
+
+        def ordered_table_lookups(current)
+          [
+            [:pricing_overrides, current.fetch(:pricing_overrides)],
+            [:prices_file, current.fetch(:file_prices)],
+            [:bundled, Registry.builtin_prices]
+          ]
+        end
+
         def current_price_tables
           cached = @prices_cache
           return cached if cached
@@ -67,8 +69,7 @@ module LlmCostTracker
 
             config = LlmCostTracker.configuration
             file_prices = Registry.file_prices(config.prices_file)
-            overrides = Registry.normalize_price_table(config.pricing_overrides)
-            value = { pricing_overrides: overrides, file_prices: file_prices }.freeze
+            value = { pricing_overrides: config.pricing_overrides, file_prices: file_prices }.freeze
             @prices_cache = value
             value
           end
@@ -85,7 +86,6 @@ module LlmCostTracker
         def cache_lookup(cache_key, match)
           MUTEX.synchronize do
             values = (@lookup_cache || {}).dup
-            values.clear if values.size >= MAX_LOOKUP_CACHE_ENTRIES
             values[cache_key] = match || NO_MATCH
             @lookup_cache = values.freeze
           end
@@ -135,7 +135,7 @@ module LlmCostTracker
         end
 
         def match(table:, source:, key:, matched_by:)
-          Match.new(source: source.to_s, key: key, prices: table[key], matched_by: matched_by.to_s)
+          Match.new(source: source, key: key, prices: table[key], matched_by: matched_by)
         end
 
         def snapshot_variant?(model, key)
@@ -147,14 +147,18 @@ module LlmCostTracker
 
         def sorted_price_keys(table)
           cached = @sorted_price_keys_cache
-          return cached[:keys] if cached && cached[:table].equal?(table)
+          existing = cached && cached[table]
+          return existing if existing
 
           MUTEX.synchronize do
             cached = @sorted_price_keys_cache
-            return cached[:keys] if cached && cached[:table].equal?(table)
+            existing = cached && cached[table]
+            return existing if existing
 
             keys = table.keys.sort_by { |key| -key.length }
-            @sorted_price_keys_cache = { table: table, keys: keys }.freeze
+            next_cache = cached ? cached.dup : {}.compare_by_identity
+            next_cache[table] = keys
+            @sorted_price_keys_cache = next_cache.freeze
             keys
           end
         end

data/lib/llm_cost_tracker/pricing/registry.rb

@@ -1,9 +1,8 @@
 # frozen_string_literal: true
 
-require "json"
 require "yaml"
 
-require_relative "components"
+require_relative "../billing/components"
 require_relative "../logging"
 
 module LlmCostTracker
@@ -11,42 +10,58 @@ module LlmCostTracker
     module Registry
       DEFAULT_PRICES_PATH = File.expand_path("../prices.json", __dir__)
       EMPTY_PRICES = {}.freeze
-      PRICE_KEYS = Pricing::COMPONENTS.map { |component| component.price_key.to_s }.freeze
-      METADATA_KEYS = %w[
-        _source _source_version _fetched_at _updated _notes _validator_override
-        _context_price_threshold_tokens
+      CONTEXT_THRESHOLD_KEY = :_context_price_threshold_tokens
+      PRICE_KEYS = Billing::Components::TOKEN_PRICED.map { |component| component.key.name }.freeze
+      METADATA_KEYS = [
+        "_source", "_source_version", "_fetched_at", "_updated", "_notes", "_validator_override",
+        CONTEXT_THRESHOLD_KEY.name
       ].freeze
-      MAX_FILE_BYTES = 2_097_152
       MUTEX = Mutex.new
 
       class << self
+        def reset!
+          MUTEX.synchronize do
+            @builtin_prices = nil
+            @metadata = nil
+            @raw_registry = nil
+            @file_prices_cache = nil
+          end
+        end
+
         def builtin_prices
           cached = @builtin_prices
           return cached if cached
 
-          value = normalize_price_table(raw_registry.fetch("models", {})).freeze
-          MUTEX.synchronize { @builtin_prices ||= value }
+          MUTEX.synchronize do
+            @builtin_prices ||= begin
+              registry = @raw_registry ||= load_raw_registry
+              normalize_price_table(registry.fetch("models", {})).freeze
+            end
+          end
         end
 
         def metadata
           cached = @metadata
           return cached if cached
 
-          value = raw_registry.fetch("metadata", {}).freeze
-          MUTEX.synchronize { @metadata ||= value }
+          MUTEX.synchronize do
+            @metadata ||= begin
+              registry = @raw_registry ||= load_raw_registry
+              registry.fetch("metadata", {}).freeze
+            end
+          end
         end
 
         def file_metadata(path)
           return {} unless path
 
-          registry = load_price_file(path.to_s)
-          raise ArgumentError, "prices_file must be a hash" unless registry.is_a?(Hash)
+          registry = YAML.safe_load_file(path, aliases: false) || {}
 
           metadata = registry.fetch("metadata", {})
           raise ArgumentError, "prices_file metadata must be a hash" unless metadata.is_a?(Hash)
 
           metadata
-        rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError => e
+        rescue Errno::ENOENT, Psych::Exception, ArgumentError, TypeError => e
           raise Error, "Unable to load prices_file #{path.inspect}: #{e.message}"
         end
 
@@ -57,8 +72,7 @@ module LlmCostTracker
         def file_prices(path)
           return EMPTY_PRICES unless path
 
-          path = path.to_s
-          cache_key = [path, File.mtime(path).to_f]
+          cache_key = [path, File.mtime(path)]
           cached = @file_prices_cache
           return cached[:value] if cached && cached[:key] == cache_key
 
@@ -66,11 +80,12 @@ module LlmCostTracker
             cached = @file_prices_cache
             return cached[:value] if cached && cached[:key] == cache_key
 
-            value = normalize_price_entries(price_file_models(load_price_file(path)), context: path).freeze
+            registry = YAML.safe_load_file(path, aliases: false) || {}
+            value = normalize_price_entries(registry.fetch("models", registry), context: path).freeze
             @file_prices_cache = { key: cache_key, value: value }.freeze
             value
           end
-        rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError => e
+        rescue Errno::ENOENT, Psych::Exception, ArgumentError, TypeError => e
           raise Error, "Unable to load prices_file #{path.inspect}: #{e.message}"
         end
 
@@ -80,20 +95,31 @@ module LlmCostTracker
           cached = @raw_registry
           return cached if cached
 
-          MUTEX.synchronize { @raw_registry ||= JSON.parse(File.read(DEFAULT_PRICES_PATH)).freeze }
+          MUTEX.synchronize { @raw_registry ||= load_raw_registry }
+        end
+
+        def load_raw_registry
+          YAML.safe_load_file(DEFAULT_PRICES_PATH, aliases: false).freeze
         end
 
         def normalize_price_entry(price)
           price.each_with_object({}) do |(key, value), normalized|
-            key = key.to_s
-            if price_key?(key)
-              normalized[key.to_sym] = Float(value)
-            elsif key == "_context_price_threshold_tokens"
-              normalized[key.to_sym] = Integer(value)
+            key = registry_key_for(key)
+            if key == CONTEXT_THRESHOLD_KEY
+              normalized[key] = Integer(value)
+            elsif key
+              normalized[key] = non_negative_float(key, value)
             end
           end
         end
 
+        def non_negative_float(key, value)
+          rate = Float(value)
+          raise ArgumentError, "price for #{key.inspect} must be non-negative (got #{rate})" if rate.negative?
+
+          rate
+        end
+
         def normalize_price_entries(table, context:)
           table = {} if table.nil?
           raise ArgumentError, "#{context} must be a hash of models" unless table.is_a?(Hash)
@@ -106,8 +132,8 @@ module LlmCostTracker
         end
 
         def warn_unknown_keys(model, price, path)
-          unknown_keys = price.keys.map(&:to_s).reject do |key|
-            price_key?(key) || METADATA_KEYS.include?(key)
+          unknown_keys = price.keys.reject do |key|
+            registry_key_for(key) || METADATA_KEYS.include?(key)
           end
           return if unknown_keys.empty?
 
@@ -117,31 +143,25 @@ module LlmCostTracker
           )
         end
 
-        def price_key?(key)
-          return true if PRICE_KEYS.include?(key)
-
-          PRICE_KEYS.any? do |base_key|
-            key.end_with?("_#{base_key}") && key.delete_suffix("_#{base_key}") != ""
-          end
-        end
-
-        def load_price_file(path)
-          raise ArgumentError, "prices_file exceeds #{MAX_FILE_BYTES} bytes" if File.size(path) > MAX_FILE_BYTES
+        def price_key_for(key)
+          name = key.is_a?(Symbol) ? key.name : key
+          Billing::Components::TOKEN_PRICED.each do |candidate|
+            return candidate.key if candidate.key.name == name
 
-          contents = File.read(path)
-          return YAML.safe_load(contents, aliases: false) || {} if yaml_file?(path)
+            suffix = "_#{candidate.key.name}"
+            next unless name.end_with?(suffix)
 
-          JSON.parse(contents)
-        end
+            prefix = name.delete_suffix(suffix)
+            return :"#{prefix}_#{candidate.key.name}" unless prefix.empty?
+          end
 
-        def yaml_file?(path)
-          %w[.yaml .yml].include?(File.extname(path).downcase)
+          nil
         end
 
-        def price_file_models(registry)
-          raise ArgumentError, "prices_file must be a hash" unless registry.is_a?(Hash)
+        def registry_key_for(key)
+          return CONTEXT_THRESHOLD_KEY if key == CONTEXT_THRESHOLD_KEY || key == CONTEXT_THRESHOLD_KEY.name
 
-          registry.fetch("models", registry)
+          price_key_for(key)
         end
 
         def validate_price_entry(price, model:, context:)

data/lib/llm_cost_tracker/pricing/service_charges.rb

@@ -0,0 +1,204 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/object/blank"
+require "bigdecimal"
+require "time"
+require "yaml"
+
+require_relative "../billing/components"
+require_relative "registry"
+
+module LlmCostTracker
+  module Pricing
+    module ServiceCharges
+      extend self
+
+      DEFAULT_CURRENCY = "USD"
+      EMPTY_RATES = {}.freeze
+      MUTEX = Mutex.new
+
+      def reset!
+        MUTEX.synchronize do
+          @builtin_rates = nil
+          @file_rates_cache = nil
+        end
+      end
+
+      def builtin_rates
+        cached = @builtin_rates
+        return cached if cached
+
+        MUTEX.synchronize do
+          @builtin_rates ||= begin
+            registry = YAML.safe_load_file(Registry::DEFAULT_PRICES_PATH, aliases: false) || {}
+            rates_from_registry(registry).freeze
+          end
+        end
+      end
+
+      def file_rates(path)
+        return EMPTY_RATES unless path
+
+        cache_key = [path, File.mtime(path)]
+        cached = @file_rates_cache
+        return cached[:value] if cached && cached[:key] == cache_key
+
+        MUTEX.synchronize do
+          cached = @file_rates_cache
+          return cached[:value] if cached && cached[:key] == cache_key
+
+          registry = YAML.safe_load_file(path, aliases: false) || {}
+          value = rates_from_registry(registry, context: path).freeze
+          @file_rates_cache = { key: cache_key, value: value }.freeze
+          value
+        end
+      rescue Errno::ENOENT, Psych::Exception, ArgumentError, TypeError => e
+        raise Error, "Unable to load prices_file #{path.inspect}: #{e.message}"
+      end
+
+      def rates_from_registry(registry, context: "price registry")
+        data = registry.fetch("service_charges", EMPTY_RATES)
+        raise ArgumentError, "#{context} service_charges must be a hash" unless data.is_a?(Hash)
+
+        data.each_with_object({}) do |(provider, entries), rates|
+          section_context = "#{context} service_charges.#{provider}"
+          rates[provider] = rates_from_section(entries, context: section_context)
+        end
+      end
+
+      def charge_rate(provider:, component:, pricing_mode:)
+        pricing_mode = Pricing.normalize_mode(pricing_mode)
+        match = charge_rate_match(provider: provider, component: component, pricing_mode: pricing_mode)
+        return nil unless match
+
+        rate = match.fetch(:rate)
+        {
+          amount: rate.fetch(:amount),
+          quantity: rate.fetch(:quantity),
+          currency: rate.fetch(:currency),
+          source: match.fetch(:source),
+          source_key: match.fetch(:key),
+          source_version: rate_source_version_for(match.fetch(:source))
+        }
+      end
+
+      private
+
+      def rates_from_section(entries, context:)
+        raise ArgumentError, "#{context} must be a hash" unless entries.is_a?(Hash)
+
+        entries.each_with_object({}) do |(key, amount), rates|
+          key = key.name if key.is_a?(Symbol)
+          component, tier = component_and_tier_for(key, context: context)
+          amount = amount_for(key, amount, context: context)
+
+          rate = {
+            amount: amount,
+            quantity: rate_quantity(component),
+            currency: DEFAULT_CURRENCY,
+            source_key: key
+          }
+          component_rates = rates[component.key] ||= { tiers: {} }
+          (tier ? component_rates[:tiers] : component_rates)[tier || :default] = rate
+        end
+      end
+
+      def component_and_tier_for(key, context:)
+        Billing::Components::REGISTRY.each do |component|
+          next if component.token_key
+
+          return [component, nil] if key == component.key.name
+
+          suffix = "_#{component.key.name}"
+          next unless key.end_with?(suffix)
+
+          tier = key.delete_suffix(suffix)
+          return [component, :"#{tier}"] unless tier.empty?
+        end
+
+        raise ArgumentError, "service charge price key #{key.inspect} in #{context} uses unknown billing component"
+      end
+
+      def amount_for(key, amount, context:)
+        value = BigDecimal(amount.to_s)
+        message = "service charge price amount for #{key.inspect} in #{context} must be non-negative"
+        raise ArgumentError, message if value.negative?
+
+        value
+      end
+
+      def rate_quantity(component)
+        component.unit == :request ? BigDecimal("1000") : BigDecimal("1")
+      end
+
+      def charge_rate_match(provider:, component:, pricing_mode:)
+        provider_name = provider.is_a?(Symbol) ? provider.name : provider.presence
+        return nil unless provider_name
+
+        component_key = charge_component_key(component)
+
+        table = ServiceCharges.file_rates(LlmCostTracker.configuration.prices_file)
+        provider_table = table.fetch(provider_name, EMPTY_RATES)
+        rate = rate_for(provider_table, component_key: component_key, pricing_mode: pricing_mode)
+        if rate
+          return {
+            source: :prices_file,
+            key: "service_charges.#{provider_name}.#{rate.fetch(:source_key)}",
+            rate: rate
+          }
+        end
+
+        table = ServiceCharges.builtin_rates
+        provider_table = table.fetch(provider_name, EMPTY_RATES)
+        rate = rate_for(provider_table, component_key: component_key, pricing_mode: pricing_mode)
+        return unless rate
+
+        {
+          source: :bundled,
+          key: "service_charges.#{provider_name}.#{rate.fetch(:source_key)}",
+          rate: rate
+        }
+      end
+
+      def rate_for(provider_table, component_key:, pricing_mode:)
+        component_rates = provider_table.fetch(component_key, EMPTY_RATES)
+        tier_rates = component_rates.fetch(:tiers, EMPTY_RATES)
+        if pricing_mode
+          rate = tier_rates[pricing_mode]
+          return rate if rate
+
+          name = pricing_mode.name
+          tier_rates.each do |candidate, candidate_rate|
+            return candidate_rate if tier_includes?(name, candidate.name)
+          end
+        end
+        component_rates[:default]
+      end
+
+      def tier_includes?(tier_name, candidate_name)
+        tier_name == candidate_name ||
+          tier_name.start_with?("#{candidate_name}_") ||
+          tier_name.end_with?("_#{candidate_name}") ||
+          tier_name.include?("_#{candidate_name}_")
+      end
+
+      def charge_component_key(component)
+        billing_component = Billing::Components::BY_KEY[component]
+        return billing_component.key if billing_component && billing_component.token_key.nil?
+
+        raise Error, "Unknown billing component: #{component.inspect}"
+      end
+
+      def rate_source_version_for(source)
+        return LlmCostTracker::VERSION if source == :bundled
+
+        path = LlmCostTracker.configuration.prices_file
+        return nil unless path
+
+        File.mtime(path).utc.iso8601
+      rescue Errno::ENOENT
+        nil
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/pricing/sync/fetcher.rb

@@ -7,6 +7,8 @@ require "openssl"
 require "time"
 require "uri"
 
+require_relative "../../version"
+
 module LlmCostTracker
   module Pricing
     module Sync
@@ -17,7 +19,7 @@ module LlmCostTracker
           end
         end
 
-        USER_AGENT = "llm_cost_tracker price refresh"
+        USER_AGENT = "llm_cost_tracker/#{LlmCostTracker::VERSION} price refresh".freeze
         MAX_REDIRECTS = 5
         MAX_BODY_BYTES = 2_097_152
         OPEN_TIMEOUT = 5
@@ -25,7 +27,8 @@ module LlmCostTracker
         WRITE_TIMEOUT = 10
 
         def get(url, etag: nil, redirects: 0)
-          raise Error, "Too many redirects while fetching #{url}" if redirects > MAX_REDIRECTS
+          safe_url = scrub_url(url)
+          raise Error, "Too many redirects while fetching #{safe_url}" if redirects > MAX_REDIRECTS
 
           uri = URI.parse(url)
           raise Error, "Pricing snapshot URL must use https" unless uri.scheme == "https"
@@ -38,23 +41,34 @@ module LlmCostTracker
 
           case response
           when Net::HTTPSuccess
-            build_response(response, body: body || limited_body(response), not_modified: false)
+            build_response(response, body: body, not_modified: false)
           when Net::HTTPNotModified
             build_response(response, body: nil, not_modified: true)
           when Net::HTTPRedirection
             location = response["location"]
-            raise Error, "Redirect without location while fetching #{url}" if location.blank?
+            raise Error, "Redirect without location while fetching #{safe_url}" if location.blank?
 
             get(URI.join(url, location).to_s, etag: etag, redirects: redirects + 1)
           else
-            raise Error, "Unable to fetch #{url}: HTTP #{response.code}"
+            raise Error, "Unable to fetch #{safe_url}: HTTP #{response.code}"
           end
         rescue OpenSSL::SSL::SSLError, SocketError, SystemCallError, Timeout::Error => e
-          raise Error, "Unable to fetch #{url}: #{e.class}: #{e.message}"
+          raise Error, "Unable to fetch #{scrub_url(url)}: #{e.class}: #{e.message}"
         end
 
         private
 
+        def scrub_url(url)
+          uri = URI.parse(url.to_s)
+          uri.user = nil
+          uri.password = nil
+          uri.query = nil
+          uri.fragment = nil
+          uri.to_s
+        rescue URI::InvalidURIError
+          "[invalid url]"
+        end
+
         def fetch_response(uri, request)
           body = nil
           response = Net::HTTP.start(
@@ -75,19 +89,14 @@ module LlmCostTracker
 
         def limited_body(response)
           body = +""
-          if response.respond_to?(:read_body)
-            response.read_body do |chunk|
-              chunk = chunk.to_s
-              if body.bytesize + chunk.bytesize > MAX_BODY_BYTES
-                raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes"
-              end
-
-              body << chunk
+          response.read_body do |chunk|
+            chunk = chunk.to_s
+            if body.bytesize + chunk.bytesize > MAX_BODY_BYTES
+              raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes"
             end
-          else
-            body = response.body.to_s
+
+            body << chunk
           end
-          raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes" if body.bytesize > MAX_BODY_BYTES
 
           body
         end

data/lib/llm_cost_tracker/pricing/sync/registry_diff.rb

@@ -18,8 +18,8 @@ module LlmCostTracker
           private
 
           def price_field_changes(current_entry, updated_entry)
-            current_price = comparable_price(current_entry)
-            updated_price = comparable_price(updated_entry)
+            current_price = current_entry || {}
+            updated_price = updated_entry || {}
 
             (current_price.keys | updated_price.keys).sort.each_with_object({}) do |field, changes|
               from = current_price[field]
@@ -30,21 +30,12 @@ module LlmCostTracker
             end
           end
 
-          def comparable_price(entry)
-            normalize_hash(entry).slice(*Registry::PRICE_KEYS)
-          end
-
           def normalize_models(models)
-            normalize_hash(models).transform_values { |entry| normalize_hash(entry) }
-          end
-
-          def normalize_hash(hash)
-            return {} if hash.nil?
-            raise Error, "pricing entries must be hashes" unless hash.is_a?(Hash)
-
-            hash.each_with_object({}) do |(key, value), normalized|
-              normalized[key.to_s] = value
+            Registry.normalize_price_table(models).transform_values do |price|
+              price.to_h { |key, value| [key.name, value] }
             end
+          rescue ArgumentError, TypeError => e
+            raise Error, e.message
           end
         end
       end