llm_cost_tracker 0.6.1 → 0.7.1

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/custom_backend.rb

@@ -1,32 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "registry"
-
-module LlmCostTracker
-  module Storage
-    class CustomBackend
-      class << self
-        def save(event)
-          result = LlmCostTracker.configuration.custom_storage&.call(event)
-          result == false ? false : event
-        end
-
-        def verify
-          if LlmCostTracker.configuration.custom_storage.respond_to?(:call)
-            return [
-              VerificationResult.new(
-                :ok,
-                "storage",
-                "custom storage callable configured; external sink was not invoked"
-              )
-            ]
-          end
-
-          [
-            VerificationResult.new(:error, "storage", "custom storage backend requires config.custom_storage")
-          ]
-        end
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/storage/dispatcher.rb

@@ -1,45 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "../logging"
-require_relative "registry"
-require_relative "active_record_backend"
-require_relative "custom_backend"
-require_relative "log_backend"
-
-module LlmCostTracker
-  module Storage
-    class Dispatcher
-      class << self
-        def save(event)
-          backend.save(event)
-        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
-          raise
-        rescue StandardError => e
-          handle_error(e)
-          false
-        end
-
-        private
-
-        def backend
-          Registry.fetch(LlmCostTracker.configuration.storage_backend)
-        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
-
-    Registry.register(:log, LogBackend)
-    Registry.register(:active_record, ActiveRecordBackend)
-    Registry.register(:custom, CustomBackend)
-  end
-end

data/lib/llm_cost_tracker/storage/log_backend.rb

@@ -1,38 +0,0 @@
-# frozen_string_literal: true
-
-require_relative "../logging"
-require_relative "registry"
-
-module LlmCostTracker
-  module Storage
-    class LogBackend
-      class << self
-        def save(event)
-          config = LlmCostTracker.configuration
-          message = "#{event.provider}/#{event.model} " \
-                    "tokens=#{event.total_tokens} " \
-                    "cost=#{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 verify
-          [
-            VerificationResult.new(:ok, "storage", "log backend configured; capture writes to logs only")
-          ]
-        end
-
-        private
-
-        def cost_label(event)
-          event.cost ? "$#{format('%.6f', event.cost.total_cost)}" : "unknown"
-        end
-      end
-    end
-  end
-end

data/lib/llm_cost_tracker/storage/registry.rb

@@ -1,63 +0,0 @@
-# frozen_string_literal: true
-
-require "monitor"
-
-require_relative "../errors"
-
-module LlmCostTracker
-  module Storage
-    VerificationResult = Data.define(:status, :name, :message)
-
-    module Registry
-      MUTEX = Monitor.new
-
-      class << self
-        def register(name, backend)
-          name = normalize_name(name)
-          validate_backend!(backend)
-          MUTEX.synchronize { @backends = backends.merge(name => backend).freeze }
-          backend
-        end
-
-        def fetch(name)
-          key = normalize_name(name)
-          backends.fetch(key) do
-            raise Error, "Unknown storage_backend: #{key.inspect}. Use one of: #{names.join(', ')}"
-          end
-        end
-
-        def registered?(name)
-          backends.key?(normalize_name(name))
-        end
-
-        def names
-          backends.keys
-        end
-
-        private
-
-        def backends
-          @backends || MUTEX.synchronize { @backends ||= {}.freeze }
-        end
-
-        def normalize_name(name)
-          name.to_sym
-        end
-
-        def validate_backend!(backend)
-          return if backend.respond_to?(:save)
-
-          raise ArgumentError, "storage backend must respond to save"
-        end
-      end
-    end
-
-    def self.register(name, backend)
-      Registry.register(name, backend)
-    end
-
-    def self.backends
-      Registry.names
-    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