llm_cost_tracker 0.7.3 → 0.8.0

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

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