llm_cost_tracker 0.13.0 → 0.14.1

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,10 +2,10 @@
 
 require "faraday"
 require "json"
-require "stringio"
 require "uri"
 
 require_relative "../capture/sse"
+require_relative "../capture/stream_tap"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -27,16 +27,17 @@ module LlmCostTracker
         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)
           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(
@@ -94,8 +95,8 @@ module LlmCostTracker
       end
 
       def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
-        return nil unless LlmCostTracker.configuration.auto_enable_stream_usage
-        return nil unless parser&.auto_enable_stream_usage?(request_url)
+        return nil unless LlmCostTracker.configuration.capture.request_stream_usage
+        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 +117,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
@@ -171,7 +173,7 @@ module LlmCostTracker
           metadata: metadata,
           context_tags: context_tags
         )
-      rescue LlmCostTracker::Error
+      rescue *LlmCostTracker::CALLER_ERRORS
         raise
       rescue StandardError => e
         Logging.warn("Error processing response: #{e.class}: #{e.message}")
@@ -198,32 +200,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 +233,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
@@ -295,7 +284,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

@@ -104,7 +104,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