llm_cost_tracker 0.13.0 → 0.14.1

data/lib/llm_cost_tracker/budget.rb

@@ -4,42 +4,94 @@ require "bigdecimal"
 
 require_relative "ledger"
 require_relative "pricing/estimator"
+require_relative "budget/per_tag"
+require_relative "tags/context"
 
 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, force: false)
+      def enforce!(provider: nil, model: nil, request: nil, estimate: nil, tags: nil, force: false)
         config = LlmCostTracker.configuration
         return unless config.enabled
-        return unless force || config.budget_exceeded_behavior == :block_requests
+
+        globally = force || config.budgets.exceeded_behavior == :block_requests
+        per_tag = force || PerTag.blocking?
+        return unless globally || per_tag
 
         estimate ||= estimate_cost(provider: provider, model: model, request: request)
-        raise_per_call_pre_send(estimate, config.per_call_budget) if config.per_call_budget && estimate.positive?
-
-        check_windowed({ monthly: config.monthly_budget, daily: config.daily_budget }.compact,
-                       time: Time.now.utc,
-estimate: estimate) do |budget_type, total, budget|
-          raise BudgetExceededError.new(**budget_payload(
-            budget_type: budget_type, total: total, budget: budget, last_event: nil, stage: :pre_send
-          ))
+        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
       end
 
-      def check!(event)
+      def check!(event, behavior_override: nil)
         config = LlmCostTracker.configuration
         return unless event.total_cost
 
-        check_per_call_budget(event, config)
-        check_windowed({ daily: config.daily_budget, monthly: config.monthly_budget }.compact,
+        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, last_event: event)
+          handle_exceeded(budget_type: budget_type,
+                          total: total,
+                          budget: budget,
+                          previous_total: total - event.total_cost,
+                          last_event: event,
+                          behavior: behavior_override)
+        end
+      end
+
+      def check_persisted!(events, behavior_override: nil)
+        by_rule = PerTag.rules_for_events(events.select(&:total_cost))
+        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
+            limit = rule.windows.fetch(window)
+            next if total < limit
+
+            handle_exceeded(
+              budget_type: window,
+              total: total,
+              budget: limit,
+              previous_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
       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 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
 
@@ -54,14 +106,48 @@ estimate: estimate) do |budget_type, total, budget|
         ))
       end
 
-      def check_per_call_budget(event, config)
-        budget = config.per_call_budget
+      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, last_event: event)
+        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:)
+        raise_per_call_pre_send(estimate, config.budgets.per_call) if config.budgets.per_call && estimate.positive?
+
+        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_payload(
+          budget_type: budget_type, total: total, budget: budget, last_event: nil, 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"))
@@ -82,38 +168,40 @@ estimate: estimate) do |budget_type, total, budget|
         period_for.transform_values { |period| period_totals.fetch(period) }
       end
 
-      def handle_exceeded(budget_type:, total:, budget:, last_event: nil)
+      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_payload(
           budget_type: budget_type,
           total: total,
           budget: budget,
           last_event: last_event,
-          stage: :post_spend
+          stage: :post_spend,
+          scope: scope
         )
 
-        if notify_exceeded?(config, budget_type: budget_type, total: total, budget: budget, last_event: last_event)
-          config.on_budget_exceeded&.call(payload)
-        end
-        raise BudgetExceededError.new(**payload) if %i[raise block_requests].include?(config.budget_exceeded_behavior)
+        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:)
+      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
+          stage: stage,
+          scope: scope
         }
       end
-
-      def notify_exceeded?(config, budget_type:, total:, budget:, last_event:)
-        return false unless config.on_budget_exceeded
-        return true if !last_event&.total_cost || budget_type == :per_call
-
-        total - last_event.total_cost < budget
-      end
     end
   end
 end

data/lib/llm_cost_tracker/capture/event_window.rb

@@ -0,0 +1,99 @@
+# 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].freeze
+
+      def initialize(notable: nil)
+        @notable = notable
+        @head = []
+        @kept = []
+        @tail = []
+        @bytes = 0
+        @overflowed = false
+      end
+
+      def push(data, type: nil)
+        return if @overflowed
+
+        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/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,7 @@ 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 "../timing"
 
 module LlmCostTracker
@@ -34,9 +34,7 @@ module LlmCostTracker
         @metadata = (metadata || {}).deep_dup
         @context_tags = (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
@@ -61,7 +59,7 @@ module LlmCostTracker
       def event(data, type: nil)
         @mutex.synchronize do
           ensure_open!
-          capture_event(data, type: type) unless data.nil?
+          @window.push(data, type: type) unless data.nil?
         end
       end
 
@@ -91,6 +89,12 @@ module LlmCostTracker
         return if snapshot.nil?
 
         record_snapshot(snapshot, errored: errored)
+      rescue TransactionAbortedError
+        raise
+      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 +106,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 +140,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,
@@ -207,55 +217,6 @@ module LlmCostTracker
           **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

data/lib/llm_cost_tracker/capture/stream_tap.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative "event_window"
+require_relative "sse"
+
+module LlmCostTracker
+  module Capture
+    class StreamTap
+      MAX_PENDING_BYTES = 16 * 1_048_576
+
+      def initialize(notable: nil)
+        @window = EventWindow.new(notable: notable)
+        @reader = SSE::Reader.new { |event| @window.push(event[:data], type: event[:event]) }
+        @received = false
+        @failed = false
+      end
+
+      def <<(chunk)
+        return if @failed
+
+        @received = true
+        @reader << chunk
+        fail! if @reader.pending_bytesize > MAX_PENDING_BYTES
+      rescue StandardError
+        fail!
+      end
+
+      def received?
+        @received
+      end
+
+      def failed?
+        @failed
+      end
+
+      def overflowed?
+        @window.overflowed?
+      end
+
+      def events
+        @reader&.finish
+        @failed ? [] : @window.events
+      rescue StandardError
+        fail!
+        []
+      end
+
+      private
+
+      def fail!
+        @failed = true
+        @reader = nil
+        @window = EventWindow.new
+      end
+    end
+  end
+end