llm_cost_tracker 0.1.1 → 0.1.2

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,8 +10,33 @@ module LlmCostTracker
     # Scopes for querying
     scope :by_provider, ->(provider) { where(provider: provider) }
     scope :by_model,    ->(model)    { where(model: model) }
-    scope :by_tag, lambda { |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..) }
@@ -35,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,25 +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
       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(
@@ -40,21 +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
-        return unless 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)
 
-        warn "[LlmCostTracker] Error processing response: #{e.message}"
+        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,7 +12,7 @@ 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

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,7 +12,7 @@ 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
@@ -28,7 +31,7 @@ module LlmCostTracker
           provider: "gemini",
           model: model,
           input_tokens: usage["promptTokenCount"] || 0,
-          output_tokens: usage["candidatesTokenCount"] || 0,
+          output_tokens: output_tokens(usage),
           total_tokens: usage["totalTokenCount"] || 0,
           cached_input_tokens: usage["cachedContentTokenCount"]
         }.compact
@@ -36,6 +39,10 @@ module LlmCostTracker
 
       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/([^/:]+)})

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require "json"
+require "uri"
+
+require_relative "base"
 
 module LlmCostTracker
   module Parsers
@@ -10,12 +13,12 @@ module LlmCostTracker
 
       def match?(url)
         uri = URI.parse(url.to_s)
-        HOSTS.include?(uri.host) && TRACKED_PATHS.include?(uri.path)
+        HOSTS.include?(uri.host.to_s.downcase) && TRACKED_PATHS.include?(uri.path)
       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)
@@ -25,7 +28,7 @@ module LlmCostTracker
         request = safe_json_parse(request_body)
 
         {
-          provider: "openai",
+          provider: provider_for(request_url),
           model: response["model"] || request["model"],
           input_tokens: usage["prompt_tokens"] || usage["input_tokens"] || 0,
           output_tokens: usage["completion_tokens"] || usage["output_tokens"] || 0,
@@ -36,6 +39,10 @@ module LlmCostTracker
 
       private
 
+      def provider_for(_request_url)
+        "openai"
+      end
+
       def cached_input_tokens(usage)
         details = usage["prompt_tokens_details"] || usage["input_tokens_details"] || {}
         details["cached_tokens"]

data/lib/llm_cost_tracker/parsers/openai_compatible.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require_relative "openai"
+
+module LlmCostTracker
+  module Parsers
+    class OpenaiCompatible < Openai
+      TRACKED_PATH_SUFFIXES = %w[/chat/completions /completions /embeddings /responses].freeze
+
+      def match?(url)
+        uri = URI.parse(url.to_s)
+        !provider_for_host(uri.host).nil? && tracked_path?(uri.path)
+      rescue URI::InvalidURIError
+        false
+      end
+
+      private
+
+      def provider_for(request_url)
+        uri = URI.parse(request_url.to_s)
+        provider_for_host(uri.host) || "openai_compatible"
+      rescue URI::InvalidURIError
+        "openai_compatible"
+      end
+
+      def provider_for_host(host)
+        host = host.to_s.downcase
+        provider_name = configured_providers[host] ||
+                        configured_providers.find do |configured_host, _provider|
+                          configured_host.to_s.downcase == host
+                        end&.last
+        provider_name&.to_s
+      end
+
+      def configured_providers
+        LlmCostTracker.configuration.openai_compatible_providers
+      end
+
+      def tracked_path?(path)
+        TRACKED_PATH_SUFFIXES.any? { |suffix| path == suffix || path.end_with?(suffix) }
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/parsers/registry.rb

@@ -4,16 +4,14 @@ module LlmCostTracker
   module Parsers
     class Registry
       class << self
+        PARSERS_MUTEX = Mutex.new
+
         def parsers
-          @parsers ||= [
-            Openai.new,
-            Anthropic.new,
-            Gemini.new
-          ]
+          @parsers || PARSERS_MUTEX.synchronize { @parsers ||= default_parsers }
         end
 
         def register(parser)
-          parsers.unshift(parser)
+          PARSERS_MUTEX.synchronize { parsers.unshift(parser) }
         end
 
         def find_for(url)
@@ -21,7 +19,18 @@ module LlmCostTracker
         end
 
         def reset!
-          @parsers = nil
+          PARSERS_MUTEX.synchronize { @parsers = nil }
+        end
+
+        private
+
+        def default_parsers
+          [
+            Openai.new,
+            OpenaiCompatible.new,
+            Anthropic.new,
+            Gemini.new
+          ]
         end
       end
     end

data/lib/llm_cost_tracker/price_registry.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+require "json"
+require "yaml"
+
+module LlmCostTracker
+  module PriceRegistry
+    DEFAULT_PRICES_PATH = File.expand_path("prices.json", __dir__)
+    PRICE_KEYS = %w[input cached_input output cache_read_input cache_creation_input].freeze
+    NORMALIZE_PRICE_ENTRY = lambda do |price|
+      (price || {}).each_with_object({}) do |(key, value), normalized|
+        key = key.to_s
+        normalized[key.to_sym] = Float(value) if PRICE_KEYS.include?(key)
+      end
+    end
+    NORMALIZE_PRICE_TABLE = lambda do |table|
+      (table || {}).each_with_object({}) do |(model, price), normalized|
+        normalized[model.to_s] = NORMALIZE_PRICE_ENTRY.call(price)
+      end
+    end
+    RAW_REGISTRY = JSON.parse(File.read(DEFAULT_PRICES_PATH)).freeze
+    PRICE_METADATA = RAW_REGISTRY.fetch("metadata", {}).freeze
+    BUILTIN_PRICES = NORMALIZE_PRICE_TABLE.call(RAW_REGISTRY.fetch("models", {})).freeze
+
+    class << self
+      def builtin_prices
+        BUILTIN_PRICES
+      end
+
+      def metadata
+        PRICE_METADATA
+      end
+
+      def normalize_price_table(table)
+        NORMALIZE_PRICE_TABLE.call(table)
+      end
+
+      def file_prices(path)
+        return {} unless path
+
+        path = path.to_s
+        cache_key = [path, File.mtime(path).to_f]
+        return @file_prices if @file_prices_cache_key == cache_key
+
+        @file_prices_cache_key = cache_key
+        @file_prices = normalize_price_table(price_file_models(load_price_file(path)))
+      rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError => e
+        raise Error, "Unable to load prices_file #{path.inspect}: #{e.message}"
+      end
+
+      private
+
+      def load_price_file(path)
+        contents = File.read(path)
+        return YAML.safe_load(contents, aliases: false) || {} if yaml_file?(path)
+
+        JSON.parse(contents)
+      end
+
+      def yaml_file?(path)
+        %w[.yaml .yml].include?(File.extname(path).downcase)
+      end
+
+      def price_file_models(registry)
+        registry.fetch("models", registry)
+      end
+    end
+  end
+end