llm_cost_tracker 0.7.0 → 0.7.2

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,7 +1,8 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "stream_tracker"
+require_relative "../capture/stream_collector"
+require_relative "../capture/stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -9,11 +10,17 @@ module LlmCostTracker
       extend Base
 
       class << self
-        def integration_name = :anthropic
+        def integration_name
+          :anthropic
+        end
 
-        def minimum_version = "1.36.0"
+        def minimum_version
+          "1.36.0"
+        end
 
-        def version_constant = "Anthropic::VERSION"
+        def version_constant
+          "Anthropic::VERSION"
+        end
 
         def patch_targets
           [
@@ -31,56 +38,84 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            usage = ObjectReader.first(message, :usage)
+            usage = object_value(message, :usage)
             next unless usage
 
-            input_tokens = ObjectReader.first(usage, :input_tokens)
-            output_tokens = ObjectReader.first(usage, :output_tokens)
+            input_tokens = object_value(usage, :input_tokens)
+            output_tokens = object_value(usage, :output_tokens)
             next if input_tokens.nil? && output_tokens.nil?
 
             LlmCostTracker::Tracker.record(
-              provider: "anthropic",
-              model: ObjectReader.first(message, :model) || request[:model],
-              input_tokens: ObjectReader.integer(input_tokens),
-              output_tokens: ObjectReader.integer(output_tokens),
-              latency_ms: latency_ms,
-              usage_source: :sdk_response,
-              provider_response_id: ObjectReader.first(message, :id),
-              metadata: usage_metadata(usage)
+              capture: UsageCapture.build(
+                provider: "anthropic",
+                model: object_value(message, :model) || request[:model],
+                pricing_mode: pricing_mode(message: message, request: request, usage: usage),
+                token_usage: token_usage(usage: usage, input_tokens: input_tokens, output_tokens: output_tokens),
+                usage_source: :sdk_response,
+                provider_response_id: object_value(message, :id)
+              ),
+              latency_ms: latency_ms
             )
           end
         end
 
-        def usage_metadata(usage)
-          {
-            cache_read_input_tokens: ObjectReader.integer(ObjectReader.first(usage, :cache_read_input_tokens)),
-            cache_write_input_tokens: ObjectReader.integer(ObjectReader.first(usage, :cache_creation_input_tokens)),
-            hidden_output_tokens: hidden_output_tokens(usage)
-          }
+        def token_usage(usage:, input_tokens:, output_tokens:)
+          cache_write_1h = object_dig(usage, :cache_creation, :ephemeral_1h_input_tokens).to_i
+          cache_write_5m = object_dig(usage, :cache_creation, :ephemeral_5m_input_tokens)
+          cache_write = if cache_write_5m.nil?
+                          total_cache_write = object_value(usage, :cache_creation_input_tokens)
+                          [total_cache_write.to_i - cache_write_1h, 0].max
+                        else
+                          cache_write_5m.to_i
+                        end
+          hidden_output = (
+            object_value(usage, :thinking_tokens, :thinking_output_tokens) ||
+            object_dig(usage, :output_tokens_details, :reasoning_tokens)
+          ).to_i
+
+          TokenUsage.build(
+            input_tokens: input_tokens.to_i,
+            output_tokens: output_tokens.to_i,
+            cache_read_input_tokens: object_value(usage, :cache_read_input_tokens).to_i,
+            cache_write_input_tokens: cache_write,
+            cache_write_1h_input_tokens: cache_write_1h,
+            hidden_output_tokens: hidden_output
+          )
         end
 
-        def hidden_output_tokens(usage)
-          ObjectReader.integer(
-            ObjectReader.first(usage, :thinking_tokens, :thinking_output_tokens) ||
-            ObjectReader.nested(usage, :output_tokens_details, :reasoning_tokens)
-          )
+        def pricing_mode(message:, request:, usage:)
+          modes = [
+            Pricing.normalize_mode(object_value(usage, :speed) || object_value(message, :speed) || request[:speed]),
+            Pricing.normalize_mode(
+              object_value(usage, :service_tier) || object_value(message, :service_tier) || request[:service_tier]
+            )
+          ]
+          modes << "data_residency" if inference_geo(message: message, request: request, usage: usage).to_s == "us"
+          modes = modes.compact.uniq
+          modes.empty? ? nil : modes.join("_")
+        end
+
+        def inference_geo(message:, request:, usage:)
+          object_value(usage, :inference_geo) ||
+            object_value(message, :inference_geo) ||
+            request[:inference_geo]
         end
 
         def track_stream(stream, collector:)
           return stream unless active?
 
-          StreamTracker.wrap(
-            stream,
+          LlmCostTracker::Capture::StreamTracker.new(
+            stream: stream,
             collector: collector,
             active: -> { active? },
             finish: ->(errored:) { finish_stream(collector, errored: errored) }
-          )
+          ).wrap
         end
 
         def stream_collector(request)
-          LlmCostTracker::StreamCollector.new(
+          LlmCostTracker::Capture::StreamCollector.new(
             provider: "anthropic",
-            model: request[:model] || request["model"]
+            model: request[:model]
           )
         end
 

data/lib/llm_cost_tracker/integrations/base.rb

@@ -1,12 +1,14 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/hash/indifferent_access"
+require "active_support/core_ext/object/try"
+require "active_support/core_ext/string/inflections"
+
 require_relative "../logging"
-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?
@@ -16,20 +18,22 @@ module LlmCostTracker
       def install
         validate_contract!
         patch_targets.each do |target|
-          target_class = constant(target.constant_name)
-          install_patch(target_class, target.patch) if target_class
+          target_class = target.fetch(:constant_name).to_s.safe_constantize
+          install_patch(target_class, target.fetch(:patch)) if target_class
         end
       end
 
       def status
         name = integration_name
-        problems = contract_problems
+        problems = version_problems + target_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) }
+        required_targets = patch_targets.reject { |target| target.fetch(:optional) }
+        installed = required_targets.count do |target|
+          target.fetch(:constant_name).to_s.safe_constantize&.ancestors&.include?(target.fetch(:patch))
+        end
         return Result.new(name, :ok, "#{name} integration installed") if installed == required_targets.count
 
         Result.new(name, :warn, "#{name} integration is enabled but not installed")
@@ -53,14 +57,31 @@ module LlmCostTracker
 
       def request_params(args, kwargs)
         params = args.first.is_a?(Hash) ? args.first : {}
-        params.merge(kwargs)
+        params.merge(kwargs).with_indifferent_access
+      end
+
+      def object_value(object, *keys)
+        keys.each do |key|
+          value = read_object_value(object, key)
+          return value unless value.nil?
+        end
+        nil
       end
 
-      def constant(path)
-        path.to_s.split("::").reduce(Object) do |scope, const_name|
-          return nil unless scope.const_defined?(const_name, false)
+      def object_dig(object, *path)
+        if object.respond_to?(:dig)
+          begin
+            value = object.dig(*path)
+            return value unless value.nil?
+          rescue NameError, TypeError
+            nil
+          end
+        end
+
+        path.reduce(object) do |current, key|
+          return nil if current.nil?
 
-          scope.const_get(const_name, false)
+          read_object_value(current, key)
         end
       end
 
@@ -71,41 +92,62 @@ module LlmCostTracker
       def patch_targets = []
 
       def patch_target(constant_name, with:, methods:, optional: false)
-        PatchTarget.new(constant_name, with, Array(methods), optional)
+        {
+          constant_name: constant_name,
+          patch: with,
+          method_names: Array(methods),
+          optional: optional
+        }
       end
 
+      module_function :object_value, :object_dig
+
       private
 
+      def read_object_value(object, key)
+        return nil if object.nil?
+        return object[key] if object.try(:key?, key)
+
+        string_key = key.to_s
+        return object[string_key] if object.try(:key?, string_key)
+
+        value = object.try(key)
+        return value unless value.nil?
+
+        indexed_object_value(object, key)
+      end
+
+      def indexed_object_value(object, key)
+        object.try(:[], key)
+      rescue IndexError, NameError, TypeError
+        nil
+      end
+
+      module_function :read_object_value, :indexed_object_value
+      private_class_method :read_object_value, :indexed_object_value
+
       def validate_contract!
-        problems = contract_problems
+        problems = version_problems + target_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
+        version = Gem.loaded_specs[integration_name.to_s]&.version || constant_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 = version_constant.to_s.safe_constantize
         value ? Gem::Version.new(value.to_s) : nil
       rescue ArgumentError
         nil
@@ -113,31 +155,28 @@ module LlmCostTracker
 
       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
+          constant_name = target.fetch(:constant_name)
+          target_class = constant_name.to_s.safe_constantize
+          next [] if target_class.nil? && target.fetch(:optional)
+          next ["#{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|
+        target.fetch(: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"
+          "#{target.fetch(:constant_name)}##{method_name} is not available"
         end
       end
 
       def install_patch(target, patch)
-        return if patch_installed?(target, patch)
+        return if target&.ancestors&.include?(patch)
 
         target.prepend(patch)
       end
-
-      def patch_installed?(target, patch)
-        target&.ancestors&.include?(patch)
-      end
     end
   end
 end

data/lib/llm_cost_tracker/integrations/openai.rb

@@ -1,7 +1,8 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "stream_tracker"
+require_relative "../capture/stream_collector"
+require_relative "../capture/stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -9,11 +10,17 @@ module LlmCostTracker
       extend Base
 
       class << self
-        def integration_name = :openai
+        def integration_name
+          :openai
+        end
 
-        def minimum_version = "0.59.0"
+        def minimum_version
+          "0.59.0"
+        end
 
-        def version_constant = "OpenAI::VERSION"
+        def version_constant
+          "OpenAI::VERSION"
+        end
 
         def patch_targets
           [
@@ -34,67 +41,66 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            usage = ObjectReader.first(response, :usage)
+            usage = object_value(response, :usage)
             next unless usage
 
-            input_tokens = ObjectReader.first(usage, :input_tokens, :prompt_tokens)
-            output_tokens = ObjectReader.first(usage, :output_tokens, :completion_tokens)
+            input_tokens = object_value(usage, :input_tokens, :prompt_tokens)
+            output_tokens = object_value(usage, :output_tokens, :completion_tokens)
             next if input_tokens.nil? && output_tokens.nil?
 
-            metadata = usage_metadata(usage)
+            cache_read = cache_read_input_tokens(usage)
             LlmCostTracker::Tracker.record(
-              provider: "openai",
-              model: ObjectReader.first(response, :model) || request[:model],
-              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: metadata
+              capture: UsageCapture.build(
+                provider: "openai",
+                model: object_value(response, :model) || request[:model],
+                pricing_mode: object_value(response, :service_tier) || request[:service_tier],
+                token_usage: TokenUsage.build(
+                  input_tokens: regular_input_tokens(input_tokens, cache_read),
+                  output_tokens: output_tokens.to_i,
+                  cache_read_input_tokens: cache_read,
+                  hidden_output_tokens: hidden_output_tokens(usage)
+                ),
+                usage_source: :sdk_response,
+                provider_response_id: object_value(response, :id)
+              ),
+              latency_ms: latency_ms
             )
           end
         end
 
-        def usage_metadata(usage)
-          {
-            cache_read_input_tokens: cache_read_input_tokens(usage),
-            hidden_output_tokens: hidden_output_tokens(usage)
-          }
-        end
-
         def cache_read_input_tokens(usage)
-          ObjectReader.integer(
-            ObjectReader.nested(usage, :input_tokens_details, :cached_tokens) ||
-            ObjectReader.nested(usage, :prompt_tokens_details, :cached_tokens)
-          )
+          (
+            object_dig(usage, :input_tokens_details, :cached_tokens) ||
+            object_dig(usage, :prompt_tokens_details, :cached_tokens)
+          ).to_i
         end
 
         def hidden_output_tokens(usage)
-          ObjectReader.integer(
-            ObjectReader.nested(usage, :output_tokens_details, :reasoning_tokens) ||
-            ObjectReader.nested(usage, :completion_tokens_details, :reasoning_tokens)
-          )
+          (
+            object_dig(usage, :output_tokens_details, :reasoning_tokens) ||
+            object_dig(usage, :completion_tokens_details, :reasoning_tokens)
+          ).to_i
         end
 
         def regular_input_tokens(input_tokens, cache_read)
-          [ObjectReader.integer(input_tokens) - cache_read.to_i, 0].max
+          [input_tokens.to_i - cache_read.to_i, 0].max
         end
 
         def track_stream(stream, collector:)
           return stream unless active?
 
-          StreamTracker.wrap(
-            stream,
+          LlmCostTracker::Capture::StreamTracker.new(
+            stream: stream,
             collector: collector,
             active: -> { active? },
             finish: ->(errored:) { finish_stream(collector, errored: errored) }
-          )
+          ).wrap
         end
 
         def stream_collector(request)
-          LlmCostTracker::StreamCollector.new(
+          LlmCostTracker::Capture::StreamCollector.new(
             provider: "openai",
-            model: request[:model] || request["model"]
+            model: request[:model]
           )
         end
 

data/lib/llm_cost_tracker/integrations/ruby_llm.rb

@@ -8,11 +8,17 @@ module LlmCostTracker
       extend Base
 
       class << self
-        def integration_name = :ruby_llm
+        def integration_name
+          :ruby_llm
+        end
 
-        def minimum_version = "1.14.1"
+        def minimum_version
+          "1.14.1"
+        end
 
-        def version_constant = "RubyLLM::VERSION"
+        def version_constant
+          "RubyLLM::VERSION"
+        end
 
         def patch_targets
           [
@@ -63,59 +69,63 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            input_tokens = ObjectReader.first(response, :input_tokens)
-            output_tokens = ObjectReader.first(response, :output_tokens) if output_tokens.nil?
+            input_tokens = object_value(response, :input_tokens)
+            output_tokens = object_value(response, :output_tokens) if output_tokens.nil?
             next if input_tokens.nil? && output_tokens.nil?
 
-            cache_read = ObjectReader.integer(ObjectReader.first(response, :cached_tokens))
+            cache_read = object_value(response, :cached_tokens).to_i
+            hidden_output = object_value(response, :thinking_tokens, :reasoning_tokens).to_i
 
             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)
+              capture: UsageCapture.build(
+                provider: provider,
+                model: model,
+                pricing_mode: pricing_mode(response),
+                token_usage: TokenUsage.build(
+                  input_tokens: regular_input_tokens(input_tokens, cache_read),
+                  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
+                ),
+                stream: stream,
+                usage_source: :ruby_llm,
+                provider_response_id: provider_response_id(response)
+              ),
+              latency_ms: latency_ms
             )
           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
+          [input_tokens.to_i - cache_read.to_i, 0].max
         end
 
         def provider_slug(provider)
-          ObjectReader.first(provider, :slug).to_s
+          object_value(provider, :slug).to_s
         end
 
         def model_id(object)
           return nil if object.nil?
 
-          value = ObjectReader.first(object, :id, :model_id, :model)
+          value = object_value(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 = object_value(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)
+          object_value(response, :id, :provider_response_id) || object_dig(response, :raw, :id)
+        end
+
+        def pricing_mode(response)
+          object_value(response, :pricing_mode, :service_tier) ||
+            object_dig(response, :raw, :pricing_mode) ||
+            object_dig(response, :raw, :service_tier)
         end
       end
 

data/lib/llm_cost_tracker/integrations.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative "errors"
+require_relative "integrations/openai"
+require_relative "integrations/anthropic"
+require_relative "integrations/ruby_llm"
+
+module LlmCostTracker
+  module Integrations
+    AVAILABLE = {
+      openai: Openai,
+      anthropic: Anthropic,
+      ruby_llm: RubyLlm
+    }.freeze
+
+    module_function
+
+    def install!(names = LlmCostTracker.configuration.instrumented_integrations)
+      normalize(names).each { |name| fetch(name).install }
+    end
+
+    def checks(names = LlmCostTracker.configuration.instrumented_integrations)
+      return [Base::Result.new(:integrations, :ok, "no SDK integrations enabled")] if names.empty?
+
+      normalize(names).map { |name| fetch(name).status }
+    end
+
+    def normalize(names)
+      Array(names).flatten.map(&:to_sym).uniq
+    end
+
+    def fetch(name)
+      AVAILABLE.fetch(name.to_sym) do
+        message = "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
+        raise LlmCostTracker::Error, message
+      end
+    end
+
+    def names
+      AVAILABLE.keys
+    end
+  end
+end