llm_cost_tracker 0.7.3 → 0.9.0

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      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
+            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
+            errors = []
+            missing = REQUIRED_COLUMNS - columns.keys
+            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 index on (#{required.join(', ')})"
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -15,23 +15,37 @@ 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
+          image_input_tokens
+          image_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 +62,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 +87,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/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

@@ -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 currency 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, currency, period_start"
+            end
+            errors
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -1,7 +1,11 @@
 # frozen_string_literal: true
 
+require "json"
+
 require_relative "../pricing"
+require_relative "../billing/line_item"
 require_relative "rollups"
+require_relative "tags/encoding"
 
 module LlmCostTracker
   module Ledger
@@ -14,9 +18,14 @@ module LlmCostTracker
           insertable = insertable_events(events)
 
           if insertable.any?
-            rows = insertable.map { |event| attributes_for(event) }
-            Ledger::Call.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)
+            end
+            increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
           end
           events
         end
@@ -28,40 +37,114 @@ module LlmCostTracker
             event_id: event.event_id,
             provider: event.provider,
             model: event.model,
-            tags: stored_tags(event.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(event.token_usage.stored_attributes)
+            .merge(event.token_usage.to_h)
             .merge(Pricing.stored_cost_attributes(event.cost || {}))
         end
 
-        def insertable_events(events)
-          existing_ids = Ledger::Call.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
-          seen_ids = Set.new
+        def call_ids_for(events)
+          LlmCostTracker::Call
+            .where(event_id: events.map(&:event_id))
+            .pluck(:event_id, :id)
+            .to_h
+        end
 
-          events.select do |event|
-            event_id = event.event_id
-            !existing_ids.include?(event_id) && seen_ids.add?(event_id)
+        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 stored_tags(tags)
-          (tags || {}).transform_keys(&:to_s).transform_values { |value| stored_tag_value(value) }
+        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 stored_tag_value(value)
-          if value.is_a?(Hash)
-            return value.transform_keys(&:to_s).transform_values { |nested| stored_tag_value(nested) }
+        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)
+          Tags::Encoding.encode(value)
+        end
 
-          value.to_s
+        def stored_details(details)
+          (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)
+          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
       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,8 +1,7 @@
 # frozen_string_literal: true
 
-require "json"
-
 require_relative "../schema/adapter"
+require_relative "encoding"
 
 module LlmCostTracker
   module Ledger
@@ -10,16 +9,12 @@ module LlmCostTracker
       module Query
         class << self
           def apply(tags)
-            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
-            return Ledger::Call.all if normalized_tags.empty?
-
-            connection = Ledger::Call.connection
-            json = normalized_tags.to_json
+            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values { |v| Encoding.encode(v) }
+            return LlmCostTracker::Call.all if normalized_tags.empty?
 
-            if Schema::Adapter.postgresql?(connection)
-              Ledger::Call.where("tags @> ?::jsonb", json)
-            else
-              Ledger::Call.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

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

@@ -1,31 +1,43 @@
 # frozen_string_literal: true
 
-require_relative "../schema/adapter"
 require_relative "../../tags/key"
 
 module LlmCostTracker
   module Ledger
     module Tags
       module Sql
+        UNTAGGED_LABEL = "(untagged)"
+
         class << self
-          def value_expression(key, table_name:)
-            key = LlmCostTracker::Tags::Key.validate!(key)
-            connection = Ledger::Call.connection
-            column = "#{table_name}.#{connection.quote_column_name('tags')}"
-
-            if Ledger::Schema::Adapter.postgresql?(connection)
-              "#{column}->>#{connection.quote(key)}"
-            elsif Ledger::Schema::Adapter.mysql?(connection)
-              "JSON_UNQUOTE(JSON_EXTRACT(#{column}, #{connection.quote(json_path(key))}))"
-            else
-              Ledger::Schema::Adapter.ensure_supported!(connection)
-            end
+          def join_relation(scope, key)
+            validated_key = LlmCostTracker::Tags::Key.validate!(key)
+            connection = scope.connection
+            join = "LEFT OUTER JOIN #{call_tag_table} ON " \
+                   "#{call_tag_table}.llm_cost_tracker_call_id = #{scope.quoted_table_name}.id AND " \
+                   "#{call_tag_table}.#{connection.quote_column_name('key')} = #{connection.quote(validated_key)}"
+            scope.joins(join)
+          end
+
+          def value_arel
+            Arel.sql("#{call_tag_table}.#{quote_column('value')}")
+          end
+
+          def label_sql(connection)
+            "COALESCE(NULLIF(#{raw_value_sql(connection)}, ''), #{connection.quote(UNTAGGED_LABEL)})"
+          end
+
+          def raw_value_sql(connection)
+            "#{call_tag_table}.#{connection.quote_column_name('value')}"
           end
 
           private
 
-          def json_path(key)
-            "$.\"#{key}\""
+          def call_tag_table
+            LlmCostTracker::CallTag.quoted_table_name
+          end
+
+          def quote_column(name)
+            LlmCostTracker::CallTag.connection.quote_column_name(name)
           end
         end
       end

data/lib/llm_cost_tracker/ledger.rb

@@ -2,11 +2,14 @@
 
 require_relative "ledger/schema/adapter"
 require_relative "ledger/schema/calls"
-require_relative "ledger/schema/period_totals"
+require_relative "ledger/schema/call_rollups"
+require_relative "ledger/schema/call_line_items"
+require_relative "ledger/schema/call_tags"
+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"
-require_relative "ledger/rollups/batch"
 require_relative "ledger/rollups/upsert_sql"
 require_relative "ledger/rollups"
 require_relative "ledger/store"

data/lib/llm_cost_tracker/logging.rb

@@ -20,12 +20,9 @@ module LlmCostTracker
       def log(level, message)
         message = prefixed(message)
         logger = Rails.logger
+        return Kernel.warn(message) unless logger
 
-        if logger
-          logger.try(level, message)
-        else
-          Kernel.warn(message)
-        end
+        logger.public_send(level, message)
       end
 
       private