llm_cost_tracker 0.7.0 → 0.7.2

data/lib/llm_cost_tracker/retention.rb

@@ -6,22 +6,23 @@ module LlmCostTracker
 
     class << self
       def prune(older_than:, batch_size: DEFAULT_BATCH_SIZE, now: Time.now.utc)
-        batch_size = normalized_batch_size(batch_size)
+        batch_size = batch_size.to_i
+        raise ArgumentError, "batch_size must be positive: #{batch_size.inspect}" unless batch_size.positive?
+
         cutoff = resolve_cutoff(older_than, now)
-        require_relative "storage/active_record_backend"
+        require_relative "ledger"
 
-        Storage::ActiveRecordBackend.prune(cutoff: cutoff, batch_size: batch_size)
+        deleted = 0
+        loop do
+          batch = prune_batch(cutoff, batch_size)
+          deleted += batch
+          break if batch < batch_size
+        end
+        deleted
       end
 
       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)
         cutoff = case older_than
                  when Time, DateTime then older_than.utc
@@ -46,6 +47,22 @@ module LlmCostTracker
 
         now - (days * 86_400)
       end
+
+      def prune_batch(cutoff, batch_size)
+        LlmCostTracker::Ledger::Call.transaction do
+          rows = LlmCostTracker::Ledger::Call
+                 .where(tracked_at: ...cutoff)
+                 .order(:id)
+                 .limit(batch_size)
+                 .lock
+                 .pluck(:id, :tracked_at, :total_cost)
+          next 0 if rows.empty?
+
+          deleted = LlmCostTracker::Ledger::Call.where(id: rows.map(&:first)).delete_all
+          LlmCostTracker::Ledger::Rollups.decrement!(rows) if deleted.positive?
+          deleted
+        end
+      end
     end
   end
 end

data/lib/llm_cost_tracker/tags/context.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/object/deep_dup"
+require "active_support/isolated_execution_state"
+
+module LlmCostTracker
+  module Tags
+    module Context
+      KEY = :llm_cost_tracker_tags
+
+      class << self
+        def with(tags)
+          stack = ActiveSupport::IsolatedExecutionState[KEY] || []
+          ActiveSupport::IsolatedExecutionState[KEY] = stack + [(tags || {}).deep_dup.to_h]
+          yield
+        ensure
+          ActiveSupport::IsolatedExecutionState[KEY] = stack
+        end
+
+        def tags
+          default_tags = LlmCostTracker.configuration.default_tags
+          default_tags = default_tags.call if default_tags.respond_to?(:call)
+
+          (default_tags || {}).deep_dup.to_h.merge(
+            (ActiveSupport::IsolatedExecutionState[KEY] || []).reduce({}) { |merged, tags| merged.merge(tags) }
+          )
+        end
+
+        def clear!
+          ActiveSupport::IsolatedExecutionState[KEY] = []
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/tags/key.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Tags
+    module Key
+      PATTERN = /\A[\w.-]+\z/
+
+      class << self
+        def validate!(key, error_class: ArgumentError)
+          key = key.to_s
+          return key if key.match?(PATTERN)
+
+          raise error_class, "invalid tag key: #{key.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/tags/sanitizer.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module Tags
+    module Sanitizer
+      REDACTED_VALUE = "[REDACTED]"
+
+      class << self
+        def call(tags, config: LlmCostTracker.configuration)
+          tags = (tags || {}).to_h
+          tags.first([config.max_tag_count.to_i, 0].max).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)
+          limit = [config.max_tag_value_bytesize.to_i, 0].max
+          return value if string.bytesize <= limit
+
+          string.byteslice(0, limit).to_s.encode("UTF-8", invalid: :replace, undef: :replace)
+        end
+
+        def redacted_key?(key, config)
+          normalized = normalized_key(key)
+          Array(config.redacted_tag_keys).map { |redacted_key| normalized_key(redacted_key) }.any? do |candidate|
+            redacted_key_component?(normalized, candidate)
+          end
+        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
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/token_usage.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/hash/keys"
+
+module LlmCostTracker
+  TokenUsage = Data.define(
+    :input_tokens,
+    :cache_read_input_tokens,
+    :cache_write_input_tokens,
+    :cache_write_1h_input_tokens,
+    :output_tokens,
+    :total_tokens,
+    :hidden_output_tokens
+  ) do
+    def self.build(input_tokens:, output_tokens:, cache_read_input_tokens: 0,
+                   cache_write_input_tokens: 0, cache_write_1h_input_tokens: 0,
+                   total_tokens: nil, hidden_output_tokens: 0)
+      input = input_tokens.to_i
+      output = output_tokens.to_i
+      cache_read = cache_read_input_tokens.to_i
+      cache_write = cache_write_input_tokens.to_i
+      cache_write_1h = cache_write_1h_input_tokens.to_i
+      calculated_total = input + cache_read + cache_write + cache_write_1h + output
+      total = total_tokens.nil? ? calculated_total : [total_tokens.to_i, calculated_total].max
+
+      new(
+        input_tokens: input,
+        cache_read_input_tokens: cache_read,
+        cache_write_input_tokens: cache_write,
+        cache_write_1h_input_tokens: cache_write_1h,
+        output_tokens: output,
+        total_tokens: total,
+        hidden_output_tokens: hidden_output_tokens.to_i
+      )
+    end
+
+    def self.from_hash(attributes)
+      attributes = attributes.to_h.symbolize_keys
+      values = TokenUsage::COMPONENT_TOKEN_KEYS.to_h { |key| [key, attributes[key]] }
+      build(
+        **values,
+        total_tokens: attributes[:total_tokens]
+      )
+    end
+
+    def price_quantities
+      {
+        input: input_tokens,
+        cache_read_input: cache_read_input_tokens,
+        cache_write_input: cache_write_input_tokens,
+        cache_write_1h_input: cache_write_1h_input_tokens,
+        output: output_tokens
+      }
+    end
+
+    def stored_attributes
+      to_h.slice(*self.class::STORED_KEYS)
+    end
+
+    def to_h
+      super.compact
+    end
+  end
+
+  TokenUsage::STORED_KEYS = TokenUsage.members.freeze
+  TokenUsage::COMPONENT_TOKEN_KEYS = (TokenUsage.members - %i[total_tokens]).freeze
+end

data/lib/llm_cost_tracker/tracker.rb

@@ -1,14 +1,18 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/object/blank"
 require "securerandom"
 
-require_relative "storage/writer"
+require_relative "ingestion"
+require_relative "ledger"
+require_relative "pricing"
 
 module LlmCostTracker
   class Tracker
     EVENT_NAME = "llm_request.llm_cost_tracker"
 
     USAGE_SOURCES = %i[response stream_final sdk_response ruby_llm manual unknown].freeze
+    TRACKING_METADATA_KEYS = (TokenUsage.members.map(&:to_s) + %w[pricing_mode provider_response_id]).freeze
 
     class << self
       def enforce_budget!
@@ -17,99 +21,65 @@ module LlmCostTracker
         Budget.enforce!
       end
 
-      def record(provider:, model:, input_tokens:, output_tokens:, latency_ms: nil, stream: false,
-                 usage_source: nil, provider_response_id: nil, pricing_mode: nil, metadata: {})
+      def record(capture:, latency_ms: nil, pricing_mode: nil, metadata: {}, context_tags: nil)
         return unless LlmCostTracker.configuration.enabled
 
-        model = normalize_model(model)
-        usage = usage_data(input_tokens, output_tokens, metadata, pricing_mode)
-        cost_data = cost_for_usage(provider, model, usage)
+        pricing_mode = Pricing.normalize_mode(pricing_mode) || capture.pricing_mode
+        cost_data = Pricing.cost_for(
+          provider: capture.provider,
+          model: capture.model,
+          token_usage: capture.token_usage,
+          pricing_mode: pricing_mode
+        )
 
-        UnknownPricing.handle!(model) unless cost_data
+        Pricing::Unknown.handle!(capture.model) unless cost_data
 
         event = build_event(
-          provider: provider,
-          model: model,
-          usage: usage,
+          capture: capture,
+          pricing_mode: pricing_mode,
           cost_data: cost_data,
           metadata: metadata,
           latency_ms: latency_ms,
-          stream: stream,
-          usage_source: usage_source,
-          provider_response_id: provider_response_id
+          context_tags: context_tags
         )
 
         ActiveSupport::Notifications.instrument(EVENT_NAME, event.to_h)
 
-        stored = Storage::Writer.save(event)
-        Budget.check!(event) unless stored == false
+        Ingestion::Inbox.save(event)
+        Budget.check!(event)
 
         event
       end
 
       private
 
-      def usage_data(input_tokens, output_tokens, metadata, pricing_mode)
-        metadata = metadata.merge(pricing_mode: pricing_mode) unless pricing_mode.nil?
-
-        EventMetadata.usage_data(
-          input_tokens,
-          output_tokens,
-          metadata
-        )
-      end
-
-      def cost_for_usage(provider, model, usage)
-        Pricing.cost_for(
-          provider: provider,
-          model: model,
-          input_tokens: usage[:input_tokens],
-          output_tokens: usage[:output_tokens],
-          cache_read_input_tokens: usage[:cache_read_input_tokens],
-          cache_write_input_tokens: usage[:cache_write_input_tokens],
-          pricing_mode: usage[:pricing_mode]
-        )
-      end
-
-      def normalize_model(value) = value.to_s.strip.then { |model| model.empty? ? ParsedUsage::UNKNOWN_MODEL : model }
+      def build_event(capture:, pricing_mode:, cost_data:, metadata:, latency_ms:, context_tags:)
+        usage_source = if capture.usage_source.nil?
+                         nil
+                       else
+                         symbol = capture.usage_source.to_sym
+                         USAGE_SOURCES.include?(symbol) ? symbol.to_s : nil
+                       end
+        tags = metadata.to_h.reject { |key, _value| TRACKING_METADATA_KEYS.include?(key.to_s) }
+        context_tags = context_tags.nil? ? LlmCostTracker::Tags::Context.tags : context_tags.to_h
 
-      def build_event(provider:, model:, usage:, cost_data:, metadata:, latency_ms:, stream:, usage_source:,
-                      provider_response_id:)
         Event.new(
           event_id: SecureRandom.uuid,
-          provider: provider,
-          model: model,
-          input_tokens: usage[:input_tokens],
-          output_tokens: usage[:output_tokens],
-          total_tokens: usage[:total_tokens],
-          cache_read_input_tokens: usage[:cache_read_input_tokens],
-          cache_write_input_tokens: usage[:cache_write_input_tokens],
-          hidden_output_tokens: usage[:hidden_output_tokens],
-          pricing_mode: usage[:pricing_mode],
+          provider: capture.provider,
+          model: capture.model,
+          token_usage: capture.token_usage,
+          pricing_mode: pricing_mode,
           cost: cost_data,
-          tags: sanitized_tags(metadata).freeze,
-          latency_ms: normalized_latency_ms(latency_ms),
-          stream: stream ? true : false,
-          usage_source: normalized_usage_source(usage_source),
-          provider_response_id: normalized_provider_response_id(provider_response_id),
+          tags: LlmCostTracker::Tags::Sanitizer.call(
+            context_tags.merge(tags)
+          ).freeze,
+          latency_ms: latency_ms.nil? ? nil : [latency_ms.to_i, 0].max,
+          stream: capture.stream ? true : false,
+          usage_source: usage_source,
+          provider_response_id: capture.provider_response_id.to_s.presence,
           tracked_at: Time.now.utc
         )
       end
-
-      def normalized_latency_ms(latency_ms) = latency_ms.nil? ? nil : [latency_ms.to_i, 0].max
-
-      def sanitized_tags(metadata)
-        LlmCostTracker::TagSanitizer.call(LlmCostTracker::TagContext.tags.merge(EventMetadata.tags(metadata)))
-      end
-
-      def normalized_usage_source(value)
-        return nil if value.nil?
-
-        symbol = value.to_sym
-        USAGE_SOURCES.include?(symbol) ? symbol.to_s : nil
-      end
-
-      def normalized_provider_response_id(value) = value.nil? || value.to_s.empty? ? nil : value.to_s
     end
   end
 end

data/lib/llm_cost_tracker/usage_capture.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/object/blank"
+
+require_relative "pricing"
+
+module LlmCostTracker
+  UsageCapture = Data.define(
+    :provider,
+    :model,
+    :token_usage,
+    :stream,
+    :usage_source,
+    :provider_response_id,
+    :pricing_mode
+  )
+
+  class UsageCapture
+    UNKNOWN_MODEL = "unknown"
+
+    def self.build(**attributes)
+      new(
+        provider: attributes.fetch(:provider).to_s,
+        model: attributes.fetch(:model).to_s.strip.presence || UNKNOWN_MODEL,
+        token_usage: attributes.fetch(:token_usage),
+        stream: attributes[:stream] || false,
+        usage_source: attributes[:usage_source],
+        provider_response_id: attributes[:provider_response_id],
+        pricing_mode: Pricing.normalize_mode(attributes[:pricing_mode])
+      )
+    end
+
+    def to_h
+      super.compact
+    end
+  end
+end

data/lib/llm_cost_tracker/version.rb

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

data/lib/llm_cost_tracker.rb

@@ -1,21 +1,26 @@
 # frozen_string_literal: true
 
+require "rails"
 require "active_support"
+require "active_support/core_ext/object/blank"
+require "active_support/core_ext/object/deep_dup"
+require "active_support/core_ext/object/try"
+require "active_support/core_ext/hash/indifferent_access"
+require "active_support/core_ext/string/inflections"
 require "active_support/notifications"
-require "monitor"
 
 require_relative "llm_cost_tracker/version"
 require_relative "llm_cost_tracker/configuration"
 require_relative "llm_cost_tracker/errors"
 require_relative "llm_cost_tracker/logging"
-require_relative "llm_cost_tracker/parameter_hash"
-require_relative "llm_cost_tracker/cost"
-require_relative "llm_cost_tracker/usage_breakdown"
+require_relative "llm_cost_tracker/tags/key"
+require_relative "llm_cost_tracker/tags/context"
+require_relative "llm_cost_tracker/tags/sanitizer"
+require_relative "llm_cost_tracker/token_usage"
 require_relative "llm_cost_tracker/event"
-require_relative "llm_cost_tracker/parsed_usage"
-require_relative "llm_cost_tracker/price_registry"
-require_relative "llm_cost_tracker/price_sync"
 require_relative "llm_cost_tracker/pricing"
+require_relative "llm_cost_tracker/usage_capture"
+require_relative "llm_cost_tracker/pricing/sync"
 require_relative "llm_cost_tracker/parsers/base"
 require_relative "llm_cost_tracker/parsers/openai_usage"
 require_relative "llm_cost_tracker/parsers/openai"
@@ -23,87 +28,59 @@ require_relative "llm_cost_tracker/parsers/openai_compatible"
 require_relative "llm_cost_tracker/parsers/anthropic"
 require_relative "llm_cost_tracker/parsers/gemini"
 require_relative "llm_cost_tracker/parsers/sse"
-require_relative "llm_cost_tracker/parsers/registry"
+require_relative "llm_cost_tracker/parsers"
 require_relative "llm_cost_tracker/middleware/faraday"
-require_relative "llm_cost_tracker/integrations/registry"
+require_relative "llm_cost_tracker/integrations"
 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/active_record_adapter"
-require_relative "llm_cost_tracker/tags_column"
-require_relative "llm_cost_tracker/tag_key"
-require_relative "llm_cost_tracker/tag_sql"
-require_relative "llm_cost_tracker/tag_query"
-require_relative "llm_cost_tracker/tag_accessors"
-require_relative "llm_cost_tracker/llm_api_call_metrics"
+require_relative "llm_cost_tracker/pricing/unknown"
+require_relative "llm_cost_tracker/ledger"
+require_relative "llm_cost_tracker/ingestion"
 require_relative "llm_cost_tracker/tracker"
 require_relative "llm_cost_tracker/retention"
-require_relative "llm_cost_tracker/report_data"
-require_relative "llm_cost_tracker/report_formatter"
 require_relative "llm_cost_tracker/report"
 require_relative "llm_cost_tracker/doctor"
-require_relative "llm_cost_tracker/capture_verifier"
+require_relative "llm_cost_tracker/doctor/capture_verifier"
 
 module LlmCostTracker
-  CONFIGURATION_MUTEX = Monitor.new
+  @configuration = Configuration.new
 
   class << self
-    def configuration
-      CONFIGURATION_MUTEX.synchronize { @configuration ||= Configuration.new }
-    end
-
-    def configuration_generation
-      CONFIGURATION_MUTEX.synchronize { @configuration_generation ||= 0 }
-    end
+    attr_reader :configuration
 
     def configure
-      config = CONFIGURATION_MUTEX.synchronize do
-        current = @configuration || Configuration.new
-        current = current.dup_for_configuration if current.finalized?
-        @configuration = current
-        yield(current)
-        current.openai_compatible_providers = current.openai_compatible_providers.dup
-        current.finalize!
-        @configuration_generation = @configuration_generation.to_i + 1
-        current
-      end
-      Integrations::Registry.install!
+      config = configuration
+      raise Error, "LlmCostTracker is already configured" if config.finalized?
+
+      yield(config)
+      config.openai_compatible_providers = config.openai_compatible_providers.dup
+      config.finalize!
+      Pricing::Lookup.reset!
+      Integrations.install!
       config
     end
 
     def reset_configuration!
-      Storage::ActiveRecordInbox.reset! if defined?(Storage::ActiveRecordInbox)
-      Storage::ActiveRecordIngestor.shutdown!(drain: false) if defined?(Storage::ActiveRecordIngestor)
-      CONFIGURATION_MUTEX.synchronize do
-        @configuration = Configuration.new
-        @configuration_generation = @configuration_generation.to_i + 1
-      end
-      UnknownPricing.reset! if defined?(UnknownPricing)
-      Storage::ActiveRecordStore.reset! if defined?(Storage::ActiveRecordStore)
-      Storage::ActiveRecordInbox.reset! if defined?(Storage::ActiveRecordInbox)
-      Storage::ActiveRecordIngestor.reset! if defined?(Storage::ActiveRecordIngestor)
-      TagContext.clear! if defined?(TagContext)
+      Ingestion::Worker.shutdown!(drain: false)
+      @configuration = Configuration.new
+      Pricing::Lookup.reset!
+      Pricing::Unknown.reset!
+      Ingestion::Worker.reset!
+      Tags::Context.clear!
     end
 
     def flush!(timeout: nil)
-      return true unless defined?(Storage::ActiveRecordIngestor)
-
       if timeout
-        Storage::ActiveRecordIngestor.flush!(timeout: timeout)
+        Ingestion::Worker.flush!(timeout: timeout)
       else
-        Storage::ActiveRecordIngestor.flush!
+        Ingestion::Worker.flush!
       end
     end
 
     def shutdown!(timeout: nil, drain: true)
-      return true unless defined?(Storage::ActiveRecordIngestor)
-
       if timeout
-        Storage::ActiveRecordIngestor.shutdown!(timeout: timeout, drain: drain)
+        Ingestion::Worker.shutdown!(timeout: timeout, drain: drain)
       else
-        Storage::ActiveRecordIngestor.shutdown!(drain: drain)
+        Ingestion::Worker.shutdown!(drain: drain)
       end
     end
 
@@ -113,21 +90,24 @@ module LlmCostTracker
 
     def with_tags(tags = nil, **kwargs, &)
       merged = (tags || {}).to_h.merge(kwargs)
-      TagContext.with(merged, &)
+      Tags::Context.with(merged, &)
     end
 
     def track(provider:, input_tokens:, output_tokens:, model: nil, latency_ms: nil, stream: false,
               usage_source: :manual, enforce_budget: false, provider_response_id: nil, pricing_mode: nil, **metadata)
       enforce_budget! if enforce_budget
+      token_usage = TokenUsage.from_hash(metadata.merge(input_tokens: input_tokens, output_tokens: output_tokens))
+
       Tracker.record(
-        provider: provider.to_s,
-        model: model,
-        input_tokens: input_tokens,
-        output_tokens: output_tokens,
+        capture: UsageCapture.build(
+          provider: provider,
+          model: model,
+          token_usage: token_usage,
+          stream: stream,
+          usage_source: usage_source,
+          provider_response_id: provider_response_id
+        ),
         latency_ms: latency_ms,
-        stream: stream,
-        usage_source: usage_source,
-        provider_response_id: provider_response_id,
         pricing_mode: pricing_mode,
         metadata: metadata
       )
@@ -135,9 +115,9 @@ module LlmCostTracker
 
     def track_stream(provider:, model: nil, latency_ms: nil, enforce_budget: false, provider_response_id: nil,
                      pricing_mode: nil, **metadata)
-      require_relative "llm_cost_tracker/stream_collector"
+      require_relative "llm_cost_tracker/capture/stream_collector"
       enforce_budget! if enforce_budget
-      collector = StreamCollector.new(
+      collector = Capture::StreamCollector.new(
         provider: provider.to_s,
         model: model,
         latency_ms: latency_ms,
@@ -154,12 +134,10 @@ module LlmCostTracker
   end
 end
 
-require_relative "llm_cost_tracker/railtie" if defined?(Rails::Railtie)
+require_relative "llm_cost_tracker/railtie"
 
-if defined?(Faraday)
-  Faraday::Middleware.register_middleware(
-    llm_cost_tracker: LlmCostTracker::Middleware::Faraday
-  )
-end
+Faraday::Middleware.register_middleware(
+  llm_cost_tracker: LlmCostTracker::Middleware::Faraday
+)
 
-at_exit { LlmCostTracker.shutdown!(drain: false) if defined?(LlmCostTracker) }
+at_exit { LlmCostTracker.shutdown!(drain: false) }