llm_cost_tracker 0.14.0 → 0.14.2

data/lib/llm_cost_tracker/budget.rb

@@ -37,25 +37,40 @@ module LlmCostTracker
         config = LlmCostTracker.configuration
         return unless event.total_cost
 
-        check_per_call_budget(event, config, behavior_override)
+        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|
-          handle_exceeded(budget_type: budget_type,
-                          total: total,
-                          budget: budget,
-                          previous_total: total - event.total_cost,
-                          last_event: event,
-                          behavior: behavior_override)
+          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?
+        raise_first(errors)
       end
 
       def check_persisted!(events, behavior_override: nil)
-        by_rule = PerTag.rules_for_events(events.select(&:total_cost))
+        raise_first(persisted_errors(events, behavior_override))
+      end
+
+      def notify_persisted_safely!(events)
+        check_persisted!(events, behavior_override: :notify)
+      rescue StandardError => e
+        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).each do |(key, window, bucket), scored|
-          totals = PerTag.spend_by_value(key, scored.keys.map(&:value), window, bucket)
-          scored.each do |rule, recorded|
-            total = totals.fetch(rule.value, 0).to_d
+        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
 
@@ -63,7 +78,7 @@ module LlmCostTracker
               budget_type: window,
               total: total,
               budget: limit,
-              previous_total: total - recorded.sum(&:total_cost),
+              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,
@@ -73,14 +88,11 @@ module LlmCostTracker
         end
       end
 
-      def notify_persisted_safely!(events)
-        check_persisted!(events, behavior_override: :notify)
-      rescue StandardError => e
-        Logging.warn("Per-tag budget check failed after ingest: #{e.class}: #{e.message}")
+      def raise_first(errors)
+        error = errors.compact.first
+        raise error if error
       end
 
-      private
-
       def window_buckets(by_rule)
         by_rule.each_with_object({}) do |(rule, events), grouped|
           rule.windows.each_key do |window|
@@ -98,14 +110,6 @@ module LlmCostTracker
         Pricing::Estimator.call(provider: provider, model: model, request: request) || BigDecimal("0")
       end
 
-      def raise_per_call_pre_send(estimate, budget)
-        return unless estimate >= budget
-
-        raise BudgetExceededError.new(**budget_payload(
-          budget_type: :per_call, total: estimate, budget: budget, last_event: nil, stage: :pre_send
-        ))
-      end
-
       def check_per_call_budget(event, config, behavior_override)
         budget = config.budgets.per_call
         return unless budget
@@ -122,7 +126,10 @@ module LlmCostTracker
       end
 
       def enforce_globally(config, estimate:, time:)
-        raise_per_call_pre_send(estimate, config.budgets.per_call) if config.budgets.per_call && estimate.positive?
+        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,
@@ -132,9 +139,9 @@ module LlmCostTracker
       end
 
       def raise_pre_send(budget_type:, total:, budget:, scope: nil)
-        raise BudgetExceededError.new(**budget_payload(
-          budget_type: budget_type, total: total, budget: budget, last_event: nil, stage: :pre_send, scope: scope
-        ))
+        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)
@@ -153,21 +160,14 @@ module LlmCostTracker
       def check_windowed(budgets, time:, estimate: BigDecimal("0"))
         return if budgets.empty?
 
-        totals = totals_for(budgets.keys, time: time)
+        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) + estimate
+          total = totals.fetch(BUDGET_TYPE_TO_PERIOD.fetch(budget_type)) + estimate
           yield(budget_type, total, budget) if total >= budget
         end
       end
 
-      def totals_for(budget_types, time:)
-        return {} if budget_types.empty?
-
-        period_for = budget_types.to_h { |type| [type, BUDGET_TYPE_TO_PERIOD.fetch(type)] }
-        period_totals = LlmCostTracker::Ledger::Period::Totals.call(period_for.values, time: time)
-        period_for.transform_values { |period| period_totals.fetch(period) }
-      end
-
       def handle_exceeded(budget_type:,
                           total:,
                           budget:,
@@ -179,28 +179,17 @@ module LlmCostTracker
         config = LlmCostTracker.configuration
         behavior ||= config.budgets.exceeded_behavior
         on_exceeded ||= config.budgets.on_exceeded
-        payload = budget_payload(
+        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)
-        raise BudgetExceededError.new(**payload) if %i[raise block_requests].include?(behavior)
-      end
-
-      def budget_payload(budget_type:, total:, budget:, last_event:, stage:, scope: nil)
-        {
-          budget_type: budget_type,
-          total: total,
-          budget: budget,
-          last_event: last_event,
-          stage: stage,
-          scope: scope
-        }
+        BudgetExceededError.new(**payload) if %i[raise block_requests].include?(behavior)
       end
     end
   end

data/lib/llm_cost_tracker/capture/event_window.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require_relative "sse"
+
+module LlmCostTracker
+  module Capture
+    class EventWindow
+      HEAD_EVENTS = 16
+      TAIL_EVENTS = 32
+      HEAVY_STRING_BYTES = 8 * 1024
+      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64 snapshot logprobs].freeze
+
+      def initialize(notable: nil)
+        @notable = notable
+        @head = []
+        @kept = []
+        @tail = []
+        @bytes = 0
+        @overflowed = false
+      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
+      rescue TypeError, SystemStackError
+        overflow!
+      end
+
+      def events
+        (@head + @kept + @tail).map(&:first)
+      end
+
+      def overflowed?
+        @overflowed
+      end
+
+      private
+
+      def settle(entry)
+        return @kept << entry if notable?(entry.first[:data])
+
+        @bytes -= entry.last
+      end
+
+      def notable?(data)
+        @notable&.call(data)
+      rescue StandardError
+        false
+      end
+
+      def overflow!
+        @overflowed = true
+        @head = []
+        @kept = []
+        @tail = []
+      end
+
+      def strip_heavy_payload(value)
+        case value
+        when Hash
+          value.each_with_object({}) do |(key, nested), out|
+            next if IGNORED_PAYLOAD_KEYS.include?(key.to_s)
+
+            out[key] = strip_heavy_payload(nested)
+          end
+        when Array
+          value.map { |nested| strip_heavy_payload(nested) }
+        when String
+          value.bytesize > HEAVY_STRING_BYTES ? "" : value
+        else
+          value
+        end
+      end
+
+      def approximate_bytesize(value)
+        case value
+        when Hash
+          value.sum { |key, nested| approximate_bytesize(key) + approximate_bytesize(nested) + 4 }
+        when Array
+          value.sum { |nested| approximate_bytesize(nested) + 2 }
+        when Numeric, true, false, nil
+          8
+        else
+          value.to_s.bytesize + 2
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/capture/sdk_payload.rb

@@ -3,6 +3,8 @@
 require "active_support/core_ext/object/deep_dup"
 require "active_support/core_ext/object/try"
 
+require_relative "event_window"
+
 module LlmCostTracker
   module Capture
     module SdkPayload
@@ -11,7 +13,9 @@ module LlmCostTracker
       def normalize(value)
         case value
         when Hash
-          value.each_with_object({}) { |(key, nested), out| out[key.to_s] = normalize(nested) }
+          value.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

data/lib/llm_cost_tracker/capture/sse.rb

@@ -11,52 +11,83 @@ module LlmCostTracker
 
       class << self
         def parse(body)
-          return [] if body.blank?
+          events = []
+          reader = Reader.new { |event| events << event }
+          reader << body
+          reader.finish
+          events
+        end
+      end
 
-          return parse_json_array(body) if body.match?(/\A\s*\[/)
+      class Reader
+        def initialize(&on_event)
+          @on_event = on_event
+          @pending = String.new(encoding: Encoding::BINARY)
+          @mode = nil
+          @event_name = nil
+          @data_lines = []
+          @scan_pos = 0
+          @depth = 0
+          @in_string = false
+          @object_start = nil
+        end
 
-          parse_event_stream(body)
+        def <<(chunk)
+          @pending << chunk.to_s.b
+          @mode ||= detect_mode
+          case @mode
+          when :sse then consume_lines
+          when :array then consume_array
+          end
         end
 
-        private
+        def pending_bytesize
+          @pending.bytesize
+        end
 
-        def parse_event_stream(body)
-          events = []
-          current_event = nil
-          data_lines = []
+        def finish
+          return unless @mode == :sse
 
-          body.each_line do |raw|
-            line = raw.chomp
+          consume_line(@pending) unless @pending.empty?
+          @pending = String.new(encoding: Encoding::BINARY)
+          dispatch
+        end
 
-            if line.empty?
-              events << finalize_event(current_event, data_lines) if data_lines.any?
-              current_event = nil
-              data_lines = []
-              next
-            end
+        private
 
-            next if line.start_with?(":")
+        def detect_mode
+          first = @pending[/\S/]
+          return nil unless first
 
-            field, _, value = line.partition(":")
-            value = value[1..] if value.start_with?(" ")
+          first == "[" ? :array : :sse
+        end
 
-            case field
-            when "event" then current_event = value
-            when "data"  then data_lines << value
-            end
+        def consume_lines
+          start = 0
+          while (newline = @pending.index("\n", start))
+            consume_line(@pending.byteslice(start, newline - start))
+            start = newline + 1
           end
-
-          events << finalize_event(current_event, data_lines) if data_lines.any?
-          events.compact
+          @pending = @pending.byteslice(start..)
         end
 
-        def parse_json_array(body)
-          parsed = JSON.parse(body)
-          return [] unless parsed.is_a?(Array)
+        def consume_line(raw)
+          line = raw.chomp("\r").force_encoding(Encoding::UTF_8)
+          return dispatch if line.empty?
+          return if line.start_with?(":")
 
-          parsed.map { |entry| { event: nil, data: entry } }
-        rescue JSON::ParserError
-          []
+          field, _, value = line.partition(":")
+          value = value[1..] if value.start_with?(" ")
+          case field
+          when "event" then @event_name = value
+          when "data" then @data_lines << value
+          end
+        end
+
+        def dispatch
+          emit(finalize_event(@event_name, @data_lines)) if @data_lines.any?
+          @event_name = nil
+          @data_lines = []
         end
 
         def finalize_event(event_name, data_lines)
@@ -73,6 +104,51 @@ module LlmCostTracker
         rescue JSON::ParserError
           payload
         end
+
+        def consume_array
+          pos = @scan_pos
+          while (index = @pending.index(@in_string ? /["\\]/ : /[{}"]/, pos))
+            pos = index + 1
+            pos = advance_array(@pending.getbyte(index).chr, index, pos)
+          end
+          trim_array_buffer(pos)
+        end
+
+        def advance_array(char, index, pos)
+          if @in_string
+            return pos + 1 if char == "\\"
+
+            @in_string = false
+          elsif char == '"'
+            @in_string = true
+          elsif char == "{"
+            @object_start = index if @depth.zero?
+            @depth += 1
+          elsif @depth.positive?
+            @depth -= 1
+            emit_object(index) if @depth.zero?
+          end
+          pos
+        end
+
+        def emit_object(end_index)
+          text = @pending.byteslice(@object_start, end_index - @object_start + 1)
+          @object_start = nil
+          emit({ event: nil, data: JSON.parse(text.force_encoding(Encoding::UTF_8)) })
+        rescue JSON::ParserError
+          nil
+        end
+
+        def trim_array_buffer(pos)
+          keep_from = @object_start || [pos, @pending.bytesize].min
+          @pending = @pending.byteslice(keep_from..)
+          @scan_pos = pos - keep_from
+          @object_start &&= 0
+        end
+
+        def emit(event)
+          @on_event.call(event) if event
+        end
       end
     end
   end

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -4,7 +4,8 @@ require "active_support/core_ext/object/blank"
 require "active_support/core_ext/object/deep_dup"
 require "json"
 
-require_relative "sse"
+require_relative "event_window"
+require_relative "sdk_payload"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -21,7 +22,6 @@ module LlmCostTracker
                      provider_workspace_id: nil,
                      pricing_mode: nil,
                      metadata: {},
-                     context_tags: nil,
                      request: nil)
         @provider = provider.to_s
         @model = model
@@ -32,11 +32,9 @@ module LlmCostTracker
         @provider_workspace_id = provider_workspace_id
         @pricing_mode = pricing_mode
         @metadata = (metadata || {}).deep_dup
-        @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
+        @context_tags = LlmCostTracker::Tags::Context.tags.deep_dup
         @request = request
-        @events = []
-        @captured_bytes = 0
-        @overflowed = false
+        @window = EventWindow.new(notable: Parsers.find_for_provider(@provider)&.method(:retain_stream_event?))
         @explicit_usage = nil
         @started_at = LlmCostTracker::Timing.now_monotonic
         @finished = false
@@ -59,9 +57,10 @@ module LlmCostTracker
       end
 
       def event(data, type: nil)
+        data = SdkPayload.normalize(data) if data.is_a?(Hash)
         @mutex.synchronize do
           ensure_open!
-          capture_event(data, type: type) unless data.nil?
+          @window.push(data, type: type&.to_s) unless data.nil?
         end
       end
 
@@ -91,6 +90,14 @@ module LlmCostTracker
         return if snapshot.nil?
 
         record_snapshot(snapshot, errored: errored)
+      rescue TransactionAbortedError
+        raise
+      rescue ActiveRecord::RecordNotUnique
+        nil
+      rescue StandardError => e
+        raise unless errored
+
+        Logging.warn("Recording an errored stream raised #{e.class}: #{e.message}; kept the stream's own exception")
       end
 
       private
@@ -102,8 +109,8 @@ module LlmCostTracker
           @recording = true
           pricing_mode = Pricing::Mode.normalize(@pricing_mode)
           {
-            events: @events.dup,
-            overflowed: @overflowed,
+            events: @window.events,
+            overflowed: @window.overflowed?,
             explicit_usage: @explicit_usage,
             model: @model,
             latency_ms: @latency_ms,
@@ -136,10 +143,16 @@ module LlmCostTracker
           @mutex.synchronize do
             @finished = save_succeeded
             @recording = false
+            release_buffers if save_succeeded
           end
         end
       end
 
+      def release_buffers
+        @window = EventWindow.new
+        @request = nil
+      end
+
       def capture_dimensions
         {
           provider_project_id: @provider_project_id.to_s.strip.presence,
@@ -155,20 +168,27 @@ module LlmCostTracker
       end
 
       def build_event(snapshot)
-        return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
-        return build_unknown_usage(snapshot) if snapshot[:overflowed]
+        return build_unparsed_event(snapshot) if snapshot[:explicit_usage]
 
-        event = Parsers.find_for_provider(@provider)&.parse_stream(
-          response_status: 200,
-          events: snapshot[:events],
-          request_body: request_body_for(snapshot[:request])
-        )
+        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
+
+        request_body = request_body_for(snapshot[:request])
+        events = Parsers.all_for_provider(@provider).filter_map do |parser|
+          parser.parse_stream(
+            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_unknown_usage(snapshot)
+        build_unparsed_event(snapshot)
       end
 
       def request_body_for(request)
@@ -184,78 +204,18 @@ module LlmCostTracker
         string unless string == Event::UNKNOWN_MODEL
       end
 
-      def build_from_explicit_usage(snapshot)
+      def build_unparsed_event(snapshot)
+        explicit_usage = snapshot[:explicit_usage]
         Event.build(
           provider: @provider,
           model: snapshot[:model] || Event::UNKNOWN_MODEL,
-          token_usage: snapshot[:explicit_usage],
+          token_usage: explicit_usage || Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
-          usage_source: Usage::Source::MANUAL,
+          usage_source: explicit_usage ? Usage::Source::MANUAL : Usage::Source::UNKNOWN,
           pricing_mode: snapshot[:pricing_mode],
           **snapshot.fetch(:capture_dimensions)
         )
       end
-
-      def build_unknown_usage(snapshot)
-        Event.build(
-          provider: @provider,
-          model: snapshot[:model] || Event::UNKNOWN_MODEL,
-          token_usage: Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
-          stream: true,
-          usage_source: Usage::Source::UNKNOWN,
-          pricing_mode: snapshot[:pricing_mode],
-          **snapshot.fetch(:capture_dimensions)
-        )
-      end
-
-      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64].freeze
-      private_constant :IGNORED_PAYLOAD_KEYS
-
-      HEAVY_STRING_BYTES = 8 * 1024
-      private_constant :HEAVY_STRING_BYTES
-
-      def capture_event(data, type:)
-        event = { event: type, data: strip_heavy_payload(data) }
-        size = approximate_bytesize(event)
-        if @captured_bytes + size <= Capture::SSE::LIMIT_BYTES
-          @events << event
-          @captured_bytes += size
-        else
-          @overflowed = true
-        end
-      rescue TypeError, SystemStackError
-        @overflowed = true
-      end
-
-      def strip_heavy_payload(value)
-        case value
-        when Hash
-          value.each_with_object({}) do |(key, nested), out|
-            next if IGNORED_PAYLOAD_KEYS.include?(key.to_s)
-
-            out[key] = strip_heavy_payload(nested)
-          end
-        when Array
-          value.map { |nested| strip_heavy_payload(nested) }
-        when String
-          value.bytesize > HEAVY_STRING_BYTES ? "" : value
-        else
-          value
-        end
-      end
-
-      def approximate_bytesize(value)
-        case value
-        when Hash
-          value.sum { |key, nested| approximate_bytesize(key) + approximate_bytesize(nested) + 4 }
-        when Array
-          value.sum { |nested| approximate_bytesize(nested) + 2 }
-        when Numeric, true, false, nil
-          8
-        else
-          value.to_s.bytesize + 2
-        end
-      end
     end
   end
 end