llm_cost_tracker 0.7.3 → 0.9.0

data/lib/llm_cost_tracker/report/data.rb

@@ -2,6 +2,7 @@
 
 require "active_support/core_ext/integer/time"
 
+require_relative "../billing/cost_status"
 require_relative "../ledger"
 
 module LlmCostTracker
@@ -32,17 +33,18 @@ module LlmCostTracker
           breakdown_limit = nil unless breakdown_limit.positive?
         end
         from = now - days.days
-        scope = Ledger::Call.where(tracked_at: from..now)
+        scope = LlmCostTracker::Call.where(tracked_at: from..now)
         tag_breakdowns ||= LlmCostTracker.configuration.report_tag_breakdowns || []
+        aggregate = totals(scope)
 
         new(
           days: days,
           from_time: from,
           to_time: now,
-          total_cost: scope.sum(:total_cost).to_f,
-          requests_count: scope.count,
-          average_latency_ms: average_latency_ms(scope),
-          unknown_pricing_count: scope.where(total_cost: nil).count,
+          total_cost: aggregate.total_cost.to_f,
+          requests_count: aggregate.requests_count.to_i,
+          average_latency_ms: aggregate.average_latency_ms&.to_f,
+          unknown_pricing_count: aggregate.unknown_pricing_count.to_i,
           cost_by_provider: scope.cost_by_provider(limit: breakdown_limit).to_a,
           cost_by_model: scope.cost_by_model(limit: breakdown_limit).to_a,
           cost_by_tags: cost_by_tags(scope, tag_breakdowns, limit: breakdown_limit),
@@ -50,8 +52,17 @@ module LlmCostTracker
         )
       end
 
-      def self.average_latency_ms(scope)
-        scope.average(:latency_ms)&.to_f
+      def self.totals(scope)
+        scope
+          .select(
+            "COALESCE(SUM(total_cost), 0) AS total_cost, " \
+            "COUNT(*) AS requests_count, " \
+            "AVG(latency_ms) AS average_latency_ms, " \
+            "COALESCE(SUM(CASE WHEN total_cost IS NULL " \
+            "OR cost_status IN ('#{Billing::CostStatus::UNKNOWN}', '#{Billing::CostStatus::PARTIAL}') " \
+            "THEN 1 ELSE 0 END), 0) AS unknown_pricing_count"
+          )
+          .take
       end
 
       def self.cost_by_tags(scope, keys, limit:)
@@ -66,7 +77,7 @@ module LlmCostTracker
           .to_a
       end
 
-      private_class_method :average_latency_ms, :cost_by_tags, :top_calls
+      private_class_method :cost_by_tags, :top_calls, :totals
     end
   end
 end

data/lib/llm_cost_tracker/report.rb

@@ -20,10 +20,6 @@ module LlmCostTracker
       rescue StandardError => e
         "Unable to build LLM cost report: #{e.class}: #{e.message}"
       end
-
-      def data(days: Data::DEFAULT_DAYS, now: Time.now.utc, tag_breakdowns: nil)
-        Data.build(days: days, now: now, tag_breakdowns: tag_breakdowns)
-      end
     end
   end
 end

data/lib/llm_cost_tracker/retention.rb

@@ -21,6 +21,17 @@ module LlmCostTracker
         deleted
       end
 
+      def prune_invoice_imports(older_than:, now: Time.now.utc)
+        cutoff = resolve_cutoff(older_than, now)
+        require_relative "ledger"
+        return 0 unless LlmCostTracker::ProviderInvoiceImport.table_exists?
+
+        LlmCostTracker::ProviderInvoiceImport
+          .where(state: %w[completed failed])
+          .where(finished_at: ...cutoff)
+          .delete_all
+      end
+
       private
 
       def resolve_cutoff(older_than, now)
@@ -49,20 +60,22 @@ module LlmCostTracker
       end
 
       def prune_batch(cutoff, batch_size)
-        LlmCostTracker::Ledger::Call.transaction do
-          rows = LlmCostTracker::Ledger::Call
-                 .where(tracked_at: ...cutoff)
-                 .order(:id)
-                 .limit(batch_size)
-                 .lock
-                 .pluck(:id, :tracked_at, :total_cost)
+        LlmCostTracker::Call.transaction do
+          rows = pluck_prunable(cutoff, batch_size)
           next 0 if rows.empty?
 
-          deleted = LlmCostTracker::Ledger::Call.where(id: rows.map(&:first)).delete_all
-          LlmCostTracker::Ledger::Rollups.decrement!(rows) if deleted.positive?
+          deleted = LlmCostTracker::Call.where(id: rows.map(&:first)).delete_all
+          if deleted.positive? && LlmCostTracker.configuration.cache_rollups
+            LlmCostTracker::Ledger::Rollups.decrement!(rows)
+          end
           deleted
         end
       end
+
+      def pluck_prunable(cutoff, batch_size)
+        LlmCostTracker::Call.where(tracked_at: ...cutoff).order(:id).limit(batch_size).lock
+                            .pluck(:id, :tracked_at, :total_cost, :pricing_snapshot, :provider)
+      end
     end
   end
 end

data/lib/llm_cost_tracker/tags/context.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/object/deep_dup"
 require "active_support/isolated_execution_state"
 
 module LlmCostTracker
@@ -11,7 +10,7 @@ module LlmCostTracker
       class << self
         def with(tags)
           stack = ActiveSupport::IsolatedExecutionState[KEY] || []
-          ActiveSupport::IsolatedExecutionState[KEY] = stack + [(tags || {}).deep_dup.to_h]
+          ActiveSupport::IsolatedExecutionState[KEY] = stack + [Sanitizer.call((tags || {}).to_h)]
           yield
         ensure
           ActiveSupport::IsolatedExecutionState[KEY] = stack
@@ -21,9 +20,7 @@ module LlmCostTracker
           default_tags = LlmCostTracker.configuration.default_tags
           default_tags = default_tags.call if default_tags.respond_to?(:call)
 
-          (default_tags || {}).deep_dup.to_h.merge(
-            (ActiveSupport::IsolatedExecutionState[KEY] || []).reduce({}) { |merged, tags| merged.merge(tags) }
-          )
+          Sanitizer.call(default_tags.to_h).merge(*Array(ActiveSupport::IsolatedExecutionState[KEY]))
         end
 
         def clear!

data/lib/llm_cost_tracker/tags/key.rb

@@ -4,10 +4,14 @@ module LlmCostTracker
   module Tags
     module Key
       PATTERN = /\A[\w.-]+\z/
+      MAX_BYTESIZE = 64
 
       class << self
         def validate!(key, error_class: ArgumentError)
           key = key.to_s
+          if key.bytesize > MAX_BYTESIZE
+            raise error_class, "tag key exceeds #{MAX_BYTESIZE} bytes: #{key[0, 16].inspect}..."
+          end
           return key if key.match?(PATTERN)
 
           raise error_class, "invalid tag key: #{key.inspect}"

data/lib/llm_cost_tracker/tags/sanitizer.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/string/inflections"
 require "json"
 
 module LlmCostTracker
@@ -7,42 +8,92 @@ module LlmCostTracker
     module Sanitizer
       REDACTED_VALUE = "[REDACTED]"
 
+      SECRET_VALUE_PATTERNS = [
+        /\Ask-(?:ant-|admin-|proj-|svcacct-|live-|test-)?[A-Za-z0-9_-]{16,}\z/,
+        /\AAKIA[0-9A-Z]{16}\z/,
+        /\Agh[opsur]_[A-Za-z0-9]{16,}\z/,
+        /\Agithub_pat_[A-Za-z0-9_]{20,}\z/,
+        /\Aeyj[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\z/i,
+        /\Abearer\s+[A-Za-z0-9_.-]{20,}\z/i,
+        /\Axox[abprs]-[A-Za-z0-9-]{10,}\z/,
+        /\A(?:sk|rk|pk)_(?:live|test)_[A-Za-z0-9]{20,}\z/,
+        /\AAIza[0-9A-Za-z_-]{35}\z/
+      ].freeze
+      private_constant :SECRET_VALUE_PATTERNS
+
       class << self
         def call(tags, config: LlmCostTracker.configuration)
           tags = (tags || {}).to_h
-          tags.first([config.max_tag_count.to_i, 0].max).each_with_object({}) do |(key, value), sanitized|
-            sanitized[key] = sanitized_value(key, value, config)
+          redacted = Array(config.redacted_tag_keys).map { |key| normalized_key(key) }
+          limit = [config.max_tag_value_bytesize.to_i, 0].max
+          max_count = [config.max_tag_count.to_i, 0].max
+          tags.to_a.last(max_count).each_with_object({}) do |(key, value), sanitized|
+            sanitized[key] = sanitized_value(key, value, redacted, limit)
           end
         end
 
         private
 
-        def sanitized_value(key, value, config)
-          return REDACTED_VALUE if redacted_key?(key, config)
+        def sanitized_value(key, value, redacted, limit)
+          return REDACTED_VALUE if redacted_key?(key, redacted)
 
-          string = value_string(value)
-          limit = [config.max_tag_value_bytesize.to_i, 0].max
-          return value if string.bytesize <= limit
+          scrubbed = scrub_secrets(value)
+          return REDACTED_VALUE if scrubbed.equal?(REDACTED_SENTINEL)
 
-          string.byteslice(0, limit).to_s.encode("UTF-8", invalid: :replace, undef: :replace)
+          scalar_truncate(scrubbed, limit)
         end
 
-        def redacted_key?(key, config)
-          normalized = normalized_key(key)
-          Array(config.redacted_tag_keys).map { |redacted_key| normalized_key(redacted_key) }.any? do |candidate|
-            redacted_key_component?(normalized, candidate)
+        REDACTED_SENTINEL = Object.new.freeze
+        private_constant :REDACTED_SENTINEL
+
+        def scalar_truncate(value, limit)
+          case value
+          when Hash
+            value.transform_values { |nested| scalar_truncate(nested, limit) }
+          when Array
+            value.map { |nested| scalar_truncate(nested, limit) }
+          else
+            return value if value == REDACTED_VALUE
+
+            string = value.to_s
+            return value if string.bytesize <= limit
+
+            string.byteslice(0, limit).encode("UTF-8", invalid: :replace, undef: :replace)
           end
         end
 
+        def scrub_secrets(value)
+          case value
+          when Hash
+            value.each_with_object({}) do |(key, nested), out|
+              scrubbed = scrub_secrets(nested)
+              out[key] = scrubbed.equal?(REDACTED_SENTINEL) ? REDACTED_VALUE : scrubbed
+            end
+          when Array
+            value.map do |nested|
+              scrubbed = scrub_secrets(nested)
+              scrubbed.equal?(REDACTED_SENTINEL) ? REDACTED_VALUE : scrubbed
+            end
+          else
+            secret_shaped?(value.to_s) ? REDACTED_SENTINEL : value
+          end
+        end
+
+        def secret_shaped?(string)
+          return false if string.bytesize < 16
+
+          SECRET_VALUE_PATTERNS.any? { |pattern| pattern.match?(string) }
+        end
+
+        def redacted_key?(key, redacted)
+          return false if redacted.empty?
+
+          normalized = normalized_key(key)
+          redacted.any? { |candidate| redacted_key_component?(normalized, candidate) }
+        end
+
         def normalized_key(key)
-          key.to_s
-             .gsub(/([A-Z]+)([A-Z][a-z])/, '\1_\2')
-             .gsub(/([a-z\d])([A-Z])/, '\1_\2')
-             .downcase
-             .gsub(/[^a-z0-9]+/, "_")
-             .gsub(/_+/, "_")
-             .delete_prefix("_")
-             .delete_suffix("_")
+          key.to_s.underscore.gsub(/[^a-z0-9]+/, "_").delete_prefix("_").delete_suffix("_")
         end
 
         def redacted_key_component?(key, candidate)

data/lib/llm_cost_tracker/timing.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Timing
+    module_function
+
+    def now_monotonic
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
+    def elapsed_ms(started_at)
+      ((now_monotonic - started_at) * 1000).round
+    end
+  end
+end

data/lib/llm_cost_tracker/token_usage.rb

@@ -1,67 +1,89 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/hash/keys"
+require_relative "billing/components"
+require_relative "logging"
 
 module LlmCostTracker
+  KNOWN_TOKEN_KEYS = (
+    Billing::Components::TOKEN_PRICED.map(&:key) + %i[total hidden_output]
+  ).freeze
+
   TokenUsage = Data.define(
     :input_tokens,
     :cache_read_input_tokens,
     :cache_write_input_tokens,
-    :cache_write_1h_input_tokens,
+    :cache_write_extended_input_tokens,
+    :audio_input_tokens,
+    :image_input_tokens,
     :output_tokens,
+    :audio_output_tokens,
+    :image_output_tokens,
     :total_tokens,
     :hidden_output_tokens
   ) do
-    def self.build(input_tokens:, output_tokens:, cache_read_input_tokens: 0,
-                   cache_write_input_tokens: 0, cache_write_1h_input_tokens: 0,
-                   total_tokens: nil, hidden_output_tokens: 0)
-      input = input_tokens.to_i
-      output = output_tokens.to_i
-      cache_read = cache_read_input_tokens.to_i
-      cache_write = cache_write_input_tokens.to_i
-      cache_write_1h = cache_write_1h_input_tokens.to_i
-      calculated_total = input + cache_read + cache_write + cache_write_1h + output
-      total = total_tokens.nil? ? calculated_total : [total_tokens.to_i, calculated_total].max
+    def self.build_from_tokens(tokens)
+      return tokens if tokens.is_a?(self)
+      raise ArgumentError, "tokens must be a Hash, got #{tokens.class}" unless tokens.respond_to?(:to_h)
 
-      new(
-        input_tokens: input,
-        cache_read_input_tokens: cache_read,
-        cache_write_input_tokens: cache_write,
-        cache_write_1h_input_tokens: cache_write_1h,
-        output_tokens: output,
-        total_tokens: total,
-        hidden_output_tokens: hidden_output_tokens.to_i
-      )
-    end
+      values = tokens.to_h.transform_keys { |key| key.to_s.to_sym }
+      warn_on_unknown_keys(values)
+      token_attributes = Billing::Components::TOKEN_PRICED.to_h do |component|
+        [component.token_key, values.fetch(component.key, 0)]
+      end
 
-    def self.from_hash(attributes)
-      attributes = attributes.to_h.symbolize_keys
-      values = TokenUsage::COMPONENT_TOKEN_KEYS.to_h { |key| [key, attributes[key]] }
       build(
-        **values,
-        total_tokens: attributes[:total_tokens]
+        **token_attributes,
+        total_tokens: values[:total],
+        hidden_output_tokens: values.fetch(:hidden_output, 0)
       )
     end
 
-    def price_quantities
-      {
-        input: input_tokens,
-        cache_read_input: cache_read_input_tokens,
-        cache_write_input: cache_write_input_tokens,
-        cache_write_1h_input: cache_write_1h_input_tokens,
-        output: output_tokens
-      }
+    def self.warn_on_unknown_keys(values)
+      return if values.empty?
+      return if values.keys.intersect?(KNOWN_TOKEN_KEYS)
+
+      Logging.warn(
+        "tokens hash contains no recognized keys (#{values.keys.inspect}); " \
+        "expected one of #{KNOWN_TOKEN_KEYS.inspect}. Did you pass a raw provider response?"
+      )
     end
 
-    def stored_attributes
-      to_h.slice(*self.class::STORED_KEYS)
+    def self.non_negative_int(value)
+      [value.to_i, 0].max
     end
 
-    def to_h
-      super.compact
+    def self.build(input_tokens:, output_tokens:, cache_read_input_tokens: 0,
+                   cache_write_input_tokens: 0, cache_write_extended_input_tokens: 0,
+                   audio_input_tokens: 0, audio_output_tokens: 0,
+                   image_input_tokens: 0, image_output_tokens: 0,
+                   total_tokens: nil, hidden_output_tokens: 0)
+      input = non_negative_int(input_tokens)
+      output = non_negative_int(output_tokens)
+      cache_read = non_negative_int(cache_read_input_tokens)
+      cache_write = non_negative_int(cache_write_input_tokens)
+      cache_write_extended = non_negative_int(cache_write_extended_input_tokens)
+      audio_input = non_negative_int(audio_input_tokens)
+      audio_output = non_negative_int(audio_output_tokens)
+      image_input = non_negative_int(image_input_tokens)
+      image_output = non_negative_int(image_output_tokens)
+      hidden_output = non_negative_int(hidden_output_tokens)
+      calculated_total = input + cache_read + cache_write + cache_write_extended +
+                         audio_input + image_input + output + audio_output + image_output
+      total = total_tokens ? [non_negative_int(total_tokens), calculated_total].max : calculated_total
+
+      new(
+        input_tokens: input,
+        cache_read_input_tokens: cache_read,
+        cache_write_input_tokens: cache_write,
+        cache_write_extended_input_tokens: cache_write_extended,
+        audio_input_tokens: audio_input,
+        image_input_tokens: image_input,
+        output_tokens: output,
+        audio_output_tokens: audio_output,
+        image_output_tokens: image_output,
+        total_tokens: total,
+        hidden_output_tokens: hidden_output
+      )
     end
   end
-
-  TokenUsage::STORED_KEYS = TokenUsage.members.freeze
-  TokenUsage::COMPONENT_TOKEN_KEYS = (TokenUsage.members - %i[total_tokens]).freeze
 end

data/lib/llm_cost_tracker/tracker.rb

@@ -1,19 +1,19 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/object/blank"
+require "bigdecimal"
 require "securerandom"
 
 require_relative "ingestion"
 require_relative "ledger"
+require_relative "logging"
 require_relative "pricing"
+require_relative "billing/cost_status"
 
 module LlmCostTracker
   class Tracker
     EVENT_NAME = "llm_request.llm_cost_tracker"
 
-    USAGE_SOURCES = %i[response stream_final sdk_response ruby_llm manual unknown].freeze
-    TRACKING_METADATA_KEYS = (TokenUsage.members.map(&:to_s) + %w[pricing_mode provider_response_id]).freeze
-
     class << self
       def enforce_budget!
         return unless LlmCostTracker.configuration.enabled
@@ -25,43 +25,73 @@ module LlmCostTracker
         return unless LlmCostTracker.configuration.enabled
 
         pricing_mode = Pricing.normalize_mode(pricing_mode) || capture.pricing_mode
-        cost_data = Pricing.cost_for(
+        cost_data, pricing_snapshot, priced_line_items = Pricing.calculate(
           provider: capture.provider,
           model: capture.model,
-          token_usage: capture.token_usage,
+          tokens: capture.token_usage,
+          line_items: capture.line_items,
           pricing_mode: pricing_mode
         )
 
-        Pricing::Unknown.handle!(capture.model) unless cost_data
+        Pricing::Unknown.handle!(capture.model) if cost_data.nil? && capture.token_usage.total_tokens.positive?
 
         event = build_event(
           capture: capture,
           pricing_mode: pricing_mode,
           cost_data: cost_data,
+          pricing_snapshot: pricing_snapshot,
+          line_items: priced_line_items,
           metadata: metadata,
           latency_ms: latency_ms,
           context_tags: context_tags
         )
 
-        ActiveSupport::Notifications.instrument(EVENT_NAME, event.to_h)
-
-        Ingestion::Inbox.save(event)
+        save_event(event)
+        yield :after_save if block_given?
+        notify_subscribers(event)
         Budget.check!(event)
 
         event
       end
 
+      def save_event(event)
+        if LlmCostTracker.configuration.durable_ingestion
+          Ingestion::Inbox.save(event)
+        else
+          Ingestion::Inline.save(event)
+        end
+      end
+
+      def notify_subscribers(event)
+        return unless ActiveSupport::Notifications.notifier.listening?(EVENT_NAME)
+
+        ActiveSupport::Notifications.instrument(EVENT_NAME, event.to_h)
+      rescue StandardError => e
+        Logging.warn("Subscriber raised on #{EVENT_NAME}: #{e.class}: #{e.message}")
+      end
+
       private
 
-      def build_event(capture:, pricing_mode:, cost_data:, metadata:, latency_ms:, context_tags:)
-        usage_source = if capture.usage_source.nil?
-                         nil
-                       else
-                         symbol = capture.usage_source.to_sym
-                         USAGE_SOURCES.include?(symbol) ? symbol.to_s : nil
-                       end
-        tags = metadata.to_h.reject { |key, _value| TRACKING_METADATA_KEYS.include?(key.to_s) }
-        context_tags = context_tags.nil? ? LlmCostTracker::Tags::Context.tags : context_tags.to_h
+      def token_pricing_partial?(token_usage:, cost_data:)
+        return false unless cost_data
+
+        Billing::Components::TOKEN_PRICED.any? do |component|
+          token_usage.public_send(component.token_key).positive? && cost_data[component.cost_key].nil?
+        end
+      end
+
+      def build_event(capture:, pricing_mode:, cost_data:, pricing_snapshot:, line_items:,
+                      metadata:, latency_ms:, context_tags:)
+        context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).to_h
+        cost = cost_with_service_lines(cost_data, line_items)
+        cost_status = Billing::CostStatus.call(
+          token_usage: capture.token_usage,
+          usage_source: capture.usage_source,
+          token_cost: cost_data,
+          token_pricing_partial: token_pricing_partial?(token_usage: capture.token_usage, cost_data: cost_data),
+          service_line_items: line_items.reject(&:token?),
+          total_cost: cost&.fetch(:total_cost, nil)
+        )
 
         Event.new(
           event_id: SecureRandom.uuid,
@@ -69,15 +99,55 @@ module LlmCostTracker
           model: capture.model,
           token_usage: capture.token_usage,
           pricing_mode: pricing_mode,
-          cost: cost_data,
-          tags: LlmCostTracker::Tags::Sanitizer.call(
-            context_tags.merge(tags)
-          ).freeze,
-          latency_ms: latency_ms.nil? ? nil : [latency_ms.to_i, 0].max,
-          stream: capture.stream ? true : false,
-          usage_source: usage_source,
-          provider_response_id: capture.provider_response_id.to_s.presence,
-          tracked_at: Time.now.utc
+          cost: cost,
+          tags: LlmCostTracker::Tags::Sanitizer.call(context_tags.merge(metadata.to_h)).freeze,
+          latency_ms: finite_latency_ms(latency_ms),
+          stream: capture.stream,
+          usage_source: capture.usage_source,
+          provider_response_id: capture.provider_response_id,
+          provider_project_id: capture.provider_project_id,
+          provider_api_key_id: capture.provider_api_key_id,
+          provider_workspace_id: capture.provider_workspace_id,
+          batch: capture.batch,
+          tracked_at: Time.now.utc,
+          cost_status: cost_status,
+          pricing_snapshot: pricing_snapshot,
+          line_items: line_items
+        )
+      end
+
+      def finite_latency_ms(latency_ms)
+        return nil if latency_ms.nil?
+
+        Integer(latency_ms).clamp(0, (1 << 31) - 1)
+      rescue ArgumentError, TypeError, FloatDomainError
+        nil
+      end
+
+      def cost_with_service_lines(cost_data, line_items)
+        priced_services = line_items.reject(&:token?).select(&:priced?)
+        return cost_data if priced_services.empty?
+
+        base_currency = (cost_data && cost_data[:currency]) || Billing::LineItem::USD
+        matching, mismatched = priced_services.partition { |line| line.currency.to_s == base_currency.to_s }
+        warn_currency_mismatch(mismatched, base_currency) if mismatched.any?
+
+        cost = cost_data ? cost_data.dup : {}
+        cost[:currency] ||= base_currency.to_s
+        return cost if matching.empty?
+
+        service_total = matching.sum(BigDecimal("0"), &:cost_value)
+        base_total = BigDecimal(cost.fetch(:total_cost, 0).to_s)
+        cost[:total_cost] = (base_total + service_total).round(8)
+        cost
+      end
+
+      def warn_currency_mismatch(lines, base_currency)
+        currencies = lines.map { |line| line.currency.to_s }.uniq.sort
+        Logging.warn(
+          "Service line currency mismatch: header is #{base_currency}, dropping " \
+          "#{lines.size} priced line(s) in #{currencies.join(', ')} from header total. " \
+          "Per-line costs are still recorded; header total reflects #{base_currency} only."
         )
       end
     end

data/lib/llm_cost_tracker/usage_capture.rb

@@ -3,6 +3,7 @@
 require "active_support/core_ext/object/blank"
 
 require_relative "pricing"
+require_relative "billing/line_item"
 
 module LlmCostTracker
   UsageCapture = Data.define(
@@ -12,26 +13,46 @@ module LlmCostTracker
     :stream,
     :usage_source,
     :provider_response_id,
-    :pricing_mode
+    :provider_project_id,
+    :provider_api_key_id,
+    :provider_workspace_id,
+    :batch,
+    :pricing_mode,
+    :line_items
   )
 
   class UsageCapture
     UNKNOWN_MODEL = "unknown"
 
+    def self.batch_from_pricing_mode?(pricing_mode)
+      pricing_mode.to_s.split("_").include?("batch")
+    end
+
     def self.build(**attributes)
+      pricing_mode = Pricing.normalize_mode(attributes[:pricing_mode])
+      batch = attributes[:batch]
+      batch = batch_from_pricing_mode?(pricing_mode) if batch.nil?
+
+      token_usage = attributes.fetch(:token_usage)
+      service_line_items = Array(attributes[:service_line_items]).map do |item|
+        item.is_a?(Billing::LineItem) ? item : Billing::LineItem.build(item)
+      end
+      line_items = attributes[:line_items] || (Billing::LineItem.from_token_usage(token_usage) + service_line_items)
+
       new(
         provider: attributes.fetch(:provider).to_s,
         model: attributes.fetch(:model).to_s.strip.presence || UNKNOWN_MODEL,
-        token_usage: attributes.fetch(:token_usage),
+        token_usage: token_usage,
         stream: attributes[:stream] || false,
         usage_source: attributes[:usage_source],
-        provider_response_id: attributes[:provider_response_id],
-        pricing_mode: Pricing.normalize_mode(attributes[:pricing_mode])
+        provider_response_id: attributes[:provider_response_id].to_s.strip.presence,
+        provider_project_id: attributes[:provider_project_id].to_s.strip.presence,
+        provider_api_key_id: attributes[:provider_api_key_id].to_s.strip.presence,
+        provider_workspace_id: attributes[:provider_workspace_id].to_s.strip.presence,
+        batch: batch,
+        pricing_mode: pricing_mode,
+        line_items: line_items
       )
     end
-
-    def to_h
-      super.compact
-    end
   end
 end