llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/token_usage.rb

@@ -1,67 +1,81 @@
 # 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,
     :output_tokens,
+    :audio_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,
+                   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)
+      hidden_output = non_negative_int(hidden_output_tokens)
+      calculated_total = input + cache_read + cache_write + cache_write_extended + audio_input + output + audio_output
+      total = total_tokens.nil? ? calculated_total : [non_negative_int(total_tokens), calculated_total].max
+
+      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,
+        output_tokens: output,
+        audio_output_tokens: audio_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,18 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/object/blank"
+require "bigdecimal"
 require "securerandom"
 
 require_relative "ingestion"
 require_relative "ledger"
 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,19 +24,20 @@ 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 = Pricing.cost_and_snapshot_for(
           provider: capture.provider,
           model: capture.model,
-          token_usage: capture.token_usage,
+          tokens: capture.token_usage,
           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,
           metadata: metadata,
           latency_ms: latency_ms,
           context_tags: context_tags
@@ -53,15 +53,32 @@ module LlmCostTracker
 
       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
+
+      # rubocop:disable Metrics/MethodLength
+      def build_event(capture:, pricing_mode:, cost_data:, pricing_snapshot:, metadata:, latency_ms:, context_tags:)
+        context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).to_h
+        line_items, = Pricing.price_line_items(
+          provider: capture.provider,
+          model: capture.model,
+          line_items: capture.line_items,
+          pricing_mode: pricing_mode
+        )
+        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,17 +86,43 @@ 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
+      # rubocop:enable Metrics/MethodLength
+
+      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)
+        service_lines = line_items.reject(&:token?)
+        return cost_data if service_lines.empty?
+        return cost_data if service_lines.none?(&:priced?)
+
+        service_total = service_lines.sum(BigDecimal("0"), &:cost_value)
+        cost = cost_data ? cost_data.dup : {}
+        base_total = BigDecimal(cost.fetch(:total_cost, 0).to_s)
+        cost[:total_cost] = (base_total + service_total).round(8).to_f
+        cost
+      end
     end
   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

data/lib/llm_cost_tracker/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LlmCostTracker
-  VERSION = "0.7.2"
+  VERSION = "0.8.0"
 end

data/lib/llm_cost_tracker.rb

@@ -17,6 +17,9 @@ require_relative "llm_cost_tracker/tags/key"
 require_relative "llm_cost_tracker/tags/context"
 require_relative "llm_cost_tracker/tags/sanitizer"
 require_relative "llm_cost_tracker/token_usage"
+require_relative "llm_cost_tracker/billing/components"
+require_relative "llm_cost_tracker/billing/line_item"
+require_relative "llm_cost_tracker/billing/cost_status"
 require_relative "llm_cost_tracker/event"
 require_relative "llm_cost_tracker/pricing"
 require_relative "llm_cost_tracker/usage_capture"
@@ -47,14 +50,19 @@ module LlmCostTracker
   class << self
     attr_reader :configuration
 
+    def table_name_prefix
+      "llm_cost_tracker_"
+    end
+
     def configure
       config = configuration
       raise Error, "LlmCostTracker is already configured" if config.finalized?
 
       yield(config)
-      config.openai_compatible_providers = config.openai_compatible_providers.dup
       config.finalize!
       Pricing::Lookup.reset!
+      Pricing::Registry.reset!
+      Pricing::ServiceCharges.reset!
       Integrations.install!
       config
     end
@@ -63,67 +71,60 @@ module LlmCostTracker
       Ingestion::Worker.shutdown!(drain: false)
       @configuration = Configuration.new
       Pricing::Lookup.reset!
+      Pricing::Registry.reset!
+      Pricing::ServiceCharges.reset!
       Pricing::Unknown.reset!
       Ingestion::Worker.reset!
       Tags::Context.clear!
     end
 
-    def flush!(timeout: nil)
-      if timeout
-        Ingestion::Worker.flush!(timeout: timeout)
-      else
-        Ingestion::Worker.flush!
-      end
-    end
-
-    def shutdown!(timeout: nil, drain: true)
-      if timeout
-        Ingestion::Worker.shutdown!(timeout: timeout, drain: drain)
-      else
-        Ingestion::Worker.shutdown!(drain: drain)
-      end
-    end
-
-    def enforce_budget!
-      Tracker.enforce_budget!
-    end
-
     def with_tags(tags = nil, **kwargs, &)
-      merged = (tags || {}).to_h.merge(kwargs)
-      Tags::Context.with(merged, &)
+      Tags::Context.with((tags || {}).merge(kwargs), &)
     end
 
-    def track(provider:, input_tokens:, output_tokens:, model: nil, latency_ms: nil, stream: false,
-              usage_source: :manual, enforce_budget: false, provider_response_id: nil, pricing_mode: nil, **metadata)
-      enforce_budget! if enforce_budget
-      token_usage = TokenUsage.from_hash(metadata.merge(input_tokens: input_tokens, output_tokens: output_tokens))
+    def track(provider:, tokens:, model: nil, tags: {}, latency_ms: nil, stream: false,
+              usage_source: :manual, enforce_budget: false,
+              provider_response_id: nil, provider_project_id: nil, provider_api_key_id: nil,
+              provider_workspace_id: nil, batch: nil, pricing_mode: nil, service_line_items: [])
+      Tracker.enforce_budget! if enforce_budget
 
       Tracker.record(
         capture: UsageCapture.build(
           provider: provider,
           model: model,
-          token_usage: token_usage,
+          token_usage: TokenUsage.build_from_tokens(tokens),
           stream: stream,
           usage_source: usage_source,
-          provider_response_id: provider_response_id
+          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,
+          service_line_items: service_line_items
         ),
         latency_ms: latency_ms,
         pricing_mode: pricing_mode,
-        metadata: metadata
+        metadata: tags
       )
     end
 
-    def track_stream(provider:, model: nil, latency_ms: nil, enforce_budget: false, provider_response_id: nil,
-                     pricing_mode: nil, **metadata)
+    def track_stream(provider:, model: nil, tags: {}, latency_ms: nil, enforce_budget: false,
+                     provider_response_id: nil, provider_project_id: nil, provider_api_key_id: nil,
+                     provider_workspace_id: nil, batch: nil, pricing_mode: nil)
       require_relative "llm_cost_tracker/capture/stream_collector"
-      enforce_budget! if enforce_budget
+      Tracker.enforce_budget! if enforce_budget
       collector = Capture::StreamCollector.new(
         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
+        metadata: tags
       )
       yield collector
       collector.finish!
@@ -140,4 +141,4 @@ Faraday::Middleware.register_middleware(
   llm_cost_tracker: LlmCostTracker::Middleware::Faraday
 )
 
-at_exit { LlmCostTracker.shutdown!(drain: false) }
+at_exit { LlmCostTracker::Ingestion::Worker.shutdown!(drain: false) }

data/lib/tasks/llm_cost_tracker.rake

@@ -1,9 +1,20 @@
 # frozen_string_literal: true
 
 require "fileutils"
+require "rails/generators"
+
+require_relative "../llm_cost_tracker/generators/llm_cost_tracker/install_generator"
+require_relative "../llm_cost_tracker/pricing/sync_change_printer"
 
 # rubocop:disable Metrics/BlockLength
 namespace :llm_cost_tracker do
+  desc "Install LLM Cost Tracker with dashboard and prices, migrate, and run doctor"
+  task :setup do
+    Rails::Generators.invoke("llm_cost_tracker:install", %w[--dashboard --prices])
+    Rake::Task["db:migrate"].invoke
+    Rake::Task["llm_cost_tracker:doctor"].invoke
+  end
+
   desc "Check LLM Cost Tracker setup"
   task :doctor do
     Rake::Task["environment"].invoke if Rake::Task.task_defined?("environment")
@@ -30,7 +41,7 @@ namespace :llm_cost_tracker do
     puts LlmCostTracker::Report.generate(days: days)
   end
 
-  desc "Delete llm_api_calls older than DAYS (default: 90). Use BATCH_SIZE=N to tune."
+  desc "Delete llm_cost_tracker_calls older than DAYS (default: 90). Use BATCH_SIZE=N to tune."
   task prune: :environment do
     days = (ENV["DAYS"] || 90).to_i
     batch_size = (ENV["BATCH_SIZE"] || LlmCostTracker::Retention::DEFAULT_BATCH_SIZE).to_i
@@ -102,15 +113,7 @@ end
 # rubocop:enable Metrics/BlockLength
 
 def print_changes(changes)
-  puts "  changed models: #{changes.size}"
-  return if changes.empty?
-
-  changes.each do |model, fields|
-    puts "    - #{model}"
-    fields.each do |field, values|
-      puts "      #{field}: #{values['from'].inspect} -> #{values['to'].inspect}"
-    end
-  end
+  LlmCostTracker::Pricing::SyncChangePrinter.call(changes)
 end
 
 def price_refresh_output_path
@@ -128,13 +131,15 @@ def price_explanation_from_env
     provider: provider,
     model: model,
     pricing_mode: ENV.fetch("PRICING_MODE", nil),
-    token_usage: LlmCostTracker::TokenUsage.build(
-      input_tokens: ENV.fetch("INPUT_TOKENS", 1).to_i,
-      output_tokens: ENV.fetch("OUTPUT_TOKENS", 1).to_i,
-      cache_read_input_tokens: ENV.fetch("CACHE_READ_INPUT_TOKENS", 0).to_i,
-      cache_write_input_tokens: ENV.fetch("CACHE_WRITE_INPUT_TOKENS", 0).to_i,
-      cache_write_1h_input_tokens: ENV.fetch("CACHE_WRITE_1H_INPUT_TOKENS", 0).to_i
-    )
+    tokens: {
+      input: ENV.fetch("INPUT_TOKENS", 1).to_i,
+      output: ENV.fetch("OUTPUT_TOKENS", 1).to_i,
+      cache_read_input: ENV.fetch("CACHE_READ_INPUT_TOKENS", 0).to_i,
+      cache_write_input: ENV.fetch("CACHE_WRITE_INPUT_TOKENS", 0).to_i,
+      cache_write_extended_input: ENV.fetch("CACHE_WRITE_EXTENDED_INPUT_TOKENS", 0).to_i,
+      audio_input: ENV.fetch("AUDIO_INPUT_TOKENS", 0).to_i,
+      audio_output: ENV.fetch("AUDIO_OUTPUT_TOKENS", 0).to_i
+    }
   )
 end