llm_cost_tracker 0.14.2 → 0.15.0

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

@@ -5,70 +5,78 @@ module LlmCostTracker
     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) }
+          Module.new.tap do |patch|
+            methods.each do |method_name|
+              patch.define_method(method_name) do |*args, **kwargs, &block|
+                seam = Openai.blocking_seam(@client, record_method)
+                Openai.wrap_blocking(args, kwargs, **seam) { super(*args, **kwargs, &block) }
+              end
+            end
           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
-                )
+          Module.new.tap do |patch|
+            methods.each do |method_name|
+              patch.define_method(method_name) do |*args, **kwargs|
+                Openai.wrap_stream(args, kwargs, **Openai.stream_seam(@client)) { super(*args, **kwargs) }
               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
       end
 
       module ResponsesPatch
-        include PatchBuilder.build(record_method: :record_response, methods: %i[create])
+        include PatchBuilder.build(record_method: :record_response, methods: %i[create compact])
         include PatchBuilder.build_stream(methods: %i[stream stream_raw])
 
         def retrieve(response_id, *args, **kwargs)
-          response = super
-          LlmCostTracker::Integrations::Openai.record_retrieved_response(
-            response, host: LlmCostTracker::Integrations::Openai.client_host_for(self)
-          )
-          response
+          super.tap { |response| Openai.record_retrieved_response(response, host: Openai.client_host(@client)) }
         end
 
         def retrieve_streaming(response_id, *args, **kwargs)
-          LlmCostTracker::Integrations::Openai.wrap_stream(
-            args, kwargs, **LlmCostTracker::Integrations::Openai.stream_seam(self)
-          ) do |collector|
+          Openai.wrap_stream(args, kwargs, **Openai.stream_seam(@client)) do |collector|
             collector.provider_response_id = response_id
             super(response_id, *args, **kwargs)
           end
         end
       end
 
+      module ResponsesConnectionPatch
+        def send_event(event)
+          llm_cost_tracker_capture.sending(event) { super }
+        end
+
+        def each(&block)
+          return super unless block
+
+          capture = llm_cost_tracker_capture
+          capture.reading do
+            super do |event|
+              capture.track(event)
+              block.call(event)
+            end
+          end
+        end
+
+        def receive
+          llm_cost_tracker_capture.raise_deferred
+          super.tap { |event| llm_cost_tracker_capture.track(event) }
+        end
+
+        private
+
+        def llm_cost_tracker_capture
+          @llm_cost_tracker_capture ||= WebsocketCapture.new(url)
+        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])
+      CreatePatch = 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])
@@ -78,10 +86,13 @@ module LlmCostTracker
       StreamingTranscriptionsPatch = PatchBuilder.build_stream(methods: %i[create_streaming])
 
       module BatchesPatch
+        def create(*args, **kwargs)
+          Openai.enforce_budget!(request: Openai.request_params(args, kwargs))
+          super
+        end
+
         def retrieve(batch_id, *args, **kwargs)
-          batch = super
-          LlmCostTracker::Integrations::Openai::BatchCapture.maybe_capture(batch, resource: self)
-          batch
+          super.tap { |batch| BatchCapture.capture(batch, client: @client) }
         end
       end
     end

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

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require "English"
+
+module LlmCostTracker
+  module Integrations
+    module Openai
+      class WebsocketCapture
+        TRACKED_EVENTS = %w[response.created error response.completed response.incomplete response.failed].freeze
+        private_constant :TRACKED_EVENTS
+
+        def initialize(url)
+          @url = url
+          @lanes = {}
+          @tags_by_response = {}
+          @deferred = nil
+        end
+
+        def sending(event)
+          request = budget_checked_create(event)
+          yield.tap { (@lanes[request["stream_id"]] ||= [nil]) << Tags::Context.tags if request }
+        end
+
+        def reading
+          yield
+        ensure
+          raise_deferred unless $ERROR_INFO
+        end
+
+        def track(event)
+          type = (event.type if event.respond_to?(:type)).to_s
+          return unless Openai.active? && TRACKED_EVENTS.include?(type)
+
+          Openai.record_safely { route(type, Capture::SdkPayload.normalize(event)) }
+        rescue BudgetExceededError, UnknownPricingError => e
+          @deferred ||= e
+        end
+
+        def raise_deferred
+          error = @deferred
+          @deferred = nil
+          raise error if error
+        end
+
+        private
+
+        def budget_checked_create(event)
+          return unless Openai.active?
+
+          request = Capture::SdkPayload.normalize(event)
+          return unless request.is_a?(Hash) && request["type"] == "response.create"
+
+          Openai.enforce_budget!(request: request.with_indifferent_access, provider: Openai.provider_for_host(url_host))
+          request
+        end
+
+        def route(type, data)
+          response = data["response"].to_h
+          case type
+          when "response.created" then @tags_by_response[response["id"]] = advance_lane(data["stream_id"])
+          when "error" then advance_lane(data["stream_id"])
+          else record(response)
+          end
+        end
+
+        def advance_lane(stream_id)
+          lane = @lanes[stream_id]
+          lane.shift if lane && lane.size > 1
+          lane&.first
+        end
+
+        def record(response)
+          tags = @tags_by_response.delete(response["id"])
+          host = url_host
+          event = Providers::Openai::ResponseParser.event_from_response(
+            response: response,
+            request: { "tools" => response["tools"] },
+            provider: Openai.provider_for_host(host),
+            host: host,
+            usage_source: Usage::Source::STREAM_FINAL
+          )
+          Tracker.record(event: event.with(stream: true), context_tags: tags) if event
+        end
+
+        def url_host = URI(@url.to_s).host
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/openai.rb

@@ -9,6 +9,7 @@ require_relative "../providers/openai/service_charges"
 require_relative "../providers/openai/usage_extractor"
 require_relative "openai/patches"
 require_relative "openai/batch_capture"
+require_relative "openai/websocket_capture"
 
 module LlmCostTracker
   module Integrations
@@ -18,44 +19,21 @@ module LlmCostTracker
       minimum_version "0.59.0"
 
       class << self
-        def stream_collector(request, host: nil)
-          LlmCostTracker::Capture::StreamCollector.new(
-            provider: provider_for_host(host),
-            model: request[:model],
-            # Host part only: the parser prefers the served tier over the requested one, e.g. after a downgrade.
-            pricing_mode: host_pricing_mode(host, request),
-            request: request
-          )
-        end
-
-        def stream_seam(resource)
-          host = client_host_for(resource)
-          {
-            provider: provider_for_host(host),
-            collector: ->(request) { stream_collector(request, host: host) }
-          }
-        end
-
-        def client_host_for(resource)
-          client = resource.instance_variable_get(:@client)
-          return nil unless client
-
-          URI.parse(client.base_url.to_s).host
-        rescue URI::InvalidURIError
-          nil
-        end
-
-        def provider_for_host(host)
-          return "azure_openai" if LlmCostTracker::Providers::Azure::Hosts.openai?(host)
-
-          LlmCostTracker.configuration.capture.openai_compatible_providers[host.to_s.downcase] || "openai"
-        end
-
         def patch_targets
           [
             patch_target("OpenAI::Resources::Responses", with: ResponsesPatch),
+            patch_target("OpenAI::Resources::Beta::Responses",
+                         with: ResponsesPatch,
+                         optional: true,
+                         skip_when_methods_missing: true,
+                         since: "0.69.0"),
+            patch_target("OpenAI::Responses::Connection",
+                         with: ResponsesConnectionPatch,
+                         optional: true,
+                         since: "0.86.0"),
             patch_target("OpenAI::Resources::Chat::Completions", with: ChatCompletionsPatch),
-            patch_target("OpenAI::Resources::Embeddings", with: EmbeddingsPatch, optional: true),
+            patch_target("OpenAI::Resources::Completions", with: CreatePatch, optional: true),
+            patch_target("OpenAI::Resources::Embeddings", with: CreatePatch, optional: true),
             patch_target("OpenAI::Resources::Images", with: ImagesPatch, optional: true),
             patch_target("OpenAI::Resources::Images",
                          with: StreamingImagesPatch,
@@ -73,29 +51,53 @@ module LlmCostTracker
           ]
         end
 
+        def provider_for_host(host)
+          return "azure_openai" if Providers::Azure::Hosts.openai?(host)
+
+          LlmCostTracker.configuration.capture.openai_compatible_providers[host.to_s.downcase] ||
+            (Providers::Openai::Hosts.bedrock?(host) ? "bedrock" : "openai")
+        end
+
+        def blocking_seam(client, record_method)
+          host = client_host(client)
+          {
+            provider: provider_for_host(host),
+            record: lambda do |response, request, latency_ms|
+              public_send(record_method, response, request: request, latency_ms: latency_ms, host: host)
+            end
+          }
+        end
+
+        def stream_seam(client)
+          host = client_host(client)
+          { provider: provider_for_host(host), collector: ->(request) { stream_collector(request, host: host) } }
+        end
+
+        def stream_collector(request, host: nil)
+          Capture::StreamCollector.new(
+            provider: provider_for_host(host),
+            parsed_as: "openai",
+            model: request[:model],
+            pricing_mode: host_pricing_mode(host, request),
+            request: request
+          )
+        end
+
         def record_response(response, request:, latency_ms:, host: nil)
           return unless active?
 
           record_safely do
-            normalized = LlmCostTracker::Capture::SdkPayload.normalize(response)
-            usage = normalized["usage"] || {}
-            input_tokens = usage["input_tokens"] || usage["prompt_tokens"]
-            output_tokens = usage["output_tokens"] || usage["completion_tokens"]
-            if input_tokens.nil? && output_tokens.nil?
-              unless normalized["background"]
-                Logging.warn("OpenAI response #{normalized['id']} has no usage; not recorded")
-              end
-              next
-            end
+            body = Capture::SdkPayload.normalize(response)
+            next warn_missing_usage(body) unless tokens_reported?(body)
 
-            event = LlmCostTracker::Providers::Openai::ResponseParser.event_from_response(
-              response: normalized,
+            event = Providers::Openai::ResponseParser.event_from_response(
+              response: body,
               request: request,
               provider: provider_for_host(host),
               host: host,
-              usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE
+              usage_source: Usage::Source::SDK_RESPONSE
             )
-            LlmCostTracker::Tracker.record(event: event, latency_ms: latency_ms)
+            Tracker.record(event: event, latency_ms: latency_ms) if event
           end
         end
 
@@ -103,11 +105,11 @@ module LlmCostTracker
           return unless active?
 
           record_safely do
-            event = LlmCostTracker::Providers::Openai::ResponseParser.retrieved_event(
-              response: LlmCostTracker::Capture::SdkPayload.normalize(response),
+            event = Providers::Openai::ResponseParser.retrieved_event(
+              response: Capture::SdkPayload.normalize(response),
               provider: provider_for_host(host),
               host: host,
-              usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE
+              usage_source: Usage::Source::SDK_RESPONSE
             )
             record_once(event) if event
           end
@@ -115,57 +117,39 @@ module LlmCostTracker
 
         def record_image(response, request:, latency_ms:, host: nil)
           usage = usage_hash_from(response) || {}
-          raw_input = usage[:input_tokens].to_i
-          image_input = LlmCostTracker::Providers::Openai::UsageExtractor.image_input_tokens(usage)
-          cache_read = LlmCostTracker::Providers::Openai::UsageExtractor.cache_read_input_tokens(usage)
-          image_output, text_output = LlmCostTracker::Providers::Openai::UsageExtractor.split_output(
-            output_tokens: usage[:output_tokens].to_i,
-            image_output_details: LlmCostTracker::Providers::Openai::UsageExtractor.image_output_tokens(usage),
-            text_output_details: LlmCostTracker::Providers::Openai::UsageExtractor.text_output_tokens(usage),
-            audio_output: 0,
-            default_to_image: true
-          )
+          extractor = Providers::Openai::UsageExtractor
           record_passthrough(
             model: request[:model],
             response: response,
             latency_ms: latency_ms,
             provider: provider_for_host(host),
             pricing_mode: host_pricing_mode(host, request),
-            input_tokens: [raw_input - image_input - cache_read, 0].max,
-            image_input_tokens: image_input,
-            output_tokens: text_output,
-            image_output_tokens: image_output,
-            cache_read_input_tokens: cache_read
+            service_line_items: Providers::Openai::ServiceCharges.billed_line_items(usage) +
+                                extractor.cache_read_line_items(usage),
+            **extractor.token_usage(usage, model: request[:model], default_to_image: true).to_h
           )
         end
 
         def record_transcription(response, request:, latency_ms:, host: nil)
-          usage = usage_hash_from(response)
+          usage = transcription_usage(response)
           record_passthrough(
             model: request[:model],
             response: response,
             latency_ms: latency_ms,
             provider: provider_for_host(host),
             pricing_mode: host_pricing_mode(host, request),
-            service_line_items: LlmCostTracker::Providers::Openai::ServiceCharges.transcription_line_items(usage),
-            usage_source: usage ? LlmCostTracker::Usage::Source::SDK_RESPONSE : LlmCostTracker::Usage::Source::UNKNOWN,
+            service_line_items: Providers::Openai::ServiceCharges.transcription_line_items(usage),
+            usage_source: usage ? Usage::Source::SDK_RESPONSE : Usage::Source::UNKNOWN,
             **transcription_token_attributes(usage)
           )
         end
 
-        def transcription_token_attributes(usage)
-          return { input_tokens: 0, output_tokens: 0 } unless usage && usage[:type].to_s == "tokens"
-
-          raw_input = usage[:input_tokens].to_i
-          audio_input = LlmCostTracker::Providers::Openai::UsageExtractor.audio_input_tokens(usage)
-          {
-            input_tokens: [raw_input - audio_input, 0].max,
-            audio_input_tokens: audio_input,
-            output_tokens: usage[:output_tokens].to_i
-          }
-        end
+        def record_speech(response, request:, latency_ms:, host: nil)
+          if request[:stream_format].to_s == "sse" && response.respond_to?(:string)
+            return record_speech_stream(response, request: request, latency_ms: latency_ms, host: host)
+          end
 
-        def record_speech(_response, request:, latency_ms:, host: nil)
+          line_items = Providers::Openai::ServiceCharges.speech_line_items(request)
           record_passthrough(
             model: request[:model],
             response: nil,
@@ -173,7 +157,8 @@ module LlmCostTracker
             provider: provider_for_host(host),
             input_tokens: 0,
             output_tokens: 0,
-            service_line_items: LlmCostTracker::Providers::Openai::ServiceCharges.speech_line_items(request)
+            service_line_items: line_items,
+            usage_source: line_items.empty? ? Usage::Source::UNKNOWN : Usage::Source::SDK_RESPONSE
           )
         end
 
@@ -188,8 +173,10 @@ module LlmCostTracker
           )
         end
 
+        private
+
         def host_pricing_mode(host, request)
-          LlmCostTracker::Providers::Openai::ResponseParser.combined_pricing_mode(
+          Providers::Openai::ResponseParser.combined_pricing_mode(
             provider: provider_for_host(host), host: host, model: request[:model], service_tier: nil
           )
         end
@@ -197,6 +184,50 @@ module LlmCostTracker
         def usage_hash_from(response)
           response.try(:usage)&.deep_to_h
         end
+
+        def tokens_reported?(body)
+          usage = body["usage"] || {}
+          input_tokens = usage["input_tokens"] || usage["prompt_tokens"]
+          output_tokens = usage["output_tokens"] || usage["completion_tokens"]
+          !(input_tokens.nil? && output_tokens.nil?)
+        end
+
+        def warn_missing_usage(body)
+          Logging.warn("OpenAI response #{body['id']} has no usage; not recorded") unless body["background"]
+        end
+
+        def transcription_usage(response)
+          body = Hash.try_convert(Capture::SdkPayload.normalize(response)) || {}
+          usage = body["usage"]&.deep_symbolize_keys
+          usage ||= { type: "duration", seconds: body["duration"].ceil } if body["duration"].to_f.positive?
+          usage
+        end
+
+        def transcription_token_attributes(usage)
+          return { input_tokens: 0, output_tokens: 0 } unless usage && usage[:type].to_s == "tokens"
+
+          raw_input = usage[:input_tokens].to_i
+          audio_input = Providers::Openai::UsageExtractor.audio_input_tokens(usage)
+          {
+            input_tokens: [raw_input - audio_input, 0].max,
+            audio_input_tokens: audio_input,
+            output_tokens: usage[:output_tokens].to_i
+          }
+        end
+
+        def record_speech_stream(response, request:, latency_ms:, host:)
+          return unless active?
+
+          collector = Capture::StreamCollector.new(
+            provider: provider_for_host(host),
+            model: request[:model],
+            latency_ms: latency_ms,
+            pricing_mode: host_pricing_mode(host, request),
+            request: request
+          )
+          Capture::SSE.parse(response.string).each { |event| collector.event(event[:data]) }
+          collector.finish!
+        end
       end
     end
   end

data/lib/llm_cost_tracker/integrations/ruby_llm/v1/patches.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V1
+        module ProviderPatch
+          def complete(*args, **kwargs, &)
+            seam = V1.blocking_seam(self, :record_completion, has_block: block_given?)
+            V1.wrap_blocking(args, kwargs, **seam) { super }
+          end
+
+          def embed(*args, **kwargs)
+            V1.wrap_blocking(args, kwargs, **V1.blocking_seam(self, :record_embedding)) { super }
+          end
+
+          def transcribe(*args, **kwargs)
+            V1.wrap_blocking(args, kwargs, **V1.blocking_seam(self, :record_transcription)) { super }
+          end
+
+          def paint(*args, **kwargs)
+            V1.wrap_blocking(args, kwargs, **V1.blocking_seam(self, :record_image)) { super }
+          end
+
+          def moderate(*args, **kwargs)
+            V1.wrap_blocking(args, kwargs, **V1.blocking_seam(self, :record_moderation)) { super }
+          end
+        end
+
+        module GeminiTranscriptionPatch
+          def transcribe(*args, **kwargs)
+            V1.wrap_blocking(args, kwargs, **V1.blocking_seam(self, :record_transcription)) { super }
+          end
+        end
+
+        module ResponseBodyPatch
+          def parse_transcription_response(response, **)
+            V1.keep_usage(super, response.body)
+          end
+
+          def parse_embedding_response(response, **)
+            V1.keep_usage(super, response.body)
+          end
+        end
+
+        module StreamPatch
+          private
+
+          def stream_response(...)
+            body = @llm_cost_tracker_stream_body = {}
+            super.tap { |message| message.raw.instance_variable_set(Reply::KEPT_BODY, body) }
+          end
+
+          def build_on_data_handler(*, &handler)
+            body = @llm_cost_tracker_stream_body
+            super do |data|
+              if body && data.is_a?(Hash)
+                body.deep_merge!(data.values_at("message", "response").find { |part| part.is_a?(Hash) } || data)
+              end
+              handler.call(data)
+            end
+          end
+        end
+      end
+    end
+  end
+end