llm_cost_tracker 0.7.0 → 0.7.2

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

@@ -1,14 +1,25 @@
 class UpgradeLlmApiCallCostPrecision < ActiveRecord::Migration<%= migration_version %>
-  COST_COLUMNS = %i[input_cost output_cost total_cost].freeze
+  COST_COLUMNS = %i[
+    input_cost
+    cache_read_input_cost
+    cache_write_input_cost
+    cache_write_1h_input_cost
+    output_cost
+    total_cost
+  ].freeze
 
   def up
     COST_COLUMNS.each do |column|
+      next unless column_exists?(:llm_api_calls, column)
+
       change_column :llm_api_calls, column, :decimal, precision: 20, scale: 8
     end
   end
 
   def down
     COST_COLUMNS.each do |column|
+      next unless column_exists?(:llm_api_calls, column)
+
       change_column :llm_api_calls, column, :decimal, precision: 12, scale: 8
     end
   end

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

@@ -1,4 +1,4 @@
-require "llm_cost_tracker/active_record_adapter"
+require "llm_cost_tracker/ledger/schema/adapter"
 
 class UpgradeLlmApiCallTagsToJsonb < ActiveRecord::Migration<%= migration_version %>
   def up
@@ -34,7 +34,7 @@ class UpgradeLlmApiCallTagsToJsonb < ActiveRecord::Migration<%= migration_versio
   private
 
   def postgresql?
-    LlmCostTracker::ActiveRecordAdapter.postgresql?(connection)
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
   end
 
   def tags_jsonb?

data/lib/llm_cost_tracker/{storage/active_record_inbox_batch.rb → ingestion/batch.rb}

@@ -1,12 +1,11 @@
 # frozen_string_literal: true
 
-require_relative "../inbox_event"
-require_relative "active_record_inbox"
-require_relative "active_record_store"
+require_relative "inbox"
+require_relative "../ledger/store"
 
 module LlmCostTracker
-  module Storage
-    class ActiveRecordInboxBatch
+  module Ingestion
+    class Batch
       BATCH_SIZE = 100
       LOCK_TIMEOUT_SECONDS = 30
 
@@ -27,14 +26,18 @@ module LlmCostTracker
         raise
       end
 
-      def pending? = model.where("attempts < ?", ActiveRecordInbox::MAX_ATTEMPTS).exists?
+      def pending?
+        Ingestion::Event.where("attempts < ?", Ingestion::Event::MAX_ATTEMPTS).exists?
+      end
 
-      def claimable? = claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists?
+      def claimable?
+        claimable_scope(Time.now.utc - LOCK_TIMEOUT_SECONDS).exists?
+      end
 
       def mark_failed(rows, error)
         message = "#{error.class}: #{error.message}".byteslice(0, 1_000)
         now = Time.now.utc
-        model
+        Ingestion::Event
           .where(id: rows.map(&:id), locked_by: identity)
           .update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
       rescue StandardError
@@ -48,16 +51,16 @@ module LlmCostTracker
       def claim
         now = Time.now.utc
         cutoff = now - LOCK_TIMEOUT_SECONDS
-        model.transaction do
+        Ingestion::Event.transaction do
           rows = claimable_scope(cutoff).order(:id).limit(BATCH_SIZE).lock.to_a
           ids = rows.map(&:id)
           next [] if ids.empty?
 
-          updates = model.sanitize_sql_array(
+          updates = Ingestion::Event.sanitize_sql_array(
             ["locked_at = ?, locked_by = ?, attempts = attempts + 1, updated_at = ?", now, identity, now]
           )
-          model.where(id: ids).update_all(updates)
-          model.where(id: ids, locked_by: identity).order(:id).to_a
+          Ingestion::Event.where(id: ids).update_all(updates)
+          Ingestion::Event.where(id: ids, locked_by: identity).order(:id).to_a
         end
       end
 
@@ -65,7 +68,7 @@ module LlmCostTracker
         valid_rows = []
         events = []
         rows.each do |row|
-          events << ActiveRecordInbox.event_from_row(row)
+          events << Ingestion::Inbox.event_from_row(row)
           valid_rows << row
         rescue StandardError => e
           mark_failed([row], e)
@@ -74,19 +77,17 @@ module LlmCostTracker
       end
 
       def persist(rows, events)
-        LlmCostTracker::LlmApiCall.transaction do
-          ActiveRecordStore.insert_many(events)
-          model.where(id: rows.map(&:id), locked_by: identity).delete_all
+        LlmCostTracker::Ledger::Call.transaction do
+          Ledger::Store.insert_many(events)
+          Ingestion::Event.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
       end
 
       def claimable_scope(cutoff)
-        model
-          .where("attempts < ?", ActiveRecordInbox::MAX_ATTEMPTS)
+        Ingestion::Event
+          .where("attempts < ?", Ingestion::Event::MAX_ATTEMPTS)
           .where("locked_at IS NULL OR locked_at < ?", cutoff)
       end
-
-      def model = LlmCostTracker::InboxEvent
     end
   end
 end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+require "json"
+require "time"
+
+require_relative "../event"
+require_relative "../pricing"
+
+module LlmCostTracker
+  module Ingestion
+    class Inbox
+      PAYLOAD_SCHEMA_VERSION = 1
+
+      class << self
+        def save(event)
+          insert_row(row_for(event))
+          Ingestion::Worker.ensure_started
+          event
+        end
+
+        def event_from_row(row)
+          payload = JSON.parse(row.payload)
+          schema_version = payload.fetch("schema_version", 0)
+          unless [0, PAYLOAD_SCHEMA_VERSION].include?(schema_version)
+            raise LlmCostTracker::Error, "unsupported ledger inbox payload schema version #{schema_version.inspect}"
+          end
+
+          cost = payload["cost"] && Pricing.stored_cost_attributes(payload["cost"])
+          token_usage = payload["token_usage"] || payload
+
+          LlmCostTracker::Event.new(
+            event_id: payload.fetch("event_id"),
+            provider: payload.fetch("provider"),
+            model: payload.fetch("model"),
+            token_usage: TokenUsage.from_hash(token_usage),
+            pricing_mode: payload["pricing_mode"],
+            cost: cost,
+            tags: payload.fetch("tags"),
+            latency_ms: payload["latency_ms"],
+            stream: payload.fetch("stream"),
+            usage_source: payload["usage_source"],
+            provider_response_id: payload["provider_response_id"],
+            tracked_at: Time.iso8601(payload.fetch("tracked_at"))
+          )
+        end
+
+        private
+
+        def row_for(event)
+          now = Time.now.utc
+          {
+            event_id: event.event_id,
+            total_cost: event.total_cost,
+            tracked_at: event.tracked_at,
+            payload: JSON.generate(payload_for(event)),
+            attempts: 0,
+            created_at: now,
+            updated_at: now
+          }
+        end
+
+        def payload_for(event)
+          event.to_h.merge(
+            schema_version: PAYLOAD_SCHEMA_VERSION,
+            event_id: event.event_id,
+            provider: event.provider,
+            model: event.model,
+            tracked_at: event.tracked_at.iso8601(6)
+          )
+        end
+
+        def insert_row(row)
+          connection = LlmCostTracker::Ledger::Call.connection
+          if connection.transaction_open?
+            insert_with_separate_connection(row)
+          else
+            execute_insert(connection, row)
+          end
+        rescue ActiveRecord::ConnectionTimeoutError => e
+          raise LlmCostTracker::Error,
+                "ledger inbox could not checkout a separate database connection: #{e.message}"
+        end
+
+        def insert_with_separate_connection(row)
+          pool = LlmCostTracker::Ledger::Call.connection_pool
+          connection = pool.checkout
+          begin
+            connection.transaction(requires_new: true) { execute_insert(connection, row) }
+          ensure
+            pool.checkin(connection)
+          end
+        end
+
+        def execute_insert(connection, row)
+          columns = row.keys
+          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(Event.table_name)
+
+          connection.execute("INSERT INTO #{table} (#{quoted_columns}) VALUES (#{quoted_values})")
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/{storage/active_record_ingestor_lease.rb → ingestion/lease_claim.rb}

@@ -1,10 +1,8 @@
 # frozen_string_literal: true
 
-require_relative "../ingestor_lease"
-
 module LlmCostTracker
-  module Storage
-    class ActiveRecordIngestorLease
+  module Ingestion
+    class LeaseClaim
       LEASE_NAME = "default"
 
       def initialize(identity:, seconds:)
@@ -14,9 +12,9 @@ module LlmCostTracker
 
       def acquire
         now = Time.now.utc
-        LlmCostTracker::IngestorLease.transaction do
-          lease = LlmCostTracker::IngestorLease.lock.find_by(name: LEASE_NAME)
-          lease ||= LlmCostTracker::IngestorLease.create!(name: LEASE_NAME)
+        LlmCostTracker::Ingestion::Lease.transaction do
+          lease = LlmCostTracker::Ingestion::Lease.lock.find_by(name: LEASE_NAME)
+          lease ||= LlmCostTracker::Ingestion::Lease.create!(name: LEASE_NAME)
           next false unless available?(lease, now)
 
           lease.update!(locked_by: identity, locked_until: now + seconds)

data/lib/llm_cost_tracker/{storage/active_record_ingestor.rb → ingestion/worker.rb}

@@ -1,17 +1,16 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/kernel/reporting"
 require "securerandom"
 
-require_relative "../inbox_event"
+require_relative "inbox"
+require_relative "batch"
+require_relative "lease_claim"
 require_relative "../logging"
-require_relative "active_record_connection_cleanup"
-require_relative "active_record_inbox"
-require_relative "active_record_inbox_batch"
-require_relative "active_record_ingestor_lease"
 
 module LlmCostTracker
-  module Storage
-    class ActiveRecordIngestor
+  module Ingestion
+    class Worker
       INTERVAL_SECONDS = 0.25
       IDLE_INTERVAL_SECONDS = 1.0
       MAX_IDLE_INTERVAL_SECONDS = 5.0
@@ -19,13 +18,12 @@ module LlmCostTracker
       FLUSH_TIMEOUT_SECONDS = 10
       class << self
         def ensure_started
-          return unless ActiveRecordInbox.enabled?
-
           thread = mutex.synchronize do
             reset_after_fork!
             unless @thread&.alive?
               @stop_requested = false
-              generation = next_generation
+              @generation = @generation.to_i + 1
+              generation = @generation
               @thread = Thread.new { run(generation) }
               @thread.name = "llm_cost_tracker_ingestor" if @thread.respond_to?(:name=)
               @thread.report_on_exception = false if @thread.respond_to?(:report_on_exception=)
@@ -36,23 +34,27 @@ module LlmCostTracker
         end
 
         def flush!(timeout: FLUSH_TIMEOUT_SECONDS, require_lease: false)
-          return true unless ActiveRecordInbox.enabled?
+          Ingestion.ensure_current_schema!
 
           deadline = Time.now.utc + timeout
           loop do
-            return true unless pending_events?
+            return true unless Ingestion::Batch.new(identity: identity).pending?
             return false if Time.now.utc >= deadline
 
             processed = ingest_once(require_lease: require_lease)
-            return false if processed.zero? && !sleep_until_next_flush(deadline)
+            next unless processed.zero?
+
+            duration = [INTERVAL_SECONDS, deadline - Time.now.utc].min
+            return false unless duration.positive?
+
+            sleep(duration)
           end
         end
 
         def shutdown!(timeout: FLUSH_TIMEOUT_SECONDS, drain: true)
-          ActiveRecordInbox.reset!
           thread = mutex.synchronize do
             @stop_requested = true
-            next_generation
+            @generation = @generation.to_i + 1
             @thread
           end
           wake_thread(thread)
@@ -68,7 +70,7 @@ module LlmCostTracker
         def reset!
           thread = mutex.synchronize do
             @stop_requested = true
-            next_generation
+            @generation = @generation.to_i + 1
             thread = @thread
             @thread = nil
             @pid = nil
@@ -79,11 +81,11 @@ module LlmCostTracker
         end
 
         def ingest_once(require_lease: true)
-          return 0 unless ActiveRecordInbox.enabled?
-          return 0 unless claimable_events?
-          return 0 if require_lease && !acquire_lease
+          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
 
-          inbox_batch.ingest
+          batch.ingest
         rescue StandardError => e
           handle_error(e)
           0
@@ -91,15 +93,17 @@ module LlmCostTracker
 
         private
 
-        def mutex = @mutex ||= Mutex.new
+        def mutex
+          @mutex ||= Mutex.new
+        end
 
         def run(generation)
           idle_interval = IDLE_INTERVAL_SECONDS
           loop do
-            break if stop_requested?(generation)
+            break if mutex.synchronize { @stop_requested || generation != @generation }
 
             processed = executor_wrap { ingest_once }
-            ActiveRecordConnectionCleanup.release!
+            release_connection!
             if processed.zero?
               sleep(idle_interval)
               idle_interval = [idle_interval * 2, MAX_IDLE_INTERVAL_SECONDS].min
@@ -108,21 +112,14 @@ module LlmCostTracker
             end
           rescue StandardError => e
             handle_error(e)
-            ActiveRecordConnectionCleanup.release!
+            release_connection!
             sleep(idle_interval)
           end
         ensure
-          ActiveRecordConnectionCleanup.release!
+          release_connection!
           mutex.synchronize { @thread = nil if @thread.equal?(Thread.current) }
         end
 
-        def sleep_until_next_flush(deadline)
-          duration = [INTERVAL_SECONDS, deadline - Time.now.utc].min
-          sleep(duration) if duration.positive?
-        end
-
-        def stop_requested?(generation) = mutex.synchronize { @stop_requested || generation != @generation }
-
         def reset_after_fork!
           return if @pid == Process.pid
 
@@ -131,8 +128,6 @@ module LlmCostTracker
           @identity = nil
         end
 
-        def next_generation = (@generation = @generation.to_i + 1)
-
         def wake_thread(thread)
           thread&.wakeup if thread&.alive?
         rescue ThreadError
@@ -147,26 +142,21 @@ module LlmCostTracker
         end
 
         def rails_executor
-          return unless defined?(Rails) && Rails.respond_to?(:application) && Rails.application.respond_to?(:executor)
-
-          Rails.application.executor
+          Rails.application.try(:executor)
         rescue StandardError
           nil
         end
 
-        def identity = @identity ||= "pid-#{Process.pid}-#{SecureRandom.hex(6)}"
-
-        def acquire_lease = ActiveRecordIngestorLease.new(identity: identity, seconds: LEASE_SECONDS).acquire
-
-        def pending_events? = inbox_batch.pending?
-
-        def claimable_events? = inbox_batch.claimable?
-
-        def inbox_batch = ActiveRecordInboxBatch.new(identity: identity)
+        def identity
+          @identity ||= "pid-#{Process.pid}-#{SecureRandom.hex(6)}"
+        end
 
         def handle_error(error)
-          Logging.warn("ActiveRecord ingestor failed: #{error.class}: #{error.message}") unless
-            LlmCostTracker.configuration.storage_error_behavior == :ignore
+          Logging.warn("ActiveRecord ingestor failed: #{error.class}: #{error.message}")
+        end
+
+        def release_connection!
+          suppress(StandardError) { ActiveRecord::Base.connection_handler.clear_active_connections! }
         end
       end
     end

data/lib/llm_cost_tracker/ingestion.rb

@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+require_relative "doctor/check"
+require_relative "errors"
+require_relative "ledger"
+require_relative "ingestion/lease_claim"
+require_relative "ingestion/inbox"
+require_relative "ingestion/batch"
+require_relative "ingestion/worker"
+
+module LlmCostTracker
+  module Ingestion
+    VERIFY_TAG = "llm_cost_tracker_verify"
+
+    class << self
+      def ensure_current_schema!
+        unless Ledger::Call.table_exists?
+          raise Error, "llm_api_calls table is missing; run install generator and migrate"
+        end
+
+        schema_errors = Ledger::Schema::Calls.current_schema_errors
+        message = "llm_api_calls table is not on the current schema: #{schema_errors.join('; ')}"
+        raise Error, message if schema_errors.any?
+
+        period_total_errors = Ledger::Schema::PeriodTotals.current_schema_errors
+        return if period_total_errors.empty?
+
+        message = "llm_cost_tracker_period_totals table is not on the current schema: " \
+                  "#{period_total_errors.join('; ')}; " \
+                  "run bin/rails generate llm_cost_tracker:add_period_totals && bin/rails db:migrate"
+        raise Error, message
+      end
+
+      def verify
+        unless LlmCostTracker::Ledger::Call.table_exists?
+          return [
+            LlmCostTracker::Doctor::Check.new(
+              :error,
+              "active_record",
+              "llm_api_calls table is missing; run install generator and migrate"
+            )
+          ]
+        end
+
+        [capture_check]
+      rescue StandardError => e
+        [LlmCostTracker::Doctor::Check.new(:error, "active_record", "#{e.class}: #{e.message}")]
+      end
+
+      private
+
+      def capture_check
+        provider, model = sample_priced_identity
+        response_id = "lct_verify_#{SecureRandom.hex(8)}"
+        notifications = []
+        subscription = subscribe_to_verification(response_id, notifications)
+
+        event = LlmCostTracker.track(
+          provider: provider,
+          model: model,
+          input_tokens: 1,
+          output_tokens: 1,
+          provider_response_id: response_id,
+          feature: VERIFY_TAG
+        )
+        LlmCostTracker.flush!
+        persisted = LlmCostTracker::Ledger::Call.where(provider_response_id: response_id).exists?
+
+        return capture_success if persisted && notifications.any?
+
+        LlmCostTracker::Doctor::Check.new(
+          :error,
+          "active_record capture",
+          capture_failure_message(persisted, notifications)
+        )
+      rescue LlmCostTracker::BudgetExceededError => e
+        LlmCostTracker::Doctor::Check.new(:error, "active_record capture", "blocked by budget guardrail: #{e.message}")
+      rescue LlmCostTracker::Error => e
+        LlmCostTracker::Doctor::Check.new(:error, "active_record capture", e.message)
+      rescue StandardError => e
+        LlmCostTracker::Doctor::Check.new(:error, "active_record capture", "#{e.class}: #{e.message}")
+      ensure
+        cleanup_verification_call(response_id) if response_id
+        LlmCostTracker::Ingestion::Event.where(event_id: event.event_id).delete_all if event
+        ActiveSupport::Notifications.unsubscribe(subscription) if subscription
+      end
+
+      def subscribe_to_verification(response_id, notifications)
+        ActiveSupport::Notifications.subscribe(LlmCostTracker::Tracker::EVENT_NAME) do |*, payload|
+          notifications << payload if payload[:provider_response_id] == response_id
+        end
+      end
+
+      def capture_success
+        LlmCostTracker::Doctor::Check.new(
+          :ok,
+          "active_record capture",
+          "manual event emitted and persisted through durable inbox"
+        )
+      end
+
+      def capture_failure_message(persisted, notifications)
+        missing = []
+        missing << "notification" if notifications.empty?
+        missing << "persisted row" unless persisted
+        "missing #{missing.join(' and ')} for synthetic manual event"
+      end
+
+      def cleanup_verification_call(response_id)
+        relation = LlmCostTracker::Ledger::Call.where(provider_response_id: response_id)
+        rows = relation.pluck(:id, :tracked_at, :total_cost)
+        return if rows.empty?
+
+        relation.delete_all
+        LlmCostTracker::Ledger::Rollups.decrement!(rows)
+      end
+
+      def sample_priced_identity
+        key = LlmCostTracker::Pricing::Registry.builtin_prices.find do |model_id, prices|
+          model_id.include?("/") && prices[:input] && prices[:output]
+        end&.first
+        provider, model = key.to_s.split("/", 2)
+        [provider || "openai", model || "gpt-4o-mini"]
+      end
+    end
+  end
+end