llm_cost_tracker 0.7.3 → 0.9.0

data/lib/llm_cost_tracker/pricing/mode.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Pricing
+    class Mode
+      COMPOUND_MODIFIERS = %w[data_residency].freeze
+
+      attr_reader :modifiers
+
+      def self.parse(value)
+        return value if value.is_a?(self)
+        return new([]) if value.nil?
+
+        new(tokenize(value.to_s))
+      end
+
+      def self.tokenize(value)
+        remaining = value.to_s.downcase.tr("-", "_")
+        tokens = []
+        loop do
+          break if remaining.empty?
+
+          compound = COMPOUND_MODIFIERS.find do |token|
+            remaining == token || remaining.start_with?("#{token}_")
+          end
+          if compound
+            tokens << compound.to_sym
+            remaining = remaining.delete_prefix(compound).delete_prefix("_")
+          else
+            first, _, rest = remaining.partition("_")
+            tokens << first.to_sym unless first.empty?
+            remaining = rest
+          end
+        end
+        tokens
+      end
+
+      def initialize(modifiers)
+        @modifiers = Array(modifiers).map(&:to_sym).uniq.sort
+        freeze
+      end
+
+      def empty?
+        modifiers.empty?
+      end
+
+      def include?(modifier)
+        modifiers.include?(modifier.to_sym)
+      end
+
+      def canonical
+        modifiers.join("_")
+      end
+      alias to_s canonical
+
+      def to_sym
+        empty? ? nil : canonical.to_sym
+      end
+
+      def permutations
+        return [canonical] if modifiers.size <= 1
+
+        modifiers.permutation.map { |permutation| permutation.join("_") }.uniq
+      end
+
+      def ==(other)
+        other.is_a?(self.class) && modifiers == other.modifiers
+      end
+      alias eql? ==
+
+      def hash
+        modifiers.hash
+      end
+    end
+  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,32 @@ 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 finite (got #{rate})" unless rate.finite?
+          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 +133,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 +144,26 @@ 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::REGISTRY.each do |candidate|
+            return candidate.key if candidate.key.name == name
+            next unless candidate.token_key
 
-          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,210 @@
+# 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)
+        if value.infinite? || value.nan?
+          raise ArgumentError,
+                "service charge price amount for #{key.inspect} in #{context} must be finite"
+        end
+        if value.negative?
+          raise ArgumentError,
+                "service charge price amount for #{key.inspect} in #{context} must be non-negative"
+        end
+
+        value
+      end
+
+      def rate_quantity(component)
+        BigDecimal(Billing::RATE_BASIS_QUANTITIES.fetch(component.rate_basis, 1).to_s)
+      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/registry_writer.rb

@@ -9,10 +9,12 @@ module LlmCostTracker
     module Sync
       class RegistryWriter
         YAML_EXTENSIONS = %w[.yml .yaml].freeze
+        MANUAL_SOURCE = "manual"
 
         def call(path:, registry:)
           FileUtils.mkdir_p(File.dirname(path))
-          payload = yaml_file?(path) ? YAML.dump(registry) : "#{JSON.pretty_generate(registry)}\n"
+          merged = merge_with_existing(path: path, registry: registry)
+          payload = yaml_file?(path) ? YAML.dump(merged) : "#{JSON.pretty_generate(merged)}\n"
           temp_path = "#{path}.tmp-#{Process.pid}-#{Thread.current.object_id}"
           File.write(temp_path, payload)
           File.rename(temp_path, path)
@@ -22,6 +24,53 @@ module LlmCostTracker
 
         private
 
+        def merge_with_existing(path:, registry:)
+          existing = read_existing(path)
+          return registry unless existing.is_a?(Hash)
+
+          merged = registry.dup
+          merged["models"] = merged_models(registry, existing) if existing["models"].is_a?(Hash)
+          if existing["service_charges"].is_a?(Hash)
+            merged["service_charges"] = merged_service_charges(registry, existing)
+          end
+          merged
+        end
+
+        def merged_models(registry, existing)
+          merged = registry.fetch("models", {}).dup
+          existing.fetch("models", {}).each do |model, attrs|
+            next unless attrs.is_a?(Hash) && attrs["_source"].to_s == MANUAL_SOURCE
+            next if merged.key?(model)
+
+            merged[model] = attrs
+          end
+          merged
+        end
+
+        def merged_service_charges(registry, existing)
+          remote = registry.fetch("service_charges", {})
+          existing.fetch("service_charges", {}).each_with_object(remote.dup) do |(provider, charges), merged|
+            next unless charges.is_a?(Hash)
+
+            merged[provider] = charges.merge(merged.fetch(provider, {}))
+          end
+        end
+
+        def read_existing(path)
+          return nil unless File.exist?(path)
+
+          contents = File.read(path)
+          return nil if contents.strip.empty?
+
+          if yaml_file?(path)
+            YAML.safe_load(contents, permitted_classes: [Symbol, Date, Time])
+          else
+            JSON.parse(contents)
+          end
+        rescue StandardError
+          nil
+        end
+
         def yaml_file?(path)
           YAML_EXTENSIONS.include?(File.extname(path).downcase)
         end