llm_cost_tracker 0.10.0 → 0.12.0

data/lib/llm_cost_tracker/tracker.rb

@@ -1,47 +1,46 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/object/blank"
+require "active_support/notifications"
 require "securerandom"
 
 require_relative "ingestion"
 require_relative "ledger"
-require_relative "logging"
 require_relative "pricing"
-require_relative "billing/cost_status"
 
 module LlmCostTracker
-  class Tracker
+  module Tracker
     EVENT_NAME = "llm_request.llm_cost_tracker"
 
     class << self
-      def enforce_budget!(provider: nil, model: nil, request: nil)
+      def record(event:, latency_ms: nil, pricing_mode: nil, metadata: {}, context_tags: nil, enforce_budget: false)
         return unless LlmCostTracker.configuration.enabled
 
-        Budget.enforce!(provider: provider, model: model, request: request)
-      end
-
-      def record(event:, latency_ms: nil, pricing_mode: nil, metadata: {}, context_tags: nil)
-        return unless LlmCostTracker.configuration.enabled
-
-        pricing_mode = Pricing.normalize_mode(pricing_mode) || event.pricing_mode
-        cost_data, pricing_snapshot, priced_line_items = Pricing.calculate(
+        pricing_mode ||= event.pricing_mode
+        calculation = Pricing::Calculation.for(
           provider: event.provider,
           model: event.model,
           tokens: event.token_usage,
           line_items: event.line_items,
-          pricing_mode: pricing_mode
+          pricing_mode: pricing_mode,
+          usage_source: event.usage_source
         )
 
-        if cost_data.nil? && event.token_usage.total_tokens.positive? && priced_line_items.none?(&:priced?)
+        if enforce_budget
+          Budget.enforce!(provider: event.provider,
+                          model: event.model,
+                          estimate: calculation.cost&.total,
+                          force: true)
+        end
+
+        if calculation.token_cost.nil? && event.token_usage.total_tokens.positive? &&
+           calculation.priced_line_items.none?(&:priced?)
           Pricing::Unknown.process(event.model)
         end
 
         event = build_event(
           event: event,
-          pricing_mode: pricing_mode,
-          cost_data: cost_data,
-          pricing_snapshot: pricing_snapshot,
-          line_items: priced_line_items,
+          calculation: calculation,
           metadata: metadata,
           latency_ms: latency_ms,
           context_tags: context_tags
@@ -51,7 +50,7 @@ module LlmCostTracker
           Ingestion::Inbox.save(event)
           Ingestion::Worker.ensure_started
         else
-          Ledger::Store.insert(event, skip_existence_check: true)
+          Ledger::Store.insert(event)
         end
 
         yield if block_given?
@@ -71,29 +70,18 @@ module LlmCostTracker
         Logging.warn("Subscriber raised on #{EVENT_NAME}: #{e.class}: #{e.message}")
       end
 
-      def build_event(event:, pricing_mode:, cost_data:, pricing_snapshot:, line_items:,
-                      metadata:, latency_ms:, context_tags:)
+      def build_event(event:, calculation:, metadata:, latency_ms:, context_tags:)
         context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).to_h
-        cost = Pricing.combine_with_service_lines(cost_data, line_items)
-        cost_status = Billing::CostStatus.call(
-          token_usage: event.token_usage,
-          usage_source: event.usage_source,
-          token_cost: cost_data,
-          token_pricing_partial: Pricing.token_pricing_partial?(event.token_usage, cost_data),
-          service_line_items: line_items.reject(&:token?),
-          total_cost: cost&.fetch(:total_cost, nil)
-        )
-
         event.with(
           event_id: SecureRandom.uuid,
-          pricing_mode: pricing_mode,
-          cost: cost,
+          pricing_mode: calculation.mode,
+          cost: calculation.cost,
           tags: build_tags(context_tags: context_tags, metadata: metadata),
           latency_ms: finite_latency_ms(latency_ms),
           tracked_at: Time.now.utc,
-          cost_status: cost_status,
-          pricing_snapshot: pricing_snapshot,
-          line_items: line_items
+          cost_status: calculation.cost_status,
+          pricing_snapshot: calculation.snapshot,
+          line_items: calculation.priced_line_items
         )
       end
 

data/lib/llm_cost_tracker/usage/catalog.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/module/delegation"
+require "psych"
+
+require_relative "dimension"
+
+module LlmCostTracker
+  module Usage
+    module Catalog
+      DEFINITIONS_PATH = File.expand_path("dimensions.yml", __dir__)
+
+      DEFAULT_RATE_BASIS_BY_UNIT = {
+        "token" => "per_million_tokens",
+        "character" => "per_million_characters",
+        "request" => "per_request",
+        "session" => "per_session",
+        "hour" => "per_hour",
+        "minute" => "per_minute",
+        "image" => "per_image"
+      }.freeze
+
+      class << self
+        delegate :[], :fetch, to: :index
+
+        def all
+          @all ||= load_definitions.freeze
+        end
+
+        def token_priced
+          @token_priced ||= all.select(&:token_key).freeze
+        end
+
+        def token_priced_for(kind:, direction:, cache_state:)
+          token_priced.find do |dimension|
+            dimension.kind == kind && dimension.direction == direction && dimension.cache_state == cache_state
+          end
+        end
+
+        private
+
+        def index
+          @index ||= all.to_h { |dimension| [dimension.key, dimension] }.freeze
+        end
+
+        def load_definitions
+          Psych.safe_load_file(DEFINITIONS_PATH, permitted_classes: [], symbolize_names: true)
+               .map { |attributes| build(attributes) }
+        end
+
+        def build(attributes)
+          rate_basis = attributes[:rate_basis] || DEFAULT_RATE_BASIS_BY_UNIT.fetch(attributes.fetch(:unit))
+          Dimension.new(**attributes, rate_basis: rate_basis)
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/usage/dimension.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Usage
+    Dimension = Data.define(
+      :key, :kind, :direction, :modality, :cache_state, :unit, :rate_basis
+    ) do
+      def token?
+        unit == "token"
+      end
+
+      def token_key
+        :"#{key}_tokens" if token?
+      end
+
+      def cost_key
+        :"#{key}_cost" if token?
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/{billing/components.yml → usage/dimensions.yml}

@@ -4,9 +4,6 @@
   modality: text
   cache_state: none
   unit: token
-  category: token
-  token_key: input_tokens
-  cost_key: input_cost
 
 - key: cache_read_input
   kind: text_token
@@ -14,9 +11,6 @@
   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
@@ -24,9 +18,6 @@
   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
@@ -34,9 +25,6 @@
   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
@@ -44,9 +32,6 @@
   modality: text
   cache_state: none
   unit: token
-  category: token
-  token_key: output_tokens
-  cost_key: output_cost
 
 - key: audio_input
   kind: audio_token
@@ -54,9 +39,6 @@
   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
@@ -64,9 +46,6 @@
   modality: audio
   cache_state: none
   unit: token
-  category: token
-  token_key: audio_output_tokens
-  cost_key: audio_output_cost
 
 - key: image_input
   kind: image_token
@@ -74,9 +53,6 @@
   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
@@ -84,9 +60,6 @@
   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
@@ -94,7 +67,6 @@
   modality: text
   cache_state: none
   unit: request
-  category: tool
   rate_basis: per_1k_requests
 
 - key: web_search_preview_request_reasoning
@@ -103,7 +75,6 @@
   modality: text
   cache_state: none
   unit: request
-  category: tool
   rate_basis: per_1k_requests
 
 - key: web_search_preview_request_non_reasoning
@@ -112,7 +83,6 @@
   modality: text
   cache_state: none
   unit: request
-  category: tool
   rate_basis: per_1k_requests
 
 - key: web_fetch_request
@@ -121,7 +91,6 @@
   modality: text
   cache_state: none
   unit: request
-  category: tool
   rate_basis: per_1k_requests
 
 - key: file_search_call
@@ -130,7 +99,6 @@
   modality: text
   cache_state: none
   unit: request
-  category: tool
   rate_basis: per_1k_requests
 
 - key: container_session
@@ -139,25 +107,14 @@
   modality: none
   cache_state: none
   unit: session
-  category: runtime
   rate_basis: per_session
 
-- key: code_execution_request
-  kind: code_execution_request
-  direction: neither
-  modality: none
-  cache_state: none
-  unit: request
-  category: runtime
-  rate_basis: per_1k_requests
-
 - key: code_execution_hour
   kind: code_execution_hour
   direction: neither
   modality: none
   cache_state: none
   unit: hour
-  category: runtime
   rate_basis: per_hour
 
 - key: grounding_request
@@ -166,7 +123,6 @@
   modality: text
   cache_state: none
   unit: request
-  category: tool
   rate_basis: per_1k_requests
 
 - key: text_to_speech_character
@@ -175,14 +131,36 @@
   modality: audio
   cache_state: none
   unit: character
-  category: tool
   rate_basis: per_million_characters
 
+- key: transcription_minute
+  kind: transcription_minute
+  direction: input
+  modality: audio
+  cache_state: none
+  unit: minute
+  rate_basis: per_minute
+
+- key: image_generation_call
+  kind: image_generation_call
+  direction: output
+  modality: image
+  cache_state: none
+  unit: image
+  rate_basis: per_image
+
+- key: computer_call
+  kind: computer_call
+  direction: neither
+  modality: text
+  cache_state: none
+  unit: request
+  rate_basis: per_request
+
 - 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/usage/source.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Usage
+    module Source
+      MANUAL = "manual"
+      UNKNOWN = "unknown"
+      RESPONSE = "response"
+      STREAM_FINAL = "stream_final"
+      SDK_RESPONSE = "sdk_response"
+      SDK_BATCH_RESULT = "sdk_batch_result"
+    end
+  end
+end

data/lib/llm_cost_tracker/usage/token_usage.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require_relative "catalog"
+
+module LlmCostTracker
+  module Usage
+    KNOWN_TOKEN_KEYS = (
+      Catalog.token_priced.map { |dimension| dimension.token_key.to_s } + %w[total_tokens hidden_output_tokens]
+    ).freeze
+
+    TokenUsage = Data.define(
+      :input_tokens,
+      :cache_read_input_tokens,
+      :cache_write_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 priced_quantities
+        Catalog.token_priced.to_h { |dimension| [dimension.key, public_send(dimension.token_key)] }
+      end
+
+      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)
+
+        values = tokens.to_h.transform_keys(&:to_s)
+        warn_on_unknown_keys(values)
+        token_attributes = Catalog.token_priced.to_h do |dimension|
+          [dimension.token_key, values.fetch(dimension.token_key.to_s, 0)]
+        end
+
+        build(
+          **token_attributes,
+          total_tokens: values["total_tokens"],
+          hidden_output_tokens: values.fetch("hidden_output_tokens", 0)
+        )
+      end
+
+      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 self.non_negative_int(value)
+        [value.to_i, 0].max
+      end
+
+      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
+  end
+end

data/lib/llm_cost_tracker/version.rb

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

data/lib/llm_cost_tracker.rb

@@ -1,13 +1,8 @@
 # frozen_string_literal: true
 
 require "rails"
-require "active_support"
 require "active_support/core_ext/object/blank"
-require "active_support/core_ext/object/deep_dup"
 require "active_support/core_ext/object/try"
-require "active_support/core_ext/hash/indifferent_access"
-require "active_support/core_ext/string/inflections"
-require "active_support/notifications"
 
 require_relative "llm_cost_tracker/version"
 require_relative "llm_cost_tracker/configuration"
@@ -16,32 +11,34 @@ require_relative "llm_cost_tracker/logging"
 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/masking"
-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/currency"
+require_relative "llm_cost_tracker/usage/catalog"
+require_relative "llm_cost_tracker/usage/token_usage"
+require_relative "llm_cost_tracker/usage/source"
+require_relative "llm_cost_tracker/pricing/rate"
+require_relative "llm_cost_tracker/charges/cost"
+require_relative "llm_cost_tracker/charges/line_item"
+require_relative "llm_cost_tracker/charges/cost_status"
 require_relative "llm_cost_tracker/event"
 require_relative "llm_cost_tracker/pricing"
 require_relative "llm_cost_tracker/parsers"
 require_relative "llm_cost_tracker/middleware/faraday"
 require_relative "llm_cost_tracker/integrations"
 require_relative "llm_cost_tracker/budget"
-require_relative "llm_cost_tracker/pricing/unknown"
 require_relative "llm_cost_tracker/ledger"
 require_relative "llm_cost_tracker/ingestion"
 require_relative "llm_cost_tracker/tracker"
 
 module LlmCostTracker
-  autoload :Engine,         "llm_cost_tracker/engine"
-  autoload :Reconciliation, "llm_cost_tracker/reconciliation"
-  autoload :ReconcileTasks, "llm_cost_tracker/reconcile_tasks"
-  autoload :Doctor,         "llm_cost_tracker/doctor"
-  autoload :Report,         "llm_cost_tracker/report"
-  autoload :Retention,      "llm_cost_tracker/retention"
+  autoload :Engine, "llm_cost_tracker/engine"
+  autoload :Doctor, "llm_cost_tracker/doctor"
+  autoload :CaptureVerifier, "llm_cost_tracker/capture_verifier"
+  autoload :Report, "llm_cost_tracker/report"
+  autoload :Retention, "llm_cost_tracker/retention"
 
   module Pricing
     autoload :Sync, "llm_cost_tracker/pricing/sync"
+    autoload :Unknown, "llm_cost_tracker/pricing/unknown"
   end
 
   @configuration = Configuration.new
@@ -53,72 +50,67 @@ module LlmCostTracker
       "llm_cost_tracker_"
     end
 
-    def reconciliation_enabled?
-      configuration.reconciliation_enabled
-    end
-
     def configure
       config = configuration
       raise Error, "LlmCostTracker is already configured" if config.finalized?
 
       yield(config)
       config.finalize!
-      Pricing::Lookup.reset!
       Pricing::Registry.reset!
-      Pricing::ServiceCharges.reset!
       Integrations.install!
       config
     end
 
-    def reset_configuration!
-      Ingestion::Worker.shutdown!(drain: false)
-      Ingestion::Pool.reset!
-      @configuration = Configuration.new
-      Pricing::Lookup.reset!
-      Pricing::Registry.reset!
-      Pricing::ServiceCharges.reset!
-      Pricing::Unknown.reset!
-      Ingestion::Worker.reset!
-      Tags::Context.clear!
-      Dashboard::SetupState.reset!
-    end
-
     def with_tags(tags = nil, **kwargs, &)
       Tags::Context.with((tags || {}).merge(kwargs), &)
     end
 
-    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
-
+    def track(provider:,
+              tokens:,
+              model: nil,
+              tags: {},
+              latency_ms: nil,
+              stream: false,
+              usage_source: Usage::Source::MANUAL,
+              enforce_budget: false,
+              provider_response_id: nil,
+              provider_project_id: nil,
+              provider_api_key_id: nil,
+              provider_workspace_id: nil,
+              pricing_mode: nil,
+              service_line_items: [])
       Tracker.record(
         event: Event.build(
           provider: provider,
           model: model,
-          token_usage: TokenUsage.build_from_tokens(tokens),
+          token_usage: Usage::TokenUsage.build_from_tokens(tokens),
           stream: stream,
           usage_source: usage_source,
           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: tags
+        metadata: tags,
+        enforce_budget: enforce_budget
       )
     end
 
-    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)
+    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,
+                     pricing_mode: nil)
       require_relative "llm_cost_tracker/capture/stream_collector"
-      Tracker.enforce_budget! if enforce_budget
+      Budget.enforce!(provider: provider, model: model, force: true) if enforce_budget
       collector = Capture::StreamCollector.new(
         provider: provider.to_s,
         model: model,
@@ -127,7 +119,6 @@ module LlmCostTracker
         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: tags
       )
@@ -146,4 +137,4 @@ Faraday::Middleware.register_middleware(
   llm_cost_tracker: LlmCostTracker::Middleware::Faraday
 )
 
-at_exit { LlmCostTracker::Ingestion::Worker.shutdown!(drain: false) }
+at_exit { LlmCostTracker::Ingestion::Worker.shutdown!(drain: false) if LlmCostTracker::Ingestion.async? }