llm_cost_tracker 0.14.0 → 0.14.2

data/lib/llm_cost_tracker/integrations.rb

@@ -6,19 +6,13 @@ require_relative "errors"
 
 module LlmCostTracker
   module Integrations
-    autoload :Base, "llm_cost_tracker/integrations/base"
-
     Dir.glob(File.join(__dir__, "integrations", "*.rb")).each do |path|
       basename = File.basename(path, ".rb")
-      next if basename == "base"
-
       autoload basename.camelize.to_sym, "llm_cost_tracker/integrations/#{basename}"
     end
 
     def self.install!(names = LlmCostTracker.configuration.instrumented_integrations)
-      normalized = normalize(names)
-      warn_double_instrumentation(normalized)
-      normalized.each do |name|
+      normalize(names).each do |name|
         integration = fetch(name)
         next integration.install if integration
 
@@ -41,19 +35,6 @@ module LlmCostTracker
       Array(names).flatten.uniq
     end
 
-    def self.warn_double_instrumentation(names)
-      return unless names.include?(:ruby_llm)
-
-      overlapping = names - [:ruby_llm]
-      return if overlapping.empty?
-
-      Logging.warn(
-        ":ruby_llm is enabled together with #{overlapping.map(&:inspect).join(', ')}. " \
-        "RubyLLM uses HTTP underneath, so calls routed to those providers may be recorded twice " \
-        "(once via the SDK patch, once via the Faraday parser). Pick one path per provider."
-      )
-    end
-
     def self.fetch(name)
       const_name = name.to_s.camelize
       return nil unless const_name.match?(/\A[A-Z]\w*\z/)

data/lib/llm_cost_tracker/ledger/isolation.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require_relative "../errors"
+
+module LlmCostTracker
+  module Ledger
+    module Isolation
+      class << self
+        def guard(model = LlmCostTracker::Call, &)
+          connection = model.connection
+          nested = connection.transaction_open?
+          return yield unless nested
+
+          model.transaction(requires_new: true, &)
+        rescue ActiveRecord::TransactionRollbackError => e
+          raise unless nested && connection.savepoint_errors_invalidate_transactions?
+
+          raise TransactionAbortedError, e
+        end
+      end
+    end
+  end
+end

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

@@ -2,6 +2,7 @@
 
 require "bigdecimal/util"
 
+require_relative "../isolation"
 require_relative "../period"
 
 module LlmCostTracker
@@ -22,7 +23,7 @@ module LlmCostTracker
 
           values = periods.to_h { |period| [period, BigDecimal("0")] }
           period_by_name = periods.to_h { |period| [period.to_s, period] }
-          LlmCostTracker::Call.find_by_sql(union_sql).each do |row|
+          Isolation.guard { LlmCostTracker::Call.find_by_sql(union_sql) }.each do |row|
             values[period_by_name.fetch(row.period_key)] = row.total_cost.to_d
           end
           values
@@ -38,25 +39,29 @@ module LlmCostTracker
 
         def period_select(period)
           start = Period.range_start(period, time)
-          components = ["(#{recorded_sql(period, start)})"]
+          components = ["(#{recorded_sql(start)})"]
           components << "(#{pending_sql(start)})" if Ingestion.async?
           "SELECT #{quote(period.to_s)} AS period_key, #{components.join(' + ')} AS total_cost"
         end
 
-        def recorded_sql(period, start)
-          calls = "COALESCE(#{sum_sql(LlmCostTracker::Call.between(start, time))}, 0)"
-          return calls unless Rollups.cache_active?
+        def recorded_sql(start)
+          today = Period.range_start(:day, time)
+          return calls_sql(start) unless Rollups.cache_active? && start < today
 
-          rollup = "COALESCE(#{sum_sql(rollup_scope(period))}, 0)"
-          "GREATEST(#{rollup}, #{calls})"
+          "#{completed_days_sql(start, today)} + #{calls_sql(today)}"
         end
 
-        def pending_sql(start)
-          "COALESCE(#{sum_sql(Ingestion::InboxEntry.pending.where(tracked_at: start..time))}, 0)"
+        def calls_sql(start)
+          "COALESCE(#{sum_sql(LlmCostTracker::Call.between(start, time))}, 0)"
+        end
+
+        def completed_days_sql(start, today)
+          days = LlmCostTracker::CallRollup.where(period: "day", period_start: start.to_date...today.to_date)
+          "COALESCE(#{sum_sql(days)}, 0)"
         end
 
-        def rollup_scope(period)
-          LlmCostTracker::CallRollup.where(period: period.to_s, period_start: Period.bucket(period, time))
+        def pending_sql(start)
+          "COALESCE(#{sum_sql(Ingestion::InboxEntry.pending.where(tracked_at: start..time))}, 0)"
         end
 
         def sum_sql(scope)

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -2,15 +2,23 @@
 
 require "bigdecimal/util"
 
+require_relative "isolation"
 require_relative "period"
 
 module LlmCostTracker
   module Ledger
     module Rollups
+      SOURCE_COLUMNS = %i[tracked_at total_cost pricing_snapshot provider].freeze
+
       class << self
         def cache_active?
           return false unless LlmCostTracker.configuration.budgets.totals_source == :cache
-          return true if LlmCostTracker::CallRollup.table_exists?
+
+          rollup = LlmCostTracker::CallRollup
+          return true if rollup.table_exists?
+
+          rollup.connection.schema_cache.clear_data_source_cache!(rollup.table_name)
+          return true if rollup.table_exists?
 
           warn_missing_table
           false
@@ -19,10 +27,8 @@ module LlmCostTracker
         def increment!(events)
           return unless cache_active?
 
-          events = Array(events).select(&:total_cost)
-          return if events.empty?
-
-          LlmCostTracker::CallRollup.increment_all(period_rows_for_events(events))
+          rows = rows_from_buckets(bucket_totals(Array(events)))
+          LlmCostTracker::CallRollup.increment_all(rows) if rows.any?
         end
 
         ROLLUP_INCREMENT_ATTEMPTS = 3
@@ -33,22 +39,29 @@ module LlmCostTracker
                          :REBUILD_INSERT_SLICE
 
         def increment_safely!(events)
+          return unless LlmCostTracker.configuration.budgets.totals_source == :cache
+
+          LlmCostTracker::Call.current_transaction.after_commit { increment_with_retries(events) }
+        end
+
+        def increment_with_retries(events)
+          retryable = !LlmCostTracker::Call.connection.transaction_open?
           attempt = 0
           begin
             attempt += 1
-            increment!(events)
+            Isolation.guard { increment!(events) }
+          rescue LlmCostTracker::Error
+            raise
           rescue StandardError => e
-            raise if LlmCostTracker::Call.connection.open_transactions.positive?
-
-            if attempt < ROLLUP_INCREMENT_ATTEMPTS && rolled_back?(e)
+            if retryable && 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} 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."
+              "#{e.class}: #{e.message}. Monthly budget totals under-count these calls until you run " \
+              "bin/rails llm_cost_tracker:rebuild_rollups."
             )
           end
         end
@@ -99,12 +112,7 @@ module LlmCostTracker
         end
 
         def priced_calls
-          LlmCostTracker::Call.where.not(total_cost: nil)
-                              .select(:id, :total_cost, :pricing_snapshot, :provider, :tracked_at)
-        end
-
-        def period_rows_for_events(events)
-          rows_from_buckets(bucket_totals(events))
+          LlmCostTracker::Call.where.not(total_cost: nil).select(:id, *SOURCE_COLUMNS)
         end
 
         def rows_from_buckets(buckets)

data/lib/llm_cost_tracker/ledger/schema/adapter.rb

@@ -8,8 +8,6 @@ module LlmCostTracker
       module Adapter
         MYSQL_ADAPTERS = %w[
           ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter
-          ActiveRecord::ConnectionAdapters::Mysql2Adapter
-          ActiveRecord::ConnectionAdapters::TrilogyAdapter
         ].freeze
         POSTGRESQL_ADAPTERS = %w[
           ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
@@ -36,11 +34,16 @@ module LlmCostTracker
           MYSQL_PERIOD_FORMATS = { day: "%Y-%m-%d", month: "%Y-%m" }.freeze
           private_constant :PG_PERIOD_FORMATS, :MYSQL_PERIOD_FORMATS
 
-          def period_bucket_sql(connection, period, column)
+          def period_bucket_sql(connection, period, column, time_zone: nil)
             period = period.to_sym
+            zone = time_zone&.tzinfo&.name
             if postgresql?(connection)
+              if zone && postgresql_zone?(connection, zone)
+                column = "(#{column}::timestamp AT TIME ZONE 'UTC') AT TIME ZONE '#{zone}'"
+              end
               "TO_CHAR(DATE_TRUNC('#{period}', #{column}), '#{PG_PERIOD_FORMATS.fetch(period)}')"
             elsif mysql?(connection)
+              column = "COALESCE(CONVERT_TZ(#{column}, '+00:00', '#{zone}'), #{column})" if zone
               "DATE_FORMAT(#{column}, '#{MYSQL_PERIOD_FORMATS.fetch(period)}')"
             else
               ensure_supported!(connection)
@@ -51,6 +54,14 @@ module LlmCostTracker
 
           private
 
+          def postgresql_zone?(connection, zone)
+            @postgresql_zones ||= {}
+            @postgresql_zones.fetch(zone) do
+              sql = "SELECT 1 FROM pg_timezone_names WHERE name = #{connection.quote(zone)}"
+              @postgresql_zones[zone] = !connection.select_value(sql).nil?
+            end
+          end
+
           def adapter_instance?(value, class_names)
             class_names.any? do |class_name|
               adapter_class = class_name.safe_constantize

data/lib/llm_cost_tracker/ledger/schema/base.rb

@@ -21,10 +21,11 @@ module LlmCostTracker
         end
 
         def current_schema_errors
-          connection = model.connection
-          Adapter.ensure_supported!(connection)
-          table_name = model.table_name
-          return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+          Adapter.ensure_supported!(model.connection)
+          unless model.table_exists?
+            model.connection.schema_cache.clear_data_source_cache!(model.table_name)
+            return ["#{model.table_name} table is missing"] unless model.table_exists?
+          end
 
           columns_hash = model.columns_hash
           cache = @schema_capabilities

data/lib/llm_cost_tracker/ledger/storable.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ledger
+    module Storable
+      def self.clean(value)
+        case value
+        when Hash then value.to_h { |key, nested| [clean(key), clean(nested)] }
+        when Array then value.map { |nested| clean(nested) }
+        when String then value.dup.force_encoding(Encoding::UTF_8).scrub.delete("\u0000")
+        else value
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -6,6 +6,7 @@ require_relative "../pricing"
 require_relative "rollups"
 require_relative "../budget/per_tag"
 require_relative "tags/encoding"
+require_relative "storable"
 
 module LlmCostTracker
   module Ledger
@@ -23,11 +24,13 @@ module LlmCostTracker
           events = Array(events)
           return if events.empty?
 
-          LlmCostTracker::Call.transaction do
-            rows = events.map { |event| attributes_for(event) }
-            call_ids = insert_calls_returning_ids(rows, events)
-            insert_line_items(events, call_ids)
-            insert_call_tags(events, call_ids)
+          Isolation.guard do
+            LlmCostTracker::Call.transaction do
+              rows = events.map { |event| attributes_for(event) }
+              call_ids = insert_calls_returning_ids(rows, events)
+              insert_line_items(events, call_ids)
+              insert_call_tags(events, call_ids)
+            end
           end
         end
 
@@ -62,9 +65,7 @@ module LlmCostTracker
             pricing_snapshot: event.pricing_snapshot
           }
 
-          attributes
-            .merge(event.token_usage.to_h)
-            .merge(total_cost: event.cost&.total)
+          Storable.clean(attributes.merge(event.token_usage.to_h).merge(total_cost: event.cost&.total))
         end
 
         def call_ids_for(events)
@@ -90,7 +91,7 @@ module LlmCostTracker
         end
 
         def line_item_attributes(call_id:, line_item:, position:)
-          {
+          Storable.clean(
             llm_cost_tracker_call_id: call_id,
             position: position,
             kind: line_item.kind,
@@ -112,16 +113,16 @@ module LlmCostTracker
             provider_item_id: line_item.provider_item_id,
             details: stored_details(line_item.details),
             created_at: Time.now.utc
-          }
+          )
         end
 
         def insert_call_tags(events, call_ids)
           rows = events.flat_map do |event|
-            (event.tags || {}).map do |key, value|
+            (event.tags || {}).to_h { |key, value| [key.to_s, value] }.map do |key, value|
               {
                 llm_cost_tracker_call_id: call_ids.fetch(event.event_id),
-                key: key.to_s,
-                value: Tags::Encoding.encode(value)
+                key: key,
+                value: Tags::Encoding.encode(Storable.clean(value))
               }.merge(budget_columns_for(event))
             end
           end

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

@@ -19,7 +19,7 @@ module LlmCostTracker
           end
 
           def value_arel
-            Arel.sql("#{call_tag_table}.#{quote_column('value')}")
+            Arel.sql(raw_value_sql(LlmCostTracker::CallTag.connection))
           end
 
           def label_sql(connection)
@@ -35,10 +35,6 @@ module LlmCostTracker
           def call_tag_table
             LlmCostTracker::CallTag.quoted_table_name
           end
-
-          def quote_column(name)
-            LlmCostTracker::CallTag.connection.quote_column_name(name)
-          end
         end
       end
     end

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

@@ -10,8 +10,7 @@ module LlmCostTracker
 
         def self.encode(value)
           encoded = case value
-                    when Hash then JSON.generate(normalize_hash(value))
-                    when Array then JSON.generate(normalize_array(value))
+                    when Hash, Array then JSON.generate(normalize_value(value))
                     else value.to_s
                     end
           truncate(encoded)
@@ -24,18 +23,10 @@ module LlmCostTracker
           string.byteslice(0, limit).encode("UTF-8", invalid: :replace, undef: :replace)
         end
 
-        def self.normalize_hash(hash)
-          hash.transform_keys(&:to_s).sort.to_h.transform_values { |v| normalize_value(v) }
-        end
-
-        def self.normalize_array(array)
-          array.map { |v| normalize_value(v) }
-        end
-
         def self.normalize_value(value)
           case value
-          when Hash then normalize_hash(value)
-          when Array then normalize_array(value)
+          when Hash then value.transform_keys(&:to_s).sort.to_h.transform_values { |v| normalize_value(v) }
+          when Array then value.map { |v| normalize_value(v) }
           else value.to_s
           end
         end

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

@@ -9,8 +9,6 @@ module LlmCostTracker
         class << self
           def apply(tags)
             normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values { |v| Encoding.encode(v) }
-            return LlmCostTracker::Call.all if normalized_tags.empty?
-
             normalized_tags.inject(LlmCostTracker::Call.all) do |relation, (key, value)|
               relation.where(id: LlmCostTracker::CallTag.where(key: key,
                                                                value: value).select(:llm_cost_tracker_call_id))

data/lib/llm_cost_tracker/ledger.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "ledger/isolation"
 require_relative "ledger/schema"
 require_relative "ledger/period"
 require_relative "ledger/rollups"

data/lib/llm_cost_tracker/logging.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "redaction"
+
 module LlmCostTracker
   module Logging
     class << self
@@ -9,7 +11,7 @@ module LlmCostTracker
 
       private
 
-      def prefixed(message) = "[#{LlmCostTracker.name}] #{message}"
+      def prefixed(message) = "[#{LlmCostTracker.name}] #{Redaction.text(message)}"
     end
   end
 end