llm_cost_tracker 0.14.0 → 0.14.2

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,15 +2,17 @@
 
 require "faraday"
 require "json"
-require "stringio"
 require "uri"
 
 require_relative "../capture/sse"
+require_relative "../capture/stream_tap"
 require_relative "../timing"
 
 module LlmCostTracker
   module Middleware
     class Faraday < ::Faraday::Middleware
+      MULTIPART_MODEL = /name="model"\r\n(?:[^\r\n]+\r\n)*\r\n([^\r\n]+)\r\n/
+
       def initialize(app, **options)
         super(app)
         @tags = options.fetch(:tags, {})
@@ -20,23 +22,24 @@ module LlmCostTracker
         return @app.call(request_env) unless LlmCostTracker.configuration.enabled
 
         request_url  = request_env.url.to_s
-        request_body = read_body(request_env.body)
         parser       = Parsers.find_for(request_url)
+        request_body = read_body(request_env.body) || multipart_model_body(request_env, parser)
         request_parsed = parser&.safe_json_parse(request_body)
         streaming = parser&.streaming_request?(request_url, request_parsed)
         if streaming
           request_body = inject_stream_usage_flag(request_env, parser, request_url, request_parsed) || request_body
         end
-        stream_buffer = install_stream_tap(request_env) if streaming
+        stream_buffer = install_stream_tap(request_env, parser) if streaming
 
-        if parser
+        context_tags, metadata = tag_snapshot(request_env) if parser
+        if parser && request_env.method == :post
           Budget.enforce!(
             provider: parser.provider_for(request_url),
             model: parser.model_for(request_url, request_parsed),
-            request: request_parsed
+            request: request_parsed,
+            tags: Tracker.build_tags(context_tags: context_tags, metadata: metadata)
           )
         end
-        context_tags, metadata = tag_snapshot(request_env) if parser
         started_at = LlmCostTracker::Timing.now_monotonic
 
         invoke_app_with_capture(
@@ -71,7 +74,7 @@ module LlmCostTracker
             request_url: request_url,
             request_body: request_body,
             response_env: response_env,
-            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
+            latency_ms: (LlmCostTracker::Timing.elapsed_ms(started_at) if request_env.method == :post),
             streaming: streaming,
             stream_buffer: stream_buffer,
             context_tags: context_tags,
@@ -95,7 +98,7 @@ module LlmCostTracker
 
       def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
         return nil unless LlmCostTracker.configuration.capture.request_stream_usage
-        return nil unless parser&.auto_enable_stream_usage?(request_url)
+        return nil unless parser.auto_enable_stream_usage?(request_url, request_parsed)
 
         stream_options = request_parsed["stream_options"]
         return nil if stream_options.is_a?(Hash) && stream_options.key?("include_usage")
@@ -116,21 +119,22 @@ module LlmCostTracker
         request = parser.safe_json_parse(request_body)
         event = Event.build(
           provider: parser.provider_for(request_url),
-          model: request["model"] || Event::UNKNOWN_MODEL,
+          model: parser.model_for(request_url, request) || Event::UNKNOWN_MODEL,
           token_usage: Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: Usage::Source::UNKNOWN
         )
-        merged_metadata = (metadata || {}).merge(
-          stream_interrupted: true,
-          stream_interrupted_error: "#{error.class}: #{error.message}"
-        )
+        merged_metadata = (metadata || {}).merge(stream_interrupted: true, stream_interrupted_error: error.class.name)
+        status = error.try(:response_status)
+        merged_metadata[:stream_interrupted_status] = status if status
         Tracker.record(
           event: event,
           latency_ms: latency_ms,
           metadata: merged_metadata,
           context_tags: context_tags
         )
+      rescue LlmCostTracker::TransactionAbortedError
+        raise
       rescue StandardError => e
         Logging.warn("Error recording interrupted stream: #{e.class}: #{e.message}")
       end
@@ -163,16 +167,12 @@ module LlmCostTracker
               response_env: response_env
             )
           end
-        return unless parsed
 
-        Tracker.record(
-          event: parsed,
-          latency_ms: latency_ms,
-          metadata: metadata,
-          context_tags: context_tags
-        )
-      rescue LlmCostTracker::Error
+        Tracker.record(event: parsed, latency_ms: latency_ms, metadata: metadata, context_tags: context_tags) if parsed
+      rescue *LlmCostTracker::CALLER_ERRORS
         raise
+      rescue ActiveRecord::RecordNotUnique
+        nil
       rescue StandardError => e
         Logging.warn("Error processing response: #{e.class}: #{e.message}")
       end
@@ -198,32 +198,28 @@ module LlmCostTracker
       end
 
       def parse_stream(parser:, request_url:, request_body:, response_env:, stream_buffer:)
-        overflowed = stream_buffer&.dig(:overflowed) == true
-        Logging.warn(capture_warning(request_url, stream_buffer)) if overflowed
-
-        body = stream_buffer&.dig(:buffer)&.string
-        body = read_body(response_env.body) if body.blank?
-
-        if body.blank?
-          Logging.warn(capture_warning(request_url, stream_buffer)) unless overflowed
-          return parser.parse_stream(
-            request_url: request_url,
-            request_body: request_body,
-            response_status: response_env.status,
-            response_headers: response_env.response_headers
-          )
-        end
-
-        events = overflowed ? [] : Capture::SSE.parse(body)
         parser.parse_stream(
           request_url: request_url,
           request_body: request_body,
           response_status: response_env.status,
-          events: events,
+          events: stream_events(request_url, response_env, stream_buffer),
           response_headers: response_env.response_headers
         )
       end
 
+      def stream_events(request_url, response_env, stream_buffer)
+        if stream_buffer&.received?
+          events = stream_buffer.events
+          return events unless stream_buffer.overflowed? || stream_buffer.failed?
+        else
+          body = read_body(response_env.body)
+          return Capture::SSE.parse(body) if body.present?
+        end
+
+        Logging.warn(capture_warning(request_url, stream_buffer))
+        []
+      end
+
       def forward_on_data_chunk(callable, chunk, size, env)
         arity = callable.arity
         return callable.call(chunk, size, env) if arity.negative?
@@ -235,28 +231,19 @@ module LlmCostTracker
         end
       end
 
-      def install_stream_tap(request_env)
+      def install_stream_tap(request_env, parser)
         request = request_env.request
         return nil unless request
 
         original = request.on_data
         return nil unless original
 
-        state = { buffer: StringIO.new, bytes: 0, overflowed: false }
+        tap = Capture::StreamTap.new(notable: parser.method(:retain_stream_event?))
         request.on_data = proc do |chunk, size, env|
-          chunk = chunk.to_s
-          remaining = Capture::SSE::LIMIT_BYTES - state[:bytes]
-          if chunk.bytesize <= remaining
-            state[:buffer] << chunk
-            state[:bytes] += chunk.bytesize
-          else
-            state[:buffer] << chunk.byteslice(0, remaining) if remaining.positive?
-            state[:bytes] += [remaining, 0].max
-            state[:overflowed] = true
-          end
+          tap << chunk
           forward_on_data_chunk(original, chunk, size, env)
         end
-        state
+        tap
       rescue StandardError => e
         Logging.warn("Unable to install streaming tap: #{e.class}: #{e.message}")
         nil
@@ -272,16 +259,35 @@ module LlmCostTracker
         end
       end
 
+      def multipart_model_body(request_env, parser)
+        body = request_env.body
+        multipart = request_env.request_headers["Content-Type"].to_s.start_with?("multipart/form-data")
+        return nil unless parser && multipart && body.respond_to?(:read) && body.respond_to?(:rewind)
+
+        model = multipart_model(body)
+        model && { "model" => model }.to_json
+      rescue StandardError => e
+        Logging.warn("Unable to read the model from a multipart request: #{e.class}: #{e.message}")
+        nil
+      end
+
+      def multipart_model(body)
+        window = String.new(encoding: Encoding::BINARY)
+        while (chunk = body.read(65_536))
+          window << chunk.b
+          model = window[MULTIPART_MODEL, 1]
+          return model.force_encoding(Encoding::UTF_8) if model
+
+          window = window.byteslice([window.bytesize - 512, 0].max..)
+        end
+      ensure
+        body.rewind
+      end
+
       def resolved_tags(request_env)
-        tags =
-          if @tags.respond_to?(:call)
-            @tags.arity.zero? ? @tags.call : @tags.call(request_env)
-          else
-            @tags
-          end
-        return {} if tags.nil?
+        return @tags.to_h unless @tags.respond_to?(:call)
 
-        tags.to_h
+        (@tags.arity.zero? ? @tags.call : @tags.call(request_env)).to_h
       end
 
       def tag_snapshot(request_env)
@@ -295,7 +301,7 @@ module LlmCostTracker
         suffix = "recording usage_source=#{Usage::Source::UNKNOWN}. " \
                  "Use LlmCostTracker.track_stream for manual capture."
         label = request_url_label(request_url)
-        return "Unable to capture streaming response for #{label}; #{suffix}" unless stream_buffer&.dig(:overflowed)
+        return "Unable to capture streaming response for #{label}; #{suffix}" unless stream_buffer&.overflowed?
 
         "Streaming response for #{label} exceeded #{Capture::SSE::LIMIT_BYTES} bytes; #{suffix}"
       end

data/lib/llm_cost_tracker/parsers.rb

@@ -8,50 +8,27 @@ require_relative "providers"
 
 module LlmCostTracker
   module Parsers
-    PARSER_PROVIDERS = %i[Openai Azure OpenaiCompatible Anthropic Gemini].freeze
+    PARSER_PROVIDERS = %i[Openai Azure Anthropic Gemini OpenaiCompatible].freeze
 
     def self.find_for(url)
-      instances.each do |klass, instance|
-        return instance if klass.match?(url)
-      end
-      nil
+      instances.find { |parser| parser.class.match?(url) }
     end
 
     def self.find_for_provider(provider)
-      provider_name = provider.to_s.downcase
-      instances.each do |klass, instance|
-        return instance if klass.provider_names.include?(provider_name)
-      end
-      nil
+      all_for_provider(provider).first
     end
 
-    def self.parser_classes
-      PARSER_PROVIDERS.map { |name| Providers.const_get(name)::Parser }
+    def self.all_for_provider(provider)
+      provider_name = provider.to_s.downcase
+      instances.select { |parser| parser.class.provider_names.include?(provider_name) }
     end
-    private_class_method :parser_classes
 
     def self.instances
-      @instances ||= parser_classes.to_h { |klass| [klass, klass.new] }.freeze
+      @instances ||= PARSER_PROVIDERS.map { |name| Providers.const_get(name)::Parser.new }.freeze
     end
     private_class_method :instances
 
     module UrlMatchers
-      def match_uri?(url, hosts: nil, exact_paths: nil, path_includes: nil, path_suffixes: nil, path_pattern: nil)
-        uri_matches?(url) do |uri|
-          host_match = hosts.nil? || hosts.include?(uri.host.to_s.downcase)
-          path_match = path_matches?(
-            uri,
-            exact_paths: exact_paths,
-            path_includes: path_includes,
-            path_suffixes: path_suffixes,
-            path_pattern: path_pattern
-          )
-          extra_match = block_given? ? yield(uri) : true
-
-          !!(host_match && path_match && extra_match)
-        end
-      end
-
       def uri_matches?(url)
         uri = parsed_uri(url)
         uri ? yield(uri) : false
@@ -62,16 +39,6 @@ module LlmCostTracker
       rescue URI::InvalidURIError
         nil
       end
-
-      def path_matches?(uri, exact_paths: nil, path_includes: nil, path_suffixes: nil, path_pattern: nil)
-        path = uri.path.to_s
-        matches = true
-        matches &&= exact_paths.include?(path) if exact_paths
-        matches &&= Array(path_includes).all? { |fragment| path.include?(fragment) } if path_includes
-        matches &&= path.match?(path_pattern) if path_pattern
-        matches &&= path_suffixes.any? { |suffix| path == suffix || path.end_with?(suffix) } if path_suffixes
-        matches
-      end
     end
 
     class Base
@@ -104,7 +71,11 @@ module LlmCostTracker
         nil
       end
 
-      def auto_enable_stream_usage?(_request_url)
+      def retain_stream_event?(_data)
+        false
+      end
+
+      def auto_enable_stream_usage?(_request_url, _request_parsed)
         false
       end