llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -3,34 +3,25 @@
 require "bigdecimal"
 
 require_relative "period"
-require_relative "rollups/batch"
 require_relative "rollups/upsert_sql"
 
 module LlmCostTracker
   module Ledger
     class Rollups
+      DEFAULT_CURRENCY = "USD"
+
       class << self
         def increment!(event)
           return unless event.total_cost
 
-          Period::Total.upsert_all(
-            period_rows(event),
-            on_duplicate: Ledger::Rollups::UpsertSql.call(Period::Total),
-            record_timestamps: true,
-            unique_by: unique_by(Period::Total, %i[period period_start])
-          )
+          upsert_call_rollups(period_rows(event))
         end
 
         def increment_many!(events)
           events = Array(events).select(&:total_cost)
           return if events.empty?
 
-          Period::Total.upsert_all(
-            Ledger::Rollups::Batch.rows(events),
-            on_duplicate: Ledger::Rollups::UpsertSql.call(Period::Total),
-            record_timestamps: true,
-            unique_by: unique_by(Period::Total, %i[period period_start])
-          )
+          upsert_call_rollups(period_rows_for_events(events))
         end
 
         def decrement!(call_rows)
@@ -43,43 +34,103 @@ module LlmCostTracker
         private
 
         def period_rows(event)
+          currency = currency_for(event)
           Period::PERIODS.map do |period, name|
             {
               period: name,
               period_start: Period.bucket(period, event.tracked_at),
+              currency: currency,
               total_cost: event.total_cost
             }
           end
         end
 
+        def period_rows_for_events(events)
+          call_rollups(events).map do |(period, period_start, currency), total_cost|
+            {
+              period: period,
+              period_start: period_start,
+              currency: currency,
+              total_cost: total_cost
+            }
+          end
+        end
+
+        def call_rollups(events)
+          events.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |event, totals|
+            currency = currency_for(event)
+            Period::PERIODS.each do |period, name|
+              totals[[name, Period.bucket(period, event.tracked_at), currency]] += 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 = row
+            _id, tracked_at, total_cost, pricing_snapshot = row
             next unless total_cost
 
+            currency = currency_from_snapshot(pricing_snapshot)
             Period::PERIODS.each_key do |period|
-              totals[[period, Period.bucket(period, tracked_at)]] += BigDecimal(total_cost.to_s)
+              totals[[period, Period.bucket(period, tracked_at), currency]] += 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
+          end
 
-          totals.each do |(period, period_start), amount|
-            row = Period::Total.lock.find_by(period: Period::PERIODS.fetch(period),
-                                             period_start: period_start)
-            next unless row
-
-            row.update_columns(total_cost: [BigDecimal(row.total_cost.to_s) - amount, BigDecimal("0")].max,
-                               updated_at: now)
+          conn = LlmCostTracker::CallRollup.connection
+          table = LlmCostTracker::CallRollup.quoted_table_name
+          period_col = conn.quote_column_name("period")
+          start_col = conn.quote_column_name("period_start")
+          currency_col = conn.quote_column_name("currency")
+          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|
+            case_clauses = by_start.map do |period_start, amount|
+              "WHEN #{start_col} = #{conn.quote(period_start)} THEN #{conn.quote(amount)}"
+            end.join(" ")
+            starts = by_start.keys.map { |period_start| conn.quote(period_start) }.join(", ")
+
+            conn.execute(
+              "UPDATE #{table} " \
+              "SET #{total_col} = GREATEST(0, #{total_col} - CASE #{case_clauses} ELSE 0 END), " \
+              "#{updated_col} = #{conn.quote(now)} " \
+              "WHERE #{period_col} = #{conn.quote(Period::PERIODS.fetch(period))} " \
+              "AND #{currency_col} = #{conn.quote(currency)} " \
+              "AND #{start_col} IN (#{starts})"
+            )
           end
         end
 
-        def unique_by(model, column)
-          return unless model.connection.supports_insert_conflict_target?
+        def currency_for(event)
+          snapshot = event.respond_to?(:pricing_snapshot) ? event.pricing_snapshot : nil
+          currency_from_snapshot(snapshot)
+        end
+
+        def currency_from_snapshot(snapshot)
+          (snapshot.is_a?(Hash) && (snapshot["currency"] || snapshot[:currency])) || DEFAULT_CURRENCY
+        end
+
+        def upsert_call_rollups(rows)
+          LlmCostTracker::CallRollup.upsert_all(
+            rows,
+            on_duplicate: Ledger::Rollups::UpsertSql.call,
+            record_timestamps: true,
+            unique_by: call_rollups_unique_by
+          )
+        end
+
+        def call_rollups_unique_by
+          return unless LlmCostTracker::CallRollup.connection.supports_insert_conflict_target?
 
-          column
+          %i[period period_start currency]
         end
       end
     end

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

@@ -32,6 +32,24 @@ module LlmCostTracker
             raise Error, "Unsupported database adapter: #{adapter_name(value)}. Use PostgreSQL or MySQL."
           end
 
+          def json_column_errors(column, adapter_value, column_name)
+            return [] unless column
+
+            expected_type = postgresql?(adapter_value) ? "jsonb" : "json"
+            return [] if json_column_type?(column, adapter_value)
+
+            ["#{column_name} column must use #{expected_type} (got #{column.sql_type})"]
+          end
+
+          def json_column_type?(column, adapter_value)
+            sql_type = column.sql_type.to_s.downcase
+            if postgresql?(adapter_value)
+              column.type == :jsonb || sql_type == "jsonb"
+            else
+              column.type == :json || sql_type == "json" || sql_type == "longtext"
+            end
+          end
+
           private
 
           def adapter_instance?(value, class_names)

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

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module CallLineItems
+        REQUIRED_COLUMNS = %w[
+          llm_cost_tracker_call_id
+          position
+          kind
+          direction
+          modality
+          cache_state
+          quantity
+          unit
+          rate_amount
+          rate_quantity
+          cost
+          currency
+          cost_status
+          pricing_basis
+          price_key
+          price_source
+          price_source_version
+          provider_field
+          provider_item_id
+          details
+        ].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Call.connection
+            Adapter.ensure_supported!(connection)
+            table_name = LlmCostTracker::CallLineItem.table_name
+            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+
+            columns = LlmCostTracker::CallLineItem.columns_hash
+            errors = []
+            missing = REQUIRED_COLUMNS - columns.keys
+            errors << "missing columns: #{missing.join(', ')}" if missing.any?
+            errors.concat(Adapter.json_column_errors(columns["details"], connection, "details"))
+            errors
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+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
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::CallRollup.connection
+            Adapter.ensure_supported!(connection)
+            table_name = LlmCostTracker::CallRollup.table_name
+            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+
+            errors = []
+            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"
+            end
+            errors
+          end
+
+          private
+
+          def unique_period_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/call_tags.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module CallTags
+        REQUIRED_COLUMNS = %w[llm_cost_tracker_call_id key value].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Call.connection
+            Ledger::Schema::Adapter.ensure_supported!(connection)
+            table_name = LlmCostTracker::CallTag.table_name
+            return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
+
+            columns = LlmCostTracker::CallTag.columns_hash
+            missing = REQUIRED_COLUMNS - columns.keys
+            return [] if missing.empty?
+
+            ["missing columns: #{missing.join(', ')}"]
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -15,23 +15,35 @@ module LlmCostTracker
           total_tokens
           cache_read_input_tokens
           cache_write_input_tokens
-          cache_write_1h_input_tokens
+          cache_write_extended_input_tokens
+          audio_input_tokens
+          audio_output_tokens
           hidden_output_tokens
-          input_cost
-          output_cost
           total_cost
-          cache_read_input_cost
-          cache_write_input_cost
-          cache_write_1h_input_cost
           latency_ms
           stream
           usage_source
           provider_response_id
+          provider_project_id
+          provider_api_key_id
+          provider_workspace_id
+          batch
           pricing_mode
-          tags
+          cost_status
+          pricing_snapshot
           tracked_at
         ].freeze
 
+        REQUIRED_INDEXES = [
+          { columns: :event_id, unique: true },
+          { columns: :tracked_at },
+          { columns: %i[provider tracked_at] },
+          { columns: %i[model tracked_at] },
+          { columns: :cost_status },
+          { columns: :provider_response_id }
+        ].freeze
+        private_constant :REQUIRED_INDEXES
+
         class << self
           def current_schema?
             current_schema_errors.empty?
@@ -48,8 +60,8 @@ module LlmCostTracker
           private
 
           def schema_capabilities
-            columns = Ledger::Call.columns_hash
-            adapter_name = Ledger::Call.connection.adapter_name
+            columns = LlmCostTracker::Call.columns_hash
+            adapter_name = LlmCostTracker::Call.connection.adapter_name
             cache = @schema_capabilities
 
             return cache.fetch(:values) if cache && cache.fetch(:columns).equal?(columns) &&
@@ -73,22 +85,21 @@ module LlmCostTracker
             errors = []
             missing = missing_columns_for(columns)
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
+            errors.concat(Adapter.json_column_errors(columns["pricing_snapshot"], adapter_name, "pricing_snapshot"))
+            errors.concat(missing_index_errors)
+            errors
+          end
 
-            tag_column = columns["tags"]
-            if tag_column
-              postgresql = Ledger::Schema::Adapter.postgresql?(adapter_name)
-              expected_type = postgresql ? "jsonb" : "json"
-              valid_type =
-                if postgresql
-                  tag_column.type == :jsonb || tag_column.sql_type.to_s.downcase == "jsonb"
-                else
-                  tag_column.type == :json
-                end
-
-              errors << "tags column must use #{expected_type}" unless valid_type
-            end
+          def missing_index_errors
+            connection = LlmCostTracker::Call.connection
+            REQUIRED_INDEXES.filter_map do |spec|
+              next if connection.index_exists?(LlmCostTracker::Call.table_name, spec[:columns], **spec.except(:columns))
 
-            errors
+              prefix = spec[:unique] ? "unique " : ""
+              "missing #{prefix}index: #{Array(spec[:columns]).join(', ')}"
+            end
+          rescue StandardError
+            []
           end
 
           def missing_columns_for(columns)

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

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module ProviderInvoices
+        REQUIRED_COLUMNS = %w[
+          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
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Call.connection
+            Adapter.ensure_supported!(connection)
+            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))
+            errors
+          end
+
+          private
+
+          def column_errors
+            missing = REQUIRED_COLUMNS - LlmCostTracker::ProviderInvoice.columns_hash.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)
+              errors << "missing unique index: external_id"
+            end
+            unless connection.index_exists?(table_name, SOURCE_PERIOD_INDEX_COLUMNS)
+              errors << "missing index: source, period_start"
+            end
+            errors
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
+require "json"
+
 require_relative "../pricing"
+require_relative "../billing/line_item"
 require_relative "rollups"
 
 module LlmCostTracker
@@ -11,13 +14,17 @@ module LlmCostTracker
           events = Array(events)
           return [] if events.empty?
 
-          model = LlmCostTracker::Ledger::Call
-          insertable = new_events(model, events)
+          insertable = insertable_events(events)
 
           if insertable.any?
-            rows = insertable.map { |event| attributes_for(event) }
-            model.insert_all!(rows, record_timestamps: true, returning: false)
-            Ledger::Rollups.increment_many!(insertable)
+            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)
+              Ledger::Rollups.increment_many!(insertable)
+            end
           end
           events
         end
@@ -25,32 +32,117 @@ module LlmCostTracker
         private
 
         def attributes_for(event)
-          tags = (event.tags || {}).transform_keys(&:to_s).transform_values { |value| stringify_tag_value(value) }
-          usage = event.token_usage.stored_attributes
-
           attributes = {
             event_id: event.event_id,
             provider: event.provider,
             model: event.model,
-            tags: tags,
             tracked_at: event.tracked_at,
-            pricing_mode: event.pricing_mode,
+            pricing_mode: event.pricing_mode&.name,
             latency_ms: event.latency_ms,
             stream: event.stream,
-            usage_source: event.usage_source,
-            provider_response_id: event.provider_response_id
+            usage_source: event.usage_source&.name,
+            provider_response_id: event.provider_response_id,
+            provider_project_id: event.provider_project_id,
+            provider_api_key_id: event.provider_api_key_id,
+            provider_workspace_id: event.provider_workspace_id,
+            batch: event.batch,
+            cost_status: event.cost_status,
+            pricing_snapshot: event.pricing_snapshot
           }
 
-          attributes.merge(usage).merge(Pricing.stored_cost_attributes(event.cost || {}))
+          attributes
+            .merge(event.token_usage.to_h)
+            .merge(Pricing.stored_cost_attributes(event.cost || {}))
+        end
+
+        def call_ids_for(events)
+          LlmCostTracker::Call
+            .where(event_id: events.map(&:event_id))
+            .pluck(:event_id, :id)
+            .to_h
+        end
+
+        def insert_line_items(events, call_ids)
+          rows = events.flat_map do |event|
+            (event.line_items || []).each_with_index.map do |line_item, position|
+              line_item_attributes(
+                call_id: call_ids.fetch(event.event_id),
+                line_item: line_item,
+                position: position
+              )
+            end
+          end
+          return if rows.empty?
+
+          LlmCostTracker::CallLineItem.insert_all!(rows, record_timestamps: false, returning: false)
         end
 
-        def new_events(model, events)
-          existing_ids = model.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
-          events.reject { |event| existing_ids.include?(event.event_id) }
+        def line_item_attributes(call_id:, line_item:, position:)
+          {
+            llm_cost_tracker_call_id: call_id,
+            position: position,
+            kind: line_item.kind&.to_s,
+            direction: line_item.direction&.to_s,
+            modality: line_item.modality&.to_s,
+            cache_state: line_item.cache_state&.to_s || "none",
+            quantity: line_item.quantity,
+            unit: line_item.unit&.to_s,
+            rate_amount: line_item.rate_amount,
+            rate_quantity: line_item.rate_quantity,
+            cost: line_item.cost,
+            currency: line_item.currency,
+            cost_status: line_item.cost_status,
+            pricing_basis: line_item.pricing_basis&.to_s,
+            price_key: line_item.price_key,
+            price_source: line_item.price_source&.to_s,
+            price_source_version: line_item.price_source_version,
+            provider_field: line_item.provider_field,
+            provider_item_id: line_item.provider_item_id,
+            details: stored_details(line_item.details),
+            created_at: Time.now.utc
+          }
         end
 
-        def stringify_tag_value(value)
-          return value.transform_values { |nested| stringify_tag_value(nested) } if value.is_a?(Hash)
+        def insert_call_tags(events, call_ids)
+          rows = events.flat_map do |event|
+            (event.tags || {}).map do |key, value|
+              {
+                llm_cost_tracker_call_id: call_ids.fetch(event.event_id),
+                key: key.to_s,
+                value: tag_row_value(value)
+              }
+            end
+          end
+          return if rows.empty?
+
+          LlmCostTracker::CallTag.insert_all!(rows, record_timestamps: false, returning: false)
+        end
+
+        def tag_row_value(value)
+          case value
+          when Hash, Array then JSON.generate(stored_tag_value(value))
+          else value.to_s
+          end
+        end
+
+        def stored_details(details)
+          (details || {}).transform_keys(&:to_s).transform_values { |value| stored_tag_value(value) }
+        end
+
+        def insertable_events(events)
+          existing_ids = LlmCostTracker::Call.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
+          seen_ids = Set.new
+
+          events.select do |event|
+            event_id = event.event_id
+            !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

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

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "json"
-
 require_relative "../schema/adapter"
 
 module LlmCostTracker
@@ -9,17 +7,13 @@ module LlmCostTracker
     module Tags
       module Query
         class << self
-          def apply(model, tags)
+          def apply(tags)
             normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
-            return model.all if normalized_tags.empty?
-
-            connection = model.connection
-            json = normalized_tags.to_json
+            return LlmCostTracker::Call.all if normalized_tags.empty?
 
-            if Schema::Adapter.postgresql?(connection)
-              model.where("tags @> ?::jsonb", json)
-            else
-              model.where("JSON_CONTAINS(tags, ?)", json)
+            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))
             end
           end
         end