llm_cost_tracker 0.9.0 → 0.10.0

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "../providers/anthropic/tier_classification"
 
 module LlmCostTracker
   module Integrations
@@ -111,7 +112,7 @@ module LlmCostTracker
 
           record_safely do
             LlmCostTracker::Tracker.record(
-              capture: UsageCapture.build(
+              event: Event.build(
                 provider: provider,
                 model: model,
                 token_usage: TokenUsage.build(
@@ -140,7 +141,7 @@ module LlmCostTracker
             hidden_output = object_value(response, :thinking_tokens, :reasoning_tokens).to_i
 
             LlmCostTracker::Tracker.record(
-              capture: UsageCapture.build(
+              event: Event.build(
                 provider: provider,
                 model: model,
                 pricing_mode: pricing_mode(provider: provider, response: response),
@@ -182,17 +183,14 @@ module LlmCostTracker
         end
 
         def provider_response_id(response)
-          object_value(response, :id, :provider_response_id) || object_dig(response, :raw, :id)
+          object_value(response, :id, :provider_response_id)
         end
 
-        ANTHROPIC_STANDARD_EQUIVALENT_SERVICE_TIERS = %w[standard standard_only priority].freeze
-        private_constant :ANTHROPIC_STANDARD_EQUIVALENT_SERVICE_TIERS
-
         def pricing_mode(provider:, response:)
-          raw = object_value(response, :pricing_mode, :service_tier) ||
-                object_dig(response, :raw, :pricing_mode) ||
-                object_dig(response, :raw, :service_tier)
-          return nil if provider == "anthropic" && ANTHROPIC_STANDARD_EQUIVALENT_SERVICE_TIERS.include?(raw.to_s)
+          raw = object_value(response, :pricing_mode, :service_tier)
+          if provider == "anthropic" && LlmCostTracker::Providers::Anthropic::TierClassification.standard_equivalent_tier?(raw)
+            return nil
+          end
 
           raw
         end
@@ -202,14 +200,14 @@ module LlmCostTracker
         def complete(*args, **kwargs, &)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          integration.enforce_budget!
+          integration.enforce_budget!(request: request)
           started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_completion(
             self,
             response,
             request: request,
-            latency_ms: integration.elapsed_ms(started_at),
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
             stream: integration.streaming_request?(request, has_block: block_given?)
           )
           response
@@ -218,14 +216,14 @@ module LlmCostTracker
         def embed(*args, **kwargs)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          integration.enforce_budget!
+          integration.enforce_budget!(request: request)
           started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_embedding(
             self,
             response,
             request: request,
-            latency_ms: integration.elapsed_ms(started_at)
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at)
           )
           response
         end
@@ -233,14 +231,14 @@ module LlmCostTracker
         def transcribe(*args, **kwargs)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          integration.enforce_budget!
+          integration.enforce_budget!(request: request)
           started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_transcription(
             self,
             response,
             request: request,
-            latency_ms: integration.elapsed_ms(started_at)
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at)
           )
           response
         end
@@ -248,14 +246,14 @@ module LlmCostTracker
         def paint(*args, **kwargs)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          integration.enforce_budget!
+          integration.enforce_budget!(request: request)
           started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_image(
             self,
             response,
             request: request,
-            latency_ms: integration.elapsed_ms(started_at)
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at)
           )
           response
         end
@@ -263,14 +261,14 @@ module LlmCostTracker
         def moderate(*args, **kwargs)
           integration = LlmCostTracker::Integrations::RubyLlm
           request = integration.request_params(args, kwargs)
-          integration.enforce_budget!
+          integration.enforce_budget!(request: request)
           started_at = LlmCostTracker::Timing.now_monotonic
           response = super
           integration.record_moderation(
             self,
             response,
             request: request,
-            latency_ms: integration.elapsed_ms(started_at)
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at)
           )
           response
         end

data/lib/llm_cost_tracker/integrations.rb

@@ -2,19 +2,17 @@
 
 require_relative "errors"
 require_relative "logging"
-require_relative "integrations/openai"
-require_relative "integrations/anthropic"
-require_relative "integrations/ruby_llm"
 
 module LlmCostTracker
   module Integrations
-    AVAILABLE = {
-      openai: Openai,
-      anthropic: Anthropic,
-      ruby_llm: RubyLlm
-    }.freeze
+    autoload :Base, "llm_cost_tracker/integrations/base"
+    autoload :Openai, "llm_cost_tracker/integrations/openai"
+    autoload :Anthropic, "llm_cost_tracker/integrations/anthropic"
+    autoload :RubyLlm, "llm_cost_tracker/integrations/ruby_llm"
 
+    INTEGRATION_CONSTANTS = { openai: :Openai, anthropic: :Anthropic, ruby_llm: :RubyLlm }.freeze
     DOUBLE_INSTRUMENTATION_OVERLAPS = %i[openai anthropic].freeze
+    private_constant :DOUBLE_INSTRUMENTATION_OVERLAPS
 
     module_function
 
@@ -25,7 +23,7 @@ module LlmCostTracker
     end
 
     def checks(names = LlmCostTracker.configuration.instrumented_integrations)
-      return [Base::Result.new(:integrations, :ok, "no SDK integrations enabled")] if names.empty?
+      return [Base::Result.new(:ok, "integrations", "no SDK integrations enabled")] if names.empty?
 
       normalize(names).map { |name| fetch(name).status }
     end
@@ -48,14 +46,17 @@ module LlmCostTracker
     end
 
     def fetch(name)
-      AVAILABLE.fetch(name) do
-        message = "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
-        raise LlmCostTracker::Error, message
+      const_name = INTEGRATION_CONSTANTS[name.to_sym]
+      unless const_name
+        raise LlmCostTracker::Error,
+              "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
       end
+
+      const_get(const_name)
     end
 
     def names
-      AVAILABLE.keys
+      INTEGRATION_CONSTANTS.keys
     end
   end
 end

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

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "bigdecimal"
+
 require_relative "../period"
 
 module LlmCostTracker
@@ -26,12 +28,12 @@ module LlmCostTracker
         attr_reader :periods, :time
 
         def snapshot_totals
-          values = periods.to_h { |period| [period, 0.0] }
+          values = periods.to_h { |period| [period, BigDecimal("0")] }
           period_by_name = periods.to_h { |period| [period.name, period] }
           sql = periods.map { |period| snapshot_select(period) }.join(" UNION ALL ")
           LlmCostTracker::Call.find_by_sql(sql).each do |row|
             period = period_by_name.fetch(row.period_key)
-            values[period] = row.total_cost.to_f
+            values[period] = BigDecimal(row.total_cost.to_s)
           end
           values
         end
@@ -39,7 +41,7 @@ module LlmCostTracker
         def snapshot_select(period)
           start = Period.range_start(period, time)
           components = [period_total_sql(period, start)]
-          components << pending_total_sql(start) if Ingestion.durable?
+          components << pending_total_sql(start) if Ingestion.async?
           "SELECT #{connection.quote(period.name)} AS period_key, " \
             "(#{components.join(') + (')}) AS total_cost"
         end

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -34,8 +34,8 @@ module LlmCostTracker
         private
 
         def period_rows(event)
-          currency = currency_for(event)
-          provider = provider_for(event)
+          currency = currency_from_snapshot(event.pricing_snapshot)
+          provider = event.provider.to_s
           Period::PERIODS.map do |period, name|
             {
               period: name,
@@ -61,8 +61,8 @@ module LlmCostTracker
 
         def call_rollups(events)
           events.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |event, totals|
-            currency = currency_for(event)
-            provider = provider_for(event)
+            currency = currency_from_snapshot(event.pricing_snapshot)
+            provider = event.provider.to_s
             Period::PERIODS.each do |period, name|
               key = [name, Period.bucket(period, event.tracked_at), currency, provider]
               totals[key] += BigDecimal(event.total_cost.to_s)
@@ -118,20 +118,11 @@ module LlmCostTracker
           end
         end
 
-        def currency_for(event)
-          snapshot = event.respond_to?(:pricing_snapshot) ? event.pricing_snapshot : nil
-          currency_from_snapshot(snapshot)
-        end
-
         def currency_from_snapshot(snapshot)
           value = (snapshot.is_a?(Hash) && (snapshot["currency"] || snapshot[:currency])) || DEFAULT_CURRENCY
           value.to_s.upcase
         end
 
-        def provider_for(event)
-          (event.respond_to?(:provider) ? event.provider : nil).to_s
-        end
-
         def upsert_call_rollups(rows)
           LlmCostTracker::CallRollup.upsert_all(
             rows,

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

@@ -42,6 +42,17 @@ module LlmCostTracker
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
             columns = LlmCostTracker::CallLineItem.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
             missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?

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

@@ -16,8 +16,20 @@ module LlmCostTracker
             table_name = LlmCostTracker::CallRollup.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
+            columns = LlmCostTracker::CallRollup.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
-            missing = REQUIRED_COLUMNS - LlmCostTracker::CallRollup.columns_hash.keys
+            missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             unless unique_period_index?(connection, table_name)
               errors << "missing unique index: period, period_start, currency, provider"
@@ -25,8 +37,6 @@ module LlmCostTracker
             errors
           end
 
-          private
-
           def unique_period_index?(connection, table_name)
             connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
           end

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

@@ -19,6 +19,17 @@ module LlmCostTracker
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
             columns = LlmCostTracker::CallTag.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
             missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?

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

@@ -47,10 +47,6 @@ module LlmCostTracker
         private_constant :REQUIRED_INDEXES
 
         class << self
-          def current_schema?
-            current_schema_errors.empty?
-          end
-
           def current_schema_errors
             schema_capabilities.fetch(:current_schema_errors)
           end

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

@@ -28,15 +28,25 @@ module LlmCostTracker
             table_name = LlmCostTracker::Ingestion::InboxEntry.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
+            columns = LlmCostTracker::Ingestion::InboxEntry.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
-            missing = REQUIRED_COLUMNS - LlmCostTracker::Ingestion::InboxEntry.columns_hash.keys
+            missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             errors << "missing unique index: event_id" unless event_id_unique_index?(connection, table_name)
             errors
           end
 
-          private
-
           def event_id_unique_index?(connection, table_name)
             connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
           end

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

@@ -23,15 +23,25 @@ module LlmCostTracker
             table_name = LlmCostTracker::Ingestion::Lease.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
+            columns = LlmCostTracker::Ingestion::Lease.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
-            missing = REQUIRED_COLUMNS - LlmCostTracker::Ingestion::Lease.columns_hash.keys
+            missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             errors << "missing unique index: name" unless name_unique_index?(connection, table_name)
             errors
           end
 
-          private
-
           def name_unique_index?(connection, table_name)
             connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
           end

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

@@ -7,10 +7,10 @@ module LlmCostTracker
     module Schema
       module ProviderInvoiceImports
         REQUIRED_COLUMNS = %w[
-          source cursor window_start window_end state last_error
+          source provider cursor window_start window_end state last_error
           rows_imported started_at finished_at
         ].freeze
-        SOURCE_STARTED_AT_INDEX = %i[source started_at].freeze
+        SOURCE_PROVIDER_STARTED_AT_INDEX = %i[source provider started_at].freeze
 
         class << self
           def current_schema_errors
@@ -19,25 +19,35 @@ module LlmCostTracker
             table_name = LlmCostTracker::ProviderInvoiceImport.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
-            errors = []
-            errors.concat(column_errors)
-            errors.concat(index_errors(connection, table_name))
+            columns = LlmCostTracker::ProviderInvoiceImport.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
             errors
           end
 
           private
 
-          def column_errors
-            missing = REQUIRED_COLUMNS - LlmCostTracker::ProviderInvoiceImport.columns_hash.keys
+          def compute_errors(connection, table_name, columns)
+            errors = []
+            errors.concat(column_errors(columns))
+            errors.concat(index_errors(connection, table_name))
+            errors
+          end
+
+          def column_errors(columns)
+            missing = REQUIRED_COLUMNS - columns.keys
             return [] if missing.empty?
 
             ["missing columns: #{missing.join(', ')}"]
           end
 
           def index_errors(connection, table_name)
-            return [] if connection.index_exists?(table_name, SOURCE_STARTED_AT_INDEX)
+            return [] if connection.index_exists?(table_name, SOURCE_PROVIDER_STARTED_AT_INDEX)
 
-            ["missing index: source, started_at"]
+            ["missing index: source, provider, started_at"]
           end
         end
       end

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

@@ -19,27 +19,32 @@ module LlmCostTracker
             table_name = LlmCostTracker::ProviderInvoice.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
-            errors = []
-            errors.concat(column_errors)
-            errors.concat(metadata_type_errors(connection))
-            errors.concat(index_errors(connection, table_name))
+            columns = LlmCostTracker::ProviderInvoice.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
             errors
           end
 
           private
 
-          def column_errors
-            missing = REQUIRED_COLUMNS - LlmCostTracker::ProviderInvoice.columns_hash.keys
+          def compute_errors(connection, table_name, columns)
+            errors = []
+            errors.concat(column_errors(columns))
+            errors.concat(Adapter.json_column_errors(columns["metadata"], connection, "metadata"))
+            errors.concat(index_errors(connection, table_name))
+            errors
+          end
+
+          def column_errors(columns)
+            missing = REQUIRED_COLUMNS - columns.keys
             return [] if missing.empty?
 
             ["missing columns: #{missing.join(', ')}"]
           end
 
-          def metadata_type_errors(connection)
-            metadata = LlmCostTracker::ProviderInvoice.columns_hash["metadata"]
-            Adapter.json_column_errors(metadata, connection, "metadata")
-          end
-
           def index_errors(connection, table_name)
             errors = []
             unless connection.index_exists?(table_name, UNIQUE_INDEX_COLUMNS, unique: true)
@@ -48,8 +53,18 @@ module LlmCostTracker
             unless connection.index_exists?(table_name, SOURCE_PERIOD_INDEX_COLUMNS)
               errors << "missing index: source, currency, period_start"
             end
+            if Adapter.postgresql?(connection) && !gin_metadata_index?(connection, table_name)
+              errors << "missing GIN index on metadata " \
+                        "(run bin/rails generate llm_cost_tracker:upgrade_provider_invoices_metadata_index)"
+            end
             errors
           end
+
+          def gin_metadata_index?(connection, table_name)
+            connection.indexes(table_name).any? do |index|
+              index.columns == ["metadata"] && index.using.to_s == "gin"
+            end
+          end
         end
       end
     end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -11,27 +11,34 @@ module LlmCostTracker
   module Ledger
     class Store
       class << self
-        def insert_many(events)
+        def insert(events, skip_existence_check: false)
           events = Array(events)
-          return [] if events.empty?
+          return if events.empty?
 
-          insertable = insertable_events(events)
+          insertable = skip_existence_check ? events : insertable_events(events)
+          return unless insertable.any?
 
-          if insertable.any?
-            LlmCostTracker::Call.transaction do
-              rows = insertable.map { |event| attributes_for(event) }
-              LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: false)
-              call_ids = call_ids_for(insertable)
-              insert_line_items(insertable, call_ids)
-              insert_call_tags(insertable, call_ids)
-            end
-            increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
+          LlmCostTracker::Call.transaction do
+            rows = insertable.map { |event| attributes_for(event) }
+            call_ids = insert_calls_returning_ids(rows, insertable)
+            insert_line_items(insertable, call_ids)
+            insert_call_tags(insertable, call_ids)
           end
-          events
+          increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
         end
 
         private
 
+        def insert_calls_returning_ids(rows, insertable)
+          if LlmCostTracker::Call.connection.supports_insert_returning?
+            result = LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: %i[id event_id])
+            result.rows.to_h { |id, event_id| [event_id, id] }
+          else
+            LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: false)
+            call_ids_for(insertable)
+          end
+        end
+
         def attributes_for(event)
           attributes = {
             event_id: event.event_id,
@@ -110,7 +117,7 @@ module LlmCostTracker
               {
                 llm_cost_tracker_call_id: call_ids.fetch(event.event_id),
                 key: key.to_s,
-                value: tag_row_value(value)
+                value: Tags::Encoding.encode(value)
               }
             end
           end
@@ -119,10 +126,6 @@ module LlmCostTracker
           LlmCostTracker::CallTag.insert_all!(rows, record_timestamps: false, returning: false)
         end
 
-        def tag_row_value(value)
-          Tags::Encoding.encode(value)
-        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/query.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require_relative "../schema/adapter"
 require_relative "encoding"
 
 module LlmCostTracker

data/lib/llm_cost_tracker/logging.rb

@@ -9,10 +9,6 @@ module LlmCostTracker
         log(:debug, message)
       end
 
-      def info(message)
-        log(:info, message)
-      end
-
       def warn(message)
         log(:warn, message)
       end