llm_cost_tracker 0.5.1 → 0.5.3

data/lib/llm_cost_tracker/doctor.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "price_freshness"
+require_relative "doctor/capture_check"
 
 module LlmCostTracker
   class Doctor
@@ -38,6 +39,7 @@ module LlmCostTracker
     def checks
       [
         configuration_check,
+        capture_check,
         *integration_checks,
         active_record_check,
         table_check,
@@ -51,9 +53,12 @@ module LlmCostTracker
     private
 
     def configuration_check
-      Check.new(:ok, "configuration", "storage_backend=#{LlmCostTracker.configuration.storage_backend.inspect}")
+      config = LlmCostTracker.configuration
+      Check.new(:ok, "configuration", "storage_backend=#{config.storage_backend.inspect}, enabled=#{config.enabled}")
     end
 
+    def capture_check = CaptureCheck.call(Check)
+
     def integration_checks
       LlmCostTracker::Integrations.checks.map do |check|
         Check.new(check.status, check.name.to_s, check.message)

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

@@ -39,6 +39,7 @@ module LlmCostTracker
 
         add_engine_require
         route %(mount LlmCostTracker::Engine => "/llm-costs")
+        say "Mount /llm-costs behind your app's admin auth before deploying.", :yellow
       end
 
       private

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

@@ -11,10 +11,16 @@ LlmCostTracker.configure do |config|
   # Tags are merged into every event. Use a callable for request/job-time context.
   config.default_tags = -> { { environment: Rails.env } }
 
+  # Tag guardrails keep accidental high-cardinality or sensitive values out of the ledger.
+  # config.max_tag_count = 50
+  # config.max_tag_value_bytesize = 1024
+  # config.redacted_tag_keys = %w[api_key access_token authorization credential password refresh_token secret]
+
   # Optional SDK integrations. Provider SDK gems are not installed by LLM Cost Tracker.
-  # Enable only the SDKs your app already uses.
+  # Enabled integrations are checked at boot, so enable only clients your app loads.
   # config.instrument :openai
   # config.instrument :anthropic
+  # config.instrument :ruby_llm
 
   # Budget behavior: :notify calls on_budget_exceeded, :raise raises after recording,
   # :block_requests preflights monthly/daily budgets before supported requests.

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -10,10 +11,19 @@ module LlmCostTracker
       class << self
         def integration_name = :anthropic
 
-        def target_patches
+        def minimum_version = "1.36.0"
+
+        def version_constant = "Anthropic::VERSION"
+
+        def patch_targets
           [
-            [constant("Anthropic::Resources::Messages"), MessagesPatch],
-            [constant("Anthropic::Resources::Beta::Messages"), MessagesPatch]
+            patch_target("Anthropic::Resources::Messages", with: MessagesPatch, methods: %i[create stream stream_raw]),
+            patch_target(
+              "Anthropic::Resources::Beta::Messages",
+              with: MessagesPatch,
+              methods: %i[create stream stream_raw],
+              optional: true
+            )
           ]
         end
 
@@ -55,6 +65,28 @@ module LlmCostTracker
             ObjectReader.nested(usage, :output_tokens_details, :reasoning_tokens)
           )
         end
+
+        def track_stream(stream, collector:)
+          return stream unless active?
+
+          StreamTracker.wrap(
+            stream,
+            collector: collector,
+            active: -> { active? },
+            finish: ->(errored:) { finish_stream(collector, errored: errored) }
+          )
+        end
+
+        def stream_collector(request)
+          LlmCostTracker::StreamCollector.new(
+            provider: "anthropic",
+            model: request[:model] || request["model"]
+          )
+        end
+
+        def finish_stream(collector, errored:)
+          record_safely { collector.finish!(errored: errored) }
+        end
       end
 
       module MessagesPatch
@@ -69,6 +101,22 @@ module LlmCostTracker
           )
           message
         end
+
+        def stream(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
+          LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
+        end
+
+        def stream_raw(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
+          LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
+        end
       end
     end
   end

data/lib/llm_cost_tracker/integrations/base.rb

@@ -6,6 +6,7 @@ require_relative "object_reader"
 module LlmCostTracker
   module Integrations
     module Base
+      PatchTarget = Data.define(:constant_name, :patch, :method_names, :optional)
       Result = Data.define(:name, :status, :message)
 
       def active?
@@ -13,15 +14,23 @@ module LlmCostTracker
       end
 
       def install
-        target_patches.each { |target, patch| install_patch(target, patch) }
+        validate_contract!
+        patch_targets.each do |target|
+          target_class = constant(target.constant_name)
+          install_patch(target_class, target.patch) if target_class
+        end
       end
 
       def status
         name = integration_name
-        installed = target_patches.count { |target, patch| patch_installed?(target, patch) }
-        available = target_patches.count { |target, _patch| target }
-        return Result.new(name, :ok, "#{name} integration installed") if installed.positive?
-        return Result.new(name, :warn, "#{name} SDK classes are not loaded") if available.zero?
+        problems = contract_problems
+        if problems.any?
+          return Result.new(name, :warn, "#{name} integration cannot be installed: #{problems.join('; ')}")
+        end
+
+        required_targets = patch_targets.reject(&:optional)
+        installed = required_targets.count { |target| patch_installed?(constant(target.constant_name), target.patch) }
+        return Result.new(name, :ok, "#{name} integration installed") if installed == required_targets.count
 
         Result.new(name, :warn, "#{name} integration is enabled but not installed")
       end
@@ -55,10 +64,72 @@ module LlmCostTracker
         end
       end
 
+      def minimum_version = nil
+
+      def version_constant = nil
+
+      def patch_targets = []
+
+      def patch_target(constant_name, with:, methods:, optional: false)
+        PatchTarget.new(constant_name, with, Array(methods), optional)
+      end
+
       private
 
+      def validate_contract!
+        problems = contract_problems
+        return if problems.empty?
+
+        raise Error, "#{integration_name} integration cannot be installed: #{problems.join('; ')}"
+      end
+
+      def contract_problems
+        version_problems + target_problems
+      end
+
+      def version_problems
+        return [] unless minimum_version
+
+        name = integration_name.to_s
+        version = installed_version
+        return ["#{name} >= #{minimum_version} is required, but #{name} is not loaded"] unless version
+        return [] if version >= Gem::Version.new(minimum_version)
+
+        ["#{name} >= #{minimum_version} is required, detected #{version}"]
+      end
+
+      def installed_version
+        Gem.loaded_specs[integration_name.to_s]&.version || constant_version
+      end
+
+      def constant_version
+        return nil unless version_constant
+
+        value = constant(version_constant)
+        value ? Gem::Version.new(value.to_s) : nil
+      rescue ArgumentError
+        nil
+      end
+
+      def target_problems
+        patch_targets.flat_map do |target|
+          target_class = constant(target.constant_name)
+          next [] if target_class.nil? && target.optional
+          next ["#{target.constant_name} is not loaded"] unless target_class
+
+          missing_methods(target_class, target)
+        end
+      end
+
+      def missing_methods(target_class, target)
+        target.method_names.filter_map do |method_name|
+          next if target_class.method_defined?(method_name) || target_class.private_method_defined?(method_name)
+
+          "#{target.constant_name}##{method_name} is not available"
+        end
+      end
+
       def install_patch(target, patch)
-        return unless target
         return if patch_installed?(target, patch)
 
         target.prepend(patch)

data/lib/llm_cost_tracker/integrations/object_reader.rb

@@ -48,7 +48,7 @@ module LlmCostTracker
         return unless object.respond_to?(:[])
 
         object[key]
-      rescue IndexError, TypeError, NoMethodError
+      rescue IndexError, NameError, TypeError
         nil
       end
     end

data/lib/llm_cost_tracker/integrations/openai.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -10,10 +11,22 @@ module LlmCostTracker
       class << self
         def integration_name = :openai
 
-        def target_patches
+        def minimum_version = "0.59.0"
+
+        def version_constant = "OpenAI::VERSION"
+
+        def patch_targets
           [
-            [constant("OpenAI::Resources::Responses"), ResponsesPatch],
-            [constant("OpenAI::Resources::Chat::Completions"), ChatCompletionsPatch]
+            patch_target(
+              "OpenAI::Resources::Responses",
+              with: ResponsesPatch,
+              methods: %i[create stream stream_raw retrieve_streaming]
+            ),
+            patch_target(
+              "OpenAI::Resources::Chat::Completions",
+              with: ChatCompletionsPatch,
+              methods: %i[create stream_raw]
+            )
           ]
         end
 
@@ -28,15 +41,16 @@ module LlmCostTracker
             output_tokens = ObjectReader.first(usage, :output_tokens, :completion_tokens)
             next if input_tokens.nil? && output_tokens.nil?
 
+            metadata = usage_metadata(usage)
             LlmCostTracker::Tracker.record(
               provider: "openai",
               model: ObjectReader.first(response, :model) || request[:model],
-              input_tokens: ObjectReader.integer(input_tokens),
+              input_tokens: regular_input_tokens(input_tokens, metadata[:cache_read_input_tokens]),
               output_tokens: ObjectReader.integer(output_tokens),
               latency_ms: latency_ms,
               usage_source: :sdk_response,
               provider_response_id: ObjectReader.first(response, :id),
-              metadata: usage_metadata(usage)
+              metadata: metadata
             )
           end
         end
@@ -61,6 +75,32 @@ module LlmCostTracker
             ObjectReader.nested(usage, :completion_tokens_details, :reasoning_tokens)
           )
         end
+
+        def regular_input_tokens(input_tokens, cache_read)
+          [ObjectReader.integer(input_tokens) - cache_read.to_i, 0].max
+        end
+
+        def track_stream(stream, collector:)
+          return stream unless active?
+
+          StreamTracker.wrap(
+            stream,
+            collector: collector,
+            active: -> { active? },
+            finish: ->(errored:) { finish_stream(collector, errored: errored) }
+          )
+        end
+
+        def stream_collector(request)
+          LlmCostTracker::StreamCollector.new(
+            provider: "openai",
+            model: request[:model] || request["model"]
+          )
+        end
+
+        def finish_stream(collector, errored:)
+          record_safely { collector.finish!(errored: errored) }
+        end
       end
 
       module ResponsesPatch
@@ -75,6 +115,31 @@ module LlmCostTracker
           )
           response
         end
+
+        def stream(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
+
+        def stream_raw(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
+
+        def retrieve_streaming(response_id, *args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          collector.provider_response_id = response_id
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
       end
 
       module ChatCompletionsPatch
@@ -89,6 +154,14 @@ module LlmCostTracker
           )
           response
         end
+
+        def stream_raw(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
       end
     end
   end

data/lib/llm_cost_tracker/integrations/registry.rb

@@ -1,18 +1,31 @@
 # frozen_string_literal: true
 
+require "monitor"
+
+require_relative "../errors"
 require_relative "openai"
 require_relative "anthropic"
+require_relative "ruby_llm"
 
 module LlmCostTracker
   module Integrations
     module Registry
-      INTEGRATIONS = {
+      DEFAULT_INTEGRATIONS = {
         openai: Openai,
-        anthropic: Anthropic
+        anthropic: Anthropic,
+        ruby_llm: RubyLlm
       }.freeze
+      MUTEX = Monitor.new
 
       module_function
 
+      def register(name, integration)
+        key = name.to_sym
+        validate_integration!(integration)
+        MUTEX.synchronize { @integrations = integrations.merge(key => integration).freeze }
+        integration
+      end
+
       def install!(names = LlmCostTracker.configuration.instrumented_integrations)
         normalize(names).each { |name| fetch(name).install }
       end
@@ -28,13 +41,32 @@ module LlmCostTracker
       end
 
       def fetch(name)
-        INTEGRATIONS.fetch(name.to_sym) do
-          message = "Unknown integration: #{name.inspect}. Use one of: #{INTEGRATIONS.keys.join(', ')}"
+        integrations.fetch(name.to_sym) do
+          message = "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
           raise LlmCostTracker::Error, message
         end
       end
+
+      def names
+        integrations.keys
+      end
+
+      def reset!
+        MUTEX.synchronize { @integrations = DEFAULT_INTEGRATIONS.dup.freeze }
+      end
+
+      def integrations
+        @integrations || MUTEX.synchronize { @integrations ||= DEFAULT_INTEGRATIONS.dup.freeze }
+      end
+
+      def validate_integration!(integration)
+        return if integration.respond_to?(:install) && integration.respond_to?(:status)
+
+        raise ArgumentError, "integration must respond to install and status"
+      end
     end
 
+    def self.register(name, integration) = Registry.register(name, integration)
     def self.install! = Registry.install!
     def self.checks = Registry.checks
   end

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -0,0 +1,171 @@
+# frozen_string_literal: true
+
+require_relative "base"
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      extend Base
+
+      class << self
+        def integration_name = :ruby_llm
+
+        def minimum_version = "1.14.1"
+
+        def version_constant = "RubyLLM::VERSION"
+
+        def patch_targets
+          [
+            patch_target(
+              "RubyLLM::Provider",
+              with: ProviderPatch,
+              methods: %i[slug complete embed transcribe]
+            )
+          ]
+        end
+
+        def record_completion(provider, response, request:, latency_ms:, stream:)
+          record_usage(
+            provider: provider_slug(provider),
+            model: response_model_id(response) || model_id(request[:model]),
+            response: response,
+            latency_ms: latency_ms,
+            stream: stream
+          )
+        end
+
+        def streaming_request?(request, has_block:)
+          has_block || request[:stream] == true
+        end
+
+        def record_embedding(provider, response, request:, latency_ms:)
+          record_usage(
+            provider: provider_slug(provider),
+            model: response_model_id(response) || model_id(request[:model]),
+            response: response,
+            latency_ms: latency_ms,
+            stream: false,
+            output_tokens: 0
+          )
+        end
+
+        def record_transcription(provider, response, request:, latency_ms:)
+          record_usage(
+            provider: provider_slug(provider),
+            model: response_model_id(response) || model_id(request[:model]),
+            response: response,
+            latency_ms: latency_ms,
+            stream: false
+          )
+        end
+
+        def record_usage(provider:, model:, response:, latency_ms:, stream:, output_tokens: nil)
+          return unless active?
+
+          record_safely do
+            input_tokens = ObjectReader.first(response, :input_tokens)
+            output_tokens = ObjectReader.first(response, :output_tokens) if output_tokens.nil?
+            next if input_tokens.nil? && output_tokens.nil?
+
+            cache_read = ObjectReader.integer(ObjectReader.first(response, :cached_tokens))
+
+            LlmCostTracker::Tracker.record(
+              provider: provider,
+              model: model,
+              input_tokens: regular_input_tokens(input_tokens, cache_read),
+              output_tokens: ObjectReader.integer(output_tokens),
+              latency_ms: latency_ms,
+              stream: stream,
+              usage_source: :ruby_llm,
+              provider_response_id: provider_response_id(response),
+              metadata: usage_metadata(response, cache_read)
+            )
+          end
+        end
+
+        def usage_metadata(response, cache_read)
+          {
+            cache_read_input_tokens: cache_read,
+            cache_write_input_tokens: ObjectReader.integer(ObjectReader.first(response, :cache_creation_tokens)),
+            hidden_output_tokens: ObjectReader.integer(
+              ObjectReader.first(response, :thinking_tokens, :reasoning_tokens)
+            )
+          }
+        end
+
+        def regular_input_tokens(input_tokens, cache_read)
+          [ObjectReader.integer(input_tokens) - cache_read.to_i, 0].max
+        end
+
+        def provider_slug(provider)
+          ObjectReader.first(provider, :slug).to_s
+        end
+
+        def model_id(object)
+          return nil if object.nil?
+
+          value = ObjectReader.first(object, :id, :model_id, :model)
+          value ||= object if object.is_a?(String) || object.is_a?(Symbol)
+          value&.to_s
+        end
+
+        def response_model_id(object)
+          value = ObjectReader.first(object, :model_id, :model)
+          value&.to_s
+        end
+
+        def provider_response_id(response)
+          ObjectReader.first(response, :id, :provider_response_id) || ObjectReader.nested(response, :raw, :id)
+        end
+      end
+
+      module ProviderPatch
+        def complete(*args, **kwargs, &)
+          integration = LlmCostTracker::Integrations::RubyLlm
+          request = integration.request_params(args, kwargs)
+          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          integration.enforce_budget!
+          response = super
+          integration.record_completion(
+            self,
+            response,
+            request: request,
+            latency_ms: integration.elapsed_ms(started_at),
+            stream: integration.streaming_request?(request, has_block: block_given?)
+          )
+          response
+        end
+
+        def embed(*args, **kwargs)
+          integration = LlmCostTracker::Integrations::RubyLlm
+          request = integration.request_params(args, kwargs)
+          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          integration.enforce_budget!
+          response = super
+          integration.record_embedding(
+            self,
+            response,
+            request: request,
+            latency_ms: integration.elapsed_ms(started_at)
+          )
+          response
+        end
+
+        def transcribe(*args, **kwargs)
+          integration = LlmCostTracker::Integrations::RubyLlm
+          request = integration.request_params(args, kwargs)
+          started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          integration.enforce_budget!
+          response = super
+          integration.record_transcription(
+            self,
+            response,
+            request: request,
+            latency_ms: integration.elapsed_ms(started_at)
+          )
+          response
+        end
+      end
+    end
+  end
+end