llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,8 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../capture/stream_collector"
-require_relative "../capture/stream_tracker"
+require_relative "../billing/line_item"
 
 module LlmCostTracker
   module Integrations
@@ -52,22 +51,48 @@ module LlmCostTracker
                 pricing_mode: pricing_mode(message: message, request: request, usage: usage),
                 token_usage: token_usage(usage: usage, input_tokens: input_tokens, output_tokens: output_tokens),
                 usage_source: :sdk_response,
-                provider_response_id: object_value(message, :id)
+                provider_response_id: object_value(message, :id),
+                service_line_items: service_line_items_from(usage)
               ),
               latency_ms: latency_ms
             )
           end
         end
 
+        def service_line_items_from(usage)
+          server_tool_use = object_value(usage, :server_tool_use)
+          return [] unless server_tool_use
+
+          [
+            line_item_for_server_tool(server_tool_use, :web_search_request, :web_search_requests,
+                                      "usage.server_tool_use.web_search_requests"),
+            line_item_for_server_tool(server_tool_use, :code_execution_request, :code_execution_requests,
+                                      "usage.server_tool_use.code_execution_requests")
+          ].compact
+        end
+
+        def line_item_for_server_tool(server_tool_use, component_key, count_key, provider_field)
+          quantity = object_value(server_tool_use, count_key).to_i
+          return nil if quantity.zero?
+
+          Billing::LineItem.build(
+            component_key: component_key,
+            quantity: quantity,
+            cost_status: Billing::CostStatus::UNKNOWN,
+            pricing_basis: :provider_usage,
+            provider_field: provider_field
+          )
+        end
+
         def token_usage(usage:, input_tokens:, output_tokens:)
-          cache_write_1h = object_dig(usage, :cache_creation, :ephemeral_1h_input_tokens).to_i
-          cache_write_5m = object_dig(usage, :cache_creation, :ephemeral_5m_input_tokens)
-          cache_write = if cache_write_5m.nil?
-                          total_cache_write = object_value(usage, :cache_creation_input_tokens)
-                          [total_cache_write.to_i - cache_write_1h, 0].max
-                        else
-                          cache_write_5m.to_i
-                        end
+          cache_creation = object_value(usage, :cache_creation)
+          if cache_creation
+            cache_write_default = object_value(cache_creation, :ephemeral_5m_input_tokens).to_i
+            cache_write_extended = object_value(cache_creation, :ephemeral_1h_input_tokens).to_i
+          else
+            cache_write_default = object_value(usage, :cache_creation_input_tokens).to_i
+            cache_write_extended = 0
+          end
           hidden_output = (
             object_value(usage, :thinking_tokens, :thinking_output_tokens) ||
             object_dig(usage, :output_tokens_details, :reasoning_tokens)
@@ -77,8 +102,8 @@ module LlmCostTracker
             input_tokens: input_tokens.to_i,
             output_tokens: output_tokens.to_i,
             cache_read_input_tokens: object_value(usage, :cache_read_input_tokens).to_i,
-            cache_write_input_tokens: cache_write,
-            cache_write_1h_input_tokens: cache_write_1h,
+            cache_write_input_tokens: cache_write_default,
+            cache_write_extended_input_tokens: cache_write_extended,
             hidden_output_tokens: hidden_output
           )
         end
@@ -95,39 +120,21 @@ module LlmCostTracker
           modes.empty? ? nil : modes.join("_")
         end
 
+        def stream_pricing_mode(request)
+          pricing_mode(message: nil, request: request || {}, usage: nil)
+        end
+
         def inference_geo(message:, request:, usage:)
           object_value(usage, :inference_geo) ||
             object_value(message, :inference_geo) ||
             request[:inference_geo]
         end
-
-        def track_stream(stream, collector:)
-          return stream unless active?
-
-          LlmCostTracker::Capture::StreamTracker.new(
-            stream: stream,
-            collector: collector,
-            active: -> { active? },
-            finish: ->(errored:) { finish_stream(collector, errored: errored) }
-          ).wrap
-        end
-
-        def stream_collector(request)
-          LlmCostTracker::Capture::StreamCollector.new(
-            provider: "anthropic",
-            model: request[:model]
-          )
-        end
-
-        def finish_stream(collector, errored:)
-          record_safely { collector.finish!(errored: errored) }
-        end
       end
 
       module MessagesPatch
         def create(*args, **kwargs)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
           LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          started_at = LlmCostTracker::Timing.now_monotonic
           message = super
           LlmCostTracker::Integrations::Anthropic.record_message(
             message,
@@ -139,16 +146,16 @@ module LlmCostTracker
 
         def stream(*args, **kwargs)
           request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
-          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
           LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
           stream = super
           LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
         end
 
         def stream_raw(*args, **kwargs)
           request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
-          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
           LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
           stream = super
           LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
         end

data/lib/llm_cost_tracker/integrations/base.rb

@@ -1,10 +1,12 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/hash/indifferent_access"
-require "active_support/core_ext/object/try"
 require "active_support/core_ext/string/inflections"
 
 require_relative "../logging"
+require_relative "../timing"
+require_relative "../capture/stream_collector"
+require_relative "../capture/stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -30,17 +32,16 @@ module LlmCostTracker
           return Result.new(name, :warn, "#{name} integration cannot be installed: #{problems.join('; ')}")
         end
 
-        required_targets = patch_targets.reject { |target| target.fetch(:optional) }
-        installed = required_targets.count do |target|
+        installed = patch_targets.reject { |target| target.fetch(:optional) }.all? do |target|
           target.fetch(:constant_name).to_s.safe_constantize&.ancestors&.include?(target.fetch(:patch))
         end
-        return Result.new(name, :ok, "#{name} integration installed") if installed == required_targets.count
+        return Result.new(name, :ok, "#{name} integration installed") if installed
 
         Result.new(name, :warn, "#{name} integration is enabled but not installed")
       end
 
       def elapsed_ms(started_at)
-        ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000).round
+        Timing.elapsed_ms(started_at)
       end
 
       def enforce_budget!
@@ -60,6 +61,29 @@ module LlmCostTracker
         params.merge(kwargs).with_indifferent_access
       end
 
+      def track_stream(stream, collector:)
+        return stream unless active?
+
+        LlmCostTracker::Capture::StreamTracker.new(
+          stream: stream,
+          collector: collector,
+          active: -> { active? },
+          finish: ->(errored:) { record_safely { collector.finish!(errored: errored) } }
+        ).wrap
+      end
+
+      def stream_collector(request)
+        LlmCostTracker::Capture::StreamCollector.new(
+          provider: integration_name.to_s,
+          model: request[:model],
+          pricing_mode: stream_pricing_mode(request)
+        )
+      end
+
+      def stream_pricing_mode(_request)
+        nil
+      end
+
       def object_value(object, *keys)
         keys.each do |key|
           value = read_object_value(object, key)
@@ -69,15 +93,6 @@ module LlmCostTracker
       end
 
       def object_dig(object, *path)
-        if object.respond_to?(:dig)
-          begin
-            value = object.dig(*path)
-            return value unless value.nil?
-          rescue NameError, TypeError
-            nil
-          end
-        end
-
         path.reduce(object) do |current, key|
           return nil if current.nil?
 
@@ -106,25 +121,17 @@ module LlmCostTracker
 
       def read_object_value(object, key)
         return nil if object.nil?
-        return object[key] if object.try(:key?, key)
-
-        string_key = key.to_s
-        return object[string_key] if object.try(:key?, string_key)
-
-        value = object.try(key)
-        return value unless value.nil?
 
-        indexed_object_value(object, key)
-      end
+        if object.is_a?(Hash)
+          return object[key] if object.key?(key)
+          return object[key.name] if key.is_a?(Symbol) && object.key?(key.name)
+        end
 
-      def indexed_object_value(object, key)
-        object.try(:[], key)
-      rescue IndexError, NameError, TypeError
-        nil
+        object.public_send(key) if object.respond_to?(key)
       end
 
-      module_function :read_object_value, :indexed_object_value
-      private_class_method :read_object_value, :indexed_object_value
+      module_function :read_object_value
+      private_class_method :read_object_value
 
       def validate_contract!
         problems = version_problems + target_problems

data/lib/llm_cost_tracker/integrations/openai.rb

@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../capture/stream_collector"
-require_relative "../capture/stream_tracker"
+require_relative "../billing/line_item"
+require_relative "../parsers/openai_service_charges"
 
 module LlmCostTracker
   module Integrations
@@ -14,6 +14,10 @@ module LlmCostTracker
           :openai
         end
 
+        def stream_pricing_mode(request)
+          Pricing.normalize_mode((request || {})[:service_tier])
+        end
+
         def minimum_version
           "0.59.0"
         end
@@ -54,65 +58,106 @@ module LlmCostTracker
                 provider: "openai",
                 model: object_value(response, :model) || request[:model],
                 pricing_mode: object_value(response, :service_tier) || request[:service_tier],
-                token_usage: TokenUsage.build(
-                  input_tokens: regular_input_tokens(input_tokens, cache_read),
-                  output_tokens: output_tokens.to_i,
-                  cache_read_input_tokens: cache_read,
-                  hidden_output_tokens: hidden_output_tokens(usage)
-                ),
+                token_usage: token_usage(usage:, input_tokens:, output_tokens:, cache_read:),
                 usage_source: :sdk_response,
-                provider_response_id: object_value(response, :id)
+                provider_response_id: object_value(response, :id),
+                service_line_items: service_line_items_from(response)
               ),
               latency_ms: latency_ms
             )
           end
         end
 
+        def service_line_items_from(response)
+          output = object_value(response, :output)
+          return [] unless output.respond_to?(:each)
+
+          LlmCostTracker::Parsers::OpenaiServiceCharges
+            .line_items_from_output(output.map { |item| normalize_output_item(item) })
+        end
+
+        def normalize_output_item(item)
+          return item if item.is_a?(Hash)
+          return nil if item.nil?
+
+          {
+            "type" => object_value(item, :type),
+            "id" => object_value(item, :id),
+            "status" => object_value(item, :status),
+            "container_id" => object_value(item, :container_id),
+            "action" => normalize_output_action(object_value(item, :action))
+          }
+        end
+
+        def normalize_output_action(action)
+          return nil if action.nil?
+          return action if action.is_a?(Hash)
+
+          { "type" => object_value(action, :type) }
+        end
+
+        def token_usage(usage:, input_tokens:, output_tokens:, cache_read:)
+          audio_input = audio_input_tokens(usage)
+          audio_output = audio_output_tokens(usage)
+
+          TokenUsage.build(
+            input_tokens: regular_input_tokens(input_tokens, cache_read, audio_input),
+            output_tokens: regular_output_tokens(output_tokens, audio_output),
+            cache_read_input_tokens: cache_read,
+            audio_input_tokens: audio_input,
+            audio_output_tokens: audio_output,
+            hidden_output_tokens: hidden_output_tokens(usage)
+          )
+        end
+
+        INPUT_DETAIL_KEYS = %i[input_tokens_details input_token_details prompt_tokens_details].freeze
+        OUTPUT_DETAIL_KEYS = %i[output_tokens_details output_token_details completion_tokens_details].freeze
+
         def cache_read_input_tokens(usage)
-          (
-            object_dig(usage, :input_tokens_details, :cached_tokens) ||
-            object_dig(usage, :prompt_tokens_details, :cached_tokens)
-          ).to_i
+          input_detail(usage, :cached_tokens)
         end
 
         def hidden_output_tokens(usage)
-          (
-            object_dig(usage, :output_tokens_details, :reasoning_tokens) ||
-            object_dig(usage, :completion_tokens_details, :reasoning_tokens)
-          ).to_i
+          output_detail(usage, :reasoning_tokens)
         end
 
-        def regular_input_tokens(input_tokens, cache_read)
-          [input_tokens.to_i - cache_read.to_i, 0].max
+        def audio_input_tokens(usage)
+          input_detail(usage, :audio_tokens)
         end
 
-        def track_stream(stream, collector:)
-          return stream unless active?
+        def audio_output_tokens(usage)
+          output_detail(usage, :audio_tokens)
+        end
 
-          LlmCostTracker::Capture::StreamTracker.new(
-            stream: stream,
-            collector: collector,
-            active: -> { active? },
-            finish: ->(errored:) { finish_stream(collector, errored: errored) }
-          ).wrap
+        def input_detail(usage, key)
+          INPUT_DETAIL_KEYS.each do |container|
+            value = object_dig(usage, container, key)
+            return value.to_i if value
+          end
+          0
         end
 
-        def stream_collector(request)
-          LlmCostTracker::Capture::StreamCollector.new(
-            provider: "openai",
-            model: request[:model]
-          )
+        def output_detail(usage, key)
+          OUTPUT_DETAIL_KEYS.each do |container|
+            value = object_dig(usage, container, key)
+            return value.to_i if value
+          end
+          0
+        end
+
+        def regular_input_tokens(input_tokens, cache_read, audio_input)
+          [input_tokens.to_i - cache_read - audio_input, 0].max
         end
 
-        def finish_stream(collector, errored:)
-          record_safely { collector.finish!(errored: errored) }
+        def regular_output_tokens(output_tokens, audio_output)
+          [output_tokens.to_i - audio_output, 0].max
         end
       end
 
       module ResponsesPatch
         def create(*args, **kwargs)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
           LlmCostTracker::Integrations::Openai.enforce_budget!
+          started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           LlmCostTracker::Integrations::Openai.record_response(
             response,
@@ -124,25 +169,25 @@ module LlmCostTracker
 
         def stream(*args, **kwargs)
           request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
-          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           LlmCostTracker::Integrations::Openai.enforce_budget!
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           stream = super
           LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
         end
 
         def stream_raw(*args, **kwargs)
           request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
-          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           LlmCostTracker::Integrations::Openai.enforce_budget!
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           stream = super
           LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
         end
 
         def retrieve_streaming(response_id, *args, **kwargs)
           request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
           collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           collector.provider_response_id = response_id
-          LlmCostTracker::Integrations::Openai.enforce_budget!
           stream = super
           LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
         end
@@ -150,8 +195,8 @@ module LlmCostTracker
 
       module ChatCompletionsPatch
         def create(*args, **kwargs)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
           LlmCostTracker::Integrations::Openai.enforce_budget!
+          started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           LlmCostTracker::Integrations::Openai.record_response(
             response,
@@ -163,8 +208,8 @@ module LlmCostTracker
 
         def stream_raw(*args, **kwargs)
           request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
-          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           LlmCostTracker::Integrations::Openai.enforce_budget!
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
           stream = super
           LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
         end

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -89,7 +89,7 @@ module LlmCostTracker
                   hidden_output_tokens: hidden_output
                 ),
                 stream: stream,
-                usage_source: :ruby_llm,
+                usage_source: :sdk_response,
                 provider_response_id: provider_response_id(response)
               ),
               latency_ms: latency_ms
@@ -98,7 +98,7 @@ module LlmCostTracker
         end
 
         def regular_input_tokens(input_tokens, cache_read)
-          [input_tokens.to_i - cache_read.to_i, 0].max
+          [input_tokens.to_i - cache_read, 0].max
         end
 
         def provider_slug(provider)
@@ -133,8 +133,8 @@ module LlmCostTracker
         def complete(*args, **kwargs, &)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
           integration.enforce_budget!
+          started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_completion(
             self,
@@ -149,8 +149,8 @@ module LlmCostTracker
         def embed(*args, **kwargs)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
           integration.enforce_budget!
+          started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_embedding(
             self,
@@ -164,8 +164,8 @@ module LlmCostTracker
         def transcribe(*args, **kwargs)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
           integration.enforce_budget!
+          started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_transcription(
             self,

data/lib/llm_cost_tracker/integrations.rb

@@ -26,11 +26,11 @@ module LlmCostTracker
     end
 
     def normalize(names)
-      Array(names).flatten.map(&:to_sym).uniq
+      Array(names).flatten.uniq
     end
 
     def fetch(name)
-      AVAILABLE.fetch(name.to_sym) do
+      AVAILABLE.fetch(name) do
         message = "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
         raise LlmCostTracker::Error, message
       end

data/lib/llm_cost_tracker/ledger/period/totals.rb

@@ -27,38 +27,41 @@ module LlmCostTracker
 
         def snapshot_totals
           values = periods.to_h { |period| [period, 0.0] }
+          period_by_name = periods.to_h { |period| [period.name, period] }
           sql = periods.map { |period| snapshot_select(period) }.join(" UNION ALL ")
-          LlmCostTracker::Ledger::Call.find_by_sql(sql).each do |row|
-            values[row.period_key.to_sym] = row.total_cost.to_f
+          LlmCostTracker::Call.find_by_sql(sql).each do |row|
+            period = period_by_name.fetch(row.period_key)
+            values[period] = row.total_cost.to_f
           end
           values
         end
 
         def snapshot_select(period)
           start = Period.range_start(period, time)
-          "SELECT #{connection.quote(period.to_s)} AS period_key, " \
+          "SELECT #{connection.quote(period.name)} AS period_key, " \
             "(#{rollup_total_sql(period)}) + (#{pending_total_sql(start)}) AS total_cost"
         end
 
         def rollup_total_sql(period)
-          table = connection.quote_table_name("llm_cost_tracker_period_totals")
-          "COALESCE((SELECT total_cost FROM #{table} " \
+          table = connection.quote_table_name("llm_cost_tracker_call_rollups")
+          "COALESCE((SELECT SUM(total_cost) FROM #{table} " \
             "WHERE period = #{connection.quote(Period::PERIODS.fetch(period))} " \
-            "AND period_start = #{connection.quote(Period.bucket(period, time))} LIMIT 1), 0)"
+            "AND period_start = #{connection.quote(Period.bucket(period, time))} " \
+            "AND currency = #{connection.quote(Ledger::Rollups::DEFAULT_CURRENCY)}), 0)"
         end
 
         def pending_total_sql(start)
-          table = connection.quote_table_name(Ingestion::Event.table_name)
+          table = connection.quote_table_name(Ingestion::InboxEntry.table_name)
           total_cost = connection.quote_column_name("total_cost")
           tracked_at = connection.quote_column_name("tracked_at")
           attempts = connection.quote_column_name("attempts")
           "COALESCE((SELECT SUM(#{total_cost}) FROM #{table} " \
-            "WHERE #{attempts} < #{Ingestion::Event::MAX_ATTEMPTS} " \
+            "WHERE #{attempts} < #{Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE} " \
             "AND #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)}), 0)"
         end
 
         def connection
-          LlmCostTracker::Ledger::Call.connection
+          LlmCostTracker::Call.connection
         end
       end
     end

data/lib/llm_cost_tracker/ledger/period.rb

@@ -4,22 +4,22 @@ module LlmCostTracker
   module Ledger
     module Period
       PERIODS = {
-        monthly: "month",
-        daily:   "day"
+        month: "month",
+        day:   "day"
       }.freeze
 
       module_function
 
       def valid_keys(periods)
-        periods.map(&:to_sym).select { |period| PERIODS.key?(period) }
+        periods.select { |period| PERIODS.key?(period) }
       end
 
       def range_start(period, time)
         utc_time = time.to_time.utc
 
         case period
-        when :monthly then utc_time.beginning_of_month
-        when :daily then utc_time.beginning_of_day
+        when :month then utc_time.beginning_of_month
+        when :day then utc_time.beginning_of_day
         end
       end
 

data/lib/llm_cost_tracker/ledger/rollups/upsert_sql.rb

@@ -6,12 +6,8 @@ module LlmCostTracker
   module Ledger
     class Rollups
       class UpsertSql
-        def self.call(model)
-          new(model).call
-        end
-
-        def initialize(model)
-          @model = model
+        def self.call
+          new.call
         end
 
         def call
@@ -23,13 +19,11 @@ module LlmCostTracker
 
         private
 
-        attr_reader :model
-
         def postgres_sql
           total_cost = connection.quote_column_name("total_cost")
           updated_at = connection.quote_column_name("updated_at")
 
-          "#{total_cost} = #{model.quoted_table_name}.#{total_cost} + excluded.#{total_cost}, " \
+          "#{total_cost} = #{LlmCostTracker::CallRollup.quoted_table_name}.#{total_cost} + excluded.#{total_cost}, " \
             "#{updated_at} = excluded.#{updated_at}"
         end
 
@@ -38,7 +32,7 @@ module LlmCostTracker
         end
 
         def connection
-          model.connection
+          LlmCostTracker::CallRollup.connection
         end
       end
     end