llm_cost_tracker 0.7.3 → 0.9.0

data/lib/llm_cost_tracker/ingestion.rb

@@ -5,6 +5,7 @@ require "securerandom"
 require_relative "doctor/check"
 require_relative "errors"
 require_relative "ledger"
+require_relative "ingestion/inline"
 require_relative "ingestion/lease_claim"
 require_relative "ingestion/inbox"
 require_relative "ingestion/batch"
@@ -15,31 +16,59 @@ module LlmCostTracker
     VERIFY_TAG = "llm_cost_tracker_verify"
 
     class << self
+      def table_name_prefix
+        "llm_cost_tracker_ingestion_"
+      end
+
+      CORE_SCHEMA_GUARDS = [
+        ["llm_cost_tracker_calls",           Ledger::Schema::Calls],
+        ["llm_cost_tracker_call_line_items", Ledger::Schema::CallLineItems],
+        ["llm_cost_tracker_call_tags",       Ledger::Schema::CallTags]
+      ].freeze
+
+      ROLLUPS_SCHEMA_GUARD = ["llm_cost_tracker_call_rollups", Ledger::Schema::CallRollups].freeze
+
+      DURABLE_SCHEMA_GUARDS = [
+        ["llm_cost_tracker_ingestion_inbox_entries", Ledger::Schema::IngestionInboxEntries],
+        ["llm_cost_tracker_ingestion_leases",        Ledger::Schema::IngestionLeases]
+      ].freeze
+
       def ensure_current_schema!
-        unless Ledger::Call.table_exists?
-          raise Error, "llm_api_calls table is missing; run install generator and migrate"
+        unless LlmCostTracker::Call.table_exists?
+          raise Error, "llm_cost_tracker_calls table is missing; run install generator and migrate"
+        end
+
+        guards_for_current_config.each do |table_name, schema_module|
+          errors = schema_module.current_schema_errors
+          next if errors.empty?
+
+          raise Error,
+                "#{table_name} table is not on the current schema: #{errors.join('; ')}; see docs/upgrading.md"
         end
+      end
 
-        schema_errors = Ledger::Schema::Calls.current_schema_errors
-        message = "llm_api_calls table is not on the current schema: #{schema_errors.join('; ')}"
-        raise Error, message if schema_errors.any?
+      def durable?
+        LlmCostTracker.configuration.durable_ingestion
+      end
 
-        period_total_errors = Ledger::Schema::PeriodTotals.current_schema_errors
-        return if period_total_errors.empty?
+      def cache_rollups?
+        LlmCostTracker.configuration.cache_rollups
+      end
 
-        message = "llm_cost_tracker_period_totals table is not on the current schema: " \
-                  "#{period_total_errors.join('; ')}; " \
-                  "run bin/rails generate llm_cost_tracker:add_period_totals && bin/rails db:migrate"
-        raise Error, message
+      def guards_for_current_config
+        guards = CORE_SCHEMA_GUARDS.dup
+        guards << ROLLUPS_SCHEMA_GUARD if cache_rollups?
+        guards += DURABLE_SCHEMA_GUARDS if durable?
+        guards
       end
 
       def verify
-        unless LlmCostTracker::Ledger::Call.table_exists?
+        unless LlmCostTracker::Call.table_exists?
           return [
             LlmCostTracker::Doctor::Check.new(
               :error,
               "active_record",
-              "llm_api_calls table is missing; run install generator and migrate"
+              "llm_cost_tracker_calls table is missing; run install generator and migrate"
             )
           ]
         end
@@ -60,13 +89,12 @@ module LlmCostTracker
         event = LlmCostTracker.track(
           provider: provider,
           model: model,
-          input_tokens: 1,
-          output_tokens: 1,
+          tokens: { input: 1, output: 1 },
           provider_response_id: response_id,
-          feature: VERIFY_TAG
+          tags: { feature: VERIFY_TAG }
         )
-        LlmCostTracker.flush!
-        persisted = LlmCostTracker::Ledger::Call.where(provider_response_id: response_id).exists?
+        LlmCostTracker::Ingestion::Worker.flush! if durable?
+        persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
 
         return capture_success if persisted && notifications.any?
 
@@ -83,7 +111,7 @@ module LlmCostTracker
         LlmCostTracker::Doctor::Check.new(:error, "active_record capture", "#{e.class}: #{e.message}")
       ensure
         cleanup_verification_call(response_id) if response_id
-        LlmCostTracker::Ingestion::Event.where(event_id: event.event_id).delete_all if event
+        cleanup_verification_inbox(event: event, response_id: response_id)
         ActiveSupport::Notifications.unsubscribe(subscription) if subscription
       end
 
@@ -94,10 +122,11 @@ module LlmCostTracker
       end
 
       def capture_success
+        path = durable? ? "durable inbox" : "inline writer"
         LlmCostTracker::Doctor::Check.new(
           :ok,
           "active_record capture",
-          "manual event emitted and persisted through durable inbox"
+          "manual event emitted and persisted through #{path}"
         )
       end
 
@@ -109,14 +138,29 @@ module LlmCostTracker
       end
 
       def cleanup_verification_call(response_id)
-        relation = LlmCostTracker::Ledger::Call.where(provider_response_id: response_id)
-        rows = relation.pluck(:id, :tracked_at, :total_cost)
+        relation = LlmCostTracker::Call.where(provider_response_id: response_id)
+        rows = relation.pluck(:id, :tracked_at, :total_cost, :pricing_snapshot, :provider)
         return if rows.empty?
 
         relation.delete_all
+        return unless cache_rollups?
+
         LlmCostTracker::Ledger::Rollups.decrement!(rows)
       end
 
+      def cleanup_verification_inbox(event:, response_id:)
+        return unless durable? && LlmCostTracker::Ingestion::InboxEntry.table_exists?
+
+        if event
+          LlmCostTracker::Ingestion::InboxEntry.where(event_id: event.event_id).delete_all
+        elsif response_id
+          escaped = ActiveRecord::Base.sanitize_sql_like(response_id)
+          LlmCostTracker::Ingestion::InboxEntry
+            .where("payload LIKE ?", "%\"provider_response_id\":\"#{escaped}\"%")
+            .delete_all
+        end
+      end
+
       def sample_priced_identity
         key = LlmCostTracker::Pricing::Registry.builtin_prices.find do |model_id, prices|
           model_id.include?("/") && prices[:input] && prices[:output]

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,58 @@ 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, :web_fetch_request, :web_fetch_requests,
+                                      "usage.server_tool_use.web_fetch_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 = server_tool_count(server_tool_use, count_key)
+          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 server_tool_count(server_tool_use, count_key)
+          direct = object_value(server_tool_use, count_key).to_i
+          return direct if direct.positive?
+          return 0 unless server_tool_use.respond_to?(:to_h)
+
+          server_tool_use.to_h[count_key].to_i
+        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,57 +112,48 @@ 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
 
+        DATA_RESIDENCY_GEOS = %w[us].freeze
+        # Anthropic Priority Tier is committed throughput (tokens/min capacity), not a per-token
+        # surcharge. Treat it as standard pricing so cost_status doesn't fall to :unknown.
+        STANDARD_EQUIVALENT_SERVICE_TIERS = %w[standard standard_only priority].freeze
+
         def pricing_mode(message:, request:, usage:)
+          service_tier = object_value(usage, :service_tier) ||
+                         object_value(message, :service_tier) ||
+                         request[:service_tier]
+          service_tier = nil if STANDARD_EQUIVALENT_SERVICE_TIERS.include?(service_tier.to_s)
+
           modes = [
             Pricing.normalize_mode(object_value(usage, :speed) || object_value(message, :speed) || request[:speed]),
-            Pricing.normalize_mode(
-              object_value(usage, :service_tier) || object_value(message, :service_tier) || request[:service_tier]
-            )
+            Pricing.normalize_mode(service_tier)
           ]
-          modes << "data_residency" if inference_geo(message: message, request: request, usage: usage).to_s == "us"
+          geo = inference_geo(message: message, request: request, usage: usage).to_s.downcase
+          modes << "data_residency" if DATA_RESIDENCY_GEOS.include?(geo)
           modes = modes.compact.uniq
           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 +165,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!
@@ -56,8 +57,45 @@ module LlmCostTracker
       end
 
       def request_params(args, kwargs)
-        params = args.first.is_a?(Hash) ? args.first : {}
+        params =
+          case args.first
+          when Hash then args.first
+          when nil then {}
+          else args.first.respond_to?(:to_h) ? args.first.to_h : {}
+          end
         params.merge(kwargs).with_indifferent_access
+      rescue StandardError
+        kwargs.to_h.with_indifferent_access
+      end
+
+      def normalize_sdk_args(args, kwargs)
+        return args if args.any? || kwargs.empty?
+
+        [kwargs]
+      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),
+          request: request
+        )
+      end
+
+      def stream_pricing_mode(_request)
+        nil
       end
 
       def object_value(object, *keys)
@@ -69,15 +107,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?
 
@@ -91,12 +120,13 @@ module LlmCostTracker
 
       def patch_targets = []
 
-      def patch_target(constant_name, with:, methods:, optional: false)
+      def patch_target(constant_name, with:, methods:, optional: false, skip_when_methods_missing: false)
         {
           constant_name: constant_name,
           patch: with,
           method_names: Array(methods),
-          optional: optional
+          optional: optional,
+          skip_when_methods_missing: skip_when_methods_missing
         }
       end
 
@@ -106,25 +136,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?
+        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
 
-        indexed_object_value(object, key)
+        object.public_send(key) if object.respond_to?(key)
       end
 
-      def indexed_object_value(object, key)
-        object.try(:[], key)
-      rescue IndexError, NameError, TypeError
-        nil
-      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
@@ -165,6 +187,8 @@ module LlmCostTracker
       end
 
       def missing_methods(target_class, target)
+        return [] if target[:skip_when_methods_missing]
+
         target.fetch(:method_names).filter_map do |method_name|
           next if target_class.method_defined?(method_name) || target_class.private_method_defined?(method_name)