llm_cost_tracker 0.9.0 → 0.11.0

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -16,29 +16,31 @@ module LlmCostTracker
       MAX_IDLE_INTERVAL_SECONDS = 5.0
       LEASE_SECONDS = 10
       FLUSH_TIMEOUT_SECONDS = 10
+      MUTEX = Mutex.new
+
       class << self
         def ensure_started
-          return unless Ingestion.durable?
+          return unless Ingestion.async?
 
-          thread = mutex.synchronize do
+          thread = MUTEX.synchronize do
             reset_after_fork!
-            unless @thread&.alive?
-              @stop_requested = false
-              @generation = @generation.to_i + 1
-              generation = @generation
-              @thread = Thread.new { run(generation) }
-              @thread.name = "llm_cost_tracker_ingestor"
-              @thread.report_on_exception = false
-            end
+            break @thread if @stop_requested || @thread&.alive?
+
+            @generation = @generation.to_i + 1
+            generation = @generation
+            @thread = Thread.new { run(generation) }
+            @thread.name = "llm_cost_tracker_ingestor"
+            @thread.report_on_exception = false
             @thread
           end
           wake_thread(thread)
         end
 
         def flush!(timeout: nil, require_lease: false)
-          return true unless Ingestion.durable?
+          return true unless Ingestion.async?
 
           Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
 
           deadline = Time.now.utc + flush_timeout_seconds(timeout)
           loop do
@@ -56,10 +58,10 @@ module LlmCostTracker
         end
 
         def shutdown!(timeout: nil, drain: true)
-          return true unless Ingestion.durable?
+          return true unless Ingestion.async?
 
           timeout ||= FLUSH_TIMEOUT_SECONDS
-          thread = mutex.synchronize do
+          thread = MUTEX.synchronize do
             @stop_requested = true
             @generation = @generation.to_i + 1
             @thread
@@ -71,14 +73,14 @@ module LlmCostTracker
           handle_error(e)
           false
         ensure
-          mutex.synchronize do
+          MUTEX.synchronize do
             @thread = nil if @thread.equal?(thread) && !thread&.alive?
           end
         end
 
         def reset!
-          thread = mutex.synchronize do
-            @stop_requested = true
+          thread = MUTEX.synchronize do
+            @stop_requested = false
             @generation = @generation.to_i + 1
             thread = @thread
             @thread = nil
@@ -98,6 +100,7 @@ module LlmCostTracker
 
         def ingest_once(require_lease: true)
           Ingestion.ensure_current_schema!
+          MUTEX.synchronize { reset_after_fork! }
           batch = Ingestion::Batch.new(identity: identity)
           return 0 unless batch.claimable?
           return 0 if require_lease && !Ingestion::LeaseClaim.new(identity: identity, seconds: LEASE_SECONDS).acquire
@@ -110,16 +113,12 @@ module LlmCostTracker
 
         private
 
-        def mutex
-          @mutex ||= Mutex.new
-        end
-
         def run(generation)
           idle_interval = IDLE_INTERVAL_SECONDS
           loop do
-            break if mutex.synchronize { @stop_requested || generation != @generation }
+            break if MUTEX.synchronize { @stop_requested || generation != @generation }
 
-            processed = executor_wrap { ingest_once }
+            processed = Rails.application.executor.wrap { ingest_once }
             release_connection!
             if processed.zero?
               sleep(idle_interval)
@@ -134,7 +133,7 @@ module LlmCostTracker
           end
         ensure
           release_connection!
-          mutex.synchronize { @thread = nil if @thread.equal?(Thread.current) }
+          MUTEX.synchronize { @thread = nil if @thread.equal?(Thread.current) }
         end
 
         def reset_after_fork!
@@ -151,19 +150,6 @@ module LlmCostTracker
           nil
         end
 
-        def executor_wrap(&)
-          executor = rails_executor
-          return yield unless executor
-
-          executor.wrap(&)
-        end
-
-        def rails_executor
-          Rails.application.try(:executor)
-        rescue StandardError
-          nil
-        end
-
         def identity
           @identity ||= "pid-#{Process.pid}-#{SecureRandom.hex(6)}"
         end

data/lib/llm_cost_tracker/ingestion.rb

@@ -2,11 +2,10 @@
 
 require "securerandom"
 
-require_relative "doctor/check"
 require_relative "errors"
 require_relative "ledger"
-require_relative "ingestion/inline"
 require_relative "ingestion/lease_claim"
+require_relative "ingestion/pool"
 require_relative "ingestion/inbox"
 require_relative "ingestion/batch"
 require_relative "ingestion/worker"
@@ -28,7 +27,7 @@ module LlmCostTracker
 
       ROLLUPS_SCHEMA_GUARD = ["llm_cost_tracker_call_rollups", Ledger::Schema::CallRollups].freeze
 
-      DURABLE_SCHEMA_GUARDS = [
+      ASYNC_SCHEMA_GUARDS = [
         ["llm_cost_tracker_ingestion_inbox_entries", Ledger::Schema::IngestionInboxEntries],
         ["llm_cost_tracker_ingestion_leases",        Ledger::Schema::IngestionLeases]
       ].freeze
@@ -47,8 +46,8 @@ module LlmCostTracker
         end
       end
 
-      def durable?
-        LlmCostTracker.configuration.durable_ingestion
+      def async?
+        LlmCostTracker.configuration.ingestion == :async
       end
 
       def cache_rollups?
@@ -58,7 +57,7 @@ module LlmCostTracker
       def guards_for_current_config
         guards = CORE_SCHEMA_GUARDS.dup
         guards << ROLLUPS_SCHEMA_GUARD if cache_rollups?
-        guards += DURABLE_SCHEMA_GUARDS if durable?
+        guards += ASYNC_SCHEMA_GUARDS if async?
         guards
       end
 
@@ -93,7 +92,7 @@ module LlmCostTracker
           provider_response_id: response_id,
           tags: { feature: VERIFY_TAG }
         )
-        LlmCostTracker::Ingestion::Worker.flush! if durable?
+        LlmCostTracker::Ingestion::Worker.flush! if async?
         persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
 
         return capture_success if persisted && notifications.any?
@@ -122,7 +121,7 @@ module LlmCostTracker
       end
 
       def capture_success
-        path = durable? ? "durable inbox" : "inline writer"
+        path = async? ? "async inbox" : "inline writer"
         LlmCostTracker::Doctor::Check.new(
           :ok,
           "active_record capture",
@@ -149,7 +148,7 @@ module LlmCostTracker
       end
 
       def cleanup_verification_inbox(event:, response_id:)
-        return unless durable? && LlmCostTracker::Ingestion::InboxEntry.table_exists?
+        return unless async? && LlmCostTracker::Ingestion::InboxEntry.table_exists?
 
         if event
           LlmCostTracker::Ingestion::InboxEntry.where(event_id: event.event_id).delete_all

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -2,6 +2,8 @@
 
 require_relative "base"
 require_relative "../billing/line_item"
+require_relative "../providers/anthropic/server_tools"
+require_relative "../providers/anthropic/tier_classification"
 
 module LlmCostTracker
   module Integrations
@@ -45,10 +47,10 @@ module LlmCostTracker
             next if input_tokens.nil? && output_tokens.nil?
 
             LlmCostTracker::Tracker.record(
-              capture: UsageCapture.build(
+              event: Event.build(
                 provider: "anthropic",
                 model: object_value(message, :model) || request[:model],
-                pricing_mode: pricing_mode(message: message, request: request, usage: usage),
+                pricing_mode: pricing_mode(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),
@@ -63,35 +65,18 @@ module LlmCostTracker
           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)
+          Providers::Anthropic::ServerTools::LINE_ITEMS.filter_map do |component_key, count_key|
+            quantity = object_value(server_tool_use, count_key).to_i
+            next if quantity.zero?
 
-          server_tool_use.to_h[count_key].to_i
+            Billing::LineItem.build(
+              component_key: component_key,
+              quantity: quantity,
+              cost_status: Billing::CostStatus::UNKNOWN,
+              pricing_basis: :provider_usage,
+              provider_field: "usage.server_tool_use.#{count_key}"
+            )
+          end
         end
 
         def token_usage(usage:, input_tokens:, output_tokens:)
@@ -118,65 +103,58 @@ module LlmCostTracker
           )
         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)
+        def pricing_mode(request:, usage:)
+          service_tier = object_value(usage, :service_tier) || request[:service_tier]
+          tier = Providers::Anthropic::TierClassification
+          service_tier = nil if tier.standard_equivalent_tier?(service_tier)
 
           modes = [
-            Pricing.normalize_mode(object_value(usage, :speed) || object_value(message, :speed) || request[:speed]),
-            Pricing.normalize_mode(service_tier)
+            Pricing::Mode.normalize(object_value(usage, :speed) || request[:speed]),
+            Pricing::Mode.normalize(service_tier)
           ]
-          geo = inference_geo(message: message, request: request, usage: usage).to_s.downcase
-          modes << "data_residency" if DATA_RESIDENCY_GEOS.include?(geo)
+          geo = inference_geo(request: request, usage: usage).to_s.downcase
+          modes << "data_residency" if tier.data_residency_geo?(geo)
           modes = modes.compact.uniq
           modes.empty? ? nil : modes.join("_")
         end
 
         def stream_pricing_mode(request)
-          pricing_mode(message: nil, request: request || {}, usage: nil)
+          pricing_mode(request: request || {}, usage: nil)
         end
 
-        def inference_geo(message:, request:, usage:)
-          object_value(usage, :inference_geo) ||
-            object_value(message, :inference_geo) ||
-            request[:inference_geo]
+        def inference_geo(request:, usage:)
+          object_value(usage, :inference_geo) || request[:inference_geo]
+        end
+
+        def wrap_stream_call(args, kwargs)
+          request = request_params(args, kwargs)
+          enforce_budget!(request: request)
+          collector = stream_collector(request)
+          stream = yield
+          track_stream(stream, collector: collector)
+        end
+
+        def wrap_blocking_call(args, kwargs)
+          request = request_params(args, kwargs)
+          enforce_budget!(request: request)
+          started_at = LlmCostTracker::Timing.now_monotonic
+          message = yield
+          record_message(message, request: request, latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at))
+          message
         end
       end
 
       module MessagesPatch
         def create(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.enforce_budget!
-          started_at = LlmCostTracker::Timing.now_monotonic
-          message = super
-          LlmCostTracker::Integrations::Anthropic.record_message(
-            message,
-            request: LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs),
-            latency_ms: LlmCostTracker::Integrations::Anthropic.elapsed_ms(started_at)
-          )
-          message
+          LlmCostTracker::Integrations::Anthropic.wrap_blocking_call(args, kwargs) { super }
         end
 
         def stream(*args, **kwargs)
-          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
-          LlmCostTracker::Integrations::Anthropic.enforce_budget!
-          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
-          stream = super
-          LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
+          LlmCostTracker::Integrations::Anthropic.wrap_stream_call(args, kwargs) { super }
         end
 
         def stream_raw(*args, **kwargs)
-          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
-          LlmCostTracker::Integrations::Anthropic.enforce_budget!
-          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
-          stream = super
-          LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
+          LlmCostTracker::Integrations::Anthropic.wrap_stream_call(args, kwargs) { super }
         end
       end
     end

data/lib/llm_cost_tracker/integrations/base.rb

@@ -3,6 +3,7 @@
 require "active_support/core_ext/hash/indifferent_access"
 require "active_support/core_ext/string/inflections"
 
+require_relative "../doctor/check"
 require_relative "../logging"
 require_relative "../timing"
 require_relative "../capture/stream_collector"
@@ -11,7 +12,7 @@ require_relative "../capture/stream_tracker"
 module LlmCostTracker
   module Integrations
     module Base
-      Result = Data.define(:name, :status, :message)
+      Result = LlmCostTracker::Doctor::Check
 
       def active?
         LlmCostTracker.configuration.instrumented?(integration_name)
@@ -26,26 +27,28 @@ module LlmCostTracker
       end
 
       def status
-        name = integration_name
+        name = integration_name.to_s
         problems = version_problems + target_problems
         if problems.any?
-          return Result.new(name, :warn, "#{name} integration cannot be installed: #{problems.join('; ')}")
+          return Result.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
         end
 
         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
+        return Result.new(:ok, name, "#{name} integration installed") if installed
 
-        Result.new(name, :warn, "#{name} integration is enabled but not installed")
+        Result.new(:warn, name, "#{name} integration is enabled but not installed")
       end
 
-      def elapsed_ms(started_at)
-        Timing.elapsed_ms(started_at)
-      end
+      def enforce_budget!(request:)
+        return unless active?
 
-      def enforce_budget!
-        LlmCostTracker::Tracker.enforce_budget! if active?
+        LlmCostTracker::Tracker.enforce_budget!(
+          provider: integration_name.to_s,
+          model: request[:model],
+          request: request
+        )
       end
 
       def record_safely
@@ -61,7 +64,7 @@ module LlmCostTracker
           case args.first
           when Hash then args.first
           when nil then {}
-          else args.first.respond_to?(:to_h) ? args.first.to_h : {}
+          else args.first.to_h
           end
         params.merge(kwargs).with_indifferent_access
       rescue StandardError