llm_cost_tracker 0.8.0 → 0.10.0

data/lib/llm_cost_tracker/ingestion.rb

@@ -2,10 +2,10 @@
 
 require "securerandom"
 
-require_relative "doctor/check"
 require_relative "errors"
 require_relative "ledger"
 require_relative "ingestion/lease_claim"
+require_relative "ingestion/pool"
 require_relative "ingestion/inbox"
 require_relative "ingestion/batch"
 require_relative "ingestion/worker"
@@ -19,11 +19,17 @@ module LlmCostTracker
         "llm_cost_tracker_ingestion_"
       end
 
-      WRITE_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],
-        ["llm_cost_tracker_call_rollups",     Ledger::Schema::CallRollups]
+      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
+
+      ASYNC_SCHEMA_GUARDS = [
+        ["llm_cost_tracker_ingestion_inbox_entries", Ledger::Schema::IngestionInboxEntries],
+        ["llm_cost_tracker_ingestion_leases",        Ledger::Schema::IngestionLeases]
       ].freeze
 
       def ensure_current_schema!
@@ -31,7 +37,7 @@ module LlmCostTracker
           raise Error, "llm_cost_tracker_calls table is missing; run install generator and migrate"
         end
 
-        WRITE_SCHEMA_GUARDS.each do |table_name, schema_module|
+        guards_for_current_config.each do |table_name, schema_module|
           errors = schema_module.current_schema_errors
           next if errors.empty?
 
@@ -40,6 +46,21 @@ module LlmCostTracker
         end
       end
 
+      def async?
+        LlmCostTracker.configuration.ingestion == :async
+      end
+
+      def cache_rollups?
+        LlmCostTracker.configuration.cache_rollups
+      end
+
+      def guards_for_current_config
+        guards = CORE_SCHEMA_GUARDS.dup
+        guards << ROLLUPS_SCHEMA_GUARD if cache_rollups?
+        guards += ASYNC_SCHEMA_GUARDS if async?
+        guards
+      end
+
       def verify
         unless LlmCostTracker::Call.table_exists?
           return [
@@ -71,7 +92,7 @@ module LlmCostTracker
           provider_response_id: response_id,
           tags: { feature: VERIFY_TAG }
         )
-        LlmCostTracker::Ingestion::Worker.flush!
+        LlmCostTracker::Ingestion::Worker.flush! if async?
         persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
 
         return capture_success if persisted && notifications.any?
@@ -89,7 +110,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::InboxEntry.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
 
@@ -100,10 +121,11 @@ module LlmCostTracker
       end
 
       def capture_success
+        path = async? ? "async 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
 
@@ -116,13 +138,28 @@ module LlmCostTracker
 
       def cleanup_verification_call(response_id)
         relation = LlmCostTracker::Call.where(provider_response_id: response_id)
-        rows = relation.pluck(:id, :tracked_at, :total_cost, :pricing_snapshot)
+        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 async? && 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

@@ -2,6 +2,7 @@
 
 require_relative "base"
 require_relative "../billing/line_item"
+require_relative "../providers/anthropic/tier_classification"
 
 module LlmCostTracker
   module Integrations
@@ -45,10 +46,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),
@@ -66,6 +67,8 @@ module LlmCostTracker
           [
             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
@@ -108,56 +111,58 @@ module LlmCostTracker
           )
         end
 
-        def pricing_mode(message:, request:, usage:)
+        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(
-              object_value(usage, :service_tier) || object_value(message, :service_tier) || request[:service_tier]
-            )
+            Pricing.normalize_mode(object_value(usage, :speed) || request[:speed]),
+            Pricing.normalize_mode(service_tier)
           ]
-          modes << "data_residency" if inference_geo(message: message, request: request, usage: usage).to_s == "us"
+          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(request:, usage:)
+          object_value(usage, :inference_geo) || request[:inference_geo]
         end
 
-        def inference_geo(message:, request:, usage:)
-          object_value(usage, :inference_geo) ||
-            object_value(message, :inference_geo) ||
-            request[:inference_geo]
+        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
       end
 
       module MessagesPatch
         def create(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
+          LlmCostTracker::Integrations::Anthropic.enforce_budget!(request: request)
           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)
+            request: request,
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at)
           )
           message
         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
@@ -57,8 +60,21 @@ 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.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:)
@@ -68,7 +84,7 @@ module LlmCostTracker
           stream: stream,
           collector: collector,
           active: -> { active? },
-          finish: ->(errored:) { record_safely { collector.finish!(errored: errored) } }
+          finish: ->(errored) { record_safely { collector.finish!(errored: errored) } }
         ).wrap
       end
 
@@ -76,7 +92,8 @@ module LlmCostTracker
         LlmCostTracker::Capture::StreamCollector.new(
           provider: integration_name.to_s,
           model: request[:model],
-          pricing_mode: stream_pricing_mode(request)
+          pricing_mode: stream_pricing_mode(request),
+          request: request
         )
       end
 
@@ -106,12 +123,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
 
@@ -172,6 +190,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)