llm_cost_tracker 0.8.0 → 0.9.0

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

@@ -38,16 +38,32 @@ 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?
           "SELECT #{connection.quote(period.name)} AS period_key, " \
-            "(#{rollup_total_sql(period)}) + (#{pending_total_sql(start)}) AS total_cost"
+            "(#{components.join(') + (')}) AS total_cost"
         end
 
-        def rollup_total_sql(period)
+        def period_total_sql(period, start)
+          if LlmCostTracker.configuration.cache_rollups
+            "GREATEST(COALESCE(#{rollup_sum_sql(period)}, 0), COALESCE(#{calls_sum_sql(start)}, 0))"
+          else
+            "COALESCE(#{calls_sum_sql(start)}, 0)"
+          end
+        end
+
+        def rollup_sum_sql(period)
           table = connection.quote_table_name("llm_cost_tracker_call_rollups")
-          "COALESCE((SELECT SUM(total_cost) FROM #{table} " \
+          "(SELECT SUM(total_cost) FROM #{table} " \
             "WHERE period = #{connection.quote(Period::PERIODS.fetch(period))} " \
-            "AND period_start = #{connection.quote(Period.bucket(period, time))} " \
-            "AND currency = #{connection.quote(Ledger::Rollups::DEFAULT_CURRENCY)}), 0)"
+            "AND period_start = #{connection.quote(Period.bucket(period, time))})"
+        end
+
+        def calls_sum_sql(start)
+          table = connection.quote_table_name("llm_cost_tracker_calls")
+          tracked_at = connection.quote_column_name("tracked_at")
+          "(SELECT SUM(total_cost) FROM #{table} " \
+            "WHERE #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)})"
         end
 
         def pending_total_sql(start)

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -35,22 +35,25 @@ module LlmCostTracker
 
         def period_rows(event)
           currency = currency_for(event)
+          provider = provider_for(event)
           Period::PERIODS.map do |period, name|
             {
               period: name,
               period_start: Period.bucket(period, event.tracked_at),
               currency: currency,
+              provider: provider,
               total_cost: event.total_cost
             }
           end
         end
 
         def period_rows_for_events(events)
-          call_rollups(events).map do |(period, period_start, currency), total_cost|
+          call_rollups(events).map do |(period, period_start, currency, provider), total_cost|
             {
               period: period,
               period_start: period_start,
               currency: currency,
+              provider: provider,
               total_cost: total_cost
             }
           end
@@ -59,29 +62,33 @@ 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)
             Period::PERIODS.each do |period, name|
-              totals[[name, Period.bucket(period, event.tracked_at), currency]] += BigDecimal(event.total_cost.to_s)
+              key = [name, Period.bucket(period, event.tracked_at), currency, provider]
+              totals[key] += BigDecimal(event.total_cost.to_s)
             end
           end
         end
 
         def period_decrement_totals(call_rows)
           call_rows.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |row, totals|
-            _id, tracked_at, total_cost, pricing_snapshot = row
+            _id, tracked_at, total_cost, pricing_snapshot, provider = row
             next unless total_cost
 
             currency = currency_from_snapshot(pricing_snapshot)
+            provider_key = provider.to_s
             Period::PERIODS.each_key do |period|
-              totals[[period, Period.bucket(period, tracked_at), currency]] += total_cost
+              totals[[period, Period.bucket(period, tracked_at), currency, provider_key]] += total_cost
             end
           end
         end
 
         def apply_decrements(totals)
           now = Time.now.utc
-          buckets_by_period = totals.each_with_object({}) do |((period, period_start, currency), amount), grouped|
-            grouped[[period, currency]] ||= {}
-            grouped[[period, currency]][period_start] = amount
+          buckets_by_period = totals.each_with_object({}) do |(key, amount), grouped|
+            period, period_start, currency, provider = key
+            grouped[[period, currency, provider]] ||= {}
+            grouped[[period, currency, provider]][period_start] = amount
           end
 
           conn = LlmCostTracker::CallRollup.connection
@@ -89,10 +96,11 @@ module LlmCostTracker
           period_col = conn.quote_column_name("period")
           start_col = conn.quote_column_name("period_start")
           currency_col = conn.quote_column_name("currency")
+          provider_col = conn.quote_column_name("provider")
           total_col = conn.quote_column_name("total_cost")
           updated_col = conn.quote_column_name("updated_at")
 
-          buckets_by_period.each do |(period, currency), by_start|
+          buckets_by_period.each do |(period, currency, provider), by_start|
             case_clauses = by_start.map do |period_start, amount|
               "WHEN #{start_col} = #{conn.quote(period_start)} THEN #{conn.quote(amount)}"
             end.join(" ")
@@ -104,6 +112,7 @@ module LlmCostTracker
               "#{updated_col} = #{conn.quote(now)} " \
               "WHERE #{period_col} = #{conn.quote(Period::PERIODS.fetch(period))} " \
               "AND #{currency_col} = #{conn.quote(currency)} " \
+              "AND #{provider_col} = #{conn.quote(provider)} " \
               "AND #{start_col} IN (#{starts})"
             )
           end
@@ -115,7 +124,12 @@ module LlmCostTracker
         end
 
         def currency_from_snapshot(snapshot)
-          (snapshot.is_a?(Hash) && (snapshot["currency"] || snapshot[:currency])) || DEFAULT_CURRENCY
+          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)
@@ -130,7 +144,7 @@ module LlmCostTracker
         def call_rollups_unique_by
           return unless LlmCostTracker::CallRollup.connection.supports_insert_conflict_target?
 
-          %i[period period_start currency]
+          %i[period period_start currency provider]
         end
       end
     end

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

@@ -27,6 +27,11 @@ module LlmCostTracker
           provider_field
           provider_item_id
           details
+          created_at
+        ].freeze
+
+        REQUIRED_INDEX_COLUMNS = [
+          %w[llm_cost_tracker_call_id position]
         ].freeze
 
         class << self
@@ -41,7 +46,31 @@ module LlmCostTracker
             missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             errors.concat(Adapter.json_column_errors(columns["details"], connection, "details"))
-            errors
+            errors.concat(missing_index_errors(connection, table_name))
+            errors << missing_fk_error(connection, table_name) if missing_fk?(connection, table_name)
+            errors.compact
+          end
+
+          def missing_index_errors(connection, table_name)
+            existing = connection.indexes(table_name).map { |index| Array(index.columns).map(&:to_s) }
+            REQUIRED_INDEX_COLUMNS.filter_map do |required|
+              next if existing.any? { |columns| columns == required }
+
+              "missing index on (#{required.join(', ')})"
+            end
+          end
+
+          def missing_fk?(connection, table_name)
+            connection.foreign_keys(table_name).none? do |fk|
+              fk.column.to_s == "llm_cost_tracker_call_id" &&
+                fk.to_table.to_s == "llm_cost_tracker_calls"
+            end
+          rescue NotImplementedError, NoMethodError
+            false
+          end
+
+          def missing_fk_error(_connection, _table_name)
+            "missing foreign key on llm_cost_tracker_call_id referencing llm_cost_tracker_calls"
           end
         end
       end

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

@@ -6,8 +6,8 @@ module LlmCostTracker
   module Ledger
     module Schema
       module CallRollups
-        REQUIRED_COLUMNS = %w[period period_start currency total_cost].freeze
-        UNIQUE_COLUMNS = %i[period period_start currency].freeze
+        REQUIRED_COLUMNS = %w[period period_start currency provider total_cost created_at updated_at].freeze
+        UNIQUE_COLUMNS = %i[period period_start currency provider].freeze
 
         class << self
           def current_schema_errors
@@ -20,7 +20,7 @@ module LlmCostTracker
             missing = REQUIRED_COLUMNS - LlmCostTracker::CallRollup.columns_hash.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             unless unique_period_index?(connection, table_name)
-              errors << "missing unique index: period, period_start, currency"
+              errors << "missing unique index: period, period_start, currency, provider"
             end
             errors
           end

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

@@ -6,6 +6,11 @@ module LlmCostTracker
       module CallTags
         REQUIRED_COLUMNS = %w[llm_cost_tracker_call_id key value].freeze
 
+        REQUIRED_INDEX_COLUMNS = [
+          %w[key value],
+          %w[llm_cost_tracker_call_id]
+        ].freeze
+
         class << self
           def current_schema_errors
             connection = LlmCostTracker::Call.connection
@@ -14,10 +19,20 @@ module LlmCostTracker
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
             columns = LlmCostTracker::CallTag.columns_hash
+            errors = []
             missing = REQUIRED_COLUMNS - columns.keys
-            return [] if missing.empty?
+            errors << "missing columns: #{missing.join(', ')}" if missing.any?
+            errors.concat(missing_index_errors(connection, table_name))
+            errors
+          end
+
+          def missing_index_errors(connection, table_name)
+            existing = connection.indexes(table_name).map { |index| Array(index.columns).map(&:to_s) }
+            REQUIRED_INDEX_COLUMNS.filter_map do |required|
+              next if existing.any? { |columns| (required - columns).empty? }
 
-            ["missing columns: #{missing.join(', ')}"]
+              "missing index on (#{required.join(', ')})"
+            end
           end
         end
       end

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

@@ -18,6 +18,8 @@ module LlmCostTracker
           cache_write_extended_input_tokens
           audio_input_tokens
           audio_output_tokens
+          image_input_tokens
+          image_output_tokens
           hidden_output_tokens
           total_cost
           latency_ms

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

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module IngestionInboxEntries
+        REQUIRED_COLUMNS = %w[
+          event_id
+          total_cost
+          tracked_at
+          payload
+          locked_at
+          locked_by
+          attempts
+          last_error
+          created_at
+          updated_at
+        ].freeze
+
+        UNIQUE_COLUMNS = %i[event_id].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Ingestion::InboxEntry.connection
+            Adapter.ensure_supported!(connection)
+            table_name = LlmCostTracker::Ingestion::InboxEntry.table_name
+            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+
+            errors = []
+            missing = REQUIRED_COLUMNS - LlmCostTracker::Ingestion::InboxEntry.columns_hash.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
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module IngestionLeases
+        REQUIRED_COLUMNS = %w[
+          name
+          locked_by
+          locked_until
+          created_at
+          updated_at
+        ].freeze
+
+        UNIQUE_COLUMNS = %i[name].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Ingestion::Lease.connection
+            Adapter.ensure_supported!(connection)
+            table_name = LlmCostTracker::Ingestion::Lease.table_name
+            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+
+            errors = []
+            missing = REQUIRED_COLUMNS - LlmCostTracker::Ingestion::Lease.columns_hash.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
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module ProviderInvoiceImports
+        REQUIRED_COLUMNS = %w[
+          source cursor window_start window_end state last_error
+          rows_imported started_at finished_at
+        ].freeze
+        SOURCE_STARTED_AT_INDEX = %i[source started_at].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Call.connection
+            Adapter.ensure_supported!(connection)
+            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))
+            errors
+          end
+
+          private
+
+          def column_errors
+            missing = REQUIRED_COLUMNS - LlmCostTracker::ProviderInvoiceImport.columns_hash.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)
+
+            ["missing index: source, started_at"]
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -10,7 +10,7 @@ module LlmCostTracker
           source period_start period_end external_id billed_amount currency metadata imported_at
         ].freeze
         UNIQUE_INDEX_COLUMNS = %i[external_id].freeze
-        SOURCE_PERIOD_INDEX_COLUMNS = %i[source period_start].freeze
+        SOURCE_PERIOD_INDEX_COLUMNS = %i[source currency period_start].freeze
 
         class << self
           def current_schema_errors
@@ -46,7 +46,7 @@ module LlmCostTracker
               errors << "missing unique index: external_id"
             end
             unless connection.index_exists?(table_name, SOURCE_PERIOD_INDEX_COLUMNS)
-              errors << "missing index: source, period_start"
+              errors << "missing index: source, currency, period_start"
             end
             errors
           end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -5,6 +5,7 @@ require "json"
 require_relative "../pricing"
 require_relative "../billing/line_item"
 require_relative "rollups"
+require_relative "tags/encoding"
 
 module LlmCostTracker
   module Ledger
@@ -23,8 +24,8 @@ module LlmCostTracker
               call_ids = call_ids_for(insertable)
               insert_line_items(insertable, call_ids)
               insert_call_tags(insertable, call_ids)
-              Ledger::Rollups.increment_many!(insertable)
             end
+            increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
           end
           events
         end
@@ -119,14 +120,21 @@ module LlmCostTracker
         end
 
         def tag_row_value(value)
-          case value
-          when Hash, Array then JSON.generate(stored_tag_value(value))
-          else value.to_s
-          end
+          Tags::Encoding.encode(value)
         end
 
         def stored_details(details)
-          (details || {}).transform_keys(&:to_s).transform_values { |value| stored_tag_value(value) }
+          (details || {}).transform_keys(&:to_s).transform_values { |value| Tags::Encoding.normalize_value(value) }
+        end
+
+        def increment_rollups_safely(events)
+          Ledger::Rollups.increment_many!(events)
+        rescue StandardError => e
+          raise if LlmCostTracker::Call.connection.open_transactions.positive?
+
+          LlmCostTracker::Logging.warn(
+            "Rollup increment failed for #{events.size} events after ledger commit: #{e.class}: #{e.message}"
+          )
         end
 
         def insertable_events(events)
@@ -138,14 +146,6 @@ module LlmCostTracker
             !existing_ids.include?(event_id) && seen_ids.add?(event_id)
           end
         end
-
-        def stored_tag_value(value)
-          if value.is_a?(Hash)
-            return value.transform_keys(&:to_s).transform_values { |nested| stored_tag_value(nested) }
-          end
-
-          value.to_s
-        end
       end
     end
   end

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

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module Ledger
+    module Tags
+      module Encoding
+        module_function
+
+        def 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
+        end
+
+        def normalize_hash(hash)
+          hash.transform_keys(&:to_s).sort.to_h.transform_values { |v| normalize_value(v) }
+        end
+
+        def normalize_array(array)
+          array.map { |v| normalize_value(v) }
+        end
+
+        def normalize_value(value)
+          case value
+          when Hash then normalize_hash(value)
+          when Array then normalize_array(value)
+          else value.to_s
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "../schema/adapter"
+require_relative "encoding"
 
 module LlmCostTracker
   module Ledger
@@ -8,7 +9,7 @@ module LlmCostTracker
       module Query
         class << self
           def apply(tags)
-            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
+            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)|

data/lib/llm_cost_tracker/ledger.rb

@@ -5,7 +5,8 @@ require_relative "ledger/schema/calls"
 require_relative "ledger/schema/call_rollups"
 require_relative "ledger/schema/call_line_items"
 require_relative "ledger/schema/call_tags"
-require_relative "ledger/schema/provider_invoices"
+require_relative "ledger/schema/ingestion_inbox_entries"
+require_relative "ledger/schema/ingestion_leases"
 require_relative "ledger/tags/query"
 require_relative "ledger/tags/sql"
 require_relative "ledger/period"

data/lib/llm_cost_tracker/masking.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Masking
+    SENSITIVE_KEYS = %i[
+      provider_api_key_id provider_workspace_id provider_organization_id provider_project_id
+    ].to_set.freeze
+    MASK_TAIL_LENGTH = 4
+
+    module_function
+
+    def mask_value(key, value)
+      string = value.to_s
+      return string unless SENSITIVE_KEYS.include?(key.to_sym)
+      return string if string.length <= MASK_TAIL_LENGTH
+
+      "***#{string[-MASK_TAIL_LENGTH, MASK_TAIL_LENGTH]}"
+    end
+
+    def format_attribution(attribution, separator: ", ")
+      return "" if attribution.nil? || attribution.empty?
+
+      attribution.map { |key, value| "#{key}=#{mask_value(key, value)}" }.join(separator)
+    end
+
+    def mask_hash(hash)
+      return hash unless hash.is_a?(Hash)
+
+      hash.each_with_object({}) do |(key, value), masked|
+        masked[key] = case value
+                      when Hash then mask_hash(value)
+                      when Array then value.map { |entry| entry.is_a?(Hash) ? mask_hash(entry) : entry }
+                      else
+                        mask_value(key, value)
+                      end
+      end
+    end
+  end
+end