llm_cost_tracker 0.11.0 → 0.13.0

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,37 +1,23 @@
 # frozen_string_literal: true
 
 require_relative "base"
-require_relative "../billing/line_item"
-require_relative "../providers/anthropic/server_tools"
-require_relative "../providers/anthropic/tier_classification"
+require_relative "../providers/anthropic/usage_extractor"
+require_relative "../providers/anthropic/response_parser"
 
 module LlmCostTracker
   module Integrations
     module Anthropic
       extend Base
 
-      class << self
-        def integration_name
-          :anthropic
-        end
-
-        def minimum_version
-          "1.36.0"
-        end
-
-        def version_constant
-          "Anthropic::VERSION"
-        end
+      minimum_version "1.36.0"
 
+      class << self
         def patch_targets
           [
-            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
-            )
+            patch_target("Anthropic::Resources::Messages", with: MessagesPatch),
+            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)
           ]
         end
 
@@ -39,122 +25,123 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            usage = object_value(message, :usage)
+            usage = message.usage
             next unless usage
+            next if usage.input_tokens.nil? && usage.output_tokens.nil?
 
-            input_tokens = object_value(usage, :input_tokens)
-            output_tokens = object_value(usage, :output_tokens)
-            next if input_tokens.nil? && output_tokens.nil?
+            usage_hash = usage.deep_to_h
 
             LlmCostTracker::Tracker.record(
-              event: Event.build(
-                provider: "anthropic",
-                model: object_value(message, :model) || request[:model],
-                pricing_mode: pricing_mode(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),
-                service_line_items: service_line_items_from(usage)
+              event: Providers::Anthropic::ResponseParser.event_from_usage(
+                usage: usage_hash,
+                model: message.model || request[:model],
+                provider_response_id: message.id,
+                usage_source: Usage::Source::SDK_RESPONSE,
+                request: request
               ),
               latency_ms: latency_ms
             )
           end
         end
 
-        def service_line_items_from(usage)
-          server_tool_use = object_value(usage, :server_tool_use)
-          return [] unless server_tool_use
+        def record_batch_result(response)
+          return unless active?
 
-          Providers::Anthropic::ServerTools::LINE_ITEMS.filter_map do |component_key, count_key|
-            quantity = object_value(server_tool_use, count_key).to_i
-            next if quantity.zero?
+          record_safely do
+            next unless response.respond_to?(:result) && response.result
 
-            Billing::LineItem.build(
-              component_key: component_key,
-              quantity: quantity,
-              cost_status: Billing::CostStatus::UNKNOWN,
-              pricing_basis: :provider_usage,
-              provider_field: "usage.server_tool_use.#{count_key}"
+            result = response.result
+            next unless result.respond_to?(:type) && result.type.to_s == "succeeded"
+
+            message = result.respond_to?(:message) ? result.message : nil
+            next unless message
+            next if LlmCostTracker::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?
+
+            usage_hash = usage.deep_to_h
+            LlmCostTracker::Tracker.record(
+              event: Providers::Anthropic::ResponseParser.event_from_usage(
+                usage: usage_hash,
+                model: message.model,
+                provider_response_id: message.id,
+                usage_source: Usage::Source::SDK_BATCH_RESULT,
+                pricing_mode: "batch"
+              )
             )
           end
         end
 
-        def token_usage(usage:, input_tokens:, output_tokens:)
-          cache_creation = object_value(usage, :cache_creation)
-          if cache_creation
-            cache_write_default = object_value(cache_creation, :ephemeral_5m_input_tokens).to_i
-            cache_write_extended = object_value(cache_creation, :ephemeral_1h_input_tokens).to_i
-          else
-            cache_write_default = object_value(usage, :cache_creation_input_tokens).to_i
-            cache_write_extended = 0
-          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_default,
-            cache_write_extended_input_tokens: cache_write_extended,
-            hidden_output_tokens: hidden_output
-          )
+        def stream_pricing_mode(request)
+          Providers::Anthropic::UsageExtractor.pricing_mode(request: request || {}, usage: nil)
         end
+      end
 
-        def pricing_mode(request:, usage:)
-          service_tier = object_value(usage, :service_tier) || request[:service_tier]
-          tier = Providers::Anthropic::TierClassification
-          service_tier = nil if tier.standard_equivalent_tier?(service_tier)
-
-          modes = [
-            Pricing::Mode.normalize(object_value(usage, :speed) || request[:speed]),
-            Pricing::Mode.normalize(service_tier)
-          ]
-          geo = inference_geo(request: request, usage: usage).to_s.downcase
-          modes << "data_residency" if tier.data_residency_geo?(geo)
-          modes = modes.compact.uniq
-          modes.empty? ? nil : modes.join("_")
+      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_pricing_mode(request)
-          pricing_mode(request: request || {}, usage: nil)
+        def stream(*args, **kwargs)
+          LlmCostTracker::Integrations::Anthropic.wrap_stream(
+            args,
+            kwargs,
+            collector: ->(request) { LlmCostTracker::Integrations::Anthropic.stream_collector(request) }
+          ) { super }
         end
 
-        def inference_geo(request:, usage:)
-          object_value(usage, :inference_geo) || request[:inference_geo]
+        def stream_raw(*args, **kwargs)
+          LlmCostTracker::Integrations::Anthropic.wrap_stream(
+            args,
+            kwargs,
+            collector: ->(request) { LlmCostTracker::Integrations::Anthropic.stream_collector(request) }
+          ) { super }
         end
+      end
 
-        def wrap_stream_call(args, kwargs)
-          request = request_params(args, kwargs)
-          enforce_budget!(request: request)
-          collector = stream_collector(request)
-          stream = yield
-          track_stream(stream, collector: collector)
-        end
+      module BatchesPatch
+        def results_streaming(*args, **kwargs)
+          raw = super
+          return raw unless LlmCostTracker::Integrations::Anthropic.active?
 
-        def wrap_blocking_call(args, kwargs)
-          request = request_params(args, kwargs)
-          enforce_budget!(request: request)
-          started_at = LlmCostTracker::Timing.now_monotonic
-          message = yield
-          record_message(message, request: request, latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at))
-          message
+          BatchResultsCapture.new(raw)
         end
       end
 
-      module MessagesPatch
-        def create(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_blocking_call(args, kwargs) { super }
+      class BatchResultsCapture
+        include Enumerable
+
+        def initialize(raw_stream)
+          @raw_stream = raw_stream
         end
 
-        def stream(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_stream_call(args, kwargs) { super }
+        def each(&block)
+          return enum_for(:each) unless block
+
+          @raw_stream.each do |response|
+            LlmCostTracker::Integrations::Anthropic.record_batch_result(response)
+            block.call(response)
+          end
         end
 
-        def stream_raw(*args, **kwargs)
-          LlmCostTracker::Integrations::Anthropic.wrap_stream_call(args, kwargs) { super }
+        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, ...)
         end
       end
     end

data/lib/llm_cost_tracker/integrations/base.rb

@@ -3,7 +3,7 @@
 require "active_support/core_ext/hash/indifferent_access"
 require "active_support/core_ext/string/inflections"
 
-require_relative "../doctor/check"
+require_relative "../check"
 require_relative "../logging"
 require_relative "../timing"
 require_relative "../capture/stream_collector"
@@ -12,7 +12,14 @@ require_relative "../capture/stream_tracker"
 module LlmCostTracker
   module Integrations
     module Base
-      Result = LlmCostTracker::Doctor::Check
+      def integration_name
+        @integration_name ||= name.demodulize.underscore.to_sym
+      end
+
+      def provider(value = nil)
+        @provider = value.to_s if value
+        @provider ||= integration_name.to_s
+      end
 
       def active?
         LlmCostTracker.configuration.instrumented?(integration_name)
@@ -30,22 +37,22 @@ module LlmCostTracker
         name = integration_name.to_s
         problems = version_problems + target_problems
         if problems.any?
-          return Result.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
+          return Check.new(:warn, name, "#{name} integration cannot be installed: #{problems.join('; ')}")
         end
 
         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 Result.new(:ok, name, "#{name} integration installed") if installed
+        return Check.new(:ok, name, "#{name} integration installed") if installed
 
-        Result.new(:warn, name, "#{name} integration is enabled but not installed")
+        Check.new(:warn, name, "#{name} integration is enabled but not installed")
       end
 
-      def enforce_budget!(request:)
+      def enforce_budget!(request:, provider: self.provider)
         return unless active?
 
-        LlmCostTracker::Tracker.enforce_budget!(
-          provider: integration_name.to_s,
+        LlmCostTracker::Budget.enforce!(
+          provider: provider,
           model: request[:model],
           request: request
         )
@@ -71,10 +78,21 @@ module LlmCostTracker
         kwargs.to_h.with_indifferent_access
       end
 
-      def normalize_sdk_args(args, kwargs)
-        return args if args.any? || kwargs.empty?
+      def wrap_blocking(args, kwargs, record:, provider: self.provider)
+        request = request_params(args, kwargs)
+        enforce_budget!(request: request, provider: provider)
+        started_at = LlmCostTracker::Timing.now_monotonic
+        response = yield
+        record_safely { record.call(response, request, LlmCostTracker::Timing.elapsed_ms(started_at)) }
+        response
+      end
 
-        [kwargs]
+      def wrap_stream(args, kwargs, collector:, provider: self.provider)
+        request = request_params(args, kwargs)
+        enforce_budget!(request: request, provider: provider)
+        built = collector.call(request)
+        stream = yield(built)
+        track_stream(stream, collector: built)
       end
 
       def track_stream(stream, collector:)
@@ -88,9 +106,9 @@ module LlmCostTracker
         ).wrap
       end
 
-      def stream_collector(request)
+      def stream_collector(request, provider: self.provider)
         LlmCostTracker::Capture::StreamCollector.new(
-          provider: integration_name.to_s,
+          provider: provider,
           model: request[:model],
           pricing_mode: stream_pricing_mode(request),
           request: request
@@ -101,56 +119,29 @@ module LlmCostTracker
         nil
       end
 
-      def object_value(object, *keys)
-        keys.each do |key|
-          value = read_object_value(object, key)
-          return value unless value.nil?
-        end
-        nil
+      def minimum_version(value = nil)
+        @minimum_version = value if value
+        @minimum_version
       end
 
-      def object_dig(object, *path)
-        path.reduce(object) do |current, key|
-          return nil if current.nil?
-
-          read_object_value(current, key)
-        end
+      def gem_version
+        Gem.loaded_specs[integration_name.to_s]&.version
       end
 
-      def minimum_version = nil
-
-      def version_constant = nil
-
       def patch_targets = []
 
-      def patch_target(constant_name, with:, methods:, optional: false, skip_when_methods_missing: false)
+      def patch_target(constant_name, with:, optional: false, skip_when_methods_missing: false)
         {
           constant_name: constant_name,
           patch: with,
-          method_names: Array(methods),
+          method_names: with.instance_methods,
           optional: optional,
           skip_when_methods_missing: skip_when_methods_missing
         }
       end
 
-      module_function :object_value, :object_dig
-
       private
 
-      def read_object_value(object, key)
-        return nil if object.nil?
-
-        if object.is_a?(Hash)
-          return object[key] if object.key?(key)
-          return object[key.name] if key.is_a?(Symbol) && object.key?(key.name)
-        end
-
-        object.public_send(key) if object.respond_to?(key)
-      end
-
-      module_function :read_object_value
-      private_class_method :read_object_value
-
       def validate_contract!
         problems = version_problems + target_problems
         return if problems.empty?
@@ -162,22 +153,13 @@ module LlmCostTracker
         return [] unless minimum_version
 
         name = integration_name.to_s
-        version = Gem.loaded_specs[integration_name.to_s]&.version || constant_version
+        version = gem_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 constant_version
-        return nil unless version_constant
-
-        value = version_constant.to_s.safe_constantize
-        value ? Gem::Version.new(value.to_s) : nil
-      rescue ArgumentError
-        nil
-      end
-
       def target_problems
         patch_targets.flat_map do |target|
           constant_name = target.fetch(:constant_name)

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

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module Integrations
+    module Openai
+      module BatchCapture
+        DEDUP_LIMIT = 1024
+        MUTEX = Mutex.new
+        private_constant :DEDUP_LIMIT, :MUTEX
+
+        class << self
+          def maybe_capture(batch, resource:)
+            return unless Openai.active?
+            return unless batch.respond_to?(:status) && batch.status.to_s == "completed"
+
+            output_file_id = batch.respond_to?(:output_file_id) ? batch.output_file_id : nil
+            return unless output_file_id
+
+            batch_id = batch.respond_to?(:id) ? batch.id : nil
+            return unless batch_id && claim(batch_id)
+
+            client = resource.instance_variable_get(:@client)
+            host = Openai.client_host_for(resource)
+            Openai.record_safely do
+              io = client.files.content(output_file_id)
+              capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+            end
+          end
+
+          private
+
+          def claim(batch_id)
+            MUTEX.synchronize do
+              @dedup ||= Set.new
+              next false if @dedup.include?(batch_id)
+
+              @dedup.clear if @dedup.size >= DEDUP_LIMIT
+              @dedup.add(batch_id)
+              true
+            end
+          end
+
+          def capture_jsonl(jsonl, host:)
+            jsonl.each_line do |line|
+              line = line.strip
+              next if line.empty?
+
+              entry = parse_line(line)
+              next unless entry
+
+              response = entry.dig("response", "body")
+              next unless response.is_a?(Hash) && response["usage"]
+
+              record_result(response, host: host)
+            end
+          end
+
+          def parse_line(line)
+            JSON.parse(line)
+          rescue JSON::ParserError
+            nil
+          end
+
+          def record_result(response, host:)
+            provider = Openai.provider_for_host(host)
+            return if LlmCostTracker::Call.already_recorded?(provider: provider, provider_response_id: response["id"])
+
+            event = LlmCostTracker::Providers::Openai::ResponseParser.event_from_response(
+              response: response,
+              request: {},
+              provider: provider,
+              host: host,
+              usage_source: LlmCostTracker::Usage::Source::SDK_BATCH_RESULT,
+              pricing_mode: "batch"
+            )
+            LlmCostTracker::Tracker.record(event: event) if event
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module Openai
+      module PatchBuilder
+        def self.build(record_method:, methods:)
+          Module.new.tap do |mod|
+            methods.each { |method_name| define_blocking_method(mod, method_name, record_method) }
+          end
+        end
+
+        def self.build_stream(methods:)
+          Module.new.tap do |mod|
+            methods.each { |method_name| define_stream_method(mod, method_name) }
+          end
+        end
+
+        def self.define_blocking_method(mod, method_name, record_method)
+          mod.define_method(method_name) do |*args, **kwargs, &block|
+            host = LlmCostTracker::Integrations::Openai.client_host_for(self)
+            LlmCostTracker::Integrations::Openai.wrap_blocking(
+              args,
+              kwargs,
+              provider: LlmCostTracker::Integrations::Openai.provider_for_host(host),
+              record: lambda do |response, request, latency_ms|
+                LlmCostTracker::Integrations::Openai.public_send(
+                  record_method, response, request: request, latency_ms: latency_ms, host: host
+                )
+              end
+            ) { super(*args, **kwargs, &block) }
+          end
+        end
+
+        def self.define_stream_method(mod, method_name)
+          mod.define_method(method_name) do |*args, **kwargs|
+            LlmCostTracker::Integrations::Openai.wrap_stream(
+              args, kwargs, **LlmCostTracker::Integrations::Openai.stream_seam(self)
+            ) { super(*args, **kwargs) }
+          end
+        end
+      end
+
+      module ResponsesPatch
+        include PatchBuilder.build(record_method: :record_response, methods: %i[create])
+        include PatchBuilder.build_stream(methods: %i[stream stream_raw])
+
+        def retrieve_streaming(response_id, *args, **kwargs)
+          LlmCostTracker::Integrations::Openai.wrap_stream(
+            args, kwargs, **LlmCostTracker::Integrations::Openai.stream_seam(self)
+          ) do |collector|
+            collector.provider_response_id = response_id
+            super(response_id, *args, **kwargs)
+          end
+        end
+      end
+
+      module ChatCompletionsPatch
+        include PatchBuilder.build(record_method: :record_response, methods: %i[create])
+        include PatchBuilder.build_stream(methods: %i[stream stream_raw])
+      end
+
+      EmbeddingsPatch = PatchBuilder.build(record_method: :record_response, methods: %i[create])
+      ImagesPatch = PatchBuilder.build(record_method: :record_image, methods: %i[generate edit create_variation])
+      TranscriptionsPatch = PatchBuilder.build(record_method: :record_transcription, methods: %i[create])
+      TranslationsPatch = PatchBuilder.build(record_method: :record_transcription, methods: %i[create])
+      SpeechPatch = PatchBuilder.build(record_method: :record_speech, methods: %i[create])
+      ModerationsPatch = PatchBuilder.build(record_method: :record_moderation, methods: %i[create])
+      StreamingImagesPatch = PatchBuilder.build_stream(methods: %i[generate_stream_raw edit_stream_raw])
+      StreamingTranscriptionsPatch = PatchBuilder.build_stream(methods: %i[create_streaming])
+
+      module BatchesPatch
+        def retrieve(batch_id, *args, **kwargs)
+          batch = super
+          LlmCostTracker::Integrations::Openai::BatchCapture.maybe_capture(batch, resource: self)
+          batch
+        end
+      end
+    end
+  end
+end