llm_cost_tracker 0.7.2 → 0.8.0

data/lib/llm_cost_tracker/errors.rb

@@ -6,16 +6,12 @@ module LlmCostTracker
   class InvalidFilterError < Error; end
 
   class BudgetExceededError < Error
-    attr_reader :monthly_total, :daily_total, :call_cost, :total, :budget, :budget_type, :last_event
-
-    def initialize(budget:, last_event: nil, budget_type: nil, total: nil, monthly_total: nil, daily_total: nil,
-                   call_cost: nil)
-      @monthly_total = monthly_total
-      @daily_total = daily_total
-      @call_cost = call_cost
-      @total = total || monthly_total || daily_total || call_cost
+    attr_reader :total, :budget, :budget_type, :last_event
+
+    def initialize(budget:, budget_type:, total:, last_event: nil)
+      @total = total
       @budget = budget
-      @budget_type = budget_type || inferred_budget_type
+      @budget_type = budget_type
       @last_event = last_event
 
       super(
@@ -23,16 +19,6 @@ module LlmCostTracker
         "$#{format('%.6f', @total)} / $#{format('%.6f', budget)}"
       )
     end
-
-    private
-
-    def inferred_budget_type
-      return :monthly if monthly_total
-      return :daily if daily_total
-      return :per_call if call_cost
-
-      :unknown
-    end
   end
 
   class UnknownPricingError < Error

data/lib/llm_cost_tracker/event.rb

@@ -13,7 +13,14 @@ module LlmCostTracker
     :stream,
     :usage_source,
     :provider_response_id,
-    :tracked_at
+    :provider_project_id,
+    :provider_api_key_id,
+    :provider_workspace_id,
+    :batch,
+    :tracked_at,
+    :cost_status,
+    :pricing_snapshot,
+    :line_items
   ) do
     def total_cost
       cost&.fetch(:total_cost, nil)
@@ -23,7 +30,8 @@ module LlmCostTracker
       super.merge(
         token_usage: token_usage.to_h,
         cost: cost&.to_h,
-        tags: tags ? tags.to_h : {}
+        tags: tags ? tags.to_h : {},
+        line_items: (line_items || []).map(&:to_h)
       )
     end
   end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/install_generator.rb

@@ -2,6 +2,8 @@
 
 require "rails/generators"
 require "rails/generators/active_record"
+require "llm_cost_tracker/billing/components"
+require "llm_cost_tracker/billing/cost_status"
 require "llm_cost_tracker/pricing"
 require "llm_cost_tracker/token_usage"
 
@@ -18,8 +20,8 @@ module LlmCostTracker
 
       def create_migration_file
         migration_template(
-          "create_llm_api_calls.rb.erb",
-          "db/migrate/create_llm_api_calls.rb"
+          "create_llm_cost_tracker_calls.rb.erb",
+          "db/migrate/create_llm_cost_tracker_calls.rb"
         )
       end
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/prices_generator.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 require "rails/generators"
+require "yaml"
 
 require_relative "../../pricing/registry"
-require_relative "../../pricing/sync/registry_loader"
 require_relative "../../pricing/sync/registry_writer"
 
 module LlmCostTracker
@@ -12,13 +12,9 @@ module LlmCostTracker
       desc "Creates a local LLM Cost Tracker price snapshot"
 
       def create_prices_file
-        registry = LlmCostTracker::Pricing::Sync::RegistryLoader.new.call(
-          path: LlmCostTracker::Pricing::Registry::DEFAULT_PRICES_PATH,
-          seed_path: LlmCostTracker::Pricing::Registry::DEFAULT_PRICES_PATH
-        )
         LlmCostTracker::Pricing::Sync::RegistryWriter.new.call(
           path: File.join(destination_root, "config/llm_cost_tracker_prices.yml"),
-          registry: registry
+          registry: YAML.safe_load_file(LlmCostTracker::Pricing::Registry::DEFAULT_PRICES_PATH, aliases: false) || {}
         )
       end
     end

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

@@ -0,0 +1,157 @@
+require "llm_cost_tracker/billing/components"
+require "llm_cost_tracker/billing/cost_status"
+require "llm_cost_tracker/ledger/schema/adapter"
+
+class CreateLlmCostTrackerCalls < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :llm_cost_tracker_calls do |t|
+      t.string  :event_id,      null: false
+      t.string  :provider,      null: false
+      t.string  :model,         null: false
+<% LlmCostTracker::TokenUsage.members.each do |column| -%>
+      t.integer :<%= column %>, null: false, default: 0
+<% end -%>
+      t.decimal :total_cost, precision: 20, scale: 8
+      t.integer :latency_ms
+      t.boolean :stream,        null: false, default: false
+      t.string  :usage_source
+      t.string  :provider_response_id
+      t.string  :provider_project_id
+      t.string  :provider_api_key_id
+      t.string  :provider_workspace_id
+      t.boolean :batch, null: false, default: false
+      t.string  :pricing_mode
+      t.string  :cost_status, null: false, default: LlmCostTracker::Billing::CostStatus::UNKNOWN
+      if postgresql?
+        t.jsonb :pricing_snapshot
+      elsif mysql?
+        t.json :pricing_snapshot
+      else
+        raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+      end
+      t.datetime :tracked_at,   null: false
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_call_rollups do |t|
+      t.string :period, null: false
+      t.date :period_start, null: false
+      t.string :currency, null: false, default: "USD"
+      t.decimal :total_cost, precision: 20, scale: 8, null: false, default: 0
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_ingestion_inbox_entries do |t|
+      t.string   :event_id, null: false
+      t.decimal  :total_cost, precision: 20, scale: 8
+      t.datetime :tracked_at, null: false
+      t.text     :payload, null: false
+      t.datetime :locked_at
+      t.string   :locked_by
+      t.integer  :attempts, null: false, default: 0
+      t.text     :last_error
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_ingestion_leases do |t|
+      t.string   :name, null: false
+      t.string   :locked_by
+      t.datetime :locked_until
+
+      t.timestamps
+    end
+
+    create_table :llm_cost_tracker_call_line_items do |t|
+      t.references :llm_cost_tracker_call,
+                   null: false,
+                   index: false,
+                   foreign_key: { to_table: :llm_cost_tracker_calls, on_delete: :cascade }
+      t.integer :position, null: false, default: 0, limit: 2
+      t.string :kind, null: false
+      t.string :direction, null: false
+      t.string :modality, null: false
+      t.string :cache_state, null: false, default: "none"
+      t.decimal :quantity, precision: 30, scale: 10, null: false
+      t.string :unit, null: false
+      t.decimal :rate_amount, precision: 20, scale: 8
+      t.decimal :rate_quantity, precision: 30, scale: 10, null: false, default: 1
+      t.decimal :cost, precision: 20, scale: 8
+      t.string :currency, null: false, default: "USD"
+      t.string :cost_status, null: false, default: LlmCostTracker::Billing::CostStatus::UNKNOWN
+      t.string :pricing_basis
+      t.string :price_key
+      t.string :price_source
+      t.string :price_source_version
+      t.string :provider_field
+      t.string :provider_item_id
+      if postgresql?
+        t.jsonb :details, null: false, default: {}
+      elsif mysql?
+        t.json :details, null: false
+      else
+        raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+      end
+
+      t.datetime :created_at, null: false
+    end
+
+    create_table :llm_cost_tracker_call_tags do |t|
+      t.references :llm_cost_tracker_call,
+                   null: false,
+                   index: false,
+                   foreign_key: { to_table: :llm_cost_tracker_calls, on_delete: :cascade }
+      t.string :key, null: false
+      t.text :value, null: false
+    end
+
+    create_table :llm_cost_tracker_provider_invoices do |t|
+      t.string :source, null: false
+      t.date :period_start, null: false
+      t.date :period_end, null: false
+      t.string :external_id, null: false
+      t.decimal :billed_amount, precision: 20, scale: 8
+      t.string :currency, null: false, default: "USD"
+      if postgresql?
+        t.jsonb :metadata, null: false, default: {}
+      elsif mysql?
+        t.json :metadata, null: false
+      else
+        raise "LLM Cost Tracker supports PostgreSQL and MySQL only"
+      end
+      t.datetime :imported_at, null: false
+
+      t.timestamps
+    end
+
+    add_index :llm_cost_tracker_calls, :event_id, unique: true
+    add_index :llm_cost_tracker_calls, :tracked_at
+    add_index :llm_cost_tracker_calls, [:provider, :tracked_at]
+    add_index :llm_cost_tracker_calls, [:model, :tracked_at]
+    add_index :llm_cost_tracker_calls, :cost_status
+    add_index :llm_cost_tracker_calls, :provider_response_id
+    add_index :llm_cost_tracker_call_line_items, [:llm_cost_tracker_call_id, :position]
+    add_index :llm_cost_tracker_call_line_items, :kind
+    add_index :llm_cost_tracker_call_tags, :llm_cost_tracker_call_id
+    add_index :llm_cost_tracker_call_tags, :key
+    add_index :llm_cost_tracker_call_rollups, [:period, :period_start, :currency], unique: true
+    add_index :llm_cost_tracker_ingestion_inbox_entries, :event_id, unique: true
+    add_index :llm_cost_tracker_ingestion_inbox_entries, [:tracked_at, :attempts]
+    add_index :llm_cost_tracker_ingestion_inbox_entries, [:locked_at, :id]
+    add_index :llm_cost_tracker_ingestion_leases, :name, unique: true
+    add_index :llm_cost_tracker_provider_invoices, :external_id, unique: true
+    add_index :llm_cost_tracker_provider_invoices, [:source, :period_start]
+  end
+
+  private
+
+  def postgresql?
+    LlmCostTracker::Ledger::Schema::Adapter.postgresql?(connection)
+  end
+
+  def mysql?
+    LlmCostTracker::Ledger::Schema::Adapter.mysql?(connection)
+  end
+end

data/lib/llm_cost_tracker/ingestion/batch.rb

@@ -27,7 +27,7 @@ module LlmCostTracker
       end
 
       def pending?
-        Ingestion::Event.where("attempts < ?", Ingestion::Event::MAX_ATTEMPTS).exists?
+        Ingestion::InboxEntry.where("attempts < ?", Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE).exists?
       end
 
       def claimable?
@@ -37,7 +37,7 @@ module LlmCostTracker
       def mark_failed(rows, error)
         message = "#{error.class}: #{error.message}".byteslice(0, 1_000)
         now = Time.now.utc
-        Ingestion::Event
+        Ingestion::InboxEntry
           .where(id: rows.map(&:id), locked_by: identity)
           .update_all(last_error: message, locked_at: now, locked_by: nil, updated_at: now)
       rescue StandardError
@@ -51,16 +51,15 @@ module LlmCostTracker
       def claim
         now = Time.now.utc
         cutoff = now - LOCK_TIMEOUT_SECONDS
-        Ingestion::Event.transaction do
+        Ingestion::InboxEntry.transaction do
           rows = claimable_scope(cutoff).order(:id).limit(BATCH_SIZE).lock.to_a
-          ids = rows.map(&:id)
-          next [] if ids.empty?
+          next [] if rows.empty?
 
-          updates = Ingestion::Event.sanitize_sql_array(
+          updates = Ingestion::InboxEntry.sanitize_sql_array(
             ["locked_at = ?, locked_by = ?, attempts = attempts + 1, updated_at = ?", now, identity, now]
           )
-          Ingestion::Event.where(id: ids).update_all(updates)
-          Ingestion::Event.where(id: ids, locked_by: identity).order(:id).to_a
+          Ingestion::InboxEntry.where(id: rows.map(&:id)).update_all(updates)
+          rows
         end
       end
 
@@ -77,15 +76,15 @@ module LlmCostTracker
       end
 
       def persist(rows, events)
-        LlmCostTracker::Ledger::Call.transaction do
+        LlmCostTracker::Call.transaction do
           Ledger::Store.insert_many(events)
-          Ingestion::Event.where(id: rows.map(&:id), locked_by: identity).delete_all
+          Ingestion::InboxEntry.where(id: rows.map(&:id), locked_by: identity).delete_all
         end
       end
 
       def claimable_scope(cutoff)
-        Ingestion::Event
-          .where("attempts < ?", Ingestion::Event::MAX_ATTEMPTS)
+        Ingestion::InboxEntry
+          .where("attempts < ?", Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE)
           .where("locked_at IS NULL OR locked_at < ?", cutoff)
       end
     end

data/lib/llm_cost_tracker/ingestion/inbox.rb

@@ -5,11 +5,12 @@ require "time"
 
 require_relative "../event"
 require_relative "../pricing"
+require_relative "../billing/line_item"
 
 module LlmCostTracker
   module Ingestion
     class Inbox
-      PAYLOAD_SCHEMA_VERSION = 1
+      PAYLOAD_SCHEMA_VERSION = 2
 
       class << self
         def save(event)
@@ -19,32 +20,47 @@ module LlmCostTracker
         end
 
         def event_from_row(row)
-          payload = JSON.parse(row.payload)
-          schema_version = payload.fetch("schema_version", 0)
-          unless [0, PAYLOAD_SCHEMA_VERSION].include?(schema_version)
+          payload = JSON.parse(row.payload, symbolize_names: true)
+          schema_version = payload[:schema_version]
+          unless schema_version == PAYLOAD_SCHEMA_VERSION
             raise LlmCostTracker::Error, "unsupported ledger inbox payload schema version #{schema_version.inspect}"
           end
 
-          cost = payload["cost"] && Pricing.stored_cost_attributes(payload["cost"])
-          token_usage = payload["token_usage"] || payload
+          LlmCostTracker::Event.new(**event_attributes_from(payload))
+        end
+
+        private
 
-          LlmCostTracker::Event.new(
-            event_id: payload.fetch("event_id"),
-            provider: payload.fetch("provider"),
-            model: payload.fetch("model"),
-            token_usage: TokenUsage.from_hash(token_usage),
-            pricing_mode: payload["pricing_mode"],
+        def event_attributes_from(payload)
+          cost = payload[:cost] && Pricing.stored_cost_attributes(payload[:cost])
+          token_usage = TokenUsage.build(**payload.fetch(:token_usage).slice(*TokenUsage.members))
+
+          {
+            event_id: payload.fetch(:event_id),
+            provider: payload.fetch(:provider),
+            model: payload.fetch(:model),
+            token_usage: token_usage,
+            pricing_mode: Pricing.normalize_mode(payload[:pricing_mode]),
             cost: cost,
-            tags: payload.fetch("tags"),
-            latency_ms: payload["latency_ms"],
-            stream: payload.fetch("stream"),
-            usage_source: payload["usage_source"],
-            provider_response_id: payload["provider_response_id"],
-            tracked_at: Time.iso8601(payload.fetch("tracked_at"))
-          )
+            tags: payload.fetch(:tags),
+            latency_ms: payload[:latency_ms],
+            stream: payload.fetch(:stream),
+            usage_source: payload[:usage_source]&.to_sym,
+            provider_response_id: payload[:provider_response_id],
+            provider_project_id: payload[:provider_project_id],
+            provider_api_key_id: payload[:provider_api_key_id],
+            provider_workspace_id: payload[:provider_workspace_id],
+            batch: payload.fetch(:batch),
+            tracked_at: Time.iso8601(payload.fetch(:tracked_at)),
+            cost_status: payload.fetch(:cost_status),
+            pricing_snapshot: payload[:pricing_snapshot],
+            line_items: line_items_from(payload)
+          }
         end
 
-        private
+        def line_items_from(payload)
+          (payload[:line_items] || []).map { |attributes| Billing::LineItem.build(attributes) }
+        end
 
         def row_for(event)
           now = Time.now.utc
@@ -70,7 +86,7 @@ module LlmCostTracker
         end
 
         def insert_row(row)
-          connection = LlmCostTracker::Ledger::Call.connection
+          connection = LlmCostTracker::Call.connection
           if connection.transaction_open?
             insert_with_separate_connection(row)
           else
@@ -82,7 +98,7 @@ module LlmCostTracker
         end
 
         def insert_with_separate_connection(row)
-          pool = LlmCostTracker::Ledger::Call.connection_pool
+          pool = LlmCostTracker::Call.connection_pool
           connection = pool.checkout
           begin
             connection.transaction(requires_new: true) { execute_insert(connection, row) }
@@ -95,7 +111,7 @@ module LlmCostTracker
           columns = row.keys
           quoted_columns = columns.map { |column| connection.quote_column_name(column) }.join(", ")
           quoted_values = columns.map { |column| connection.quote(row.fetch(column)) }.join(", ")
-          table = connection.quote_table_name(Event.table_name)
+          table = connection.quote_table_name(InboxEntry.table_name)
 
           connection.execute("INSERT INTO #{table} (#{quoted_columns}) VALUES (#{quoted_values})")
         end

data/lib/llm_cost_tracker/ingestion/worker.rb

@@ -25,18 +25,18 @@ module LlmCostTracker
               @generation = @generation.to_i + 1
               generation = @generation
               @thread = Thread.new { run(generation) }
-              @thread.name = "llm_cost_tracker_ingestor" if @thread.respond_to?(:name=)
-              @thread.report_on_exception = false if @thread.respond_to?(:report_on_exception=)
+              @thread.name = "llm_cost_tracker_ingestor"
+              @thread.report_on_exception = false
             end
             @thread
           end
           wake_thread(thread)
         end
 
-        def flush!(timeout: FLUSH_TIMEOUT_SECONDS, require_lease: false)
+        def flush!(timeout: nil, require_lease: false)
           Ingestion.ensure_current_schema!
 
-          deadline = Time.now.utc + timeout
+          deadline = Time.now.utc + flush_timeout_seconds(timeout)
           loop do
             return true unless Ingestion::Batch.new(identity: identity).pending?
             return false if Time.now.utc >= deadline
@@ -51,7 +51,8 @@ module LlmCostTracker
           end
         end
 
-        def shutdown!(timeout: FLUSH_TIMEOUT_SECONDS, drain: true)
+        def shutdown!(timeout: nil, drain: true)
+          timeout ||= FLUSH_TIMEOUT_SECONDS
           thread = mutex.synchronize do
             @stop_requested = true
             @generation = @generation.to_i + 1
@@ -80,7 +81,15 @@ module LlmCostTracker
           wake_thread(thread)
         end
 
+        def flush_timeout_seconds(timeout)
+          numeric = Float(timeout, exception: false)
+          return FLUSH_TIMEOUT_SECONDS unless numeric&.finite? && numeric.positive?
+
+          numeric
+        end
+
         def ingest_once(require_lease: true)
+          Ingestion.ensure_current_schema!
           batch = Ingestion::Batch.new(identity: identity)
           return 0 unless batch.claimable?
           return 0 if require_lease && !Ingestion::LeaseClaim.new(identity: identity, seconds: LEASE_SECONDS).acquire

data/lib/llm_cost_tracker/ingestion.rb

@@ -15,31 +15,38 @@ module LlmCostTracker
     VERIFY_TAG = "llm_cost_tracker_verify"
 
     class << self
+      def table_name_prefix
+        "llm_cost_tracker_ingestion_"
+      end
+
+      WRITE_SCHEMA_GUARDS = [
+        ["llm_cost_tracker_calls",            Ledger::Schema::Calls],
+        ["llm_cost_tracker_call_line_items",  Ledger::Schema::CallLineItems],
+        ["llm_cost_tracker_call_tags",        Ledger::Schema::CallTags],
+        ["llm_cost_tracker_call_rollups",     Ledger::Schema::CallRollups]
+      ].freeze
+
       def ensure_current_schema!
-        unless Ledger::Call.table_exists?
-          raise Error, "llm_api_calls table is missing; run install generator and migrate"
+        unless LlmCostTracker::Call.table_exists?
+          raise Error, "llm_cost_tracker_calls table is missing; run install generator and migrate"
         end
 
-        schema_errors = Ledger::Schema::Calls.current_schema_errors
-        message = "llm_api_calls table is not on the current schema: #{schema_errors.join('; ')}"
-        raise Error, message if schema_errors.any?
+        WRITE_SCHEMA_GUARDS.each do |table_name, schema_module|
+          errors = schema_module.current_schema_errors
+          next if errors.empty?
 
-        period_total_errors = Ledger::Schema::PeriodTotals.current_schema_errors
-        return if period_total_errors.empty?
-
-        message = "llm_cost_tracker_period_totals table is not on the current schema: " \
-                  "#{period_total_errors.join('; ')}; " \
-                  "run bin/rails generate llm_cost_tracker:add_period_totals && bin/rails db:migrate"
-        raise Error, message
+          raise Error,
+                "#{table_name} table is not on the current schema: #{errors.join('; ')}; see docs/upgrading.md"
+        end
       end
 
       def verify
-        unless LlmCostTracker::Ledger::Call.table_exists?
+        unless LlmCostTracker::Call.table_exists?
           return [
             LlmCostTracker::Doctor::Check.new(
               :error,
               "active_record",
-              "llm_api_calls table is missing; run install generator and migrate"
+              "llm_cost_tracker_calls table is missing; run install generator and migrate"
             )
           ]
         end
@@ -60,13 +67,12 @@ module LlmCostTracker
         event = LlmCostTracker.track(
           provider: provider,
           model: model,
-          input_tokens: 1,
-          output_tokens: 1,
+          tokens: { input: 1, output: 1 },
           provider_response_id: response_id,
-          feature: VERIFY_TAG
+          tags: { feature: VERIFY_TAG }
         )
-        LlmCostTracker.flush!
-        persisted = LlmCostTracker::Ledger::Call.where(provider_response_id: response_id).exists?
+        LlmCostTracker::Ingestion::Worker.flush!
+        persisted = LlmCostTracker::Call.where(provider_response_id: response_id).exists?
 
         return capture_success if persisted && notifications.any?
 
@@ -83,7 +89,7 @@ module LlmCostTracker
         LlmCostTracker::Doctor::Check.new(:error, "active_record capture", "#{e.class}: #{e.message}")
       ensure
         cleanup_verification_call(response_id) if response_id
-        LlmCostTracker::Ingestion::Event.where(event_id: event.event_id).delete_all if event
+        LlmCostTracker::Ingestion::InboxEntry.where(event_id: event.event_id).delete_all if event
         ActiveSupport::Notifications.unsubscribe(subscription) if subscription
       end
 
@@ -109,8 +115,8 @@ module LlmCostTracker
       end
 
       def cleanup_verification_call(response_id)
-        relation = LlmCostTracker::Ledger::Call.where(provider_response_id: response_id)
-        rows = relation.pluck(:id, :tracked_at, :total_cost)
+        relation = LlmCostTracker::Call.where(provider_response_id: response_id)
+        rows = relation.pluck(:id, :tracked_at, :total_cost, :pricing_snapshot)
         return if rows.empty?
 
         relation.delete_all