llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/billing/components.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+require "psych"
+
+require_relative "../errors"
+
+module LlmCostTracker
+  module Billing
+    module Components
+      Component = Data.define(
+        :key,
+        :kind,
+        :direction,
+        :modality,
+        :cache_state,
+        :unit,
+        :category,
+        :token_key,
+        :cost_key
+      )
+
+      REQUIRED_FIELDS = %i[key kind direction modality cache_state unit category].freeze
+      DEFINITIONS_PATH = File.expand_path("components.yml", __dir__)
+
+      def self.load_registry
+        Psych.safe_load_file(DEFINITIONS_PATH, permitted_classes: [], symbolize_names: true)
+             .map { |attributes| build(attributes) }
+             .freeze
+      end
+
+      def self.build(attributes)
+        missing = REQUIRED_FIELDS - attributes.keys
+        raise Error, "components.yml entry missing #{missing.join(', ')}: #{attributes.inspect}" if missing.any?
+
+        Component.new(
+          key: attributes.fetch(:key).to_sym,
+          kind: attributes.fetch(:kind).to_sym,
+          direction: attributes.fetch(:direction).to_sym,
+          modality: attributes.fetch(:modality).to_sym,
+          cache_state: attributes.fetch(:cache_state).to_sym,
+          unit: attributes.fetch(:unit).to_sym,
+          category: attributes.fetch(:category).to_sym,
+          token_key: attributes[:token_key]&.to_sym,
+          cost_key: attributes[:cost_key]&.to_sym
+        )
+      end
+
+      REGISTRY = load_registry
+      BY_KEY = REGISTRY.to_h { |component| [component.key, component] }.freeze
+      TOKEN_PRICED = REGISTRY.select { |component| component.token_key && component.cost_key }.freeze
+    end
+  end
+end

data/lib/llm_cost_tracker/billing/components.yml

@@ -0,0 +1,117 @@
+- key: input
+  kind: text_token
+  direction: input
+  modality: text
+  cache_state: none
+  unit: token
+  category: token
+  token_key: input_tokens
+  cost_key: input_cost
+
+- key: cache_read_input
+  kind: text_token
+  direction: input
+  modality: text
+  cache_state: read
+  unit: token
+  category: token
+  token_key: cache_read_input_tokens
+  cost_key: cache_read_input_cost
+
+- key: cache_write_input
+  kind: text_token
+  direction: input
+  modality: text
+  cache_state: write_default
+  unit: token
+  category: token
+  token_key: cache_write_input_tokens
+  cost_key: cache_write_input_cost
+
+- key: cache_write_extended_input
+  kind: text_token
+  direction: input
+  modality: text
+  cache_state: write_extended
+  unit: token
+  category: token
+  token_key: cache_write_extended_input_tokens
+  cost_key: cache_write_extended_input_cost
+
+- key: output
+  kind: text_token
+  direction: output
+  modality: text
+  cache_state: none
+  unit: token
+  category: token
+  token_key: output_tokens
+  cost_key: output_cost
+
+- key: audio_input
+  kind: audio_token
+  direction: input
+  modality: audio
+  cache_state: none
+  unit: token
+  category: token
+  token_key: audio_input_tokens
+  cost_key: audio_input_cost
+
+- key: audio_output
+  kind: audio_token
+  direction: output
+  modality: audio
+  cache_state: none
+  unit: token
+  category: token
+  token_key: audio_output_tokens
+  cost_key: audio_output_cost
+
+- key: web_search_request
+  kind: web_search_request
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool
+
+- key: file_search_call
+  kind: file_search_call
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool
+
+- key: container_session
+  kind: container_session
+  direction: neither
+  modality: none
+  cache_state: none
+  unit: session
+  category: runtime
+
+- key: code_execution_request
+  kind: code_execution_request
+  direction: neither
+  modality: none
+  cache_state: none
+  unit: request
+  category: runtime
+
+- key: code_execution_hour
+  kind: code_execution_hour
+  direction: neither
+  modality: none
+  cache_state: none
+  unit: hour
+  category: runtime
+
+- key: grounding_request
+  kind: grounding_request
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  category: tool

data/lib/llm_cost_tracker/billing/cost_status.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require_relative "components"
+
+module LlmCostTracker
+  module Billing
+    module CostStatus
+      COMPLETE = "complete"
+      FREE = "free"
+      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
+
+          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?
+
+            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
+
+          total_cost.nil? || total_cost.zero? ? FREE : COMPLETE
+        end
+        # rubocop:enable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/billing/line_item.rb

@@ -0,0 +1,189 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "components"
+require_relative "cost_status"
+
+module LlmCostTracker
+  module Billing
+    LineItem = Data.define(
+      :kind,
+      :direction,
+      :modality,
+      :cache_state,
+      :quantity,
+      :unit,
+      :rate_amount,
+      :rate_quantity,
+      :cost,
+      :currency,
+      :cost_status,
+      :pricing_basis,
+      :price_key,
+      :price_source,
+      :price_source_version,
+      :provider_field,
+      :provider_item_id,
+      :details
+    )
+
+    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 = {
+          kind: symbol_or_nil(attributes[:kind]) || component&.kind,
+          direction: symbol_or_nil(attributes[:direction]) || component&.direction,
+          modality: symbol_or_nil(attributes[:modality]) || component&.modality,
+          cache_state: symbol_or_nil(attributes[:cache_state]) || component&.cache_state,
+          quantity: decimal_or_zero(attributes[:quantity]),
+          unit: symbol_or_nil(attributes[:unit]) || component&.unit,
+          rate_amount: decimal_or_nil(attributes[:rate_amount]),
+          rate_quantity: decimal_or_nil(attributes[:rate_quantity]) || BigDecimal("1"),
+          cost: decimal_or_nil(attributes[:cost]),
+          currency: attributes[:currency] || USD,
+          cost_status: cost_status_for(attributes),
+          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)
+          next unless quantity.positive?
+
+          new(
+            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: {}
+          )
+        end
+      end
+
+      def self.cost_status_for(attributes)
+        explicit = attributes[:cost_status]
+        return explicit.to_s if explicit
+
+        cost = decimal_or_nil(attributes[:cost])
+        return CostStatus::UNKNOWN if cost.nil?
+
+        cost.zero? ? CostStatus::FREE : CostStatus::COMPLETE
+      end
+
+      def self.component_for(attributes)
+        component_key = attributes[:component_key] || attributes[:price_key]
+        return nil unless component_key
+
+        Components::BY_KEY[component_key.to_sym]
+      end
+
+      def self.symbol_or_nil(value)
+        return nil if value.nil?
+
+        value.is_a?(Symbol) ? value : value.to_s.to_sym
+      end
+
+      def self.decimal_or_nil(value)
+        return nil if value.nil? || value == ""
+
+        BigDecimal(value.to_s)
+      end
+
+      def self.decimal_or_zero(value)
+        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
+
+      def billable?
+        quantity.positive?
+      end
+
+      def priced?
+        [CostStatus::COMPLETE, CostStatus::FREE].include?(cost_status)
+      end
+
+      def unpriced?
+        cost_status == CostStatus::UNKNOWN
+      end
+
+      def token?
+        unit == :token
+      end
+
+      def cost_value
+        cost || BigDecimal("0")
+      end
+
+      def apply_rate(rate)
+        rate_amount = rate.fetch(:amount)
+        rate_quantity = rate.fetch(:quantity)
+        applied_cost = (quantity / rate_quantity) * rate_amount
+        with(
+          rate_amount: rate_amount,
+          rate_quantity: rate_quantity,
+          cost: applied_cost,
+          currency: rate.fetch(:currency),
+          cost_status: applied_cost.zero? ? CostStatus::FREE : CostStatus::COMPLETE,
+          price_key: rate.fetch(:source_key),
+          price_source: rate.fetch(:source),
+          price_source_version: rate.fetch(:source_version)
+        )
+      end
+
+      def to_h
+        super.transform_values do |value|
+          value.is_a?(BigDecimal) ? value.to_s("F") : value
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/budget.rb

@@ -5,20 +5,22 @@ require_relative "ledger"
 
 module LlmCostTracker
   class Budget
+    BUDGET_TYPE_TO_PERIOD = { monthly: :month, daily: :day }.freeze
+
     class << self
       def enforce!
         config = LlmCostTracker.configuration
         return unless config.budget_exceeded_behavior == :block_requests
 
-        budgets = enforce_period_budgets(config)
+        budgets = { monthly: config.monthly_budget, daily: config.daily_budget }.compact
         return if budgets.empty?
 
-        totals = LlmCostTracker::Ledger::Period::Totals.call(budgets.keys, time: Time.now.utc)
+        totals = totals_for(budgets.keys, time: Time.now.utc)
 
-        budgets.each do |period, budget|
-          total = totals.fetch(period)
+        budgets.each do |budget_type, budget|
+          total = totals.fetch(budget_type)
 
-          handle_exceeded(budget_type: period, total: total, budget: budget) if total >= budget
+          handle_exceeded(budget_type: budget_type, total: total, budget: budget) if total >= budget
         end
       end
 
@@ -27,13 +29,13 @@ module LlmCostTracker
         return unless event.total_cost
 
         check_per_call_budget(event, config)
-        budgets = check_period_budgets(config)
-        totals = totals_for_check(event, budgets)
+        budgets = { daily: config.daily_budget, monthly: config.monthly_budget }.compact
+        totals = totals_for(budgets.keys, time: event.tracked_at)
 
-        budgets.each do |period, budget|
-          total = totals.fetch(period)
+        budgets.each do |budget_type, budget|
+          total = totals.fetch(budget_type)
 
-          handle_exceeded(budget_type: period, total: total, budget: budget, last_event: event) if total >= budget
+          handle_exceeded(budget_type: budget_type, total: total, budget: budget, last_event: event) if total >= budget
         end
       end
 
@@ -43,30 +45,20 @@ module LlmCostTracker
         budget = config.per_call_budget
         return unless budget
 
-        call_cost = event.total_cost
-        return unless call_cost >= budget
+        total = event.total_cost
+        return unless total >= budget
 
-        handle_exceeded(budget_type: :per_call, total: call_cost, budget: budget, last_event: event)
+        handle_exceeded(budget_type: :per_call, total: total, budget: budget, last_event: event)
       end
 
-      def enforce_period_budgets(config)
-        {
-          monthly: config.monthly_budget,
-          daily: config.daily_budget
-        }.compact
-      end
+      def totals_for(budget_types, time:)
+        return {} if budget_types.empty?
 
-      def check_period_budgets(config)
-        {
-          daily: config.daily_budget,
-          monthly: config.monthly_budget
-        }.compact
-      end
-
-      def totals_for_check(event, budgets)
-        return {} if budgets.empty?
-
-        LlmCostTracker::Ledger::Period::Totals.call(budgets.keys, time: event.tracked_at)
+        periods = budget_types.map { |type| BUDGET_TYPE_TO_PERIOD.fetch(type) }
+        period_totals = LlmCostTracker::Ledger::Period::Totals.call(periods, time: time)
+        BUDGET_TYPE_TO_PERIOD.each_with_object({}) do |(budget_type, period), totals|
+          totals[budget_type] = period_totals[period] if period_totals.key?(period)
+        end
       end
 
       def handle_exceeded(budget_type:, total:, budget:, last_event: nil)
@@ -85,16 +77,12 @@ module LlmCostTracker
       end
 
       def budget_payload(budget_type:, total:, budget:, last_event:)
-        payload = {
+        {
           budget_type: budget_type,
           total: total,
           budget: budget,
           last_event: last_event
         }
-        payload[:monthly_total] = total if budget_type == :monthly
-        payload[:daily_total] = total if budget_type == :daily
-        payload[:call_cost] = total if budget_type == :per_call
-        payload
       end
 
       def notify_exceeded?(config, budget_type:, total:, budget:, last_event:)

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -2,20 +2,27 @@
 
 require "active_support/core_ext/object/blank"
 require "active_support/core_ext/object/deep_dup"
+require "json"
 
 require_relative "stream"
+require_relative "../timing"
 
 module LlmCostTracker
   module Capture
     class StreamCollector
       attr_reader :provider
 
-      def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, pricing_mode: nil, metadata: {},
-                     context_tags: nil)
+      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)
         @provider = provider.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
+        @batch = batch
         @pricing_mode = pricing_mode
         @metadata = (metadata || {}).deep_dup
         @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
@@ -23,7 +30,7 @@ module LlmCostTracker
         @captured_bytes = 0
         @overflowed = false
         @explicit_usage = nil
-        @started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        @started_at = LlmCostTracker::Timing.now_monotonic
         @finished = false
         @mutex = Mutex.new
       end
@@ -66,10 +73,16 @@ module LlmCostTracker
         @mutex.synchronize do
           ensure_open!
           @provider_response_id = extra.delete(:provider_response_id) || @provider_response_id
-          @explicit_usage = TokenUsage.from_hash(extra.merge(
-                                                   input_tokens: input_tokens.to_i,
-                                                   output_tokens: output_tokens.to_i
-                                                 ))
+          @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
+          batch = extra.delete(:batch)
+          @batch = batch unless batch.nil?
+          @explicit_usage = TokenUsage.build(
+            **extra.slice(*TokenUsage.members),
+            input_tokens: input_tokens,
+            output_tokens: output_tokens
+          )
         end
         self
       end
@@ -79,6 +92,7 @@ module LlmCostTracker
           return if @finished
 
           @finished = true
+          pricing_mode = Pricing.normalize_mode(@pricing_mode)
           {
             events: @events.dup,
             overflowed: @overflowed,
@@ -86,7 +100,8 @@ module LlmCostTracker
             model: @model,
             latency_ms: @latency_ms,
             provider_response_id: @provider_response_id,
-            pricing_mode: @pricing_mode,
+            capture_dimensions: capture_dimensions(pricing_mode),
+            pricing_mode: pricing_mode,
             metadata: @metadata.deep_dup,
             context_tags: @context_tags.deep_dup
           }
@@ -98,8 +113,7 @@ module LlmCostTracker
 
         Tracker.record(
           capture: capture,
-          latency_ms: snapshot[:latency_ms] ||
-            ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at) * 1000).round,
+          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]
@@ -108,6 +122,16 @@ module LlmCostTracker
 
       private
 
+      def capture_dimensions(pricing_mode)
+        batch = @batch.nil? ? UsageCapture.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,
+          provider_workspace_id: @provider_workspace_id.to_s.strip.presence,
+          batch: batch
+        }.compact
+      end
+
       def ensure_open!
         return unless @finished
 
@@ -116,15 +140,14 @@ module LlmCostTracker
 
       def build_usage_capture(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(
           response_status: 200,
           events: snapshot[:events]
         )
-        if capture
+        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)
+          return capture.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
         end
 
         build_unknown_usage(snapshot)
@@ -145,7 +168,9 @@ module LlmCostTracker
           model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
           token_usage: snapshot[:explicit_usage],
           stream: true,
-          usage_source: :manual
+          usage_source: :manual,
+          pricing_mode: snapshot[:pricing_mode],
+          **snapshot.fetch(:capture_dimensions)
         )
       end
 
@@ -155,34 +180,23 @@ module LlmCostTracker
           model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
-          usage_source: :unknown
+          usage_source: :unknown,
+          pricing_mode: snapshot[:pricing_mode],
+          **snapshot.fetch(:capture_dimensions)
         )
       end
 
       def capture_event(data, type:)
-        size = type.to_s.bytesize + estimated_bytes(data) + 32
+        event = { event: type, data: data }
+        size = JSON.generate(event).bytesize
         if @captured_bytes + size <= Capture::Stream::LIMIT_BYTES
-          @events << { event: type, data: data.deep_dup }
+          @events << event.deep_dup
           @captured_bytes += size
         else
           @overflowed = true
-          @events.clear
-        end
-      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
+      rescue JSON::JSONError, TypeError
+        @overflowed = true
       end
     end
   end