llm_cost_tracker 0.5.0 → 0.5.2

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

@@ -10,10 +10,14 @@ 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: :create),
+            patch_target("OpenAI::Resources::Chat::Completions", with: ChatCompletionsPatch, methods: :create)
           ]
         end
 
@@ -28,15 +32,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 +66,10 @@ 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
       end
 
       module ResponsesPatch

data/lib/llm_cost_tracker/integrations/registry.rb

@@ -2,13 +2,15 @@
 
 require_relative "openai"
 require_relative "anthropic"
+require_relative "ruby_llm"
 
 module LlmCostTracker
   module Integrations
     module Registry
       INTEGRATIONS = {
         openai: Openai,
-        anthropic: Anthropic
+        anthropic: Anthropic,
+        ruby_llm: RubyLlm
       }.freeze
 
       module_function

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

data/lib/llm_cost_tracker/llm_api_call.rb

@@ -73,12 +73,15 @@ module LlmCostTracker
     end
 
     def self.group_by_tag(key)
-      group(Arel.sql(tag_group_expression(key)))
+      group(Arel.sql(tag_value_expression(key)))
     end
 
-    def self.cost_by_tag(key)
-      costs = group_by_tag(key).sum(:total_cost).each_with_object(Hash.new(0.0)) do |(tag_value, cost), grouped|
-        grouped[tag_label(tag_value)] += cost.to_f
+    def self.cost_by_tag(key, limit: nil)
+      relation = group_by_tag(key).order(Arel.sql("COALESCE(SUM(total_cost), 0) DESC"))
+      relation = relation.limit(limit) if limit
+
+      costs = relation.sum(:total_cost).each_with_object(Hash.new(0.0)) do |(tag_value, cost), grouped|
+        grouped[tag_value_label(tag_value)] += cost.to_f
       end
       costs.sort_by { |_label, cost| -cost }.to_h
     end
@@ -101,14 +104,13 @@ module LlmCostTracker
       group(:provider).average(:latency_ms).transform_values(&:to_f)
     end
 
-    def self.tag_label(value)
+    def self.tag_value_label(value)
       value.nil? || value == "" ? "(untagged)" : value.to_s
     end
-    private_class_method :tag_label
 
-    def self.tag_group_expression(key)
+    def self.tag_value_expression(key, table_name: quoted_table_name)
       key = validated_tag_key(key)
-      column = "#{quoted_table_name}.#{connection.quote_column_name('tags')}"
+      column = "#{table_name}.#{connection.quote_column_name('tags')}"
 
       case connection.adapter_name
       when /postgres/i
@@ -120,7 +122,6 @@ module LlmCostTracker
         "json_extract(#{column}, #{connection.quote(json_path(key))})"
       end
     end
-    private_class_method :tag_group_expression
 
     def self.validated_tag_key(key)
       TagKey.validate!(key)

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -5,12 +5,11 @@ require "json"
 
 require_relative "../logging"
 require_relative "../request_url"
+require_relative "../stream_capture"
 
 module LlmCostTracker
   module Middleware
     class Faraday < ::Faraday::Middleware
-      STREAM_CAPTURE_LIMIT_BYTES = 1_048_576
-
       def initialize(app, **options)
         super(app)
         @tags = options.fetch(:tags, {})
@@ -78,8 +77,8 @@ module LlmCostTracker
         unless response_body
           Logging.warn(
             "Unable to read response body for #{RequestUrl.label(request_url)}; " \
-            "streaming responses are captured automatically for OpenAI/Anthropic/Gemini " \
-            "or via LlmCostTracker.track_stream for custom clients."
+            "known streaming responses are captured automatically, or via LlmCostTracker.track_stream " \
+            "for custom clients."
           )
           return nil
         end
@@ -88,6 +87,11 @@ module LlmCostTracker
       end
 
       def parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+        if stream_buffer&.dig(:overflowed)
+          Logging.warn(capture_warning(request_url, stream_buffer))
+          return parser.parse_stream(request_url, request_body, response_env.status, [])
+        end
+
         body = stream_buffer&.dig(:buffer)&.string
         body = read_body(response_env.body) if body.nil? || body.empty?
 
@@ -110,7 +114,7 @@ module LlmCostTracker
         request_env.request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
           unless state[:overflowed]
-            if state[:bytes] + chunk.bytesize <= STREAM_CAPTURE_LIMIT_BYTES
+            if state[:bytes] + chunk.bytesize <= StreamCapture::LIMIT_BYTES
               state[:buffer] << chunk
               state[:bytes] += chunk.bytesize
             else
@@ -161,7 +165,7 @@ module LlmCostTracker
                  "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
         end
 
-        "Streaming response for #{RequestUrl.label(request_url)} exceeded #{STREAM_CAPTURE_LIMIT_BYTES} bytes; " \
+        "Streaming response for #{RequestUrl.label(request_url)} exceeded #{StreamCapture::LIMIT_BYTES} bytes; " \
           "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
       end
     end

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -72,7 +72,7 @@ module LlmCostTracker
           model: extract_model_from_url(request_url),
           input_tokens: [usage["promptTokenCount"].to_i - cache_read, 0].max,
           output_tokens: output_tokens(usage),
-          total_tokens: usage["totalTokenCount"].to_i,
+          total_tokens: total_tokens(usage, cache_read),
           cache_read_input_tokens: usage["cachedContentTokenCount"],
           hidden_output_tokens: usage["thoughtsTokenCount"],
           stream: stream,
@@ -92,6 +92,13 @@ module LlmCostTracker
         usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
       end
 
+      def total_tokens(usage, cache_read)
+        total = usage["totalTokenCount"]
+        return total.to_i unless total.nil?
+
+        [usage["promptTokenCount"].to_i - cache_read, 0].max + cache_read + output_tokens(usage)
+      end
+
       def stream_response_id(events)
         find_event_value(events) { |data| data["responseId"] }
       end

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -21,7 +21,7 @@ module LlmCostTracker
           model: response["model"] || request["model"],
           input_tokens: regular_input_tokens(usage, cache_read),
           output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
-          total_tokens: usage["total_tokens"].to_i,
+          total_tokens: total_tokens(usage, cache_read),
           cache_read_input_tokens: cache_read,
           hidden_output_tokens: hidden_output_tokens(usage),
           usage_source: :response
@@ -44,7 +44,7 @@ module LlmCostTracker
             model: model,
             input_tokens: regular_input_tokens(usage, cache_read),
             output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
-            total_tokens: usage["total_tokens"].to_i,
+            total_tokens: total_tokens(usage, cache_read),
             cache_read_input_tokens: cache_read,
             hidden_output_tokens: hidden_output_tokens(usage),
             stream: true,
@@ -87,6 +87,15 @@ module LlmCostTracker
         details = usage["completion_tokens_details"] || usage["output_tokens_details"] || {}
         details["reasoning_tokens"]
       end
+
+      def total_tokens(usage, cache_read)
+        total = usage["total_tokens"]
+        return total.to_i unless total.nil?
+
+        regular_input_tokens(usage, cache_read) +
+          cache_read.to_i +
+          (usage["completion_tokens"] || usage["output_tokens"]).to_i
+      end
     end
   end
 end

data/lib/llm_cost_tracker/price_freshness.rb

@@ -17,20 +17,20 @@ module LlmCostTracker
 
         [:ok, "updated_at=#{updated_at}"]
       rescue Date::Error
-        [:warn, "metadata.updated_at=#{updated_at.inspect} is invalid; run bin/rails llm_cost_tracker:prices:sync"]
+        [:warn, "metadata.updated_at=#{updated_at.inspect} is invalid; run bin/rails llm_cost_tracker:prices:refresh"]
       end
 
       private
 
       def missing
-        [:warn, "metadata.updated_at missing; run bin/rails llm_cost_tracker:prices:sync"]
+        [:warn, "metadata.updated_at missing; run bin/rails llm_cost_tracker:prices:refresh"]
       end
 
       def stale(updated_at)
         [
           :warn,
           "updated_at=#{updated_at} is older than #{STALE_AFTER_DAYS} days; " \
-          "run bin/rails llm_cost_tracker:prices:sync"
+          "run bin/rails llm_cost_tracker:prices:refresh"
         ]
       end
     end

data/lib/llm_cost_tracker/price_registry.rb

@@ -12,6 +12,7 @@ module LlmCostTracker
     EMPTY_PRICES = {}.freeze
     PRICE_KEYS = %w[input output cache_read_input cache_write_input].freeze
     METADATA_KEYS = %w[_source _source_version _fetched_at _updated _notes _validator_override].freeze
+    MAX_FILE_BYTES = 2_097_152
     MUTEX = Monitor.new
 
     class << self
@@ -114,6 +115,8 @@ module LlmCostTracker
       end
 
       def load_price_file(path)
+        raise ArgumentError, "prices_file exceeds #{MAX_FILE_BYTES} bytes" if File.size(path) > MAX_FILE_BYTES
+
         contents = File.read(path)
         return YAML.safe_load(contents, aliases: false) || {} if yaml_file?(path)
 

data/lib/llm_cost_tracker/price_sync/fetcher.rb

@@ -15,8 +15,9 @@ module LlmCostTracker
         end
       end
 
-      USER_AGENT = "llm_cost_tracker price sync"
+      USER_AGENT = "llm_cost_tracker price refresh"
       MAX_REDIRECTS = 5
+      MAX_BODY_BYTES = 2_097_152
       OPEN_TIMEOUT = 5
       READ_TIMEOUT = 10
       WRITE_TIMEOUT = 10
@@ -25,24 +26,17 @@ module LlmCostTracker
         raise Error, "Too many redirects while fetching #{url}" if redirects > MAX_REDIRECTS
 
         uri = URI.parse(url)
+        raise Error, "Pricing snapshot URL must use https" unless uri.scheme == "https"
+
         request = Net::HTTP::Get.new(uri)
         request["User-Agent"] = USER_AGENT
         request["If-None-Match"] = etag if etag
 
-        response = Net::HTTP.start(
-          uri.host,
-          uri.port,
-          use_ssl: uri.scheme == "https",
-          open_timeout: OPEN_TIMEOUT,
-          read_timeout: READ_TIMEOUT,
-          write_timeout: WRITE_TIMEOUT
-        ) do |http|
-          http.request(request)
-        end
+        response, body = fetch_response(uri, request)
 
         case response
         when Net::HTTPSuccess
-          build_response(response, not_modified: false)
+          build_response(response, body: body || limited_body(response), not_modified: false)
         when Net::HTTPNotModified
           build_response(response, body: nil, not_modified: true)
         when Net::HTTPRedirection
@@ -59,6 +53,43 @@ module LlmCostTracker
 
       private
 
+      def fetch_response(uri, request)
+        body = nil
+        response = Net::HTTP.start(
+          uri.host,
+          uri.port,
+          use_ssl: uri.scheme == "https",
+          open_timeout: OPEN_TIMEOUT,
+          read_timeout: READ_TIMEOUT,
+          write_timeout: WRITE_TIMEOUT
+        ) do |http|
+          http.request(request) do |streamed_response|
+            body = limited_body(streamed_response) if streamed_response.is_a?(Net::HTTPSuccess)
+          end
+        end
+
+        [response, body]
+      end
+
+      def limited_body(response)
+        body = +""
+        if response.respond_to?(:read_body)
+          response.read_body do |chunk|
+            chunk = chunk.to_s
+            if body.bytesize + chunk.bytesize > MAX_BODY_BYTES
+              raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes"
+            end
+
+            body << chunk
+          end
+        else
+          body = response.body.to_s
+        end
+        raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes" if body.bytesize > MAX_BODY_BYTES
+
+        body
+      end
+
       def build_response(response, not_modified:, body: response.body)
         Response.new(
           body: body,