llm_cost_tracker 0.8.0 → 0.10.0

data/lib/llm_cost_tracker/billing/components.yml

@@ -68,6 +68,26 @@
   token_key: audio_output_tokens
   cost_key: audio_output_cost
 
+- key: image_input
+  kind: image_token
+  direction: input
+  modality: image
+  cache_state: none
+  unit: token
+  category: token
+  token_key: image_input_tokens
+  cost_key: image_input_cost
+
+- key: image_output
+  kind: image_token
+  direction: output
+  modality: image
+  cache_state: none
+  unit: token
+  category: token
+  token_key: image_output_tokens
+  cost_key: image_output_cost
+
 - key: web_search_request
   kind: web_search_request
   direction: neither
@@ -75,6 +95,34 @@
   cache_state: none
   unit: request
   category: tool
+  rate_basis: per_1k_requests
+
+- key: web_search_preview_request_reasoning
+  kind: web_search_preview_request_reasoning
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool
+  rate_basis: per_1k_requests
+
+- key: web_search_preview_request_non_reasoning
+  kind: web_search_preview_request_non_reasoning
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool
+  rate_basis: per_1k_requests
+
+- key: web_fetch_request
+  kind: web_fetch_request
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool
+  rate_basis: per_1k_requests
 
 - key: file_search_call
   kind: file_search_call
@@ -83,6 +131,7 @@
   cache_state: none
   unit: request
   category: tool
+  rate_basis: per_1k_requests
 
 - key: container_session
   kind: container_session
@@ -91,6 +140,7 @@
   cache_state: none
   unit: session
   category: runtime
+  rate_basis: per_session
 
 - key: code_execution_request
   kind: code_execution_request
@@ -99,6 +149,7 @@
   cache_state: none
   unit: request
   category: runtime
+  rate_basis: per_1k_requests
 
 - key: code_execution_hour
   kind: code_execution_hour
@@ -107,6 +158,7 @@
   cache_state: none
   unit: hour
   category: runtime
+  rate_basis: per_hour
 
 - key: grounding_request
   kind: grounding_request
@@ -115,3 +167,22 @@
   cache_state: none
   unit: request
   category: tool
+  rate_basis: per_1k_requests
+
+- key: text_to_speech_character
+  kind: text_to_speech_character
+  direction: output
+  modality: audio
+  cache_state: none
+  unit: character
+  category: tool
+  rate_basis: per_million_characters
+
+- key: mcp_call
+  kind: mcp_call
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool
+  rate_basis: per_request

data/lib/llm_cost_tracker/billing/cost_status.rb

@@ -10,36 +10,32 @@ module LlmCostTracker
       PARTIAL = "partial"
       UNKNOWN = "unknown"
 
-      class << self
-        # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
-        def call(token_usage:, usage_source:, token_cost:, service_line_items:, total_cost:,
-                 token_pricing_partial: false)
-          return UNKNOWN if usage_source == :unknown
+      # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+      def self.call(token_usage:, usage_source:, token_cost:, service_line_items:, total_cost:,
+                    token_pricing_partial: false)
+        return UNKNOWN if usage_source == :unknown
 
-          token_billable = Components::TOKEN_PRICED.any? do |component|
-            token_usage.public_send(component.token_key).positive?
-          end
-          service_billable = false
-          service_priced = false
-          service_unpriced = false
-          service_line_items.each do |line_item|
-            next unless line_item.billable?
+        token_billable = token_usage.priced_quantities.any? { |_key, quantity| quantity.positive? }
+        service_billable = false
+        service_priced = false
+        service_unpriced = false
+        service_line_items.each do |line_item|
+          next unless line_item.billable?
 
-            service_billable = true
-            service_priced ||= line_item.priced?
-            service_unpriced ||= line_item.unpriced?
-            break if service_priced && service_unpriced
-          end
+          service_billable = true
+          service_priced ||= line_item.priced?
+          service_unpriced ||= line_item.unpriced?
+          break if service_priced && service_unpriced
+        end
 
-          priced = (token_billable && !token_cost.nil?) || service_priced || (!token_billable && !service_billable)
-          unpriced = (token_billable && (token_cost.nil? || token_pricing_partial)) || service_unpriced
-          return UNKNOWN if unpriced && !priced
-          return PARTIAL if unpriced
+        priced = (token_billable && !token_cost.nil?) || service_priced || (!token_billable && !service_billable)
+        unpriced = (token_billable && (token_cost.nil? || token_pricing_partial)) || service_unpriced
+        return UNKNOWN if unpriced && !priced
+        return PARTIAL if unpriced
 
-          total_cost.nil? || total_cost.zero? ? FREE : COMPLETE
-        end
-        # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+        total_cost.nil? || total_cost.zero? ? FREE : COMPLETE
       end
+      # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
     end
   end
 end

data/lib/llm_cost_tracker/billing/line_item.rb

@@ -30,28 +30,11 @@ module LlmCostTracker
 
     class LineItem
       USD = "USD"
-      OPTIONAL_ATTRIBUTES = %i[
-        pricing_basis
-        price_key
-        price_source
-        price_source_version
-        provider_field
-        provider_item_id
-      ].freeze
-      SYMBOL_ATTRIBUTES = %i[
-        kind
-        direction
-        modality
-        cache_state
-        unit
-        pricing_basis
-        price_source
-      ].freeze
 
       def self.build(attributes)
         attributes = attributes.to_h
         component = component_for(attributes)
-        normalized = {
+        new(
           kind: symbol_or_nil(attributes[:kind]) || component&.kind,
           direction: symbol_or_nil(attributes[:direction]) || component&.direction,
           modality: symbol_or_nil(attributes[:modality]) || component&.modality,
@@ -63,38 +46,30 @@ module LlmCostTracker
           cost: decimal_or_nil(attributes[:cost]),
           currency: attributes[:currency] || USD,
           cost_status: cost_status_for(attributes),
+          pricing_basis: symbol_or_nil(attributes[:pricing_basis]),
+          price_key: attributes[:price_key],
+          price_source: symbol_or_nil(attributes[:price_source]),
+          price_source_version: attributes[:price_source_version],
+          provider_field: attributes[:provider_field],
+          provider_item_id: attributes[:provider_item_id],
           details: attributes[:details] || {}
-        }.merge(optional_attributes_for(attributes))
-
-        new(**normalized)
+        )
       end
 
       def self.from_token_usage(token_usage)
         return [] unless token_usage
 
-        Components::TOKEN_PRICED.filter_map do |component|
-          quantity = token_usage.public_send(component.token_key)
+        token_usage.priced_quantities.filter_map do |key, quantity|
           next unless quantity.positive?
 
-          new(
+          component = Components::BY_KEY.fetch(key)
+          build(
             kind: component.kind,
             direction: component.direction,
             modality: component.modality,
             cache_state: component.cache_state,
-            quantity: BigDecimal(quantity.to_s),
-            unit: component.unit,
-            rate_amount: nil,
-            rate_quantity: BigDecimal("1"),
-            cost: nil,
-            currency: USD,
-            cost_status: CostStatus::UNKNOWN,
-            pricing_basis: nil,
-            price_key: nil,
-            price_source: nil,
-            price_source_version: nil,
-            provider_field: nil,
-            provider_item_id: nil,
-            details: {}
+            quantity: quantity,
+            unit: component.unit
           )
         end
       end
@@ -119,7 +94,7 @@ module LlmCostTracker
       def self.symbol_or_nil(value)
         return nil if value.nil?
 
-        value.is_a?(Symbol) ? value : value.to_s.to_sym
+        value.to_s.to_sym
       end
 
       def self.decimal_or_nil(value)
@@ -132,16 +107,7 @@ module LlmCostTracker
         decimal_or_nil(value) || BigDecimal("0")
       end
 
-      def self.optional_attributes_for(attributes)
-        OPTIONAL_ATTRIBUTES.to_h do |key|
-          value = attributes[key]
-          value = value.to_sym if value.is_a?(String) && SYMBOL_ATTRIBUTES.include?(key)
-          [key, value]
-        end
-      end
-
-      private_class_method :cost_status_for, :component_for, :symbol_or_nil, :decimal_or_nil, :decimal_or_zero,
-                           :optional_attributes_for
+      private_class_method :cost_status_for, :component_for, :symbol_or_nil, :decimal_or_nil, :decimal_or_zero
 
       def billable?
         quantity.positive?
@@ -163,7 +129,7 @@ module LlmCostTracker
         cost || BigDecimal("0")
       end
 
-      def apply_rate(rate)
+      def with_rate(rate)
         rate_amount = rate.fetch(:amount)
         rate_quantity = rate.fetch(:quantity)
         applied_cost = (quantity / rate_quantity) * rate_amount

data/lib/llm_cost_tracker/budget.rb

@@ -1,26 +1,35 @@
 # frozen_string_literal: true
 
+require "bigdecimal"
+
 require_relative "logging"
 require_relative "ledger"
+require_relative "pricing/estimator"
 
 module LlmCostTracker
   class Budget
     BUDGET_TYPE_TO_PERIOD = { monthly: :month, daily: :day }.freeze
 
     class << self
-      def enforce!
+      def enforce!(provider: nil, model: nil, request: nil)
         config = LlmCostTracker.configuration
         return unless config.budget_exceeded_behavior == :block_requests
 
+        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?
+
         budgets = { monthly: config.monthly_budget, daily: config.daily_budget }.compact
         return if budgets.empty?
 
         totals = totals_for(budgets.keys, time: Time.now.utc)
 
         budgets.each do |budget_type, budget|
-          total = totals.fetch(budget_type)
+          total = totals.fetch(budget_type) + estimate
+          next unless total >= budget
 
-          handle_exceeded(budget_type: budget_type, total: total, budget: budget) if total >= budget
+          raise BudgetExceededError.new(**budget_payload(
+            budget_type: budget_type, total: total, budget: budget, last_event: nil, stage: :pre_send
+          ))
         end
       end
 
@@ -41,6 +50,20 @@ module LlmCostTracker
 
       private
 
+      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 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)
         budget = config.per_call_budget
         return unless budget
@@ -67,7 +90,8 @@ module LlmCostTracker
           budget_type: budget_type,
           total: total,
           budget: budget,
-          last_event: last_event
+          last_event: last_event,
+          stage: :post_spend
         )
 
         if notify_exceeded?(config, budget_type: budget_type, total: total, budget: budget, last_event: last_event)
@@ -76,18 +100,18 @@ module LlmCostTracker
         raise BudgetExceededError.new(**payload) if %i[raise block_requests].include?(config.budget_exceeded_behavior)
       end
 
-      def budget_payload(budget_type:, total:, budget:, last_event:)
+      def budget_payload(budget_type:, total:, budget:, last_event:, stage:)
         {
           budget_type: budget_type,
           total: total,
           budget: budget,
-          last_event: last_event
+          last_event: last_event,
+          stage: stage
         }
       end
 
       def notify_exceeded?(config, budget_type:, total:, budget:, last_event:)
         return false unless config.on_budget_exceeded
-        return true unless config.budget_exceeded_behavior == :notify
         return true unless last_event&.total_cost
         return true if budget_type == :per_call
 

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -5,6 +5,7 @@ require "active_support/core_ext/object/deep_dup"
 require "json"
 
 require_relative "stream"
+require_relative "../pricing/mode"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -14,7 +15,7 @@ module LlmCostTracker
 
       def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, provider_project_id: nil,
                      provider_api_key_id: nil, provider_workspace_id: nil, batch: nil, pricing_mode: nil,
-                     metadata: {}, context_tags: nil)
+                     metadata: {}, context_tags: nil, request: nil)
         @provider = provider.to_s
         @model = model
         @latency_ms = latency_ms
@@ -26,12 +27,14 @@ module LlmCostTracker
         @pricing_mode = pricing_mode
         @metadata = (metadata || {}).deep_dup
         @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
+        @request = request
         @events = []
         @captured_bytes = 0
         @overflowed = false
         @explicit_usage = nil
         @started_at = LlmCostTracker::Timing.now_monotonic
         @finished = false
+        @recording = false
         @mutex = Mutex.new
       end
 
@@ -66,7 +69,6 @@ module LlmCostTracker
           ensure_open!
           capture_event(data, type: type) unless data.nil?
         end
-        self
       end
 
       def usage(input_tokens:, output_tokens:, **extra)
@@ -84,14 +86,22 @@ module LlmCostTracker
             output_tokens: output_tokens
           )
         end
-        self
       end
 
       def finish!(errored: false)
-        snapshot = @mutex.synchronize do
-          return if @finished
+        snapshot = claim_recording_slot
+        return if snapshot.nil?
 
-          @finished = true
+        record_snapshot(snapshot, errored: errored)
+      end
+
+      private
+
+      def claim_recording_slot
+        @mutex.synchronize do
+          return nil if @finished || @recording
+
+          @recording = true
           pricing_mode = Pricing.normalize_mode(@pricing_mode)
           {
             events: @events.dup,
@@ -103,27 +113,50 @@ module LlmCostTracker
             capture_dimensions: capture_dimensions(pricing_mode),
             pricing_mode: pricing_mode,
             metadata: @metadata.deep_dup,
-            context_tags: @context_tags.deep_dup
+            context_tags: @context_tags.deep_dup,
+            request: @request
           }
         end
+      end
 
-        capture = build_usage_capture(snapshot)
-        provider_response_id = capture.provider_response_id || snapshot[:provider_response_id]
-        capture = capture.with(provider_response_id: provider_response_id)
+      def record_snapshot(snapshot, errored:)
+        save_succeeded = false
+        begin
+          event = build_event(snapshot)
+          provider_response_id = event.provider_response_id || snapshot[:provider_response_id]
+          event = event.with(provider_response_id: provider_response_id)
 
-        Tracker.record(
-          capture: capture,
-          latency_ms: snapshot[:latency_ms] || LlmCostTracker::Timing.elapsed_ms(@started_at),
-          pricing_mode: snapshot[:pricing_mode],
-          metadata: (errored ? { stream_errored: true } : {}).merge(snapshot[:metadata]),
-          context_tags: snapshot[:context_tags]
-        )
+          Tracker.record(
+            event: event,
+            latency_ms: snapshot[:latency_ms] || LlmCostTracker::Timing.elapsed_ms(@started_at),
+            pricing_mode: merge_pricing_modes(event.pricing_mode, snapshot[:pricing_mode]),
+            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
+          end
+        end
       end
 
-      private
+      HOST_DERIVED_MODE_TOKENS = %i[data_residency].freeze
+      private_constant :HOST_DERIVED_MODE_TOKENS
+
+      def merge_pricing_modes(provider_mode, request_mode)
+        return Pricing.normalize_mode(request_mode) if provider_mode.to_s.strip.empty?
+
+        provider_tokens = Pricing::Mode.tokenize(provider_mode) - Pricing::STANDARD_MODE_VALUES
+        request_host_tokens = Pricing::Mode.tokenize(request_mode || "") & HOST_DERIVED_MODE_TOKENS
+        combined = provider_tokens | request_host_tokens
+        return nil if combined.empty?
+
+        Pricing.normalize_mode(combined.join("_"))
+      end
 
       def capture_dimensions(pricing_mode)
-        batch = @batch.nil? ? UsageCapture.batch_from_pricing_mode?(pricing_mode).presence : @batch
+        batch = @batch.nil? ? Event.batch_from_pricing_mode?(pricing_mode).presence : @batch
         {
           provider_project_id: @provider_project_id.to_s.strip.presence,
           provider_api_key_id: @provider_api_key_id.to_s.strip.presence,
@@ -138,34 +171,44 @@ module LlmCostTracker
         raise FrozenError, "can't modify finished LlmCostTracker::Capture::StreamCollector"
       end
 
-      def build_usage_capture(snapshot)
+      def build_event(snapshot)
         return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
+        return build_unknown_usage(snapshot) if snapshot[:overflowed]
 
-        capture = Parsers.find_for_provider(@provider)&.parse_stream(
+        event = Parsers.find_for_provider(@provider)&.parse_stream(
           response_status: 200,
-          events: snapshot[:events]
+          events: snapshot[:events],
+          request_body: request_body_for(snapshot[:request])
         )
-        if capture && (capture.usage_source != :unknown || !snapshot[:overflowed])
-          model = present_model(capture.model) || present_model(snapshot[:model]) || UsageCapture::UNKNOWN_MODEL
-          return capture.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
+        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)
       end
 
+      def request_body_for(request)
+        return nil unless request
+
+        JSON.generate(request)
+      rescue StandardError
+        nil
+      end
+
       def present_model(value)
         return nil if value.nil?
 
         string = value.to_s.presence
-        return nil if string.nil? || string == "unknown"
+        return nil if string.nil? || string == Event::UNKNOWN_MODEL
 
         string
       end
 
       def build_from_explicit_usage(snapshot)
-        UsageCapture.build(
+        Event.build(
           provider: @provider,
-          model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
+          model: snapshot[:model] || Event::UNKNOWN_MODEL,
           token_usage: snapshot[:explicit_usage],
           stream: true,
           usage_source: :manual,
@@ -175,9 +218,9 @@ module LlmCostTracker
       end
 
       def build_unknown_usage(snapshot)
-        UsageCapture.build(
+        Event.build(
           provider: @provider,
-          model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
+          model: snapshot[:model] || Event::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: :unknown,
@@ -186,18 +229,54 @@ module LlmCostTracker
         )
       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: data }
-        size = JSON.generate(event).bytesize
+        event = { event: type, data: strip_heavy_payload(data) }
+        size = approximate_bytesize(event)
         if @captured_bytes + size <= Capture::Stream::LIMIT_BYTES
-          @events << event.deep_dup
+          @events << event
           @captured_bytes += size
         else
           @overflowed = true
         end
-      rescue JSON::JSONError, TypeError
+      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