llm_cost_tracker 0.8.0 → 0.9.0

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_call_tags_key_value_index_generator.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeCallTagsKeyValueIndexGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a (key, value) composite index on llm_cost_tracker_call_tags " \
+           "so high-cardinality tag filters use an index lookup instead of a key-only scan."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_call_tags_key_value_index.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_call_tags_key_value_index.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_image_tokens_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeImageTokensGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds image_input_tokens and image_output_tokens columns to llm_cost_tracker_calls."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_image_tokens.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_image_tokens.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -112,7 +112,6 @@ module LlmCostTracker
           quoted_columns = columns.map { |column| connection.quote_column_name(column) }.join(", ")
           quoted_values = columns.map { |column| connection.quote(row.fetch(column)) }.join(", ")
           table = connection.quote_table_name(InboxEntry.table_name)
-
           connection.execute("INSERT INTO #{table} (#{quoted_columns}) VALUES (#{quoted_values})")
         end
       end

data/lib/llm_cost_tracker/ingestion/inline.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative "../ledger/store"
+
+module LlmCostTracker
+  module Ingestion
+    module Inline
+      class << self
+        def save(event)
+          persist(event)
+          event
+        end
+
+        private
+
+        def persist(event)
+          Ledger::Store.insert_many([event])
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -18,6 +18,8 @@ module LlmCostTracker
       FLUSH_TIMEOUT_SECONDS = 10
       class << self
         def ensure_started
+          return unless Ingestion.durable?
+
           thread = mutex.synchronize do
             reset_after_fork!
             unless @thread&.alive?
@@ -34,6 +36,8 @@ module LlmCostTracker
         end
 
         def flush!(timeout: nil, require_lease: false)
+          return true unless Ingestion.durable?
+
           Ingestion.ensure_current_schema!
 
           deadline = Time.now.utc + flush_timeout_seconds(timeout)
@@ -52,6 +56,8 @@ module LlmCostTracker
         end
 
         def shutdown!(timeout: nil, drain: true)
+          return true unless Ingestion.durable?
+
           timeout ||= FLUSH_TIMEOUT_SECONDS
           thread = mutex.synchronize do
             @stop_requested = true
@@ -59,13 +65,15 @@ module LlmCostTracker
             @thread
           end
           wake_thread(thread)
-          thread&.join([timeout, 1].min)
+          thread&.join(timeout)
           drain ? flush!(timeout: timeout, require_lease: true) : true
         rescue StandardError => e
           handle_error(e)
           false
         ensure
-          mutex.synchronize { @thread = nil if @thread.equal?(thread) }
+          mutex.synchronize do
+            @thread = nil if @thread.equal?(thread) && !thread&.alive?
+          end
         end
 
         def reset!

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"
@@ -19,11 +20,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
+
+      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!
@@ -31,7 +38,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 +47,21 @@ module LlmCostTracker
         end
       end
 
+      def durable?
+        LlmCostTracker.configuration.durable_ingestion
+      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 += DURABLE_SCHEMA_GUARDS if durable?
+        guards
+      end
+
       def verify
         unless LlmCostTracker::Call.table_exists?
           return [
@@ -71,7 +93,7 @@ module LlmCostTracker
           provider_response_id: response_id,
           tags: { feature: VERIFY_TAG }
         )
-        LlmCostTracker::Ingestion::Worker.flush!
+        LlmCostTracker::Ingestion::Worker.flush! if durable?
         persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
 
         return capture_success if persisted && notifications.any?
@@ -89,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::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 +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
 
@@ -116,13 +139,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 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

@@ -66,13 +66,15 @@ 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
         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
+          quantity = server_tool_count(server_tool_use, count_key)
           return nil if quantity.zero?
 
           Billing::LineItem.build(
@@ -84,6 +86,14 @@ module LlmCostTracker
           )
         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_creation = object_value(usage, :cache_creation)
           if cache_creation
@@ -108,14 +118,23 @@ 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)
+
           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

data/lib/llm_cost_tracker/integrations/base.rb

@@ -57,8 +57,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.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:)
@@ -68,7 +81,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 +89,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 +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
 
@@ -172,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)