llm_cost_tracker 0.14.2 → 0.15.0

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -2,6 +2,8 @@
 
 require_relative "base"
 require_relative "../providers/anthropic/response_parser"
+require_relative "anthropic/patches"
+require_relative "anthropic/batch_results_capture"
 
 module LlmCostTracker
   module Integrations
@@ -17,28 +19,47 @@ module LlmCostTracker
             patch_target("Anthropic::Resources::Beta::Messages", with: MessagesPatch, optional: true),
             patch_target("Anthropic::Resources::Messages::Batches", with: BatchesPatch, optional: true),
             patch_target("Anthropic::Resources::Beta::Messages::Batches", with: BatchesPatch, optional: true),
-            patch_target("Anthropic::BetaRefusalFallbackMiddleware", with: FallbackMiddlewarePatch, optional: true)
+            patch_target("Anthropic::BetaRefusalFallbackMiddleware",
+                         with: FallbackMiddlewarePatch,
+                         optional: true,
+                         since: "1.54.0")
           ]
         end
 
-        def record_message(message, request:, latency_ms:)
+        def blocking_seam(client)
+          host = client_host(client)
+          {
+            record: lambda do |message, request, latency_ms|
+              record_message(message, request: request, latency_ms: latency_ms, host: host)
+            end
+          }
+        end
+
+        def stream_seam(client)
+          host = client_host(client)
+          { collector: ->(request) { stream_collector(request, host: host) } }
+        end
+
+        def stream_collector(request, host: nil)
+          regional = Providers::Anthropic::UsageExtractor.regional_host?(host, request[:model])
+          super(request, pricing_mode: ("data_residency" if regional))
+        end
+
+        def record_message(message, request:, latency_ms:, host: nil)
           return unless active?
 
           record_safely do
             usage = message.usage
-            next unless usage
-            next if usage.input_tokens.nil? && usage.output_tokens.nil?
-
-            usage_hash = usage.deep_to_h
+            next unless tokens_reported?(usage)
 
-            LlmCostTracker::Tracker.record(
-              event: Providers::Anthropic::ResponseParser.event_from_usage(
-                usage: usage_hash,
+            Tracker.record(
+              event: message_event(
+                message,
+                usage: usage.deep_to_h,
                 model: message.model || request[:model],
-                provider_response_id: message.id,
                 usage_source: Usage::Source::SDK_RESPONSE,
                 request: request,
-                **response_fields(message)
+                host: host
               ),
               latency_ms: latency_ms
             )
@@ -49,26 +70,19 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            next unless response.respond_to?(:result) && response.result
-
-            result = response.result
-            next unless result.respond_to?(:type) && result.type.to_s == "succeeded"
-
-            message = result.respond_to?(:message) ? result.message : nil
+            message = succeeded_message(response)
             next unless message
-            next if LlmCostTracker::Call.already_recorded?(provider: "anthropic", provider_response_id: message.id)
+            next if Call.already_recorded?(provider: "anthropic", provider_response_id: message.id)
 
             usage = message.usage
-            next unless usage
-            next if usage.input_tokens.nil? && usage.output_tokens.nil?
+            next unless tokens_reported?(usage)
 
             record_once(
-              Providers::Anthropic::ResponseParser.event_from_usage(
+              message_event(
+                message,
                 usage: usage.deep_to_h.merge(service_tier: "batch"),
                 model: message.model,
-                provider_response_id: message.id,
-                usage_source: Usage::Source::SDK_BATCH_RESULT,
-                **response_fields(message)
+                usage_source: Usage::Source::SDK_BATCH_RESULT
               )
             )
           end
@@ -76,7 +90,7 @@ module LlmCostTracker
 
         def record_refused_hop(request, response)
           record_safely { record_message(response.parse, request: request.body, latency_ms: nil) }
-        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
+        rescue BudgetExceededError, UnknownPricingError
           nil
         end
 
@@ -84,123 +98,48 @@ module LlmCostTracker
           return unless active?
 
           usage = hop.dig(:refused, :usage).to_h.deep_symbolize_keys
-          output = usage[:output_tokens].to_i.positive?
-          usage = usage.slice(:server_tool_use) if output
-          return if output && Providers::Anthropic::UsageExtractor.service_line_items(usage).empty?
+          produced_output = usage[:output_tokens].to_i.positive?
+          usage = usage.slice(:server_tool_use) if produced_output
+          return if produced_output && Providers::Anthropic::UsageExtractor.service_line_items(usage).empty?
 
           record_safely do
-            LlmCostTracker::Tracker.record(
+            Tracker.record(
               event: Providers::Anthropic::ResponseParser.event_from_usage(
                 usage: usage,
                 model: hop[:model],
                 provider_response_id: nil,
                 usage_source: Usage::Source::STREAM_FINAL,
                 stream: true,
-                stop_reason: ("refusal" unless output),
+                stop_reason: ("refusal" unless produced_output),
                 refusal_category: hop.dig(:refused, :stop_details, "category")
               )
             )
           end
-        rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError
+        rescue BudgetExceededError, UnknownPricingError
           nil
         end
 
-        def response_fields(message)
-          { stop_reason: message.stop_reason, refusal_category: message.stop_details&.category,
-            content: message.deep_to_h[:content] }
-        end
-      end
-
-      module MessagesPatch
-        def create(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_blocking(
-            args,
-            kwargs,
-            record: lambda do |message, request, latency_ms|
-              LlmCostTracker::Integrations::Anthropic.record_message(
-                message, request: request, latency_ms: latency_ms
-              )
-            end
-          ) { super }
-        end
-
-        def stream(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_stream(
-            args,
-            kwargs,
-            collector: ->(request) { LlmCostTracker::Integrations::Anthropic.stream_collector(request) }
-          ) { super }
-        end
-
-        def stream_raw(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_stream(
-            args,
-            kwargs,
-            collector: ->(request) { LlmCostTracker::Integrations::Anthropic.stream_collector(request) }
-          ) { super }
-        end
-      end
-
-      module FallbackMiddlewarePatch
-        def call(req, nxt)
-          return super if req.streaming? || !LlmCostTracker::Integrations::Anthropic.active?
-
-          hops = []
-          response = super(req, ->(hop_req) { nxt.call(hop_req).tap { |hop| hops << [hop_req, hop] } })
-          answered = hops.select { |_hop_req, hop| hop.status < 300 }
-          answered[...-1].each { |hop| LlmCostTracker::Integrations::Anthropic.record_refused_hop(*hop) }
-          response
-        end
-
         private
 
-        def consume_hop(*args, **kwargs)
-          refused = Thread.current[:llm_cost_tracker_refused_hop]
-          Thread.current[:llm_cost_tracker_refused_hop] = nil
-          LlmCostTracker::Integrations::Anthropic.record_refused_stream_hop(refused) if refused && kwargs[:splice]
-          super.tap { |hop| Thread.current[:llm_cost_tracker_refused_hop] = hop if hop[:refused] }
-        end
-      end
-
-      module BatchesPatch
-        def results_streaming(*args, **kwargs)
-          raw = super
-          return raw unless LlmCostTracker::Integrations::Anthropic.active?
+        def succeeded_message(response)
+          result = response.result if response.respond_to?(:result)
+          return unless result.respond_to?(:type) && result.type.to_s == "succeeded"
 
-          BatchResultsCapture.new(raw)
+          result.message if result.respond_to?(:message)
         end
-      end
-
-      class BatchResultsCapture
-        include Enumerable
-
-        def initialize(raw_stream)
-          @raw_stream = raw_stream
-        end
-
-        def each(&block)
-          return enum_for(:each) unless block
 
-          deferred = nil
-          @raw_stream.each do |response|
-            begin
-              LlmCostTracker::Integrations::Anthropic.record_batch_result(response)
-            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
-              deferred ||= e
-            end
-            block.call(response)
-          end
-          raise deferred if deferred
+        def tokens_reported?(usage)
+          usage && !(usage.input_tokens.nil? && usage.output_tokens.nil?)
         end
 
-        def respond_to_missing?(name, include_private = false)
-          @raw_stream.respond_to?(name, include_private) || super
-        end
-
-        def method_missing(name, ...)
-          return super unless @raw_stream.respond_to?(name)
-
-          @raw_stream.public_send(name, ...)
+        def message_event(message, **attributes)
+          Providers::Anthropic::ResponseParser.event_from_usage(
+            provider_response_id: message.id,
+            stop_reason: message.stop_reason,
+            refusal_category: message.stop_details&.category,
+            content: message.deep_to_h[:content],
+            **attributes
+          )
         end
       end
     end

data/lib/llm_cost_tracker/integrations/base/patch_target.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/string/inflections"
+
+module LlmCostTracker
+  module Integrations
+    module Base
+      PatchTarget = Data.define(:constant_name, :patch, :optional, :skip_when_methods_missing, :since) do
+        def target_class = constant_name.to_s.safe_constantize
+
+        def installed? = target_class&.ancestors&.include?(patch)
+
+        def missing?(version) = target_class.nil? && (since.nil? || version.nil? || version >= Gem::Version.new(since))
+
+        def install
+          target = target_class
+          target.prepend(patch) unless target.nil? || target.ancestors.include?(patch)
+        end
+
+        def problems
+          target = target_class
+          return [] if target.nil? && optional
+          return ["#{constant_name} is not loaded"] unless target
+          return [] if skip_when_methods_missing
+
+          patch.instance_methods.filter_map do |method_name|
+            next if target.method_defined?(method_name) || target.private_method_defined?(method_name)
+
+            "#{constant_name}##{method_name} is not available"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/base.rb

@@ -8,6 +8,7 @@ require_relative "../logging"
 require_relative "../timing"
 require_relative "../capture/stream_collector"
 require_relative "../capture/stream_tracker"
+require_relative "base/patch_target"
 
 module LlmCostTracker
   module Integrations
@@ -22,39 +23,57 @@ module LlmCostTracker
         LlmCostTracker.configuration.enabled && LlmCostTracker.configuration.instrumented?(integration_name)
       end
 
+      def minimum_version(value = nil)
+        @minimum_version = value if value
+        @minimum_version
+      end
+
+      def maximum_version(value = nil)
+        @maximum_version = value if value
+        @maximum_version
+      end
+
+      def gem_version
+        Gem.loaded_specs[integration_name.to_s]&.version
+      end
+
+      def patch_targets = []
+
+      def patch_target(constant_name, with:, optional: false, skip_when_methods_missing: false, since: nil)
+        PatchTarget.new(constant_name:, patch: with, optional:, skip_when_methods_missing:, since:)
+      end
+
       def install
         validate_contract!
         Logging.warn(untested_version_message) if untested_version?
-        patch_targets.each do |target|
-          target_class = target.fetch(:constant_name).to_s.safe_constantize
-          patch = target.fetch(:patch)
-          target_class.prepend(patch) if target_class && !target_class.ancestors.include?(patch)
-        end
+        patch_targets.each(&:install)
       end
 
       def status
         name = integration_name.to_s
-        problems = version_problems + target_problems
-        if problems.any?
-          return Check.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
-        end
+        violation = contract_violation
+        return Check.new(:warn, name, violation) if violation
         return Check.new(:warn, name, untested_version_message) if untested_version?
+        return Check.new(:warn, name, "#{name} integration is enabled but not installed") unless installed?
 
-        installed = patch_targets.reject { |target| target.fetch(:optional) }.all? do |target|
-          target.fetch(:constant_name).to_s.safe_constantize&.ancestors&.include?(target.fetch(:patch))
-        end
-        return Check.new(:ok, name, "#{name} integration installed") if installed
+        missing = missing_targets
+        return Check.new(:ok, name, "#{name} integration installed") if missing.empty?
 
-        Check.new(:warn, name, "#{name} integration is enabled but not installed")
+        Check.new(
+          :warn,
+          name,
+          "#{name} integration installed without #{missing.join(', ')}; their calls are not recorded"
+        )
       end
 
-      def enforce_budget!(request:, provider: self.provider)
+      def enforce_budget!(request:, provider: self.provider, tags: nil)
         return unless active?
 
         LlmCostTracker::Budget.enforce!(
           provider: provider,
           model: request[:model],
-          request: request
+          request: request,
+          tags: tags
         )
       end
 
@@ -92,21 +111,23 @@ module LlmCostTracker
         end
       end
 
-      def record_once(event)
-        LlmCostTracker::Tracker.record(event: event.keyed_by_response_id)
+      def record_once(event, **)
+        LlmCostTracker::Tracker.record(event: event.keyed_by_response_id, **)
       rescue ActiveRecord::RecordNotUnique
         nil
       end
 
-      def provider_response_id_for(response) = response&.try(:id)
+      def provider_response_id_for(response) = (response.id if response.respond_to?(:id))
+
+      def client_host(client)
+        URI.parse(client.base_url.to_s).host if client
+      rescue URI::InvalidURIError
+        nil
+      end
 
       def request_params(args, kwargs)
-        params =
-          case args.first
-          when Hash then args.first
-          when nil then {}
-          else args.first.to_h
-          end
+        params = args.first
+        params = params.to_h unless params.is_a?(Hash)
         params.merge(kwargs).with_indifferent_access
       rescue StandardError
         kwargs.to_h.with_indifferent_access
@@ -129,58 +150,36 @@ module LlmCostTracker
         track_stream(stream, collector: built)
       end
 
-      def track_stream(stream, collector:)
-        return stream unless active?
-
-        LlmCostTracker::Capture::StreamTracker.new(
-          stream: stream,
-          collector: collector,
-          active: -> { active? },
-          finish: ->(errored) { record_safely { collector.finish!(errored: errored) } }
-        ).wrap
-      end
-
-      def stream_collector(request)
+      def stream_collector(request, pricing_mode: nil)
         LlmCostTracker::Capture::StreamCollector.new(
           provider: provider,
           model: request[:model],
+          pricing_mode: pricing_mode,
           request: request
         )
       end
 
-      def minimum_version(value = nil)
-        @minimum_version = value if value
-        @minimum_version
-      end
+      private
 
-      def maximum_version(value = nil)
-        @maximum_version = value if value
-        @maximum_version
-      end
+      def track_stream(stream, collector:)
+        return stream unless active?
 
-      def gem_version
-        Gem.loaded_specs[integration_name.to_s]&.version
+        LlmCostTracker::Capture::StreamTracker.new(
+          stream: stream,
+          collector: collector,
+          active: -> { active? },
+          finish: ->(errored) { record_safely { collector.finish!(errored: errored) } }
+        ).wrap
       end
 
-      def patch_targets = []
-
-      def patch_target(constant_name, with:, optional: false, skip_when_methods_missing: false)
-        {
-          constant_name: constant_name,
-          patch: with,
-          method_names: with.instance_methods,
-          optional: optional,
-          skip_when_methods_missing: skip_when_methods_missing
-        }
+      def validate_contract!
+        violation = contract_violation
+        raise Error, violation if violation
       end
 
-      private
-
-      def validate_contract!
+      def contract_violation
         problems = version_problems + target_problems
-        return if problems.empty?
-
-        raise Error, "#{integration_name} integration cannot be installed: #{problems.join('; ')}"
+        "#{integration_name} integration cannot be installed: #{problems.join('; ')}" if problems.any?
       end
 
       def version_problems
@@ -203,26 +202,11 @@ module LlmCostTracker
           "its calls may not be recorded"
       end
 
-      def target_problems
-        patch_targets.flat_map do |target|
-          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)
-        return [] if target[:skip_when_methods_missing]
+      def target_problems = patch_targets.flat_map(&:problems)
 
-        target.fetch(:method_names).filter_map do |method_name|
-          next if target_class.method_defined?(method_name) || target_class.private_method_defined?(method_name)
+      def missing_targets = patch_targets.select { |target| target.missing?(gem_version) }.map(&:constant_name).uniq
 
-          "#{target.fetch(:constant_name)}##{method_name} is not available"
-        end
-      end
+      def installed? = patch_targets.reject(&:optional).all?(&:installed?)
     end
   end
 end

data/lib/llm_cost_tracker/integrations/openai/batch_capture.rb

@@ -6,19 +6,16 @@ module LlmCostTracker
   module Integrations
     module Openai
       module BatchCapture
+        FINISHED_STATUSES = %w[completed expired cancelled].freeze
         DEDUP_LIMIT = 1024
         MUTEX = Mutex.new
-        private_constant :DEDUP_LIMIT, :MUTEX
+        private_constant :FINISHED_STATUSES, :DEDUP_LIMIT, :MUTEX
 
         class << self
-          def maybe_capture(batch, resource:)
-            return unless Openai.active?
-            return unless %w[completed expired cancelled].include?(batch.status.to_s)
-            return unless batch.output_file_id && batch.id
-            return if captured?(batch.id)
+          def capture(batch, client:)
+            return unless Openai.active? && capturable?(batch)
 
-            client = resource.instance_variable_get(:@client)
-            host = Openai.client_host_for(resource)
+            host = Openai.client_host(client)
             Openai.record_safely do
               jsonl = client.files.content(batch.output_file_id).read
               deferred = capture_jsonl(jsonl, host: host, model: batch.model)
@@ -29,6 +26,10 @@ module LlmCostTracker
 
           private
 
+          def capturable?(batch)
+            FINISHED_STATUSES.include?(batch.status.to_s) && batch.output_file_id && batch.id && !captured?(batch.id)
+          end
+
           def captured?(batch_id)
             MUTEX.synchronize { @dedup&.include?(batch_id) || false }
           end
@@ -51,7 +52,7 @@ module LlmCostTracker
               next unless response.is_a?(Hash) && response["usage"]
 
               record_result({ "id" => entry["id"] }.merge(response), host: host, model: model)
-            rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
+            rescue BudgetExceededError, UnknownPricingError => e
               deferred ||= e
             end
             deferred
@@ -65,24 +66,24 @@ module LlmCostTracker
 
           def record_result(response, host:, model:)
             provider = Openai.provider_for_host(host)
-            return if LlmCostTracker::Call.already_recorded?(provider: provider, provider_response_id: response["id"])
+            return if Call.already_recorded?(provider: provider, provider_response_id: response["id"])
 
-            parser = LlmCostTracker::Providers::Openai::ResponseParser
-            event = parser.event_from_response(
+            event = Providers::Openai::ResponseParser.event_from_response(
               response: response,
               request: { "model" => model },
               provider: provider,
               host: host,
-              usage_source: LlmCostTracker::Usage::Source::SDK_BATCH_RESULT,
-              # /v1/batches runs regional processing only on the us and eu hosts.
-              pricing_mode: parser.combined_pricing_mode(
-                host: (host if host.to_s.match?(/\A(?:us|eu)\./i)),
-                model: response["model"] || model,
-                service_tier: "batch"
-              )
+              usage_source: Usage::Source::SDK_BATCH_RESULT,
+              pricing_mode: batch_pricing_mode(host, response["model"] || model)
             )
             Openai.record_once(event)
           end
+
+          def batch_pricing_mode(host, model)
+            Providers::Openai::ResponseParser.combined_pricing_mode(
+              host: (host if host.to_s.match?(/\A(?:us|eu)\./i)), model: model, service_tier: "batch"
+            )
+          end
         end
       end
     end