llm_cost_tracker 0.7.0 → 0.7.2

data/lib/llm_cost_tracker/storage/active_record_store.rb

@@ -1,145 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "active_record_inbox"
-require_relative "active_record_period_totals"
-require_relative "active_record_rollups"
-
-module LlmCostTracker
-  module Storage
-    class ActiveRecordStore
-      class << self
-        def reset!
-          ActiveRecordRollups.reset!
-        end
-
-        def save(event)
-          model = LlmCostTracker::LlmApiCall
-          attributes = attributes_for(event, model)
-
-          model.transaction do
-            call = model.create!(attributes)
-            ActiveRecordRollups.increment!(event)
-            call
-          end
-        end
-
-        def insert_many(events)
-          events = Array(events)
-          return [] if events.empty?
-
-          model = LlmCostTracker::LlmApiCall
-          insertable = new_events(model, events)
-
-          if insertable.any?
-            rows = insertable.map { |event| attributes_for(event, model) }
-            model.insert_all!(rows, **insert_options)
-            ActiveRecordRollups.increment_many!(insertable)
-          end
-          events
-        end
-
-        def attributes_for(event, model = LlmCostTracker::LlmApiCall)
-          tags = stringify_tags(event.tags || {})
-          columns = model.columns_hash
-
-          attributes = {
-            provider:      event.provider,
-            model:         event.model,
-            input_tokens:  event.input_tokens,
-            output_tokens: event.output_tokens,
-            total_tokens:  event.total_tokens,
-            input_cost:    event.cost&.input_cost,
-            output_cost:   event.cost&.output_cost,
-            total_cost:    event.cost&.total_cost,
-            tags:          tags_for_storage(tags, model),
-            tracked_at:    event.tracked_at
-          }
-          attributes[:event_id] = event.event_id if columns.key?("event_id")
-          optional_attributes(event).each do |name, value|
-            attributes[name] = value if columns.key?(name.to_s)
-          end
-          attributes[:latency_ms] = event.latency_ms if columns.key?("latency_ms")
-          attributes[:stream] = event.stream if columns.key?("stream")
-          attributes[:usage_source] = event.usage_source if columns.key?("usage_source")
-          attributes[:provider_response_id] = event.provider_response_id if columns.key?("provider_response_id")
-
-          attributes
-        end
-
-        def monthly_total(time: Time.now.utc)
-          period_totals(%i[monthly], time: time).fetch(:monthly)
-        end
-
-        def daily_total(time: Time.now.utc)
-          period_totals(%i[daily], time: time).fetch(:daily)
-        end
-
-        def period_totals(periods, time: Time.now.utc)
-          ActiveRecordPeriodTotals.call(periods, time: time)
-        end
-
-        def prune(cutoff:, batch_size:)
-          deleted = 0
-          loop do
-            batch = prune_batch(cutoff, batch_size)
-            deleted += batch
-            break if batch < batch_size
-          end
-          deleted
-        end
-
-        private
-
-        def new_events(model, events)
-          return events unless model.columns_hash.key?("event_id")
-
-          existing_ids = model.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
-          events.reject { |event| existing_ids.include?(event.event_id) }
-        end
-
-        def insert_options = { record_timestamps: true, returning: false }
-
-        def prune_batch(cutoff, batch_size)
-          LlmCostTracker::LlmApiCall.transaction do
-            rows = LlmCostTracker::LlmApiCall
-                   .where(tracked_at: ...cutoff)
-                   .order(:id)
-                   .limit(batch_size)
-                   .lock
-                   .pluck(:id, :tracked_at, :total_cost)
-            next 0 if rows.empty?
-
-            deleted = LlmCostTracker::LlmApiCall.where(id: rows.map(&:first)).delete_all
-            ActiveRecordRollups.decrement!(rows) if deleted.positive?
-            deleted
-          end
-        end
-
-        def stringify_tags(tags)
-          tags.transform_keys(&:to_s).transform_values { |value| stringify_tag_value(value) }
-        end
-
-        def tags_for_storage(tags, model)
-          model.tags_json_column? ? tags : tags.to_json
-        end
-
-        def optional_attributes(event)
-          {
-            cache_read_input_tokens: event.cache_read_input_tokens,
-            cache_write_input_tokens: event.cache_write_input_tokens,
-            hidden_output_tokens: event.hidden_output_tokens,
-            cache_read_input_cost: event.cost&.cache_read_input_cost,
-            cache_write_input_cost: event.cost&.cache_write_input_cost,
-            pricing_mode: event.pricing_mode
-          }
-        end
-
-        def stringify_tag_value(value)
-          return value.transform_values { |nested| stringify_tag_value(nested) } if value.is_a?(Hash)
-
-          value.to_s
-        end
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/storage/writer.rb

@@ -1,35 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "../errors"
-require_relative "../logging"
-require_relative "active_record_backend"
-
-module LlmCostTracker
-  module Storage
-    class Writer
-      class << self
-        def save(event)
-          ActiveRecordBackend.save(event)
-        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
-          raise
-        rescue StandardError => e
-          handle_error(e)
-          false
-        end
-
-        private
-
-        def handle_error(error)
-          case LlmCostTracker.configuration.storage_error_behavior
-          when :ignore
-            nil
-          when :warn
-            Logging.warn("ActiveRecord ledger write failed: #{error.class}: #{error.message}")
-          when :raise
-            raise StorageError, error
-          end
-        end
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/stream_capture.rb

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

data/lib/llm_cost_tracker/stream_collector.rb

@@ -1,199 +0,0 @@
-# frozen_string_literal: true
-
-require "monitor"
-
-require_relative "stream_capture"
-require_relative "value_helpers"
-
-module LlmCostTracker
-  class StreamCollector
-    attr_reader :provider
-
-    def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, pricing_mode: nil, metadata: {})
-      @provider = provider.to_s
-      @model = model
-      @latency_ms = latency_ms
-      @provider_response_id = provider_response_id
-      @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
-      @monitor = Monitor.new
-    end
-
-    def model = @monitor.synchronize { @model }
-
-    def metadata = @monitor.synchronize { ValueHelpers.deep_dup(@metadata) }
-
-    def provider_response_id = @monitor.synchronize { @provider_response_id }
-
-    def model=(value)
-      @monitor.synchronize do
-        ensure_open!
-        @model = value
-      end
-    end
-
-    def provider_response_id=(value)
-      @monitor.synchronize do
-        ensure_open!
-        @provider_response_id = value
-      end
-    end
-
-    def event(data, type: nil)
-      @monitor.synchronize do
-        ensure_open!
-        capture_event(data, type: type) unless data.nil?
-      end
-      self
-    end
-    alias chunk event
-
-    def usage(input_tokens:, output_tokens:, **extra)
-      @monitor.synchronize do
-        ensure_open!
-        @explicit_usage = ValueHelpers.deep_dup(
-          extra.merge(
-            input_tokens: input_tokens.to_i,
-            output_tokens: output_tokens.to_i
-          )
-        )
-      end
-      self
-    end
-
-    def finish!(errored: false)
-      snapshot = @monitor.synchronize do
-        return if @finished
-
-        @finished = true
-        {
-          events: @events.dup,
-          overflowed: @overflowed,
-          explicit_usage: ValueHelpers.deep_dup(@explicit_usage),
-          model: @model,
-          latency_ms: @latency_ms,
-          provider_response_id: @provider_response_id,
-          pricing_mode: @pricing_mode,
-          metadata: ValueHelpers.deep_dup(@metadata)
-        }
-      end
-
-      parsed = build_parsed_usage(snapshot)
-      Tracker.record(
-        provider: parsed.provider,
-        model: parsed.model,
-        input_tokens: parsed.input_tokens,
-        output_tokens: parsed.output_tokens,
-        latency_ms: snapshot[:latency_ms] || elapsed_ms,
-        stream: true,
-        usage_source: parsed.usage_source,
-        provider_response_id: parsed.provider_response_id || snapshot[:provider_response_id],
-        pricing_mode: snapshot[:pricing_mode],
-        metadata: error_metadata(errored).merge(snapshot[:metadata]).merge(parsed.metadata)
-      )
-    end
-
-    private
-
-    def ensure_open!
-      return unless @finished
-
-      raise FrozenError, "can't modify finished LlmCostTracker::StreamCollector"
-    end
-
-    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
-
-      build_unknown_usage(snapshot)
-    end
-
-    def finalize(parsed, snapshot)
-      parsed.with(
-        provider: @provider,
-        model: present_model(parsed.model) || present_model(snapshot[:model]) || ParsedUsage::UNKNOWN_MODEL
-      )
-    end
-
-    def present_model(value)
-      return nil if value.nil?
-
-      string = value.to_s
-      return nil if string.empty? || string == "unknown"
-
-      string
-    end
-
-    def build_from_explicit_usage(snapshot)
-      explicit = snapshot[:explicit_usage]
-      input = explicit[:input_tokens]
-      output = explicit[:output_tokens]
-      extras = explicit.except(:input_tokens, :output_tokens)
-
-      ParsedUsage.build(
-        provider: @provider,
-        model: snapshot[:model] || ParsedUsage::UNKNOWN_MODEL,
-        input_tokens: input,
-        output_tokens: output,
-        stream: true,
-        usage_source: :manual,
-        **extras
-      )
-    end
-
-    def build_unknown_usage(snapshot)
-      ParsedUsage.build(
-        provider: @provider,
-        model: snapshot[:model] || ParsedUsage::UNKNOWN_MODEL,
-        input_tokens: 0,
-        output_tokens: 0,
-        total_tokens: 0,
-        stream: true,
-        usage_source: :unknown
-      )
-    end
-
-    def capture_event(data, type:)
-      size = event_bytes(data, type)
-      if @captured_bytes + size <= StreamCapture::LIMIT_BYTES
-        @events << { event: type, data: ValueHelpers.deep_dup(data) }
-        @captured_bytes += size
-      else
-        @overflowed = true
-        @events.clear
-      end
-    end
-
-    def event_bytes(data, type)
-      type.to_s.bytesize + estimated_bytes(data) + 32
-    end
-
-    def estimated_bytes(value)
-      case value
-      when Hash
-        value.sum { |key, nested| estimated_bytes(key) + estimated_bytes(nested) + 4 }
-      when Array
-        value.sum { |nested| estimated_bytes(nested) + 2 }
-      when String
-        value.bytesize + 2
-      when Numeric, true, false, nil
-        value.to_s.bytesize
-      else
-        value.to_s.bytesize + 2
-      end
-    end
-
-    def error_metadata(errored) = errored ? { stream_errored: true } : {}
-
-    def elapsed_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at) * 1000).round
-  end
-end

data/lib/llm_cost_tracker/tag_accessors.rb

@@ -1,15 +0,0 @@
-# frozen_string_literal: true
-
-require "json"
-
-module LlmCostTracker
-  module TagAccessors
-    def parsed_tags
-      return tags.transform_keys(&:to_s) if tags.is_a?(Hash)
-
-      JSON.parse(tags || "{}")
-    rescue JSON::ParserError
-      {}
-    end
-  end
-end

data/lib/llm_cost_tracker/tag_context.rb

@@ -1,52 +0,0 @@
-# frozen_string_literal: true
-
-require "active_support/isolated_execution_state"
-
-require_relative "value_helpers"
-
-module LlmCostTracker
-  module TagContext
-    KEY = :llm_cost_tracker_tags
-
-    class << self
-      def with(tags)
-        stack = current_stack
-        ActiveSupport::IsolatedExecutionState[KEY] = stack + [normalize(tags)]
-        yield
-      ensure
-        ActiveSupport::IsolatedExecutionState[KEY] = stack
-      end
-
-      def tags
-        config_tags.merge(scoped_tags)
-      end
-
-      def clear!
-        ActiveSupport::IsolatedExecutionState[KEY] = []
-      end
-
-      private
-
-      def config_tags
-        normalize(resolve_default_tags)
-      end
-
-      def resolve_default_tags
-        tags = LlmCostTracker.configuration.default_tags
-        tags.respond_to?(:call) ? tags.call : tags
-      end
-
-      def scoped_tags
-        current_stack.reduce({}) { |merged, tags| merged.merge(tags) }
-      end
-
-      def current_stack
-        ActiveSupport::IsolatedExecutionState[KEY] || []
-      end
-
-      def normalize(tags)
-        ValueHelpers.deep_dup(tags || {}).to_h
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/tag_key.rb

@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-
-module LlmCostTracker
-  module TagKey
-    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

data/lib/llm_cost_tracker/tag_query.rb

@@ -1,43 +0,0 @@
-# frozen_string_literal: true
-
-require "json"
-
-module LlmCostTracker
-  module TagQuery
-    class << self
-      def apply(model, tags)
-        normalized_tags = normalize_tags(tags)
-        return model.all if normalized_tags.empty?
-
-        return postgres_json_query(model, normalized_tags) if model.tags_jsonb_column?
-        return mysql_json_query(model, normalized_tags) if model.tags_mysql_json_column?
-
-        text_query(model, normalized_tags)
-      end
-
-      def normalize_tags(tags)
-        (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
-      end
-
-      private
-
-      def postgres_json_query(model, tags)
-        model.where("tags @> ?::jsonb", tags.to_json)
-      end
-
-      def mysql_json_query(model, tags)
-        model.where("JSON_CONTAINS(tags, ?)", tags.to_json)
-      end
-
-      def text_query(model, tags)
-        tags.reduce(model.all) do |relation, (key, value)|
-          relation.where("tags LIKE ? ESCAPE '\\'", "%#{model.sanitize_sql_like(json_tag_fragment(key, value))}%")
-        end
-      end
-
-      def json_tag_fragment(key, value)
-        JSON.generate(key => value).delete_prefix("{").delete_suffix("}")
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/tag_sanitizer.rb

@@ -1,81 +0,0 @@
-# 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/tag_sql.rb

@@ -1,34 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "active_record_adapter"
-require_relative "tag_key"
-
-module LlmCostTracker
-  module TagSql
-    class << self
-      def value_expression(model, key, table_name:)
-        key = TagKey.validate!(key)
-        column = "#{table_name}.#{model.connection.quote_column_name('tags')}"
-
-        if ActiveRecordAdapter.postgresql?(model.connection)
-          json_column = model.tags_jsonb_column? ? column : "(#{column})::jsonb"
-          "#{json_column}->>#{model.connection.quote(key)}"
-        elsif ActiveRecordAdapter.mysql?(model.connection)
-          "JSON_UNQUOTE(JSON_EXTRACT(#{column}, #{model.connection.quote(json_path(key))}))"
-        else
-          ActiveRecordAdapter.ensure_supported!(model.connection)
-        end
-      end
-
-      def value_label(value)
-        value.nil? || value == "" ? "(untagged)" : value.to_s
-      end
-
-      private
-
-      def json_path(key)
-        "$.\"#{key}\""
-      end
-    end
-  end
-end