llm_cost_tracker 0.14.2 → 0.15.0

data/lib/llm_cost_tracker/budget.rb

@@ -1,18 +1,13 @@
 # frozen_string_literal: true
 
-require "bigdecimal"
-
-require_relative "ledger"
-require_relative "pricing/estimator"
-require_relative "budget/per_tag"
-require_relative "tags/context"
+require_relative "budget/estimate"
+require_relative "budget/global"
+require_relative "budget/tagged"
 
 module LlmCostTracker
   module Budget
-    BUDGET_TYPE_TO_PERIOD = { monthly: :month, daily: :day }.freeze
-
     class << self
-      def enforce!(provider: nil, model: nil, request: nil, estimate: nil, tags: nil, force: false)
+      def enforce!(provider: nil, model: nil, request: nil, tags: nil, force: false)
         config = LlmCostTracker.configuration
         return unless config.enabled
 
@@ -20,39 +15,20 @@ module LlmCostTracker
         per_tag = force || PerTag.blocking?
         return unless globally || per_tag
 
-        estimate ||= estimate_cost(provider: provider, model: model, request: request)
+        estimate = Estimate.for(provider: provider, model: model, request: request)
         now = Time.now.utc
-        enforce_globally(config, estimate: estimate, time: now) if globally
-        return unless per_tag
-
-        check_per_tag(tags || Tags::Context.tags,
-                      time: now,
-                      estimate: estimate,
-                      blocking_only: !force) do |rule, window, total, limit|
-          raise_pre_send(budget_type: window, total: total, budget: limit, scope: scope_for(rule))
-        end
+        Global.enforce!(estimate, time: now) if globally
+        Tagged.enforce!(tags, estimate.total, time: now, blocking_only: !force) if per_tag
       end
 
       def check!(event, behavior_override: nil)
-        config = LlmCostTracker.configuration
-        return unless event.total_cost
-
-        errors = [check_per_call_budget(event, config, behavior_override)]
-        check_windowed({ daily: config.budgets.daily, monthly: config.budgets.monthly }.compact,
-                       time: event.tracked_at) do |budget_type, total, budget|
-          errors << handle_exceeded(budget_type: budget_type,
-                                    total: total,
-                                    budget: budget,
-                                    previous_total: total - event.total_cost,
-                                    last_event: event,
-                                    behavior: behavior_override)
-        end
-        errors.concat(persisted_errors([event], behavior_override)) unless Ingestion.async?
+        errors = event.total_cost ? Global.post_spend_errors(event, behavior_override) : []
+        errors.concat(Tagged.post_spend_errors([event], behavior_override)) unless Ingestion.async?
         raise_first(errors)
       end
 
       def check_persisted!(events, behavior_override: nil)
-        raise_first(persisted_errors(events, behavior_override))
+        raise_first(Tagged.post_spend_errors(events, behavior_override))
       end
 
       def notify_persisted_safely!(events)
@@ -61,136 +37,22 @@ module LlmCostTracker
         Logging.warn("Per-tag budget check failed after ingest: #{e.class}: #{e.message}")
       end
 
-      private
-
-      def persisted_errors(events, behavior_override)
-        by_rule = PerTag.rules_for_events(events)
-        by_rule = by_rule.reject { |rule, _| rule.on_exceeded.nil? } if behavior_override == :notify
-        window_buckets(by_rule).flat_map do |(key, window, bucket), scored|
-          upto = scored.values.flatten.map(&:tracked_at).max unless Ingestion.async?
-          totals = PerTag.spend_by_value(key, scored.keys.map(&:value), window, bucket, upto)
-          scored.filter_map do |rule, recorded|
-            total, upto_total = totals.fetch(rule.value, [0, 0])
-            limit = rule.windows.fetch(window)
-            next if total < limit
+      def notify_repriced_safely!(changes)
+        return if changes.empty?
 
-            handle_exceeded(
-              budget_type: window,
-              total: total,
-              budget: limit,
-              previous_total: (upto_total >= limit ? upto_total : total) - recorded.sum(&:total_cost),
-              last_event: recorded.last,
-              scope: scope_for(rule),
-              behavior: behavior_override || rule.behavior,
-              on_exceeded: rule.on_exceeded
-            )
-          end
-        end
+        now = Time.now.utc
+        Global.notify_repriced(changes, time: now)
+        Tagged.notify_repriced(changes, time: now)
+      rescue StandardError => e
+        Logging.warn("Budget check failed after repricing: #{e.class}: #{e.message}")
       end
 
+      private
+
       def raise_first(errors)
         error = errors.compact.first
         raise error if error
       end
-
-      def window_buckets(by_rule)
-        by_rule.each_with_object({}) do |(rule, events), grouped|
-          rule.windows.each_key do |window|
-            events.group_by { |event| PerTag.window_start(window, event.tracked_at) }
-                  .each do |bucket, bucket_events|
-              (grouped[[rule.key, window, bucket]] ||= {})[rule] = bucket_events
-            end
-          end
-        end
-      end
-
-      def estimate_cost(provider:, model:, request:)
-        return BigDecimal("0") unless provider && model && request
-
-        Pricing::Estimator.call(provider: provider, model: model, request: request) || BigDecimal("0")
-      end
-
-      def check_per_call_budget(event, config, behavior_override)
-        budget = config.budgets.per_call
-        return unless budget
-
-        total = event.total_cost
-        return unless total >= budget
-
-        handle_exceeded(budget_type: :per_call,
-                        total: total,
-                        budget: budget,
-                        previous_total: nil,
-                        last_event: event,
-                        behavior: behavior_override)
-      end
-
-      def enforce_globally(config, estimate:, time:)
-        per_call = config.budgets.per_call
-        if per_call && estimate.positive? && estimate >= per_call
-          raise_pre_send(budget_type: :per_call, total: estimate, budget: per_call)
-        end
-
-        check_windowed({ monthly: config.budgets.monthly, daily: config.budgets.daily }.compact,
-                       time: time,
-                       estimate: estimate) do |budget_type, total, budget|
-          raise_pre_send(budget_type: budget_type, total: total, budget: budget)
-        end
-      end
-
-      def raise_pre_send(budget_type:, total:, budget:, scope: nil)
-        raise BudgetExceededError.new(
-          budget_type: budget_type, total: total, budget: budget, stage: :pre_send, scope: scope
-        )
-      end
-
-      def check_per_tag(tags, time:, estimate: BigDecimal("0"), blocking_only: false)
-        PerTag.rules_for(tags, blocking_only: blocking_only).each do |rule|
-          rule.windows.each do |window, limit|
-            total = PerTag.spend(rule.key, rule.value, window, time: time) + estimate
-            yield(rule, window, total, limit) if total >= limit
-          end
-        end
-      end
-
-      def scope_for(rule)
-        { key: rule.key, value: rule.value }
-      end
-
-      def check_windowed(budgets, time:, estimate: BigDecimal("0"))
-        return if budgets.empty?
-
-        periods = budgets.keys.map { |type| BUDGET_TYPE_TO_PERIOD.fetch(type) }
-        totals = LlmCostTracker::Ledger::Period::Totals.call(periods, time: time)
-        budgets.each do |budget_type, budget|
-          total = totals.fetch(BUDGET_TYPE_TO_PERIOD.fetch(budget_type)) + estimate
-          yield(budget_type, total, budget) if total >= budget
-        end
-      end
-
-      def handle_exceeded(budget_type:,
-                          total:,
-                          budget:,
-                          previous_total:,
-                          last_event: nil,
-                          scope: nil,
-                          behavior: nil,
-                          on_exceeded: nil)
-        config = LlmCostTracker.configuration
-        behavior ||= config.budgets.exceeded_behavior
-        on_exceeded ||= config.budgets.on_exceeded
-        payload = {
-          budget_type: budget_type,
-          total: total,
-          budget: budget,
-          last_event: last_event,
-          stage: :post_spend,
-          scope: scope
-        }
-
-        on_exceeded.call(payload) if on_exceeded && (previous_total.nil? || previous_total < budget)
-        BudgetExceededError.new(**payload) if %i[raise block_requests].include?(behavior)
-      end
     end
   end
 end

data/lib/llm_cost_tracker/capture/event_window.rb

@@ -8,10 +8,11 @@ module LlmCostTracker
       HEAD_EVENTS = 16
       TAIL_EVENTS = 32
       HEAVY_STRING_BYTES = 8 * 1024
-      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64 snapshot logprobs].freeze
+      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64 snapshot logprobs embedding].freeze
 
-      def initialize(notable: nil)
+      def initialize(notable: nil, trim: nil)
         @notable = notable
+        @trim = trim
         @head = []
         @kept = []
         @tail = []
@@ -20,19 +21,10 @@ module LlmCostTracker
       end
 
       def push(data, type: nil)
-        # The openai gem's chat stream helper resends all logprobs so far in logprobs.* events; none carry usage.
         return if @overflowed || type&.start_with?("logprobs.")
 
-        event = { event: type, data: strip_heavy_payload(data) }
-        size = approximate_bytesize(event)
-        if @head.size < HEAD_EVENTS
-          @head << [event, size]
-        else
-          @tail << [event, size]
-          settle(@tail.shift) if @tail.size > TAIL_EVENTS
-        end
-        @bytes += size
-        overflow! if @bytes > SSE::LIMIT_BYTES
+        event = { event: type, data: strip_heavy_payload(@trim ? @trim.call(data) : data) }
+        append([event, approximate_bytesize(event)])
       rescue TypeError, SystemStackError
         overflow!
       end
@@ -47,6 +39,17 @@ module LlmCostTracker
 
       private
 
+      def append(entry)
+        if @head.size < HEAD_EVENTS
+          @head << entry
+        else
+          @tail << entry
+          settle(@tail.shift) if @tail.size > TAIL_EVENTS
+        end
+        @bytes += entry.last
+        overflow! if @bytes > SSE::LIMIT_BYTES
+      end
+
       def settle(entry)
         return @kept << entry if notable?(entry.first[:data])
 

data/lib/llm_cost_tracker/capture/sdk_payload.rb

@@ -8,30 +8,35 @@ require_relative "event_window"
 module LlmCostTracker
   module Capture
     module SdkPayload
-      module_function
+      class << self
+        def normalize(value)
+          case value
+          when Hash then normalize_hash(value)
+          when Array then value.map { |nested| normalize(nested) }
+          when Symbol then value.to_s
+          when NilClass then nil
+          else normalize_object(value)
+          end
+        end
+
+        private
 
-      def normalize(value)
-        case value
-        when Hash
-          value.each_with_object({}) do |(key, nested), out|
+        def normalize_hash(hash)
+          hash.each_with_object({}) do |(key, nested), out|
             out[key.to_s] = normalize(nested) unless EventWindow::IGNORED_PAYLOAD_KEYS.include?(key.to_s)
           end
-        when Array
-          value.map { |nested| normalize(nested) }
-        when Symbol
-          value.to_s
-        when NilClass
-          nil
-        else
+        end
+
+        def normalize_object(value)
           converted = container_for(value)
           converted ? normalize(converted) : value.deep_dup
         end
-      end
 
-      def container_for(value)
-        value.try(:deep_to_h) || value.try(:to_h)
-      rescue StandardError
-        nil
+        def container_for(value)
+          value.try(:to_h)
+        rescue StandardError
+          nil
+        end
       end
     end
   end

data/lib/llm_cost_tracker/capture/sse.rb

@@ -64,11 +64,12 @@ module LlmCostTracker
 
         def consume_lines
           start = 0
-          while (newline = @pending.index("\n", start))
+          while (newline = @pending.index("\n", [start, @scan_pos].max))
             consume_line(@pending.byteslice(start, newline - start))
             start = newline + 1
           end
-          @pending = @pending.byteslice(start..)
+          @pending.slice!(0, start) if start.positive?
+          @scan_pos = @pending.bytesize
         end
 
         def consume_line(raw)
@@ -140,9 +141,10 @@ module LlmCostTracker
         end
 
         def trim_array_buffer(pos)
-          keep_from = @object_start || [pos, @pending.bytesize].min
-          @pending = @pending.byteslice(keep_from..)
-          @scan_pos = pos - keep_from
+          size = @pending.bytesize
+          keep_from = @object_start || [pos, size].min
+          @pending.slice!(0, keep_from) if keep_from.positive?
+          @scan_pos = [pos, size].max - keep_from
           @object_start &&= 0
         end
 

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -11,10 +11,16 @@ require_relative "../timing"
 module LlmCostTracker
   module Capture
     class StreamCollector
+      DIMENSIONS = %i[provider_project_id provider_api_key_id provider_workspace_id].freeze
+      Snapshot = Data.define(*%i[events overflowed explicit_usage model latency_ms provider_response_id
+                                 capture_dimensions pricing_mode metadata context_tags request])
+      private_constant :DIMENSIONS, :Snapshot
+
       attr_reader :provider
 
       def initialize(provider:,
                      model:,
+                     parsed_as: provider,
                      latency_ms: nil,
                      provider_response_id: nil,
                      provider_project_id: nil,
@@ -24,44 +30,33 @@ module LlmCostTracker
                      metadata: {},
                      request: nil)
         @provider = provider.to_s
+        @parsed_as = parsed_as.to_s
         @model = model
         @latency_ms = latency_ms
         @provider_response_id = provider_response_id
-        @provider_project_id = provider_project_id
-        @provider_api_key_id = provider_api_key_id
-        @provider_workspace_id = provider_workspace_id
+        @dimensions = { provider_project_id:, provider_api_key_id:, provider_workspace_id: }
         @pricing_mode = pricing_mode
         @metadata = (metadata || {}).deep_dup
         @context_tags = LlmCostTracker::Tags::Context.tags.deep_dup
         @request = request
-        @window = EventWindow.new(notable: Parsers.find_for_provider(@provider)&.method(:retain_stream_event?))
+        @window = parser_window
         @explicit_usage = nil
         @started_at = LlmCostTracker::Timing.now_monotonic
-        @finished = false
-        @recording = false
+        @state = :open
         @mutex = Mutex.new
       end
 
       def model=(value)
-        @mutex.synchronize do
-          ensure_open!
-          @model = value
-        end
+        modify { @model = value }
       end
 
       def provider_response_id=(value)
-        @mutex.synchronize do
-          ensure_open!
-          @provider_response_id = value
-        end
+        modify { @provider_response_id = value }
       end
 
       def event(data, type: nil)
         data = SdkPayload.normalize(data) if data.is_a?(Hash)
-        @mutex.synchronize do
-          ensure_open!
-          @window.push(data, type: type&.to_s) unless data.nil?
-        end
+        modify { @window.push(data, type: type&.to_s) unless data.nil? }
       end
 
       def usage(input_tokens:, output_tokens:, **extra)
@@ -71,17 +66,10 @@ module LlmCostTracker
                 "pass `pricing_mode: :batch` to track_stream"
         end
 
-        @mutex.synchronize do
-          ensure_open!
+        modify do
           @provider_response_id = extra.delete(:provider_response_id) || @provider_response_id
-          @provider_project_id = extra.delete(:provider_project_id) || @provider_project_id
-          @provider_api_key_id = extra.delete(:provider_api_key_id) || @provider_api_key_id
-          @provider_workspace_id = extra.delete(:provider_workspace_id) || @provider_workspace_id
-          @explicit_usage = Usage::TokenUsage.build(
-            **extra,
-            input_tokens: input_tokens,
-            output_tokens: output_tokens
-          )
+          DIMENSIONS.each { |key| @dimensions[key] = extra.delete(key) || @dimensions[key] }
+          @explicit_usage = Usage::TokenUsage.build(**extra, input_tokens: input_tokens, output_tokens: output_tokens)
         end
       end
 
@@ -102,49 +90,57 @@ module LlmCostTracker
 
       private
 
+      def modify
+        @mutex.synchronize do
+          raise FrozenError, "can't modify finished LlmCostTracker::Capture::StreamCollector" if @state == :finished
+
+          yield
+        end
+      end
+
+      def parser_window
+        parser = Parsers.find_for_provider(@parsed_as)
+        EventWindow.new(notable: parser&.method(:retain_stream_event?), trim: parser&.method(:trim_stream_event))
+      end
+
       def claim_recording_slot
         @mutex.synchronize do
-          return nil if @finished || @recording
-
-          @recording = true
-          pricing_mode = Pricing::Mode.normalize(@pricing_mode)
-          {
-            events: @window.events,
-            overflowed: @window.overflowed?,
-            explicit_usage: @explicit_usage,
-            model: @model,
-            latency_ms: @latency_ms,
-            provider_response_id: @provider_response_id,
-            capture_dimensions: capture_dimensions,
-            pricing_mode: pricing_mode,
-            metadata: @metadata.deep_dup,
-            context_tags: @context_tags.deep_dup,
-            request: @request
-          }
+          return nil unless @state == :open
+
+          @state = :recording
+          snapshot
         end
       end
 
-      def record_snapshot(snapshot, errored:)
-        save_succeeded = false
-        begin
-          event = build_event(snapshot)
-          event = event.with(
-            provider_response_id: event.provider_response_id || snapshot[:provider_response_id],
-            pricing_mode: Pricing::Mode.merge(event.pricing_mode, snapshot[:pricing_mode])
-          )
+      def snapshot
+        pricing_mode = Pricing::Mode.normalize(@pricing_mode)
+        Snapshot.new(
+          events: @window.events,
+          overflowed: @window.overflowed?,
+          explicit_usage: @explicit_usage,
+          model: @model,
+          latency_ms: @latency_ms,
+          provider_response_id: @provider_response_id,
+          capture_dimensions: @dimensions.transform_values { |value| value.to_s.strip.presence }.compact,
+          pricing_mode: pricing_mode,
+          metadata: @metadata.deep_dup,
+          context_tags: @context_tags.deep_dup,
+          request: @request
+        )
+      end
 
-          Tracker.record(
-            event: event,
-            latency_ms: snapshot[:latency_ms] || LlmCostTracker::Timing.elapsed_ms(@started_at),
-            metadata: (errored ? { stream_errored: true } : {}).merge(snapshot[:metadata]),
-            context_tags: snapshot[:context_tags]
-          ) { save_succeeded = true }
-        ensure
-          @mutex.synchronize do
-            @finished = save_succeeded
-            @recording = false
-            release_buffers if save_succeeded
-          end
+      def record_snapshot(snapshot, errored:)
+        saved = false
+        Tracker.record(
+          event: event_for(snapshot),
+          latency_ms: snapshot.latency_ms || LlmCostTracker::Timing.elapsed_ms(@started_at),
+          metadata: (errored ? { stream_errored: true } : {}).merge(snapshot.metadata),
+          context_tags: snapshot.context_tags
+        ) { saved = true }
+      ensure
+        @mutex.synchronize do
+          @state = saved ? :finished : :open
+          release_buffers if saved
         end
       end
 
@@ -153,42 +149,39 @@ module LlmCostTracker
         @request = nil
       end
 
-      def capture_dimensions
-        {
-          provider_project_id: @provider_project_id.to_s.strip.presence,
-          provider_api_key_id: @provider_api_key_id.to_s.strip.presence,
-          provider_workspace_id: @provider_workspace_id.to_s.strip.presence
-        }.compact
+      def event_for(snapshot)
+        event = build_event(snapshot)
+        event.with(
+          provider_response_id: event.provider_response_id || snapshot.provider_response_id,
+          pricing_mode: Pricing::Mode.merge(event.pricing_mode, snapshot.pricing_mode)
+        )
       end
 
-      def ensure_open!
-        return unless @finished
+      def build_event(snapshot)
+        return build_unparsed_event(snapshot) if snapshot.explicit_usage
+        return overflowed_event(snapshot) if snapshot.overflowed
 
-        raise FrozenError, "can't modify finished LlmCostTracker::Capture::StreamCollector"
-      end
+        parsed = parse_events(snapshot)
+        return build_unparsed_event(snapshot) unless parsed
 
-      def build_event(snapshot)
-        return build_unparsed_event(snapshot) if snapshot[:explicit_usage]
+        model = present_model(parsed.model) || present_model(snapshot.model) || Event::UNKNOWN_MODEL
+        parsed.with(provider: @provider, model: model, **snapshot.capture_dimensions)
+      end
 
-        if snapshot[:overflowed]
-          Logging.warn("#{@provider} stream events exceeded #{SSE::LIMIT_BYTES} bytes; " \
-                       "recording usage_source=#{Usage::Source::UNKNOWN}.")
-          return build_unparsed_event(snapshot)
-        end
+      def overflowed_event(snapshot)
+        Logging.warn("#{@provider} stream events exceeded #{SSE::LIMIT_BYTES} bytes; " \
+                     "recording usage_source=#{Usage::Source::UNKNOWN}.")
+        build_unparsed_event(snapshot)
+      end
 
-        request_body = request_body_for(snapshot[:request])
-        events = Parsers.all_for_provider(@provider).filter_map do |parser|
+      def parse_events(snapshot)
+        request_body = request_body_for(snapshot.request)
+        events = Parsers.all_for_provider(@parsed_as).filter_map do |parser|
           parser.parse_stream(
-            response_status: 200, events: snapshot[:events], request_body: request_body, model: snapshot[:model]
+            response_status: 200, events: snapshot.events, request_body: request_body, model: snapshot.model
           )
         end
-        event = events.find { |parsed| parsed.usage_source != Usage::Source::UNKNOWN } || events.first
-        if event
-          model = present_model(event.model) || present_model(snapshot[:model]) || Event::UNKNOWN_MODEL
-          return event.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
-        end
-
-        build_unparsed_event(snapshot)
+        events.find { |parsed| parsed.usage_source != Usage::Source::UNKNOWN } || events.first
       end
 
       def request_body_for(request)
@@ -205,15 +198,15 @@ module LlmCostTracker
       end
 
       def build_unparsed_event(snapshot)
-        explicit_usage = snapshot[:explicit_usage]
+        explicit_usage = snapshot.explicit_usage
         Event.build(
           provider: @provider,
-          model: snapshot[:model] || Event::UNKNOWN_MODEL,
+          model: snapshot.model || Event::UNKNOWN_MODEL,
           token_usage: explicit_usage || Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: explicit_usage ? Usage::Source::MANUAL : Usage::Source::UNKNOWN,
-          pricing_mode: snapshot[:pricing_mode],
-          **snapshot.fetch(:capture_dimensions)
+          pricing_mode: snapshot.pricing_mode,
+          **snapshot.capture_dimensions
         )
       end
     end

data/lib/llm_cost_tracker/capture/stream_tap.rb

@@ -8,8 +8,8 @@ module LlmCostTracker
     class StreamTap
       MAX_PENDING_BYTES = 16 * 1_048_576
 
-      def initialize(notable: nil)
-        @window = EventWindow.new(notable: notable)
+      def initialize(notable: nil, trim: nil)
+        @window = EventWindow.new(notable: notable, trim: trim)
         @reader = SSE::Reader.new { |event| @window.push(event[:data], type: event[:event]) }
         @received = false
         @failed = false