llm_cost_tracker 0.1.0 → 0.1.2

data/lib/llm_cost_tracker/configuration.rb

@@ -2,15 +2,31 @@
 
 module LlmCostTracker
   class Configuration
+    # Hostname => provider name for OpenAI-compatible APIs.
+    OPENAI_COMPATIBLE_PROVIDERS = {
+      "openrouter.ai" => "openrouter",
+      "api.deepseek.com" => "deepseek"
+    }.freeze
+
+    BUDGET_EXCEEDED_BEHAVIORS = %i[notify raise block_requests].freeze
+    STORAGE_ERROR_BEHAVIORS = %i[ignore warn raise].freeze
+    UNKNOWN_PRICING_BEHAVIORS = %i[ignore warn raise].freeze
+
     attr_accessor :enabled,
                   :storage_backend,    # :log, :active_record, :custom
                   :custom_storage,     # callable object for :custom backend
                   :default_tags,       # Hash of default tags added to every event
                   :on_budget_exceeded, # callable, receives event hash
                   :monthly_budget,     # Float, in USD — nil means no limit
+                  :budget_exceeded_behavior, # :notify, :raise, :block_requests
+                  :storage_error_behavior, # :ignore, :warn, :raise
+                  :unknown_pricing_behavior, # :ignore, :warn, :raise
                   :log_level,          # :debug, :info, :warn
+                  :prices_file,        # JSON/YAML file that overrides built-in prices
                   :pricing_overrides   # Hash to override built-in pricing
 
+    attr_reader :openai_compatible_providers
+
     def initialize
       @enabled            = true
       @storage_backend    = :log
@@ -18,8 +34,21 @@ module LlmCostTracker
       @default_tags       = {}
       @on_budget_exceeded = nil
       @monthly_budget     = nil
+      @budget_exceeded_behavior = :notify
+      @storage_error_behavior = :warn
+      @unknown_pricing_behavior = :warn
       @log_level          = :info
+      @prices_file        = nil
       @pricing_overrides  = {}
+      self.openai_compatible_providers = OPENAI_COMPATIBLE_PROVIDERS
+    end
+
+    def openai_compatible_providers=(providers)
+      @openai_compatible_providers = normalize_openai_compatible_providers(providers)
+    end
+
+    def normalize_openai_compatible_providers!
+      self.openai_compatible_providers = openai_compatible_providers
     end
 
     def active_record?
@@ -29,5 +58,13 @@ module LlmCostTracker
     def log?
       storage_backend == :log
     end
+
+    private
+
+    def normalize_openai_compatible_providers(providers)
+      (providers || {}).each_with_object({}) do |(host, provider), normalized|
+        normalized[host.to_s.downcase] = provider.to_s
+      end
+    end
   end
 end

data/lib/llm_cost_tracker/errors.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  class Error < StandardError; end
+
+  class BudgetExceededError < Error
+    attr_reader :monthly_total, :budget, :last_event
+
+    def initialize(monthly_total:, budget:, last_event: nil)
+      @monthly_total = monthly_total
+      @budget = budget
+      @last_event = last_event
+
+      super("LLM monthly budget exceeded: $#{format('%.6f', monthly_total)} / $#{format('%.6f', budget)}")
+    end
+  end
+
+  class UnknownPricingError < Error
+    attr_reader :model
+
+    def initialize(model:)
+      @model = model
+
+      super("No pricing configured for LLM model: #{model.inspect}")
+    end
+  end
+
+  class StorageError < Error
+    attr_reader :original_error
+
+    def initialize(original_error)
+      @original_error = original_error
+
+      super("Failed to store LLM cost event: #{original_error.class}: #{original_error.message}")
+    end
+  end
+end

data/lib/llm_cost_tracker/event_metadata.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module EventMetadata
+    INTERNAL_TAG_KEYS = %w[
+      cache_creation_input_tokens
+      cache_creation_tokens
+      cache_read_input_tokens
+      cache_read_tokens
+      cached_input_tokens
+      input_tokens
+      output_tokens
+      reasoning_tokens
+      total_tokens
+    ].freeze
+
+    class << self
+      def usage_data(input_tokens, output_tokens, metadata)
+        cache_read_input_tokens = integer_metadata(metadata, :cache_read_input_tokens, :cache_read_tokens)
+        cache_creation_input_tokens = integer_metadata(
+          metadata,
+          :cache_creation_input_tokens,
+          :cache_creation_tokens
+        )
+        cached_input_tokens = integer_metadata(metadata, :cached_input_tokens)
+
+        {
+          input_tokens: input_tokens.to_i,
+          output_tokens: output_tokens.to_i,
+          cached_input_tokens: cached_input_tokens,
+          cache_read_input_tokens: cache_read_input_tokens,
+          cache_creation_input_tokens: cache_creation_input_tokens,
+          total_tokens: input_tokens.to_i + output_tokens.to_i +
+            cache_read_input_tokens + cache_creation_input_tokens
+        }
+      end
+
+      def tags(metadata)
+        metadata.reject { |key, _value| INTERNAL_TAG_KEYS.include?(key.to_s) }
+      end
+
+      private
+
+      def integer_metadata(metadata, *keys)
+        keys.each do |key|
+          value = metadata[key] || metadata[key.to_s]
+          return value.to_i unless value.nil?
+        end
+
+        0
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/add_latency_ms_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class AddLatencyMsGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates a migration to add llm_api_calls.latency_ms"
+
+      def create_migration_file
+        migration_template(
+          "add_latency_ms_to_llm_api_calls.rb.erb",
+          "db/migrate/add_latency_ms_to_llm_api_calls.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

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

@@ -0,0 +1,9 @@
+class AddLatencyMsToLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
+  def up
+    add_column :llm_api_calls, :latency_ms, :integer unless column_exists?(:llm_api_calls, :latency_ms)
+  end
+
+  def down
+    remove_column :llm_api_calls, :latency_ms if column_exists?(:llm_api_calls, :latency_ms)
+  end
+end

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

@@ -6,10 +6,15 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
       t.integer :input_tokens,  null: false, default: 0
       t.integer :output_tokens, null: false, default: 0
       t.integer :total_tokens,  null: false, default: 0
-      t.decimal :input_cost,    precision: 12, scale: 8
-      t.decimal :output_cost,   precision: 12, scale: 8
-      t.decimal :total_cost,    precision: 12, scale: 8
-      t.text    :tags
+      t.decimal :input_cost,    precision: 20, scale: 8
+      t.decimal :output_cost,   precision: 20, scale: 8
+      t.decimal :total_cost,    precision: 20, scale: 8
+      t.integer :latency_ms
+      if postgresql?
+        t.jsonb :tags, null: false, default: {}
+      else
+        t.text :tags
+      end
       t.datetime :tracked_at,   null: false
 
       t.timestamps
@@ -19,5 +24,12 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
     add_index :llm_api_calls, :model
     add_index :llm_api_calls, :tracked_at
     add_index :llm_api_calls, [:provider, :tracked_at]
+    add_index :llm_api_calls, :tags, using: :gin if postgresql?
+  end
+
+  private
+
+  def postgresql?
+    connection.adapter_name.downcase.include?("postgres")
   end
 end

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

@@ -12,6 +12,13 @@ LlmCostTracker.configure do |config|
 
   # Monthly budget in USD. Set to nil to disable budget alerts.
   # config.monthly_budget = 100.00
+  # config.budget_exceeded_behavior = :notify # :notify, :raise, or :block_requests
+
+  # What to do when storage fails.
+  # config.storage_error_behavior = :warn # :ignore, :warn, or :raise
+
+  # What to do when a model has no built-in price and no pricing_overrides entry.
+  # config.unknown_pricing_behavior = :warn # :ignore, :warn, or :raise
 
   # Callback when monthly budget is exceeded.
   # config.on_budget_exceeded = ->(data) {
@@ -20,8 +27,14 @@ LlmCostTracker.configure do |config|
   #   # Or send a Slack notification, email, etc.
   # }
 
-  # Override built-in pricing for specific models (per 1M tokens, USD)
+  # Load a local JSON/YAML price table that overrides built-in pricing.
+  # config.prices_file = Rails.root.join("config/llm_cost_tracker_prices.json")
+
+  # Override pricing for specific models in Ruby (per 1M tokens, USD).
   # config.pricing_overrides = {
   #   "my-custom-model" => { input: 1.00, output: 2.00 }
   # }
+
+  # OpenAI-compatible APIs. OpenRouter and DeepSeek are included by default.
+  # config.openai_compatible_providers["llm.my-company.com"] = "internal_gateway"
 end

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

@@ -0,0 +1,15 @@
+class UpgradeLlmApiCallCostPrecision < ActiveRecord::Migration<%= migration_version %>
+  COST_COLUMNS = %i[input_cost output_cost total_cost].freeze
+
+  def up
+    COST_COLUMNS.each do |column|
+      change_column :llm_api_calls, column, :decimal, precision: 20, scale: 8
+    end
+  end
+
+  def down
+    COST_COLUMNS.each do |column|
+      change_column :llm_api_calls, column, :decimal, precision: 12, scale: 8
+    end
+  end
+end

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

@@ -0,0 +1,41 @@
+class UpgradeLlmApiCallTagsToJsonb < ActiveRecord::Migration<%= migration_version %>
+  def up
+    unless postgresql?
+      say "Skipping llm_api_calls.tags JSONB upgrade: database adapter is #{connection.adapter_name}."
+      return
+    end
+
+    return if tags_jsonb?
+
+    remove_index :llm_api_calls, :tags if index_exists?(:llm_api_calls, :tags)
+
+    change_column(
+      :llm_api_calls,
+      :tags,
+      :jsonb,
+      using: "CASE WHEN tags IS NULL OR tags = '' THEN '{}'::jsonb ELSE tags::jsonb END",
+      default: {},
+      null: false
+    )
+
+    add_index :llm_api_calls, :tags, using: :gin unless index_exists?(:llm_api_calls, :tags)
+  end
+
+  def down
+    return unless postgresql?
+
+    remove_index :llm_api_calls, :tags if index_exists?(:llm_api_calls, :tags)
+    change_column :llm_api_calls, :tags, :text, using: "tags::text"
+  end
+
+  private
+
+  def postgresql?
+    connection.adapter_name.downcase.include?("postgres")
+  end
+
+  def tags_jsonb?
+    column = connection.columns(:llm_api_calls).find { |candidate| candidate.name == "tags" }
+    column&.sql_type.to_s.downcase == "jsonb"
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_cost_precision_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeCostPrecisionGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates a migration to widen llm_api_calls cost decimal precision"
+
+      def create_migration_file
+        migration_template(
+          "upgrade_llm_api_call_cost_precision.rb.erb",
+          "db/migrate/upgrade_llm_api_call_cost_precision.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/upgrade_tags_to_jsonb_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class UpgradeTagsToJsonbGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates a migration to upgrade llm_api_calls.tags to PostgreSQL JSONB"
+
+      def create_migration_file
+        migration_template(
+          "upgrade_llm_api_call_tags_to_jsonb.rb.erb",
+          "db/migrate/upgrade_llm_api_call_tags_to_jsonb.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/llm_api_call.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "active_record"
+require "json"
 
 module LlmCostTracker
   class LlmApiCall < ActiveRecord::Base
@@ -9,7 +10,34 @@ module LlmCostTracker
     # Scopes for querying
     scope :by_provider, ->(provider) { where(provider: provider) }
     scope :by_model,    ->(model)    { where(model: model) }
-    scope :by_tag,      ->(key, value) { where("tags LIKE ?", "%\"#{key}\":\"#{value}\"%") }
+    scope :by_tag, ->(key, value) { by_tags(key => value) }
+    scope :by_tags, lambda { |tags|
+      normalized_tags = normalize_tags(tags)
+
+      if normalized_tags.empty?
+        all
+      elsif tags_json_column?
+        where("tags @> ?::jsonb", normalized_tags.to_json)
+      else
+        normalized_tags.reduce(all) do |relation, (key, value)|
+          relation.where("tags LIKE ? ESCAPE '\\'", "%#{sanitize_sql_like(json_tag_fragment(key, value))}%")
+        end
+      end
+    }
+    scope :by_user,    ->(user_id) { by_tag("user_id", user_id) }
+    scope :by_feature, ->(feature) { by_tag("feature", feature) }
+    scope :with_cost, -> { where.not(total_cost: nil) }
+    scope :without_cost, -> { where(total_cost: nil) }
+    scope :unknown_pricing, -> { without_cost }
+    scope :with_latency, -> { latency_column? ? where.not(latency_ms: nil) : none }
+
+    scope :with_json_tags, lambda {
+      if tags_json_column?
+        where.not(tags: {})
+      else
+        where.not(tags: [nil, "", "{}"])
+      end
+    }
 
     scope :today,       -> { where(tracked_at: Time.now.utc.beginning_of_day..) }
     scope :this_week,   -> { where(tracked_at: Time.now.utc.beginning_of_week..) }
@@ -33,13 +61,53 @@ module LlmCostTracker
       group(:provider).sum(:total_cost)
     end
 
+    def self.average_latency_ms
+      return nil unless latency_column?
+
+      average(:latency_ms)&.to_f
+    end
+
+    def self.latency_by_model
+      return {} unless latency_column?
+
+      group(:model).average(:latency_ms).transform_values(&:to_f)
+    end
+
+    def self.latency_by_provider
+      return {} unless latency_column?
+
+      group(:provider).average(:latency_ms).transform_values(&:to_f)
+    end
+
     def self.daily_costs(days: 30)
       where(tracked_at: days.days.ago..)
         .group("DATE(tracked_at)")
         .sum(:total_cost)
+        .transform_keys(&:to_s)
+    end
+
+    def self.tags_json_column?
+      column = columns_hash["tags"]
+      return false unless column
+
+      %i[json jsonb].include?(column.type) || column.sql_type.to_s.downcase == "jsonb"
+    end
+
+    def self.latency_column?
+      columns_hash.key?("latency_ms")
+    end
+
+    def self.normalize_tags(tags)
+      (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
+    end
+
+    def self.json_tag_fragment(key, value)
+      JSON.generate(key => value).delete_prefix("{").delete_suffix("}")
     end
 
     def parsed_tags
+      return tags.transform_keys(&:to_s) if tags.is_a?(Hash)
+
       JSON.parse(tags || "{}")
     rescue JSON::ParserError
       {}

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "faraday"
+require "json"
 
 module LlmCostTracker
   module Middleware
@@ -14,28 +15,23 @@ module LlmCostTracker
         return @app.call(request_env) unless LlmCostTracker.configuration.enabled
 
         request_url  = request_env.url.to_s
-        request_body = read_body(request_env.body)
+        request_body = read_body(request_env.body) || ""
+
+        enforce_budget_before_request(request_url)
+        started_at = monotonic_time
 
         @app.call(request_env).on_complete do |response_env|
-          process(request_url, request_body, response_env)
+          process(request_url, request_body, response_env, elapsed_ms(started_at))
         end
-      rescue StandardError => e
-        # Never break the actual request — log and re-raise
-        raise e
       end
 
       private
 
-      def process(request_url, request_body, response_env)
+      def process(request_url, request_body, response_env, latency_ms)
         parser = Parsers::Registry.find_for(request_url)
         return unless parser
 
-        parsed = parser.parse(
-          request_url,
-          request_body,
-          response_env.status,
-          read_body(response_env.body)
-        )
+        parsed = parse_response(parser, request_url, request_body, response_env)
         return unless parsed
 
         Tracker.record(
@@ -43,19 +39,60 @@ module LlmCostTracker
           model: parsed[:model],
           input_tokens: parsed[:input_tokens],
           output_tokens: parsed[:output_tokens],
+          latency_ms: latency_ms,
           metadata: @tags.merge(parsed.except(:provider, :model, :input_tokens, :output_tokens, :total_tokens))
         )
+      rescue LlmCostTracker::Error
+        raise
       rescue StandardError => e
-        warn "[LlmCostTracker] Error processing response: #{e.message}" if LlmCostTracker.configuration.log_level == :debug
+        log_warning("Error processing response: #{e.class}: #{e.message}")
+      end
+
+      def parse_response(parser, request_url, request_body, response_env)
+        response_body = read_body(response_env.body)
+        unless response_body
+          log_warning(
+            "Unable to read response body for #{request_url}; streaming/SSE responses require manual tracking."
+          )
+          return nil
+        end
+
+        parser.parse(request_url, request_body, response_env.status, response_body)
+      end
+
+      def enforce_budget_before_request(request_url)
+        return unless Parsers::Registry.find_for(request_url)
+
+        Tracker.enforce_budget!
       end
 
       def read_body(body)
         case body
         when String then body
         when nil then ""
-        else body.to_s
+        when Hash, Array then body.to_json
+        else
+          body.respond_to?(:to_str) ? body.to_str : nil
+        end
+      end
+
+      def log_warning(message)
+        message = "[LlmCostTracker] #{message}"
+
+        if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
+          Rails.logger.warn(message)
+        else
+          warn message
         end
       end
+
+      def monotonic_time
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+
+      def elapsed_ms(started_at)
+        ((monotonic_time - started_at) * 1000).round
+      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require "json"
+require "uri"
+
+require_relative "base"
 
 module LlmCostTracker
   module Parsers
@@ -9,12 +12,12 @@ module LlmCostTracker
 
       def match?(url)
         uri = URI.parse(url.to_s)
-        HOSTS.include?(uri.host) && uri.path.include?("/v1/messages")
+        HOSTS.include?(uri.host.to_s.downcase) && uri.path.include?("/v1/messages")
       rescue URI::InvalidURIError
         false
       end
 
-      def parse(request_url, request_body, response_status, response_body)
+      def parse(_request_url, request_body, response_status, response_body)
         return nil unless response_status == 200
 
         response = safe_json_parse(response_body)
@@ -28,9 +31,11 @@ module LlmCostTracker
           model: response["model"] || request["model"],
           input_tokens: usage["input_tokens"] || 0,
           output_tokens: usage["output_tokens"] || 0,
-          total_tokens: (usage["input_tokens"] || 0) + (usage["output_tokens"] || 0),
-          cache_read_tokens: usage["cache_read_input_tokens"],
-          cache_creation_tokens: usage["cache_creation_input_tokens"]
+          total_tokens: (usage["input_tokens"] || 0) + (usage["output_tokens"] || 0) +
+            (usage["cache_read_input_tokens"] || 0) +
+            (usage["cache_creation_input_tokens"] || 0),
+          cache_read_input_tokens: usage["cache_read_input_tokens"],
+          cache_creation_input_tokens: usage["cache_creation_input_tokens"]
         }.compact
       end
     end

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require "json"
+require "uri"
+
+require_relative "base"
 
 module LlmCostTracker
   module Parsers
@@ -9,12 +12,12 @@ module LlmCostTracker
 
       def match?(url)
         uri = URI.parse(url.to_s)
-        HOSTS.include?(uri.host)
+        HOSTS.include?(uri.host.to_s.downcase)
       rescue URI::InvalidURIError
         false
       end
 
-      def parse(request_url, request_body, response_status, response_body)
+      def parse(request_url, _request_body, response_status, response_body)
         return nil unless response_status == 200
 
         response = safe_json_parse(response_body)
@@ -28,13 +31,18 @@ module LlmCostTracker
           provider: "gemini",
           model: model,
           input_tokens: usage["promptTokenCount"] || 0,
-          output_tokens: usage["candidatesTokenCount"] || 0,
-          total_tokens: usage["totalTokenCount"] || 0
-        }
+          output_tokens: output_tokens(usage),
+          total_tokens: usage["totalTokenCount"] || 0,
+          cached_input_tokens: usage["cachedContentTokenCount"]
+        }.compact
       end
 
       private
 
+      def output_tokens(usage)
+        (usage["candidatesTokenCount"] || 0) + (usage["thoughtsTokenCount"] || 0)
+      end
+
       def extract_model_from_url(url)
         uri = URI.parse(url.to_s)
         match = uri.path.match(%r{/models/([^/:]+)})