llm_cost_tracker 0.10.0 → 0.12.0

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.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