llm_cost_tracker 0.14.2 → 0.15.0

data/lib/llm_cost_tracker/ledger/store.rb

@@ -11,6 +11,15 @@ require_relative "storable"
 module LlmCostTracker
   module Ledger
     module Store
+      CALL_FIELDS = %i[
+        event_id provider model tracked_at pricing_mode latency_ms stream usage_source
+        provider_response_id provider_project_id provider_api_key_id provider_workspace_id
+      ].freeze
+      LINE_ITEM_FIELDS = %i[
+        kind direction modality cache_state quantity unit rate_amount rate_quantity cost currency cost_status
+        pricing_basis price_key price_source price_source_version provider_field provider_item_id
+      ].freeze
+
       class << self
         def insert(events)
           events = Array(events)
@@ -37,7 +46,7 @@ module LlmCostTracker
         private
 
         def insert_calls_returning_ids(rows, insertable)
-          if LlmCostTracker::Call.connection.supports_insert_returning?
+          if LlmCostTracker::Call.lease_connection.supports_insert_returning?
             result = LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: %i[id event_id])
             result.rows.to_h { |id, event_id| [event_id, id] }
           else
@@ -47,25 +56,14 @@ module LlmCostTracker
         end
 
         def attributes_for(event)
-          attributes = {
-            event_id: event.event_id,
-            provider: event.provider,
-            model: event.model,
-            tracked_at: event.tracked_at,
-            pricing_mode: event.pricing_mode,
-            latency_ms: event.latency_ms,
-            stream: event.stream,
-            usage_source: event.usage_source,
-            provider_response_id: event.provider_response_id,
-            provider_project_id: event.provider_project_id,
-            provider_api_key_id: event.provider_api_key_id,
-            provider_workspace_id: event.provider_workspace_id,
+          Storable.clean(
+            **CALL_FIELDS.to_h { |field| [field, event.public_send(field)] },
             batch: event.batch?,
             cost_status: event.cost_status,
-            pricing_snapshot: event.pricing_snapshot
-          }
-
-          Storable.clean(attributes.merge(event.token_usage.to_h).merge(total_cost: event.cost&.total))
+            pricing_snapshot: event.pricing_snapshot,
+            **event.token_usage.to_h,
+            total_cost: event.cost&.total
+          )
         end
 
         def call_ids_for(events)
@@ -94,23 +92,7 @@ module LlmCostTracker
           Storable.clean(
             llm_cost_tracker_call_id: call_id,
             position: position,
-            kind: line_item.kind,
-            direction: line_item.direction,
-            modality: line_item.modality,
-            cache_state: line_item.cache_state,
-            quantity: line_item.quantity,
-            unit: line_item.unit,
-            rate_amount: line_item.rate_amount,
-            rate_quantity: line_item.rate_quantity,
-            cost: line_item.cost,
-            currency: line_item.currency,
-            cost_status: line_item.cost_status,
-            pricing_basis: line_item.pricing_basis,
-            price_key: line_item.price_key,
-            price_source: line_item.price_source,
-            price_source_version: line_item.price_source_version,
-            provider_field: line_item.provider_field,
-            provider_item_id: line_item.provider_item_id,
+            **LINE_ITEM_FIELDS.to_h { |field| [field, line_item.public_send(field)] },
             details: stored_details(line_item.details),
             created_at: Time.now.utc
           )

data/lib/llm_cost_tracker/ledger/tags/breakdown.rb

@@ -11,7 +11,7 @@ module LlmCostTracker
         class << self
           def join_relation(scope, key)
             validated_key = LlmCostTracker::Tags::Key.validate!(key)
-            connection = scope.connection
+            connection = scope.lease_connection
             join = "LEFT OUTER JOIN #{call_tag_table} ON " \
                    "#{call_tag_table}.llm_cost_tracker_call_id = #{scope.quoted_table_name}.id AND " \
                    "#{call_tag_table}.#{connection.quote_column_name('key')} = #{connection.quote(validated_key)}"
@@ -19,7 +19,7 @@ module LlmCostTracker
           end
 
           def value_arel
-            Arel.sql(raw_value_sql(LlmCostTracker::CallTag.connection))
+            Arel.sql(raw_value_sql(LlmCostTracker::CallTag.lease_connection))
           end
 
           def label_sql(connection)

data/lib/llm_cost_tracker/middleware/faraday/body.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "faraday"
+require "json"
+
+module LlmCostTracker
+  module Middleware
+    class Faraday < ::Faraday::Middleware
+      module Body
+        MULTIPART_BOUNDARY = %r{\Amultipart/form-data\b.*?\bboundary="?([^";\s]+)"?}i
+        MODEL_PART = "\r\nContent-Disposition: form-data; name=\"model\"\r\n(?:[^\r\n]+\r\n)*\r\n([^\r\n]+)\r\n"
+
+        class << self
+          def read(body)
+            case body
+            when String then body
+            when nil then ""
+            when Hash, Array then body.to_json
+            else body.to_str if body.respond_to?(:to_str)
+            end
+          end
+
+          def multipart_model_json(env, parser)
+            body = env.body
+            boundary = env.request_headers["Content-Type"].to_s[MULTIPART_BOUNDARY, 1]
+            return nil unless parser && boundary && body.respond_to?(:read) && body.respond_to?(:rewind)
+
+            model = scan_model(body, model_part(boundary))
+            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
+
+          private
+
+          def model_part(boundary)
+            Regexp.new("--#{Regexp.escape(boundary)}#{MODEL_PART}", Regexp::NOENCODING)
+          end
+
+          def scan_model(io, pattern)
+            window = String.new(encoding: Encoding::BINARY)
+            while (chunk = io.read(65_536))
+              window << chunk.b
+              model = window[pattern, 1]
+              return model.force_encoding(Encoding::UTF_8) if model
+
+              window = window.byteslice([window.bytesize - 512, 0].max..)
+            end
+          ensure
+            io.rewind
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/middleware/faraday/exchange.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+require "faraday"
+require "json"
+
+require_relative "../../timing"
+require_relative "body"
+require_relative "response_reader"
+require_relative "stream_tee"
+
+module LlmCostTracker
+  module Middleware
+    class Faraday < ::Faraday::Middleware
+      class Exchange
+        def initialize(env, tags:)
+          @env = env
+          @tags = tags
+          @url = env.url.to_s
+          @parser = Parsers.find_for(@url)
+          read_request if @parser
+        end
+
+        def call(app)
+          prepare_stream if @streaming
+          if tracked?
+            @context_tags, @metadata = tag_snapshot
+            enforce_budget if post?
+          end
+          @request = nil
+          perform(app, LlmCostTracker::Timing.now_monotonic)
+        end
+
+        private
+
+        def read_request
+          @body = Body.read(@env.body) || Body.multipart_model_json(@env, @parser)
+          @request = @parser.safe_json_parse(@body)
+          @model = @parser.model_for(@url, @request)
+          @streaming = @parser.streaming_request?(@url, @request)
+        rescue StandardError => e
+          @request = {}
+          Logging.warn("Unable to read the request body: #{e.class}: #{e.message}")
+        end
+
+        def prepare_stream
+          @body = request_stream_usage || @body
+          @stream_tap = StreamTee.install(@env, @parser)
+        end
+
+        def request_stream_usage
+          return unless LlmCostTracker.configuration.capture.request_stream_usage
+          return unless @parser.auto_enable_stream_usage?(@url, @request)
+
+          options = @request["stream_options"]
+          return if options.is_a?(Hash) && options.key?("include_usage")
+
+          @request["stream_options"] = (options || {}).merge("include_usage" => true)
+          @env.body = @request.to_json
+        end
+
+        def tracked?
+          @parser || (post? && Parsers.batch_submission?(@env.url))
+        end
+
+        def post? = @env.method == :post
+
+        def tag_snapshot
+          [tags_or_empty { LlmCostTracker::Tags::Context.tags }, tags_or_empty { resolved_tags }]
+        end
+
+        def tags_or_empty
+          yield
+        rescue StandardError => e
+          Logging.warn("Error resolving request tags: #{e.class}: #{e.message}")
+          {}
+        end
+
+        def resolved_tags
+          return @tags.to_h unless @tags.respond_to?(:call)
+
+          (@tags.arity.zero? ? @tags.call : @tags.call(@env)).to_h
+        end
+
+        def enforce_budget
+          Budget.enforce!(
+            provider: @parser&.provider_for(@url),
+            model: @model,
+            request: @request,
+            tags: Tracker.build_tags(context_tags: @context_tags, metadata: @metadata)
+          )
+        end
+
+        def perform(app, started_at)
+          received = false
+          app.call(@env).on_complete do |response_env|
+            received = true
+            record_response(response_env, (LlmCostTracker::Timing.elapsed_ms(started_at) if post?))
+          end
+        rescue StandardError => e
+          record_interruption(e, LlmCostTracker::Timing.elapsed_ms(started_at)) if @streaming && !received
+          raise
+        end
+
+        def record_response(response_env, latency_ms)
+          return unless @parser
+
+          reader = ResponseReader.new(parser: @parser, request_url: @url, request_body: @body, stream_tap: @stream_tap)
+          event = reader.call(response_env, streaming: @streaming)
+          return unless event
+
+          Tracker.record(event: event, latency_ms: latency_ms, metadata: @metadata, context_tags: @context_tags)
+        rescue *LlmCostTracker::CALLER_ERRORS
+          raise
+        rescue ActiveRecord::RecordNotUnique
+          nil
+        rescue StandardError => e
+          Logging.warn("Error processing #{@parser.provider_for(@url)} response for model #{@model || 'unknown'}: " \
+                       "#{e.class}: #{e.message} (#{Array(e.backtrace).first})")
+        end
+
+        def record_interruption(error, latency_ms)
+          Tracker.record(
+            event: interrupted_event,
+            latency_ms: latency_ms,
+            metadata: interruption_metadata(error),
+            context_tags: @context_tags
+          )
+        rescue LlmCostTracker::TransactionAbortedError
+          raise
+        rescue StandardError => e
+          Logging.warn("Error recording interrupted stream: #{e.class}: #{e.message}")
+        end
+
+        def interrupted_event
+          request = @parser.safe_json_parse(@body)
+          Event.build(
+            provider: @parser.provider_for(@url),
+            model: @parser.model_for(@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
+          )
+        end
+
+        def interruption_metadata(error)
+          metadata = (@metadata || {}).merge(stream_interrupted: true, stream_interrupted_error: error.class.name)
+          status = error.response_status if error.respond_to?(:response_status)
+          status ? metadata.merge(stream_interrupted_status: status) : metadata
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/middleware/faraday/response_reader.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/object/blank"
+require "faraday"
+
+require_relative "../../capture/sse"
+require_relative "body"
+
+module LlmCostTracker
+  module Middleware
+    class Faraday < ::Faraday::Middleware
+      ResponseReader = ::Data.define(:parser, :request_url, :request_body, :stream_tap)
+
+      class ResponseReader
+        def call(response_env, streaming:)
+          streaming ? parse_stream(response_env) : parse_response(response_env)
+        end
+
+        private
+
+        def parse_response(response_env)
+          response_body = Body.read(response_env.body)
+          unless response_body
+            Logging.warn(
+              "Unable to read response body for #{Redaction.url(request_url)}; " \
+              "known streaming responses are captured automatically, or via LlmCostTracker.track_stream " \
+              "for custom clients."
+            )
+            return nil
+          end
+
+          parser.parse(
+            request_url: request_url,
+            request_body: request_body,
+            response_status: response_env.status,
+            response_body: response_body,
+            response_headers: response_env.response_headers
+          )
+        end
+
+        def parse_stream(response_env)
+          parser.parse_stream(
+            request_url: request_url,
+            request_body: request_body,
+            response_status: response_env.status,
+            events: stream_events(response_env),
+            response_headers: response_env.response_headers
+          )
+        end
+
+        def stream_events(response_env)
+          if stream_tap&.received?
+            events = stream_tap.events
+            return events unless stream_tap.overflowed? || stream_tap.failed?
+          else
+            body = Body.read(response_env.body)
+            return Capture::SSE.parse(body) if body.present?
+          end
+
+          Logging.warn(capture_warning)
+          []
+        end
+
+        def capture_warning
+          suffix = "recording usage_source=#{Usage::Source::UNKNOWN}. " \
+                   "Use LlmCostTracker.track_stream for manual capture."
+          label = Redaction.url(request_url)
+          return "Unable to capture streaming response for #{label}; #{suffix}" unless stream_tap&.overflowed?
+
+          "Streaming response for #{label} exceeded #{Capture::SSE::LIMIT_BYTES} bytes; #{suffix}"
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/middleware/faraday/stream_tee.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require "faraday"
+
+require_relative "../../capture/stream_tap"
+
+module LlmCostTracker
+  module Middleware
+    class Faraday < ::Faraday::Middleware
+      module StreamTee
+        class << self
+          def install(env, parser)
+            request = env.request
+            return nil unless request
+
+            original = request.on_data
+            return nil unless original
+
+            tap = Capture::StreamTap.new(notable: parser.method(:retain_stream_event?),
+                                         trim: parser.method(:trim_stream_event))
+            request.on_data = proc do |chunk, size, chunk_env|
+              tap << chunk
+              forward(original, chunk, size, chunk_env)
+            end
+            tap
+          rescue StandardError => e
+            Logging.warn("Unable to install streaming tap: #{e.class}: #{e.message}")
+            nil
+          end
+
+          private
+
+          def forward(callable, chunk, size, env)
+            arity = callable.arity
+            return callable.call(chunk, size, env) if arity.negative?
+
+            case arity
+            when 0, 1 then callable.call(chunk)
+            when 2 then callable.call(chunk, size)
+            else callable.call(chunk, size, env)
+            end
+          end
+        end
+      end
+    end
+  end
+end