llm_cost_tracker 0.5.0 → 0.5.2

data/lib/llm_cost_tracker/pricing.rb

@@ -32,27 +32,20 @@ module LlmCostTracker
       end
 
       def lookup(provider:, model:)
-        table = prices
         provider_name = provider.to_s
         model_name = model.to_s
         provider_model = provider_name.empty? ? model_name : "#{provider_name}/#{model_name}"
         normalized_model = normalize_model_name(model_name)
+        current = current_price_tables
 
-        table[provider_model] ||
-          table[model_name] ||
-          table[normalized_model] ||
-          fuzzy_match(provider_model, normalized_model, table)
-      end
-
-      def models
-        prices.keys
+        lookup_in_table(current.fetch(:pricing_overrides), provider_model, model_name, normalized_model) ||
+          lookup_in_table(current.fetch(:file_prices), provider_model, model_name, normalized_model) ||
+          lookup_in_table(PRICES, provider_model, model_name, normalized_model)
       end
 
-      def metadata
-        PriceRegistry.metadata
-      end
+      private
 
-      def prices
+      def current_price_tables
         file_prices = PriceRegistry.file_prices(LlmCostTracker.configuration.prices_file)
         overrides = PriceRegistry.normalize_price_table(LlmCostTracker.configuration.pricing_overrides)
         cache_key = [file_prices.object_id, LlmCostTracker.configuration.pricing_overrides.hash]
@@ -64,13 +57,22 @@ module LlmCostTracker
           cached = @prices_cache
           return cached[:value] if cached && cached[:key] == cache_key
 
-          value = PRICES.merge(file_prices).merge(overrides).freeze
+          value = { pricing_overrides: overrides, file_prices: file_prices }.freeze
           @prices_cache = { key: cache_key, value: value }.freeze
           value
         end
       end
 
-      private
+      def lookup_in_table(table, provider_model, model_name, normalized_model)
+        return nil if table.empty?
+
+        table[provider_model] ||
+          table[model_name] ||
+          table[normalized_model] ||
+          unique_providerless_lookup(normalized_model, table) ||
+          fuzzy_match(provider_model, normalized_model, table) ||
+          unique_providerless_fuzzy_match(normalized_model, table)
+      end
 
       def calculate_costs(usage, prices, pricing_mode:)
         {
@@ -113,6 +115,11 @@ module LlmCostTracker
         model.to_s.split("/").last
       end
 
+      def unique_providerless_lookup(model, table)
+        matches = sorted_price_keys(table).select { |key| normalize_model_name(key) == model }
+        table[matches.first] if matches.one?
+      end
+
       def fuzzy_match(model, normalized_model, table)
         sorted_price_keys(table).each do |key|
           return table[key] if snapshot_variant?(model, key) || snapshot_variant?(normalized_model, key)
@@ -121,6 +128,11 @@ module LlmCostTracker
         nil
       end
 
+      def unique_providerless_fuzzy_match(model, table)
+        matches = sorted_price_keys(table).select { |key| snapshot_variant?(model, normalize_model_name(key)) }
+        table[matches.first] if matches.one?
+      end
+
       def snapshot_variant?(model, key)
         suffix = model.delete_prefix("#{key}-")
         return false if suffix == model

data/lib/llm_cost_tracker/report.rb

@@ -9,7 +9,14 @@ module LlmCostTracker
 
     class << self
       def generate(days: DEFAULT_DAYS, now: Time.now.utc, tag_breakdowns: nil)
-        ReportFormatter.new(data(days: days, now: now, tag_breakdowns: tag_breakdowns)).to_s
+        report_data = ReportData.build(
+          days: days,
+          now: now,
+          tag_breakdowns: tag_breakdowns,
+          breakdown_limit: ReportFormatter::TOP_LIMIT
+        )
+
+        ReportFormatter.new(report_data).to_s
       rescue LoadError => e
         "Unable to build LLM cost report: ActiveRecord storage is unavailable (#{e.message})"
       rescue StandardError => e

data/lib/llm_cost_tracker/report_data.rb

@@ -23,10 +23,11 @@ module LlmCostTracker
     DEFAULT_DAYS = 30
     TOP_LIMIT = 5
 
-    def self.build(days: DEFAULT_DAYS, now: Time.now.utc, tag_breakdowns: nil)
+    def self.build(days: DEFAULT_DAYS, now: Time.now.utc, tag_breakdowns: nil, breakdown_limit: nil)
       require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
 
       days = normalized_days(days)
+      breakdown_limit = normalized_limit(breakdown_limit)
       from = now - days.days
       scope = LlmApiCall.where(tracked_at: from..now)
       tag_breakdowns ||= LlmCostTracker.configuration.report_tag_breakdowns || []
@@ -39,9 +40,9 @@ module LlmCostTracker
         requests_count: scope.count,
         average_latency_ms: average_latency_ms(scope),
         unknown_pricing_count: scope.where(total_cost: nil).count,
-        cost_by_provider: cost_by(scope, :provider),
-        cost_by_model: cost_by(scope, :model),
-        cost_by_tags: cost_by_tags(scope, tag_breakdowns),
+        cost_by_provider: cost_by(scope, :provider, limit: breakdown_limit),
+        cost_by_model: cost_by(scope, :model, limit: breakdown_limit),
+        cost_by_tags: cost_by_tags(scope, tag_breakdowns, limit: breakdown_limit),
         top_calls: top_calls(scope)
       )
     end
@@ -51,18 +52,33 @@ module LlmCostTracker
       days.positive? ? days : DEFAULT_DAYS
     end
 
+    def self.normalized_limit(limit)
+      return nil if limit.nil?
+
+      limit = limit.to_i
+      limit.positive? ? limit : nil
+    end
+
     def self.average_latency_ms(scope)
       return nil unless LlmApiCall.latency_column?
 
       scope.average(:latency_ms)&.to_f
     end
 
-    def self.cost_by(scope, column)
-      scope.group(column).sum(:total_cost).transform_values(&:to_f).sort_by { |_name, cost| -cost }
+    def self.cost_by(scope, column, limit:)
+      relation = scope.group(column)
+                      .order(Arel.sql("COALESCE(SUM(total_cost), 0) DESC"))
+
+      relation = relation.limit(limit) if limit
+
+      relation
+        .sum(:total_cost)
+        .transform_values(&:to_f)
+        .sort_by { |_name, cost| -cost }
     end
 
-    def self.cost_by_tags(scope, keys)
-      keys.to_h { |key| [key, scope.cost_by_tag(key).to_a] }
+    def self.cost_by_tags(scope, keys, limit:)
+      keys.to_h { |key| [key, scope.cost_by_tag(key, limit: limit).to_a] }
     end
 
     def self.top_calls(scope)
@@ -73,6 +89,6 @@ module LlmCostTracker
         .map { |call| TopCall.new(provider: call.provider, model: call.model, total_cost: call.total_cost.to_f) }
     end
 
-    private_class_method :normalized_days, :average_latency_ms, :cost_by, :cost_by_tags, :top_calls
+    private_class_method :normalized_days, :normalized_limit, :average_latency_ms, :cost_by, :cost_by_tags, :top_calls
   end
 end

data/lib/llm_cost_tracker/retention.rb

@@ -6,6 +6,7 @@ module LlmCostTracker
 
     class << self
       def prune(older_than:, batch_size: DEFAULT_BATCH_SIZE, now: Time.now.utc)
+        batch_size = normalized_batch_size(batch_size)
         cutoff = resolve_cutoff(older_than, now)
         require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
 
@@ -20,14 +21,36 @@ module LlmCostTracker
 
       private
 
+      def normalized_batch_size(value)
+        value = value.to_i
+        raise ArgumentError, "batch_size must be positive: #{value.inspect}" unless value.positive?
+
+        value
+      end
+
       def resolve_cutoff(older_than, now)
-        case older_than
-        when Time, DateTime then older_than.utc
-        when ActiveSupport::Duration then now - older_than
-        when Integer then now - (older_than * 86_400)
-        else
-          raise ArgumentError, "older_than must be a Duration, Time, or Integer days: #{older_than.inspect}"
-        end
+        cutoff = case older_than
+                 when Time, DateTime then older_than.utc
+                 when ActiveSupport::Duration then duration_cutoff(older_than, now)
+                 when Integer then integer_day_cutoff(older_than, now)
+                 else
+                   raise ArgumentError, "older_than must be a Duration, Time, or Integer days: #{older_than.inspect}"
+                 end
+        raise ArgumentError, "older_than cutoff must be before now: #{cutoff.inspect}" unless cutoff < now
+
+        cutoff
+      end
+
+      def duration_cutoff(duration, now)
+        raise ArgumentError, "older_than duration must be positive: #{duration.inspect}" unless duration.to_i.positive?
+
+        now - duration
+      end
+
+      def integer_day_cutoff(days, now)
+        raise ArgumentError, "older_than days must be positive: #{days.inspect}" unless days.positive?
+
+        now - (days * 86_400)
       end
     end
   end

data/lib/llm_cost_tracker/storage/dispatcher.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require_relative "../logging"
+
+module LlmCostTracker
+  module Storage
+    class Dispatcher
+      class << self
+        def save(event)
+          config = LlmCostTracker.configuration
+          case config.storage_backend
+          when :log           then log_event(event, config)
+          when :active_record then active_record_save(event)
+          when :custom        then custom_save(event, config)
+          end
+        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
+          raise
+        rescue StandardError => e
+          handle_error(e)
+          false
+        end
+
+        private
+
+        def log_event(event, config)
+          message = "#{event.provider}/#{event.model} " \
+                    "tokens=#{event.total_tokens} " \
+                    "cost=#{log_cost_label(event)}"
+          message += " latency=#{event.latency_ms}ms" if event.latency_ms
+          message += " stream=#{event.stream}" if event.stream
+          message += " source=#{event.usage_source}" if event.usage_source
+          message += " tags=#{event.tags}" unless event.tags.empty?
+
+          Logging.log(config.log_level, message)
+          event
+        end
+
+        def log_cost_label(event) = event.cost ? "$#{format('%.6f', event.cost.total_cost)}" : "unknown"
+
+        def active_record_save(event)
+          require_relative "../llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
+          require_relative "active_record_store" unless defined?(LlmCostTracker::Storage::ActiveRecordStore)
+
+          ActiveRecordStore.save(event)
+          event
+        rescue LoadError => e
+          raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
+        end
+
+        def custom_save(event, config)
+          result = config.custom_storage&.call(event)
+          result == false ? false : event
+        end
+
+        def handle_error(error)
+          case LlmCostTracker.configuration.storage_error_behavior
+          when :ignore
+            nil
+          when :warn
+            Logging.warn("Storage failed; tracking event was not persisted: #{error.class}: #{error.message}")
+          when :raise
+            raise StorageError, error
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/stream_capture.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module StreamCapture
+    LIMIT_BYTES = 1_048_576
+  end
+end

data/lib/llm_cost_tracker/stream_collector.rb

@@ -1,7 +1,9 @@
 # frozen_string_literal: true
 
+require "json"
 require "monitor"
 
+require_relative "stream_capture"
 require_relative "value_helpers"
 
 module LlmCostTracker
@@ -16,6 +18,8 @@ module LlmCostTracker
       @pricing_mode = pricing_mode
       @metadata = ValueHelpers.deep_dup(metadata || {})
       @events = []
+      @captured_bytes = 0
+      @overflowed = false
       @explicit_usage = nil
       @started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
       @finished = false
@@ -45,7 +49,7 @@ module LlmCostTracker
     def event(data, type: nil)
       @monitor.synchronize do
         ensure_open!
-        @events << { event: type, data: ValueHelpers.deep_dup(data) } unless data.nil?
+        capture_event(data, type: type) unless data.nil?
       end
       self
     end
@@ -71,6 +75,7 @@ module LlmCostTracker
         @finished = true
         {
           events: @events.dup,
+          overflowed: @overflowed,
           explicit_usage: ValueHelpers.deep_dup(@explicit_usage),
           model: @model,
           latency_ms: @latency_ms,
@@ -105,6 +110,7 @@ module LlmCostTracker
 
     def build_parsed_usage(snapshot)
       return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
+      return build_unknown_usage(snapshot) if snapshot[:overflowed]
 
       parsed = Parsers::Registry.find_for_provider(@provider)&.parse_stream(nil, nil, 200, snapshot[:events])
       return finalize(parsed, snapshot) if parsed
@@ -157,6 +163,24 @@ module LlmCostTracker
       )
     end
 
+    def capture_event(data, type:)
+      copied = ValueHelpers.deep_dup(data)
+      size = event_bytes(copied, type)
+      if @captured_bytes + size <= StreamCapture::LIMIT_BYTES
+        @events << { event: type, data: copied }
+        @captured_bytes += size
+      else
+        @overflowed = true
+        @events.clear
+      end
+    end
+
+    def event_bytes(data, type)
+      JSON.generate(event: type, data: data).bytesize
+    rescue JSON::GeneratorError, TypeError
+      type.to_s.bytesize + data.to_s.bytesize
+    end
+
     def error_metadata(errored) = errored ? { stream_errored: true } : {}
 
     def elapsed_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at) * 1000).round

data/lib/llm_cost_tracker/tag_sanitizer.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module TagSanitizer
+    REDACTED_VALUE = "[REDACTED]"
+
+    class << self
+      def call(tags, config: LlmCostTracker.configuration)
+        tags = (tags || {}).to_h
+        tags.first(max_tag_count(config)).each_with_object({}) do |(key, value), sanitized|
+          sanitized[key] = sanitized_value(key, value, config)
+        end
+      end
+
+      private
+
+      def sanitized_value(key, value, config)
+        return REDACTED_VALUE if redacted_key?(key, config)
+
+        string = value_string(value)
+        return value if string.bytesize <= max_tag_value_bytesize(config)
+
+        truncate_bytes(string, max_tag_value_bytesize(config))
+      end
+
+      def redacted_key?(key, config)
+        normalized = normalized_key(key)
+        redacted_keys(config).any? do |candidate|
+          redacted_key_component?(normalized, candidate)
+        end
+      end
+
+      def redacted_keys(config)
+        Array(config.redacted_tag_keys).map { |key| normalized_key(key) }
+      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("_")
+      end
+
+      def redacted_key_component?(key, candidate)
+        key == candidate ||
+          key.start_with?("#{candidate}_") ||
+          key.end_with?("_#{candidate}") ||
+          key.include?("_#{candidate}_")
+      end
+
+      def value_string(value)
+        case value
+        when Hash, Array
+          JSON.generate(value)
+        else
+          value.to_s
+        end
+      rescue JSON::GeneratorError, TypeError
+        value.to_s
+      end
+
+      def truncate_bytes(string, limit)
+        string.byteslice(0, limit).to_s.encode("UTF-8", invalid: :replace, undef: :replace)
+      end
+
+      def max_tag_count(config)
+        [config.max_tag_count.to_i, 0].max
+      end
+
+      def max_tag_value_bytesize(config)
+        [config.max_tag_value_bytesize.to_i, 0].max
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/tracker.rb

@@ -1,12 +1,12 @@
 # frozen_string_literal: true
 
-require_relative "logging"
+require_relative "storage/dispatcher"
 
 module LlmCostTracker
   class Tracker
     EVENT_NAME = "llm_request.llm_cost_tracker"
 
-    USAGE_SOURCES = %i[response stream_final sdk_response manual unknown].freeze
+    USAGE_SOURCES = %i[response stream_final sdk_response ruby_llm manual unknown].freeze
 
     class << self
       def enforce_budget!
@@ -39,7 +39,7 @@ module LlmCostTracker
 
         ActiveSupport::Notifications.instrument(EVENT_NAME, event.to_h)
 
-        stored = store(event)
+        stored = Storage::Dispatcher.save(event)
         Budget.check!(event) unless stored == false
 
         event
@@ -84,7 +84,7 @@ module LlmCostTracker
           hidden_output_tokens: usage[:hidden_output_tokens],
           pricing_mode: usage[:pricing_mode],
           cost: cost_data,
-          tags: LlmCostTracker::TagContext.tags.merge(EventMetadata.tags(metadata)).freeze,
+          tags: sanitized_tags(metadata).freeze,
           latency_ms: normalized_latency_ms(latency_ms),
           stream: stream ? true : false,
           usage_source: normalized_usage_source(usage_source),
@@ -93,64 +93,12 @@ module LlmCostTracker
         )
       end
 
-      def store(event)
-        config = LlmCostTracker.configuration
-        case config.storage_backend
-        when :log            then log_event(event, config)
-        when :active_record  then active_record_save(event)
-        when :custom         then custom_save(event, config)
-        end
-      rescue BudgetExceededError, UnknownPricingError
-        raise
-      rescue StandardError => e
-        handle_storage_error(e)
-        false
-      end
-
-      def log_event(event, config)
-        message = "#{event.provider}/#{event.model} " \
-                  "tokens=#{event.total_tokens} " \
-                  "cost=#{log_cost_label(event)}"
-        message += " latency=#{event.latency_ms}ms" if event.latency_ms
-        message += " stream=#{event.stream}" if event.stream
-        message += " source=#{event.usage_source}" if event.usage_source
-        message += " tags=#{event.tags}" unless event.tags.empty?
-
-        Logging.log(config.log_level, message)
-        event
-      end
-
-      def log_cost_label(event) = event.cost ? "$#{format('%.6f', event.cost.total_cost)}" : "unknown"
-
-      def active_record_save(event)
-        require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
-        require_relative "storage/active_record_store" unless defined?(LlmCostTracker::Storage::ActiveRecordStore)
-
-        Storage::ActiveRecordStore.save(event)
-        event
-      rescue LoadError => e
-        raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
-      end
-
-      def custom_save(event, config)
-        result = config.custom_storage&.call(event)
-        result == false ? false : event
-      end
+      def normalized_latency_ms(latency_ms) = latency_ms.nil? ? nil : [latency_ms.to_i, 0].max
 
-      def handle_storage_error(error)
-        case LlmCostTracker.configuration.storage_error_behavior
-        when :ignore
-          nil
-        when :warn
-          Logging.warn("Storage failed; tracking event was not persisted: #{error.class}: #{error.message}")
-        when :raise
-          storage_error = StorageError.new(error)
-          raise storage_error
-        end
+      def sanitized_tags(metadata)
+        LlmCostTracker::TagSanitizer.call(LlmCostTracker::TagContext.tags.merge(EventMetadata.tags(metadata)))
       end
 
-      def normalized_latency_ms(latency_ms) = latency_ms.nil? ? nil : [latency_ms.to_i, 0].max
-
       def normalized_usage_source(value)
         return nil if value.nil?
 

data/lib/llm_cost_tracker/version.rb

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

data/lib/llm_cost_tracker.rb

@@ -30,6 +30,7 @@ require_relative "llm_cost_tracker/budget"
 require_relative "llm_cost_tracker/unknown_pricing"
 require_relative "llm_cost_tracker/event_metadata"
 require_relative "llm_cost_tracker/tag_context"
+require_relative "llm_cost_tracker/tag_sanitizer"
 require_relative "llm_cost_tracker/tags_column"
 require_relative "llm_cost_tracker/tag_key"
 require_relative "llm_cost_tracker/tag_query"