llm_cost_tracker 0.14.1 → 0.14.2

data/lib/llm_cost_tracker/integrations/base.rb

@@ -16,13 +16,10 @@ module LlmCostTracker
         @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 provider = integration_name.to_s
 
       def active?
-        LlmCostTracker.configuration.instrumented?(integration_name)
+        LlmCostTracker.configuration.enabled && LlmCostTracker.configuration.instrumented?(integration_name)
       end
 
       def install
@@ -30,7 +27,8 @@ module LlmCostTracker
         Logging.warn(untested_version_message) if untested_version?
         patch_targets.each do |target|
           target_class = target.fetch(:constant_name).to_s.safe_constantize
-          install_patch(target_class, target.fetch(:patch)) if target_class
+          patch = target.fetch(:patch)
+          target_class.prepend(patch) if target_class && !target_class.ancestors.include?(patch)
         end
       end
 
@@ -68,6 +66,40 @@ module LlmCostTracker
         Logging.warn("#{integration_name} integration failed to record usage: #{e.class}: #{e.message}")
       end
 
+      def record_passthrough(provider:,
+                             model:,
+                             response:,
+                             latency_ms:,
+                             service_line_items: [],
+                             usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE,
+                             pricing_mode: nil,
+                             **token_attributes)
+        return unless active?
+
+        record_safely do
+          LlmCostTracker::Tracker.record(
+            event: LlmCostTracker::Event.build(
+              provider: provider,
+              model: model,
+              token_usage: LlmCostTracker::Usage::TokenUsage.build(**token_attributes),
+              usage_source: usage_source,
+              pricing_mode: pricing_mode,
+              provider_response_id: provider_response_id_for(response),
+              service_line_items: service_line_items
+            ),
+            latency_ms: latency_ms
+          )
+        end
+      end
+
+      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 request_params(args, kwargs)
         params =
           case args.first
@@ -108,19 +140,14 @@ module LlmCostTracker
         ).wrap
       end
 
-      def stream_collector(request, provider: self.provider)
+      def stream_collector(request)
         LlmCostTracker::Capture::StreamCollector.new(
           provider: provider,
           model: request[:model],
-          pricing_mode: stream_pricing_mode(request),
           request: request
         )
       end
 
-      def stream_pricing_mode(_request)
-        nil
-      end
-
       def minimum_version(value = nil)
         @minimum_version = value if value
         @minimum_version
@@ -196,12 +223,6 @@ module LlmCostTracker
           "#{target.fetch(:constant_name)}##{method_name} is not available"
         end
       end
-
-      def install_patch(target, patch)
-        return if target&.ancestors&.include?(patch)
-
-        target.prepend(patch)
-      end
     end
   end
 end

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

@@ -13,15 +13,15 @@ module LlmCostTracker
         class << self
           def maybe_capture(batch, resource:)
             return unless Openai.active?
-            return unless batch.status.to_s == "completed"
+            return unless %w[completed expired cancelled].include?(batch.status.to_s)
             return unless batch.output_file_id && batch.id
             return if captured?(batch.id)
 
             client = resource.instance_variable_get(:@client)
             host = Openai.client_host_for(resource)
             Openai.record_safely do
-              io = client.files.content(batch.output_file_id)
-              deferred = capture_jsonl(io.respond_to?(:read) ? io.read : io.to_s, host: host)
+              jsonl = client.files.content(batch.output_file_id).read
+              deferred = capture_jsonl(jsonl, host: host, model: batch.model)
               mark_captured(batch.id)
               raise deferred if deferred
             end
@@ -41,19 +41,16 @@ module LlmCostTracker
             end
           end
 
-          def capture_jsonl(jsonl, host:)
+          def capture_jsonl(jsonl, host:, model:)
             deferred = nil
             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)
+              record_result({ "id" => entry["id"] }.merge(response), host: host, model: model)
             rescue LlmCostTracker::BudgetExceededError, LlmCostTracker::UnknownPricingError => e
               deferred ||= e
             end
@@ -66,19 +63,25 @@ module LlmCostTracker
             nil
           end
 
-          def record_result(response, host:)
+          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"])
 
-            event = LlmCostTracker::Providers::Openai::ResponseParser.event_from_response(
+            parser = LlmCostTracker::Providers::Openai::ResponseParser
+            event = parser.event_from_response(
               response: response,
-              request: {},
+              request: { "model" => model },
               provider: provider,
               host: host,
               usage_source: LlmCostTracker::Usage::Source::SDK_BATCH_RESULT,
-              pricing_mode: "batch"
+              # /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"
+              )
             )
-            LlmCostTracker::Tracker.record(event: event) if event
+            Openai.record_once(event)
           end
         end
       end

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

@@ -45,6 +45,14 @@ module LlmCostTracker
         include PatchBuilder.build(record_method: :record_response, methods: %i[create])
         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
+        end
+
         def retrieve_streaming(response_id, *args, **kwargs)
           LlmCostTracker::Integrations::Openai.wrap_stream(
             args, kwargs, **LlmCostTracker::Integrations::Openai.stream_seam(self)

data/lib/llm_cost_tracker/integrations/openai.rb

@@ -18,19 +18,12 @@ module LlmCostTracker
       minimum_version "0.59.0"
 
       class << self
-        def stream_pricing_mode(request, host: nil)
-          LlmCostTracker::Providers::Openai::ResponseParser.combined_pricing_mode(
-            host: host,
-            model: (request || {})[:model],
-            service_tier: (request || {})[:service_tier]
-          )
-        end
-
         def stream_collector(request, host: nil)
           LlmCostTracker::Capture::StreamCollector.new(
             provider: provider_for_host(host),
             model: request[:model],
-            pricing_mode: stream_pricing_mode(request, host: host),
+            # 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
@@ -53,29 +46,25 @@ module LlmCostTracker
         end
 
         def provider_for_host(host)
-          LlmCostTracker::Providers::Azure::Hosts.openai?(host) ? "azure_openai" : "openai"
+          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::Chat::Completions", with: ChatCompletionsPatch),
-            *auxiliary_patch_targets
-          ]
-        end
-
-        def auxiliary_patch_targets
-          [
             patch_target("OpenAI::Resources::Embeddings", with: EmbeddingsPatch, optional: true),
             patch_target("OpenAI::Resources::Images", with: ImagesPatch, optional: true),
             patch_target("OpenAI::Resources::Images",
                          with: StreamingImagesPatch,
-                                                                               optional: true,
+                         optional: true,
                          skip_when_methods_missing: true),
             patch_target("OpenAI::Resources::Audio::Transcriptions", with: TranscriptionsPatch, optional: true),
             patch_target("OpenAI::Resources::Audio::Transcriptions",
                          with: StreamingTranscriptionsPatch,
-                                                                                              optional: true,
+                         optional: true,
                          skip_when_methods_missing: true),
             patch_target("OpenAI::Resources::Audio::Translations", with: TranslationsPatch, optional: true),
             patch_target("OpenAI::Resources::Audio::Speech", with: SpeechPatch, optional: true),
@@ -89,11 +78,14 @@ module LlmCostTracker
 
           record_safely do
             normalized = LlmCostTracker::Capture::SdkPayload.normalize(response)
-            usage = normalized["usage"]
-            if usage
-              input_tokens = usage["input_tokens"] || usage["prompt_tokens"]
-              output_tokens = usage["output_tokens"] || usage["completion_tokens"]
-              next if input_tokens.nil? && output_tokens.nil?
+            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
 
             event = LlmCostTracker::Providers::Openai::ResponseParser.event_from_response(
@@ -103,7 +95,21 @@ module LlmCostTracker
               host: host,
               usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE
             )
-            LlmCostTracker::Tracker.record(event: event, latency_ms: latency_ms) if event
+            LlmCostTracker::Tracker.record(event: event, latency_ms: latency_ms)
+          end
+        end
+
+        def record_retrieved_response(response, host:)
+          return unless active?
+
+          record_safely do
+            event = LlmCostTracker::Providers::Openai::ResponseParser.retrieved_event(
+              response: LlmCostTracker::Capture::SdkPayload.normalize(response),
+              provider: provider_for_host(host),
+              host: host,
+              usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE
+            )
+            record_once(event) if event
           end
         end
 
@@ -123,7 +129,8 @@ module LlmCostTracker
             model: request[:model],
             response: response,
             latency_ms: latency_ms,
-            host: host,
+            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,
@@ -138,8 +145,10 @@ module LlmCostTracker
             model: request[:model],
             response: response,
             latency_ms: latency_ms,
-            host: host,
+            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,
             **transcription_token_attributes(usage)
           )
         end
@@ -161,54 +170,28 @@ module LlmCostTracker
             model: request[:model],
             response: nil,
             latency_ms: latency_ms,
-            host: host,
+            provider: provider_for_host(host),
             input_tokens: 0,
             output_tokens: 0,
-            service_line_items: speech_line_items(request)
+            service_line_items: LlmCostTracker::Providers::Openai::ServiceCharges.speech_line_items(request)
           )
         end
 
-        def speech_line_items(request)
-          input = request[:input]
-          return [] unless input.is_a?(String)
-          return [] unless LlmCostTracker::Providers::Openai::ModelFamilies.character_billed_tts?(request[:model])
-
-          [LlmCostTracker::Charges::LineItem.build(
-            dimension_key: "text_to_speech_character",
-            quantity: input.length,
-            cost_status: LlmCostTracker::Charges::CostStatus::UNKNOWN,
-            pricing_basis: "provider_usage",
-            provider_field: "request.input"
-          )]
-        end
-
         def record_moderation(response, request:, latency_ms:, host: nil)
           record_passthrough(
             model: response.model || request[:model],
             response: response,
             latency_ms: latency_ms,
-            host: host,
+            provider: provider_for_host(host),
             input_tokens: 0,
             output_tokens: 0
           )
         end
 
-        def record_passthrough(model:, response:, latency_ms:, host: nil, service_line_items: [], **token_attributes)
-          return unless active?
-
-          record_safely do
-            LlmCostTracker::Tracker.record(
-              event: Event.build(
-                provider: provider_for_host(host),
-                model: model,
-                token_usage: Usage::TokenUsage.build(**token_attributes),
-                usage_source: LlmCostTracker::Usage::Source::SDK_RESPONSE,
-                provider_response_id: response&.try(:id),
-                service_line_items: service_line_items
-              ),
-              latency_ms: latency_ms
-            )
-          end
+        def host_pricing_mode(host, request)
+          LlmCostTracker::Providers::Openai::ResponseParser.combined_pricing_mode(
+            provider: provider_for_host(host), host: host, model: request[:model], service_tier: nil
+          )
         end
 
         def usage_hash_from(response)