llm_cost_tracker 0.13.0 → 0.14.0

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_indexes.rb.erb

@@ -0,0 +1,40 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerIndexes < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  CALLS = :llm_cost_tracker_calls
+  INBOX = :llm_cost_tracker_ingestion_inbox_entries
+  UNUSED_INBOX_INDEX = %i[locked_at id].freeze
+  UNPRICED_INDEX = :index_llm_cost_tracker_calls_on_unpriced
+
+  def up
+    add_index CALLS, :id, name: UNPRICED_INDEX, where: "total_cost IS NULL", **add_index_options
+    return unless table_exists?(INBOX)
+
+    remove_index INBOX, column: UNUSED_INBOX_INDEX, **remove_index_options
+  end
+
+  def down
+    add_index INBOX, UNUSED_INBOX_INDEX, **add_index_options if table_exists?(INBOX)
+    remove_index CALLS, name: UNPRICED_INDEX, **remove_index_options
+  end
+
+  private
+
+  def add_index_options
+    return { if_not_exists: true } unless postgresql?
+
+    { if_not_exists: true, algorithm: :concurrently }
+  end
+
+  def remove_index_options
+    return { if_exists: true } unless postgresql?
+
+    { if_exists: true, algorithm: :concurrently }
+  end
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/upgrade_per_tag_budgets.rb.erb

@@ -0,0 +1,41 @@
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class UpgradeLlmCostTrackerPerTagBudgets < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  TABLE = :llm_cost_tracker_call_tags
+  INDEX = %i[key value tracked_at].freeze
+  REPLACED_INDEX = %i[key value].freeze
+
+  def up
+    add_column TABLE, :total_cost, :decimal, precision: 20, scale: 8 unless column_exists?(TABLE, :total_cost)
+    add_column TABLE, :tracked_at, :datetime unless column_exists?(TABLE, :tracked_at)
+    add_index TABLE, INDEX, **add_index_options
+    remove_index TABLE, column: REPLACED_INDEX, **remove_index_options
+  end
+
+  def down
+    add_index TABLE, REPLACED_INDEX, **add_index_options
+    remove_index TABLE, column: INDEX, **remove_index_options
+    remove_column TABLE, :tracked_at, if_exists: true
+    remove_column TABLE, :total_cost, if_exists: true
+  end
+
+  private
+
+  def add_index_options
+    return { if_not_exists: true, length: { value: 191 } } unless postgresql?
+
+    { if_not_exists: true, algorithm: :concurrently }
+  end
+
+  def remove_index_options
+    return { if_exists: true } unless postgresql?
+
+    { if_exists: true, algorithm: :concurrently }
+  end
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_indexes_generator.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeIndexesGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Adds a partial index over unpriced calls and drops the unused ingestion inbox lock index" \
+           "scanning the whole ledger."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_indexes.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_indexes.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_per_tag_budgets_generator.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradePerTagBudgetsGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Carries each call's cost and time onto its tag rows so per-tag budgets read " \
+           "llm_cost_tracker_call_tags without a join. Required when config.budgets.per_tag is set."
+
+      def create_migration_file
+        migration_template(
+          "upgrade_per_tag_budgets.rb.erb",
+          "db/migrate/upgrade_llm_cost_tracker_per_tag_budgets.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "inbox"
+require_relative "../budget"
 require_relative "../ledger/store"
 
 module LlmCostTracker
@@ -121,7 +122,8 @@ module LlmCostTracker
           Ledger::Store.persist_records(events)
           Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
-        Ledger::Rollups.increment_safely!(events) if Ingestion.cache_rollups?
+        Ledger::Rollups.increment_safely!(events)
+        Budget.notify_persisted_safely!(events)
       rescue ActiveRecord::RecordNotUnique
         raise unless retry_on_conflict
 

data/lib/llm_cost_tracker/ingestion/pool.rb

@@ -12,7 +12,14 @@ module LlmCostTracker
         delegate :with_connection, to: :pool
 
         def pool
-          @pool || MUTEX.synchronize { @pool ||= connect! }
+          return @pool if @pool && @pid == Process.pid
+
+          MUTEX.synchronize do
+            next @pool if @pool && @pid == Process.pid
+
+            @pid = Process.pid
+            @pool = connect!
+          end
         end
 
         private
@@ -27,7 +34,7 @@ module LlmCostTracker
         end
 
         def pool_size
-          configured = LlmCostTracker.configuration.ingestion_pool_size.to_i
+          configured = LlmCostTracker.configuration.ingestion.pool_size.to_i
           configured.positive? ? configured : DEFAULT_POOL_SIZE
         end
       end

data/lib/llm_cost_tracker/ingestion.rb

@@ -37,16 +37,12 @@ module LlmCostTracker
       end
 
       def async?
-        LlmCostTracker.configuration.ingestion == :async
-      end
-
-      def cache_rollups?
-        LlmCostTracker.configuration.cache_rollups
+        LlmCostTracker.configuration.ingestion.mode == :async
       end
 
       def guards_for_current_config
         guards = Ledger::Schema::CORE_SCHEMAS.dup
-        guards << Ledger::Schema::CACHE_ROLLUPS_SCHEMA if cache_rollups?
+        guards << Ledger::Schema::CACHE_ROLLUPS_SCHEMA if LlmCostTracker.configuration.budgets.totals_source == :cache
         guards += Ledger::Schema::ASYNC_SCHEMAS if async?
         guards
       end
@@ -132,7 +128,7 @@ module LlmCostTracker
         return if records.empty?
 
         relation.delete_all
-        LlmCostTracker::Ledger::Rollups.decrement!(records) if cache_rollups?
+        LlmCostTracker::Ledger::Rollups.decrement!(records)
       end
 
       def cleanup_verification_inbox(event:, response_id:)

data/lib/llm_cost_tracker/integrations/openai/batch_capture.rb

@@ -13,32 +13,30 @@ module LlmCostTracker
         class << self
           def maybe_capture(batch, resource:)
             return unless Openai.active?
-            return unless batch.respond_to?(:status) && batch.status.to_s == "completed"
-
-            output_file_id = batch.respond_to?(:output_file_id) ? batch.output_file_id : nil
-            return unless output_file_id
-
-            batch_id = batch.respond_to?(:id) ? batch.id : nil
-            return unless batch_id && claim(batch_id)
+            return unless batch.status.to_s == "completed"
+            return unless batch.output_file_id && batch.id
+            return if captured?(batch.id)
 
             client = resource.instance_variable_get(:@client)
             host = Openai.client_host_for(resource)
             Openai.record_safely do
-              io = client.files.content(output_file_id)
+              io = client.files.content(batch.output_file_id)
               capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+              mark_captured(batch.id)
             end
           end
 
           private
 
-          def claim(batch_id)
+          def captured?(batch_id)
+            MUTEX.synchronize { @dedup&.include?(batch_id) || false }
+          end
+
+          def mark_captured(batch_id)
             MUTEX.synchronize do
               @dedup ||= Set.new
-              next false if @dedup.include?(batch_id)
-
               @dedup.clear if @dedup.size >= DEDUP_LIMIT
               @dedup.add(batch_id)
-              true
             end
           end
 

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

@@ -45,7 +45,7 @@ module LlmCostTracker
 
         def recorded_sql(period, start)
           calls = "COALESCE(#{sum_sql(LlmCostTracker::Call.between(start, time))}, 0)"
-          return calls unless LlmCostTracker.configuration.cache_rollups
+          return calls unless Rollups.cache_active?
 
           rollup = "COALESCE(#{sum_sql(rollup_scope(period))}, 0)"
           "GREATEST(#{rollup}, #{calls})"

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -8,7 +8,17 @@ module LlmCostTracker
   module Ledger
     module Rollups
       class << self
+        def cache_active?
+          return false unless LlmCostTracker.configuration.budgets.totals_source == :cache
+          return true if LlmCostTracker::CallRollup.table_exists?
+
+          warn_missing_table
+          false
+        end
+
         def increment!(events)
+          return unless cache_active?
+
           events = Array(events).select(&:total_cost)
           return if events.empty?
 
@@ -18,7 +28,9 @@ module LlmCostTracker
         ROLLUP_INCREMENT_ATTEMPTS = 3
         ROLLUP_INCREMENT_BASE_DELAY_SECONDS = 0.05
         REBUILD_INSERT_SLICE = 1_000
-        private_constant :ROLLUP_INCREMENT_ATTEMPTS, :ROLLUP_INCREMENT_BASE_DELAY_SECONDS, :REBUILD_INSERT_SLICE
+        private_constant :ROLLUP_INCREMENT_ATTEMPTS,
+                         :ROLLUP_INCREMENT_BASE_DELAY_SECONDS,
+                         :REBUILD_INSERT_SLICE
 
         def increment_safely!(events)
           attempt = 0
@@ -28,19 +40,26 @@ module LlmCostTracker
           rescue StandardError => e
             raise if LlmCostTracker::Call.connection.open_transactions.positive?
 
-            if attempt < ROLLUP_INCREMENT_ATTEMPTS
+            if attempt < ROLLUP_INCREMENT_ATTEMPTS && rolled_back?(e)
               sleep(ROLLUP_INCREMENT_BASE_DELAY_SECONDS * (2**(attempt - 1)))
               retry
             end
 
             LlmCostTracker::Logging.warn(
-              "Rollup increment failed for #{events.size} events after #{attempt} attempts: " \
-              "#{e.class}: #{e.message}"
+              "Rollup increment failed for #{events.size} events after #{attempt} attempt(s): " \
+              "#{e.class}: #{e.message}. Budget reads fall back to the calls ledger; " \
+              "run bin/rails llm_cost_tracker:rebuild_rollups to resync the cache."
             )
           end
         end
 
+        def rolled_back?(error)
+          error.is_a?(ActiveRecord::Deadlocked) || error.is_a?(ActiveRecord::LockWaitTimeout)
+        end
+
         def decrement!(records)
+          return unless cache_active?
+
           buckets = bucket_totals(records)
           return if buckets.empty?
 
@@ -60,6 +79,17 @@ module LlmCostTracker
 
         private
 
+        def warn_missing_table
+          return if @missing_table_warned
+
+          @missing_table_warned = true
+          Logging.warn(
+            "config.budgets.totals_source = :cache is set but llm_cost_tracker_call_rollups is missing; " \
+            "budget and dashboard totals fall back to aggregating llm_cost_tracker_calls. " \
+            "Run the call_rollups generator and migrate, or set config.budgets.totals_source = :ledger."
+          )
+        end
+
         def accumulate_priced_calls
           accumulated = Hash.new { |totals, key| totals[key] = BigDecimal("0") }
           priced_calls.in_batches do |batch|

data/lib/llm_cost_tracker/ledger/store.rb

@@ -4,6 +4,7 @@ require "json"
 
 require_relative "../pricing"
 require_relative "rollups"
+require_relative "../budget/per_tag"
 require_relative "tags/encoding"
 
 module LlmCostTracker
@@ -15,7 +16,7 @@ module LlmCostTracker
           return if events.empty?
 
           persist_records(events)
-          Ledger::Rollups.increment_safely!(events) if LlmCostTracker.configuration.cache_rollups
+          Ledger::Rollups.increment_safely!(events)
         end
 
         def persist_records(events)
@@ -121,7 +122,7 @@ module LlmCostTracker
                 llm_cost_tracker_call_id: call_ids.fetch(event.event_id),
                 key: key.to_s,
                 value: Tags::Encoding.encode(value)
-              }
+              }.merge(budget_columns_for(event))
             end
           end
           return if rows.empty?
@@ -129,6 +130,12 @@ module LlmCostTracker
           LlmCostTracker::CallTag.insert_all!(rows, record_timestamps: false, returning: false)
         end
 
+        def budget_columns_for(event)
+          return {} unless LlmCostTracker::Budget::PerTag.columns?
+
+          { total_cost: event.cost&.total, tracked_at: event.tracked_at }
+        end
+
         def stored_details(details)
           (details || {}).transform_keys(&:to_s).transform_values { |value| Tags::Encoding.normalize_value(value) }
         end

data/lib/llm_cost_tracker/ledger/tags/encoding.rb

@@ -6,12 +6,22 @@ module LlmCostTracker
   module Ledger
     module Tags
       module Encoding
+        INDEXABLE_BYTES = 2_048
+
         def self.encode(value)
-          case value
-          when Hash then JSON.generate(normalize_hash(value))
-          when Array then JSON.generate(normalize_array(value))
-          else value.to_s
-          end
+          encoded = case value
+                    when Hash then JSON.generate(normalize_hash(value))
+                    when Array then JSON.generate(normalize_array(value))
+                    else value.to_s
+                    end
+          truncate(encoded)
+        end
+
+        def self.truncate(string)
+          limit = [LlmCostTracker.configuration.tags.max_value_bytesize.to_i, INDEXABLE_BYTES].min
+          return string if string.bytesize <= limit
+
+          string.byteslice(0, limit).encode("UTF-8", invalid: :replace, undef: :replace)
         end
 
         def self.normalize_hash(hash)

data/lib/llm_cost_tracker/logging.rb

@@ -3,15 +3,13 @@
 module LlmCostTracker
   module Logging
     class << self
-      def debug(message) = tagged.debug(message)
+      def debug(message) = Rails.logger.debug(prefixed(message))
 
-      def warn(message) = tagged.warn(message)
+      def warn(message) = Rails.logger.warn(prefixed(message))
 
       private
 
-      def tagged
-        Rails.logger.tagged(LlmCostTracker.name)
-      end
+      def prefixed(message) = "[#{LlmCostTracker.name}] #{message}"
     end
   end
 end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -94,7 +94,7 @@ module LlmCostTracker
       end
 
       def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
-        return nil unless LlmCostTracker.configuration.auto_enable_stream_usage
+        return nil unless LlmCostTracker.configuration.capture.request_stream_usage
         return nil unless parser&.auto_enable_stream_usage?(request_url)
 
         stream_options = request_parsed["stream_options"]