llm_cost_tracker 0.11.0 → 0.12.0

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -1,40 +1,23 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../providers/anthropic/tier_classification"
 
 module LlmCostTracker
   module Integrations
     module RubyLlm
       extend Base
 
-      class << self
-        def integration_name
-          :ruby_llm
-        end
-
-        def minimum_version
-          "1.14.1"
-        end
-
-        def version_constant
-          "RubyLLM::VERSION"
-        end
+      minimum_version "1.15.0"
 
+      class << self
         def patch_targets
-          [
-            patch_target(
-              "RubyLLM::Provider",
-              with: ProviderPatch,
-              methods: %i[slug complete embed transcribe paint moderate]
-            )
-          ]
+          [patch_target("RubyLLM::Provider", with: ProviderPatch)]
         end
 
         def record_completion(provider, response, request:, latency_ms:, has_block:)
           record_usage(
-            provider: provider_slug(provider),
-            model: response_model_id(response) || model_id(request[:model]),
+            provider: provider.slug.to_s,
+            model: response_model_id(response) || model_id_from_request(request[:model]),
             response: response,
             latency_ms: latency_ms,
             stream: has_block || request[:stream] == true
@@ -43,8 +26,8 @@ module LlmCostTracker
 
         def record_embedding(provider, response, request:, latency_ms:)
           record_usage(
-            provider: provider_slug(provider),
-            model: response_model_id(response) || model_id(request[:model]),
+            provider: provider.slug.to_s,
+            model: response_model_id(response) || model_id_from_request(request[:model]),
             response: response,
             latency_ms: latency_ms,
             stream: false,
@@ -54,8 +37,8 @@ module LlmCostTracker
 
         def record_transcription(provider, response, request:, latency_ms:)
           record_usage(
-            provider: provider_slug(provider),
-            model: response_model_id(response) || model_id(request[:model]),
+            provider: provider.slug.to_s,
+            model: response_model_id(response) || model_id_from_request(request[:model]),
             response: response,
             latency_ms: latency_ms,
             stream: false
@@ -63,30 +46,27 @@ module LlmCostTracker
         end
 
         def record_image(provider, response, request:, latency_ms:)
-          usage = object_value(response, :usage)
-          usage = {} unless usage.is_a?(Hash)
-          raw_input = (usage[:input_tokens] || usage["input_tokens"]).to_i
-          raw_output = (usage[:output_tokens] || usage["output_tokens"]).to_i
+          usage = response.usage.with_indifferent_access
+          raw_input = usage[:input_tokens].to_i
+          raw_output = usage[:output_tokens].to_i
           image_input = image_token_detail(usage, :input)
           image_output = image_token_detail(usage, :output)
-          text_input = [raw_input - image_input, 0].max
-          text_output = [raw_output - image_output, 0].max
           record_passthrough(
-            provider: provider_slug(provider),
-            model: response_model_id(response) || model_id(request[:model]),
+            provider: provider.slug.to_s,
+            model: response_model_id(response) || model_id_from_request(request[:model]),
             response: response,
             latency_ms: latency_ms,
-            input_tokens: text_input,
+            input_tokens: [raw_input - image_input, 0].max,
             image_input_tokens: image_input,
-            output_tokens: text_output,
+            output_tokens: [raw_output - image_output, 0].max,
             image_output_tokens: image_output
           )
         end
 
         def record_moderation(provider, response, request:, latency_ms:)
           record_passthrough(
-            provider: provider_slug(provider),
-            model: response_model_id(response) || model_id(request[:model]),
+            provider: provider.slug.to_s,
+            model: response_model_id(response) || model_id_from_request(request[:model]),
             response: response,
             latency_ms: latency_ms,
             input_tokens: 0,
@@ -96,14 +76,20 @@ module LlmCostTracker
 
         def image_token_detail(usage, direction)
           container_key = direction == :input ? :input_tokens_details : :output_tokens_details
-          details = usage[container_key] || usage[container_key.to_s] || {}
+          details = usage[container_key]
           return 0 unless details.is_a?(Hash)
 
-          (details[:image_tokens] || details["image_tokens"]).to_i
+          details.with_indifferent_access[:image_tokens].to_i
         end
 
-        def record_passthrough(provider:, model:, response:, latency_ms:, input_tokens:, output_tokens:,
-                               image_input_tokens: 0, image_output_tokens: 0)
+        def record_passthrough(provider:,
+                               model:,
+                               response:,
+                               latency_ms:,
+                               input_tokens:,
+                               output_tokens:,
+                               image_input_tokens: 0,
+                               image_output_tokens: 0)
           return unless active?
 
           record_safely do
@@ -111,14 +97,14 @@ module LlmCostTracker
               event: Event.build(
                 provider: provider,
                 model: model,
-                token_usage: TokenUsage.build(
+                token_usage: Usage::TokenUsage.build(
                   input_tokens: input_tokens,
                   output_tokens: output_tokens,
                   image_input_tokens: image_input_tokens,
                   image_output_tokens: image_output_tokens
                 ),
-                usage_source: :sdk_response,
-                provider_response_id: provider_response_id(response)
+                usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE,
+                provider_response_id: provider_response_id_for(response)
               ),
               latency_ms: latency_ms
             )
@@ -129,104 +115,106 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            input_tokens = object_value(response, :input_tokens)
-            output_tokens = object_value(response, :output_tokens) if output_tokens.nil?
+            input_tokens = response.input_tokens
+            output_tokens = response.output_tokens if output_tokens.nil?
             next if input_tokens.nil? && output_tokens.nil?
 
-            cache_read = object_value(response, :cached_tokens).to_i
-            hidden_output = object_value(response, :thinking_tokens, :reasoning_tokens).to_i
-
+            cache_write_5m, cache_write_1h = cache_creation_split(provider, response)
             LlmCostTracker::Tracker.record(
               event: Event.build(
                 provider: provider,
                 model: model,
-                pricing_mode: pricing_mode(provider: provider, response: response),
-                token_usage: TokenUsage.build(
-                  input_tokens: regular_input_tokens(input_tokens, cache_read),
+                pricing_mode: pricing_mode_for(provider: provider, response: response),
+                token_usage: Usage::TokenUsage.build(
+                  input_tokens: input_tokens.to_i,
                   output_tokens: output_tokens.to_i,
-                  cache_read_input_tokens: cache_read,
-                  cache_write_input_tokens: object_value(response, :cache_creation_tokens).to_i,
-                  hidden_output_tokens: hidden_output
+                  cache_read_input_tokens: response.try(:cached_tokens).to_i,
+                  cache_write_input_tokens: cache_write_5m,
+                  cache_write_extended_input_tokens: cache_write_1h,
+                  hidden_output_tokens: response.try(:thinking_tokens).to_i
                 ),
                 stream: stream,
-                usage_source: :sdk_response,
-                provider_response_id: provider_response_id(response)
+                usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE,
+                provider_response_id: provider_response_id_for(response)
               ),
               latency_ms: latency_ms
             )
           end
         end
 
-        def regular_input_tokens(input_tokens, cache_read)
-          [input_tokens.to_i - cache_read, 0].max
-        end
+        def cache_creation_split(provider, response)
+          return [response.try(:cache_creation_tokens).to_i, 0] unless provider == "anthropic"
+
+          cache = raw_body(response).dig("usage", "cache_creation")
+          return [response.try(:cache_creation_tokens).to_i, 0] unless cache.is_a?(Hash)
 
-        def provider_slug(provider)
-          object_value(provider, :slug).to_s
+          [cache["ephemeral_5m_input_tokens"].to_i, cache["ephemeral_1h_input_tokens"].to_i]
         end
 
-        def model_id(object)
-          return nil if object.nil?
+        def model_id_from_request(value)
+          return nil if value.nil?
+          return value.to_s if value.is_a?(String) || value.is_a?(Symbol)
 
-          value = object_value(object, :id, :model_id, :model)
-          value ||= object if object.is_a?(String) || object.is_a?(Symbol)
-          value&.to_s
+          (value.try(:id) || value.try(:model_id) || value.try(:model))&.to_s
         end
 
-        def response_model_id(object)
-          value = object_value(object, :model_id, :model)
-          value&.to_s
+        def provider_response_id_for(response)
+          body = raw_body(response)
+          body["id"] || body["responseId"]
         end
 
-        def provider_response_id(response)
-          object_value(response, :id, :provider_response_id)
+        def raw_body(response)
+          body = response.try(:raw)&.body
+          body.is_a?(Hash) ? body : {}
         end
 
-        def pricing_mode(provider:, response:)
-          raw = object_value(response, :pricing_mode, :service_tier)
-          if provider == "anthropic" && LlmCostTracker::Providers::Anthropic::TierClassification.standard_equivalent_tier?(raw)
-            return nil
+        def response_model_id(response)
+          (response.try(:model_id) || response.try(:model))&.to_s
+        end
+
+        def pricing_mode_for(provider:, response:)
+          body = raw_body(response)
+          case provider
+          when "anthropic" then body.dig("usage", "service_tier")
+          when "gemini" then body.dig("usageMetadata", "serviceTier")
+          else body["service_tier"]
           end
+        end
 
-          raw
+        def blocking_seam(resource, record_method, **extras)
+          {
+            provider: resource.slug.to_s,
+            record: lambda do |response, request, latency_ms|
+              public_send(record_method, resource, response, request: request, latency_ms: latency_ms, **extras)
+            end
+          }
         end
       end
 
       module ProviderPatch
         def complete(*args, **kwargs, &)
-          measure(args, kwargs, recorder: :record_completion, has_block: block_given?) { super }
+          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_completion, has_block: block_given?)
+          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
         end
 
         def embed(*args, **kwargs)
-          measure(args, kwargs, recorder: :record_embedding) { super }
+          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_embedding)
+          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
         end
 
         def transcribe(*args, **kwargs)
-          measure(args, kwargs, recorder: :record_transcription) { super }
+          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_transcription)
+          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
         end
 
         def paint(*args, **kwargs)
-          measure(args, kwargs, recorder: :record_image) { super }
+          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_image)
+          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
         end
 
         def moderate(*args, **kwargs)
-          measure(args, kwargs, recorder: :record_moderation) { super }
-        end
-
-        private
-
-        def measure(args, kwargs, recorder:, **extras)
-          request = RubyLlm.request_params(args, kwargs)
-          RubyLlm.enforce_budget!(request: request)
-          started_at = LlmCostTracker::Timing.now_monotonic
-          response = yield
-          RubyLlm.public_send(
-            recorder, self, response,
-            request: request,
-            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
-            **extras
-          )
-          response
+          seam = LlmCostTracker::Integrations::RubyLlm.blocking_seam(self, :record_moderation)
+          LlmCostTracker::Integrations::RubyLlm.wrap_blocking(args, kwargs, **seam) { super }
         end
       end
     end

data/lib/llm_cost_tracker/integrations.rb

@@ -1,41 +1,50 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/string/inflections"
+require_relative "check"
 require_relative "errors"
-require_relative "logging"
 
 module LlmCostTracker
   module Integrations
     autoload :Base, "llm_cost_tracker/integrations/base"
-    autoload :Openai, "llm_cost_tracker/integrations/openai"
-    autoload :Anthropic, "llm_cost_tracker/integrations/anthropic"
-    autoload :RubyLlm, "llm_cost_tracker/integrations/ruby_llm"
 
-    INTEGRATION_CONSTANTS = { openai: :Openai, anthropic: :Anthropic, ruby_llm: :RubyLlm }.freeze
-    DOUBLE_INSTRUMENTATION_OVERLAPS = %i[openai anthropic].freeze
-    private_constant :DOUBLE_INSTRUMENTATION_OVERLAPS
+    Dir.glob(File.join(__dir__, "integrations", "*.rb")).each do |path|
+      basename = File.basename(path, ".rb")
+      next if basename == "base"
 
-    module_function
+      autoload basename.camelize.to_sym, "llm_cost_tracker/integrations/#{basename}"
+    end
 
-    def install!(names = LlmCostTracker.configuration.instrumented_integrations)
+    def self.install!(names = LlmCostTracker.configuration.instrumented_integrations)
       normalized = normalize(names)
       warn_double_instrumentation(normalized)
-      normalized.each { |name| fetch(name).install }
+      normalized.each do |name|
+        integration = fetch(name)
+        next integration.install if integration
+
+        Logging.warn("Unknown integration: #{name.inspect}. Known: #{self.names.map(&:inspect).join(', ')}")
+      end
     end
 
-    def checks(names = LlmCostTracker.configuration.instrumented_integrations)
-      return [Base::Result.new(:ok, "integrations", "no SDK integrations enabled")] if names.empty?
+    def self.checks(names = LlmCostTracker.configuration.instrumented_integrations)
+      return [Check.new(:ok, "integrations", "no SDK integrations enabled")] if names.empty?
 
-      normalize(names).map { |name| fetch(name).status }
+      normalize(names).map do |name|
+        integration = fetch(name)
+        next integration.status if integration
+
+        Check.new(:warn, name.to_s, "unknown integration; check your config.instrument(...) call")
+      end
     end
 
-    def normalize(names)
+    def self.normalize(names)
       Array(names).flatten.uniq
     end
 
-    def warn_double_instrumentation(names)
+    def self.warn_double_instrumentation(names)
       return unless names.include?(:ruby_llm)
 
-      overlapping = names & DOUBLE_INSTRUMENTATION_OVERLAPS
+      overlapping = names - [:ruby_llm]
       return if overlapping.empty?
 
       Logging.warn(
@@ -45,18 +54,16 @@ module LlmCostTracker
       )
     end
 
-    def fetch(name)
-      const_name = INTEGRATION_CONSTANTS[name.to_sym]
-      unless const_name
-        raise LlmCostTracker::Error,
-              "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
-      end
+    def self.fetch(name)
+      const_name = name.to_s.camelize
+      return nil unless const_name.match?(/\A[A-Z]\w*\z/)
+      return nil unless const_defined?(const_name, false)
 
-      const_get(const_name)
+      const_get(const_name, false)
     end
 
-    def names
-      INTEGRATION_CONSTANTS.keys
+    def self.names
+      constants(false).reject { |c| c == :Base }.map { |c| c.to_s.underscore.to_sym }.sort
     end
   end
 end

data/lib/llm_cost_tracker/ledger/period/totals.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require "bigdecimal"
+require "bigdecimal/util"
 
 require_relative "../period"
 
@@ -20,66 +20,51 @@ module LlmCostTracker
         def totals
           return {} if periods.empty?
 
-          snapshot_totals
+          values = periods.to_h { |period| [period, BigDecimal("0")] }
+          period_by_name = periods.to_h { |period| [period.to_s, period] }
+          LlmCostTracker::Call.find_by_sql(union_sql).each do |row|
+            values[period_by_name.fetch(row.period_key)] = row.total_cost.to_d
+          end
+          values
         end
 
         private
 
         attr_reader :periods, :time
 
-        def snapshot_totals
-          values = periods.to_h { |period| [period, BigDecimal("0")] }
-          period_by_name = periods.to_h { |period| [period.name, period] }
-          sql = periods.map { |period| snapshot_select(period) }.join(" UNION ALL ")
-          LlmCostTracker::Call.find_by_sql(sql).each do |row|
-            period = period_by_name.fetch(row.period_key)
-            values[period] = BigDecimal(row.total_cost.to_s)
-          end
-          values
+        def union_sql
+          periods.map { |period| period_select(period) }.join(" UNION ALL ")
         end
 
-        def snapshot_select(period)
+        def period_select(period)
           start = Period.range_start(period, time)
-          components = [period_total_sql(period, start)]
-          components << pending_total_sql(start) if Ingestion.async?
-          "SELECT #{connection.quote(period.name)} AS period_key, " \
-            "(#{components.join(') + (')}) AS total_cost"
+          components = ["(#{recorded_sql(period, start)})"]
+          components << "(#{pending_sql(start)})" if Ingestion.async?
+          "SELECT #{quote(period.to_s)} AS period_key, #{components.join(' + ')} AS total_cost"
         end
 
-        def period_total_sql(period, start)
-          if LlmCostTracker.configuration.cache_rollups
-            "GREATEST(COALESCE(#{rollup_sum_sql(period)}, 0), COALESCE(#{calls_sum_sql(start)}, 0))"
-          else
-            "COALESCE(#{calls_sum_sql(start)}, 0)"
-          end
+        def recorded_sql(period, start)
+          calls = "COALESCE(#{sum_sql(LlmCostTracker::Call.between(start, time))}, 0)"
+          return calls unless LlmCostTracker.configuration.cache_rollups
+
+          rollup = "COALESCE(#{sum_sql(rollup_scope(period))}, 0)"
+          "GREATEST(#{rollup}, #{calls})"
         end
 
-        def rollup_sum_sql(period)
-          table = connection.quote_table_name("llm_cost_tracker_call_rollups")
-          "(SELECT SUM(total_cost) FROM #{table} " \
-            "WHERE period = #{connection.quote(Period::PERIODS.fetch(period))} " \
-            "AND period_start = #{connection.quote(Period.bucket(period, time))})"
+        def pending_sql(start)
+          "COALESCE(#{sum_sql(Ingestion::InboxEntry.pending.where(tracked_at: start..time))}, 0)"
         end
 
-        def calls_sum_sql(start)
-          table = connection.quote_table_name("llm_cost_tracker_calls")
-          tracked_at = connection.quote_column_name("tracked_at")
-          "(SELECT SUM(total_cost) FROM #{table} " \
-            "WHERE #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)})"
+        def rollup_scope(period)
+          LlmCostTracker::CallRollup.where(period: period.to_s, period_start: Period.bucket(period, time))
         end
 
-        def pending_total_sql(start)
-          table = connection.quote_table_name(Ingestion::InboxEntry.table_name)
-          total_cost = connection.quote_column_name("total_cost")
-          tracked_at = connection.quote_column_name("tracked_at")
-          attempts = connection.quote_column_name("attempts")
-          "COALESCE((SELECT SUM(#{total_cost}) FROM #{table} " \
-            "WHERE #{attempts} < #{Ingestion::InboxEntry::MAX_ATTEMPTS_BEFORE_QUARANTINE} " \
-            "AND #{tracked_at} BETWEEN #{connection.quote(start)} AND #{connection.quote(time)}), 0)"
+        def sum_sql(scope)
+          "(#{scope.select('SUM(total_cost)').to_sql})"
         end
 
-        def connection
-          LlmCostTracker::Call.connection
+        def quote(value)
+          LlmCostTracker::Call.connection.quote(value)
         end
       end
     end

data/lib/llm_cost_tracker/ledger/period.rb

@@ -3,18 +3,13 @@
 module LlmCostTracker
   module Ledger
     module Period
-      PERIODS = {
-        month: "month",
-        day:   "day"
-      }.freeze
+      PERIODS = %i[month day].freeze
 
-      module_function
-
-      def valid_keys(periods)
-        periods.select { |period| PERIODS.key?(period) }
+      def self.valid_keys(periods)
+        PERIODS & periods
       end
 
-      def range_start(period, time)
+      def self.range_start(period, time)
         utc_time = time.to_time.utc
 
         case period
@@ -23,7 +18,7 @@ module LlmCostTracker
         end
       end
 
-      def bucket(period, time)
+      def self.bucket(period, time)
         range_start(period, time).to_date
       end
     end