llm_cost_tracker 0.5.1 → 0.5.3

data/lib/llm_cost_tracker/integrations/stream_tracker.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+require_relative "../logging"
+require_relative "../stream_collector"
+require_relative "../value_helpers"
+require_relative "object_reader"
+
+module LlmCostTracker
+  module Integrations
+    class StreamTracker
+      def self.wrap(stream, collector:, active:, finish: nil) = new(stream, collector, active, finish).wrap
+
+      def initialize(stream, collector, active, finish)
+        @stream = stream
+        @collector = collector
+        @active = active
+        @finish = finish || proc { |errored:| @collector.finish!(errored: errored) }
+        @finished = false
+        @capture_failed = false
+        @monitor = Monitor.new
+      end
+
+      def wrap
+        return @stream unless @stream
+
+        iterator_wrapped = @stream.instance_variable_defined?(:@iterator) && wrap_iterator?
+        wrap_each if !iterator_wrapped && @stream.respond_to?(:each)
+
+        @stream
+      rescue StandardError => e
+        Logging.warn("stream integration failed to install wrapper: #{e.class}: #{e.message}")
+        @stream
+      end
+
+      private
+
+      def wrap_iterator?
+        iterator = @stream.instance_variable_get(:@iterator)
+        return false unless iterator.respond_to?(:each)
+
+        @stream.instance_variable_set(:@iterator, tracked_iterator(iterator))
+        true
+      end
+
+      def wrap_each
+        tracker = self
+        original_each = @stream.method(:each)
+        @stream.define_singleton_method(:each) do |&block|
+          next enum_for(:each) unless block
+
+          tracker.__send__(:each_from, original_each, &block)
+        end
+      end
+
+      def tracked_iterator(iterator)
+        Enumerator.new do |yielder|
+          each_from(iterator) { |event| yielder << event }
+        end
+      end
+
+      def each_from(iterable)
+        errored = false
+        iterate(iterable) do |event|
+          capture(event)
+          yield event
+        end
+      rescue StandardError
+        errored = true
+        raise
+      ensure
+        finish!(errored: errored)
+      end
+
+      def iterate(iterable, &)
+        if iterable.respond_to?(:each)
+          iterable.each(&)
+        else
+          iterable.call(&)
+        end
+      end
+
+      def capture(event)
+        payload = normalize(event_payload(event))
+        @collector.event(payload, type: event_type(event, payload))
+      rescue StandardError => e
+        warn_capture_failure(e)
+      end
+
+      def event_payload(event)
+        if event.respond_to?(:deep_to_h)
+          event.deep_to_h
+        elsif event.respond_to?(:to_h)
+          event.to_h
+        else
+          event_attributes(event)
+        end
+      end
+
+      def event_attributes(event)
+        %i[type id model usage response message].each_with_object({}) do |key, attributes|
+          value = ObjectReader.read(event, key)
+          attributes[key] = value unless value.nil?
+        end
+      end
+
+      def event_type(event, payload)
+        value = ObjectReader.first(event, :type) || payload["type"]
+        value&.to_s
+      end
+
+      def normalize(value)
+        case value
+        when Hash
+          value.each_with_object({}) do |(key, nested), normalized|
+            normalized[key.to_s] = normalize(nested)
+          end
+        when Array
+          value.map { |nested| normalize(nested) }
+        when Symbol
+          value.to_s
+        when NilClass
+          nil
+        else
+          converted = object_hash(value)
+          converted ? normalize(converted) : ValueHelpers.deep_dup(value)
+        end
+      end
+
+      def object_hash(value)
+        if value.respond_to?(:deep_to_h)
+          value.deep_to_h
+        elsif value.respond_to?(:to_h)
+          value.to_h
+        end
+      rescue StandardError
+        nil
+      end
+
+      def warn_capture_failure(error)
+        should_warn = @monitor.synchronize do
+          next false if @capture_failed
+
+          @capture_failed = true
+          true
+        end
+        return unless should_warn
+
+        Logging.warn("stream integration failed to capture event: #{error.class}: #{error.message}")
+      end
+
+      def finish!(errored:)
+        should_finish = @monitor.synchronize do
+          next false if @finished
+
+          @finished = true
+          true
+        end
+        return unless should_finish && @active.call
+
+        @finish.call(errored: errored)
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/llm_api_call.rb

@@ -2,9 +2,9 @@
 
 require "active_record"
 
+require_relative "llm_api_call_metrics"
 require_relative "period_grouping"
 require_relative "tag_accessors"
-require_relative "tag_key"
 require_relative "tag_query"
 require_relative "tags_column"
 
@@ -12,6 +12,7 @@ module LlmCostTracker
   class LlmApiCall < ActiveRecord::Base
     extend PeriodGrouping
     extend TagsColumn
+    extend LlmApiCallMetrics
     include TagAccessors
 
     self.table_name = "llm_api_calls"
@@ -55,81 +56,5 @@ module LlmCostTracker
     def self.by_tags(tags)
       TagQuery.apply(self, tags)
     end
-
-    def self.total_cost
-      sum(:total_cost).to_f
-    end
-
-    def self.total_tokens
-      sum(:total_tokens).to_i
-    end
-
-    def self.cost_by_model
-      group(:model).sum(:total_cost)
-    end
-
-    def self.cost_by_provider
-      group(:provider).sum(:total_cost)
-    end
-
-    def self.group_by_tag(key)
-      group(Arel.sql(tag_group_expression(key)))
-    end
-
-    def self.cost_by_tag(key)
-      costs = group_by_tag(key).sum(:total_cost).each_with_object(Hash.new(0.0)) do |(tag_value, cost), grouped|
-        grouped[tag_label(tag_value)] += cost.to_f
-      end
-      costs.sort_by { |_label, cost| -cost }.to_h
-    end
-
-    def self.average_latency_ms
-      return nil unless latency_column?
-
-      average(:latency_ms)&.to_f
-    end
-
-    def self.latency_by_model
-      return {} unless latency_column?
-
-      group(:model).average(:latency_ms).transform_values(&:to_f)
-    end
-
-    def self.latency_by_provider
-      return {} unless latency_column?
-
-      group(:provider).average(:latency_ms).transform_values(&:to_f)
-    end
-
-    def self.tag_label(value)
-      value.nil? || value == "" ? "(untagged)" : value.to_s
-    end
-    private_class_method :tag_label
-
-    def self.tag_group_expression(key)
-      key = validated_tag_key(key)
-      column = "#{quoted_table_name}.#{connection.quote_column_name('tags')}"
-
-      case connection.adapter_name
-      when /postgres/i
-        json_column = tags_jsonb_column? ? column : "(#{column})::jsonb"
-        "#{json_column}->>#{connection.quote(key)}"
-      when /mysql/i
-        "JSON_UNQUOTE(JSON_EXTRACT(#{column}, #{connection.quote(json_path(key))}))"
-      else
-        "json_extract(#{column}, #{connection.quote(json_path(key))})"
-      end
-    end
-    private_class_method :tag_group_expression
-
-    def self.validated_tag_key(key)
-      TagKey.validate!(key)
-    end
-    private_class_method :validated_tag_key
-
-    def self.json_path(key)
-      "$.\"#{key}\""
-    end
-    private_class_method :json_path
   end
 end

data/lib/llm_cost_tracker/llm_api_call_metrics.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require_relative "tag_sql"
+
+module LlmCostTracker
+  module LlmApiCallMetrics
+    def total_cost
+      sum(:total_cost).to_f
+    end
+
+    def total_tokens
+      sum(:total_tokens).to_i
+    end
+
+    def cost_by_model
+      group(:model).sum(:total_cost)
+    end
+
+    def cost_by_provider
+      group(:provider).sum(:total_cost)
+    end
+
+    def group_by_tag(key)
+      group(Arel.sql(tag_value_expression(key)))
+    end
+
+    def cost_by_tag(key, limit: nil)
+      relation = group_by_tag(key).order(Arel.sql("COALESCE(SUM(total_cost), 0) DESC"))
+      relation = relation.limit(limit) if limit
+
+      costs = relation.sum(:total_cost).each_with_object(Hash.new(0.0)) do |(tag_value, cost), grouped|
+        grouped[tag_value_label(tag_value)] += cost.to_f
+      end
+      costs.sort_by { |_label, cost| -cost }.to_h
+    end
+
+    def average_latency_ms
+      return nil unless latency_column?
+
+      average(:latency_ms)&.to_f
+    end
+
+    def latency_by_model
+      return {} unless latency_column?
+
+      group(:model).average(:latency_ms).transform_values(&:to_f)
+    end
+
+    def latency_by_provider
+      return {} unless latency_column?
+
+      group(:provider).average(:latency_ms).transform_values(&:to_f)
+    end
+
+    def tag_value_label(value)
+      TagSql.value_label(value)
+    end
+
+    def tag_value_expression(key, table_name: quoted_table_name)
+      TagSql.value_expression(self, key, table_name: table_name)
+    end
+  end
+end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -5,12 +5,11 @@ require "json"
 
 require_relative "../logging"
 require_relative "../request_url"
+require_relative "../stream_capture"
 
 module LlmCostTracker
   module Middleware
     class Faraday < ::Faraday::Middleware
-      STREAM_CAPTURE_LIMIT_BYTES = 1_048_576
-
       def initialize(app, **options)
         super(app)
         @tags = options.fetch(:tags, {})
@@ -88,6 +87,11 @@ module LlmCostTracker
       end
 
       def parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+        if stream_buffer&.dig(:overflowed)
+          Logging.warn(capture_warning(request_url, stream_buffer))
+          return parser.parse_stream(request_url, request_body, response_env.status, [])
+        end
+
         body = stream_buffer&.dig(:buffer)&.string
         body = read_body(response_env.body) if body.nil? || body.empty?
 
@@ -110,7 +114,7 @@ module LlmCostTracker
         request_env.request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
           unless state[:overflowed]
-            if state[:bytes] + chunk.bytesize <= STREAM_CAPTURE_LIMIT_BYTES
+            if state[:bytes] + chunk.bytesize <= StreamCapture::LIMIT_BYTES
               state[:buffer] << chunk
               state[:bytes] += chunk.bytesize
             else
@@ -161,7 +165,7 @@ module LlmCostTracker
                  "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
         end
 
-        "Streaming response for #{RequestUrl.label(request_url)} exceeded #{STREAM_CAPTURE_LIMIT_BYTES} bytes; " \
+        "Streaming response for #{RequestUrl.label(request_url)} exceeded #{StreamCapture::LIMIT_BYTES} bytes; " \
           "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
       end
     end

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -72,7 +72,7 @@ module LlmCostTracker
           model: extract_model_from_url(request_url),
           input_tokens: [usage["promptTokenCount"].to_i - cache_read, 0].max,
           output_tokens: output_tokens(usage),
-          total_tokens: usage["totalTokenCount"].to_i,
+          total_tokens: total_tokens(usage, cache_read),
           cache_read_input_tokens: usage["cachedContentTokenCount"],
           hidden_output_tokens: usage["thoughtsTokenCount"],
           stream: stream,
@@ -92,6 +92,13 @@ module LlmCostTracker
         usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
       end
 
+      def total_tokens(usage, cache_read)
+        total = usage["totalTokenCount"]
+        return total.to_i unless total.nil?
+
+        [usage["promptTokenCount"].to_i - cache_read, 0].max + cache_read + output_tokens(usage)
+      end
+
       def stream_response_id(events)
         find_event_value(events) { |data| data["responseId"] }
       end

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -21,7 +21,7 @@ module LlmCostTracker
           model: response["model"] || request["model"],
           input_tokens: regular_input_tokens(usage, cache_read),
           output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
-          total_tokens: usage["total_tokens"].to_i,
+          total_tokens: total_tokens(usage, cache_read),
           cache_read_input_tokens: cache_read,
           hidden_output_tokens: hidden_output_tokens(usage),
           usage_source: :response
@@ -44,7 +44,7 @@ module LlmCostTracker
             model: model,
             input_tokens: regular_input_tokens(usage, cache_read),
             output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
-            total_tokens: usage["total_tokens"].to_i,
+            total_tokens: total_tokens(usage, cache_read),
             cache_read_input_tokens: cache_read,
             hidden_output_tokens: hidden_output_tokens(usage),
             stream: true,
@@ -61,7 +61,7 @@ module LlmCostTracker
 
       def detect_stream_usage(events)
         find_event_value(events, reverse: true) do |data|
-          usage = data["usage"]
+          usage = data["usage"] || data.dig("response", "usage")
           usage if usage.is_a?(Hash)
         end
       end
@@ -87,6 +87,15 @@ module LlmCostTracker
         details = usage["completion_tokens_details"] || usage["output_tokens_details"] || {}
         details["reasoning_tokens"]
       end
+
+      def total_tokens(usage, cache_read)
+        total = usage["total_tokens"]
+        return total.to_i unless total.nil?
+
+        regular_input_tokens(usage, cache_read) +
+          cache_read.to_i +
+          (usage["completion_tokens"] || usage["output_tokens"]).to_i
+      end
     end
   end
 end

data/lib/llm_cost_tracker/price_registry.rb

@@ -12,6 +12,7 @@ module LlmCostTracker
     EMPTY_PRICES = {}.freeze
     PRICE_KEYS = %w[input output cache_read_input cache_write_input].freeze
     METADATA_KEYS = %w[_source _source_version _fetched_at _updated _notes _validator_override].freeze
+    MAX_FILE_BYTES = 2_097_152
     MUTEX = Monitor.new
 
     class << self
@@ -114,6 +115,8 @@ module LlmCostTracker
       end
 
       def load_price_file(path)
+        raise ArgumentError, "prices_file exceeds #{MAX_FILE_BYTES} bytes" if File.size(path) > MAX_FILE_BYTES
+
         contents = File.read(path)
         return YAML.safe_load(contents, aliases: false) || {} if yaml_file?(path)
 

data/lib/llm_cost_tracker/price_sync/fetcher.rb

@@ -17,6 +17,7 @@ module LlmCostTracker
 
       USER_AGENT = "llm_cost_tracker price refresh"
       MAX_REDIRECTS = 5
+      MAX_BODY_BYTES = 2_097_152
       OPEN_TIMEOUT = 5
       READ_TIMEOUT = 10
       WRITE_TIMEOUT = 10
@@ -25,26 +26,17 @@ module LlmCostTracker
         raise Error, "Too many redirects while fetching #{url}" if redirects > MAX_REDIRECTS
 
         uri = URI.parse(url)
-        raise Error, "Pricing snapshot URL must use http or https" unless %w[http https].include?(uri.scheme)
+        raise Error, "Pricing snapshot URL must use https" unless uri.scheme == "https"
 
         request = Net::HTTP::Get.new(uri)
         request["User-Agent"] = USER_AGENT
         request["If-None-Match"] = etag if etag
 
-        response = Net::HTTP.start(
-          uri.host,
-          uri.port,
-          use_ssl: uri.scheme == "https",
-          open_timeout: OPEN_TIMEOUT,
-          read_timeout: READ_TIMEOUT,
-          write_timeout: WRITE_TIMEOUT
-        ) do |http|
-          http.request(request)
-        end
+        response, body = fetch_response(uri, request)
 
         case response
         when Net::HTTPSuccess
-          build_response(response, not_modified: false)
+          build_response(response, body: body || limited_body(response), not_modified: false)
         when Net::HTTPNotModified
           build_response(response, body: nil, not_modified: true)
         when Net::HTTPRedirection
@@ -61,6 +53,43 @@ module LlmCostTracker
 
       private
 
+      def fetch_response(uri, request)
+        body = nil
+        response = Net::HTTP.start(
+          uri.host,
+          uri.port,
+          use_ssl: uri.scheme == "https",
+          open_timeout: OPEN_TIMEOUT,
+          read_timeout: READ_TIMEOUT,
+          write_timeout: WRITE_TIMEOUT
+        ) do |http|
+          http.request(request) do |streamed_response|
+            body = limited_body(streamed_response) if streamed_response.is_a?(Net::HTTPSuccess)
+          end
+        end
+
+        [response, body]
+      end
+
+      def limited_body(response)
+        body = +""
+        if response.respond_to?(:read_body)
+          response.read_body do |chunk|
+            chunk = chunk.to_s
+            if body.bytesize + chunk.bytesize > MAX_BODY_BYTES
+              raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes"
+            end
+
+            body << chunk
+          end
+        else
+          body = response.body.to_s
+        end
+        raise Error, "Pricing snapshot response exceeds #{MAX_BODY_BYTES} bytes" if body.bytesize > MAX_BODY_BYTES
+
+        body
+      end
+
       def build_response(response, not_modified:, body: response.body)
         Response.new(
           body: body,

data/lib/llm_cost_tracker/price_sync/registry_loader.rb

@@ -3,6 +3,8 @@
 require "json"
 require "yaml"
 
+require_relative "../price_registry"
+
 module LlmCostTracker
   module PriceSync
     class RegistryLoader
@@ -18,6 +20,10 @@ module LlmCostTracker
       private
 
       def load_registry_file(path)
+        if File.size(path) > PriceRegistry::MAX_FILE_BYTES
+          raise ArgumentError, "pricing registry exceeds #{PriceRegistry::MAX_FILE_BYTES} bytes"
+        end
+
         contents = File.read(path)
         registry = yaml_file?(path) ? (YAML.safe_load(contents, aliases: false) || {}) : JSON.parse(contents)
         raise ArgumentError, "pricing registry must be a hash" unless registry.is_a?(Hash)

data/lib/llm_cost_tracker/pricing/effective_prices.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Pricing
+    EffectivePriceSet = Data.define(:input, :cache_read_input, :cache_write_input, :output) do
+      def to_h
+        {
+          input: input,
+          cache_read_input: cache_read_input,
+          cache_write_input: cache_write_input,
+          output: output
+        }
+      end
+
+      def complete?
+        missing_keys.empty?
+      end
+
+      def missing_keys
+        to_h.filter_map { |key, value| key if value.nil? }
+      end
+    end
+
+    module EffectivePrices
+      class << self
+        def call(usage:, prices:, pricing_mode:)
+          EffectivePriceSet.new(
+            input: price_for_usage(usage.input_tokens, prices, :input, pricing_mode),
+            cache_read_input: price_for_cache_usage(
+              usage.cache_read_input_tokens,
+              prices,
+              :cache_read_input,
+              pricing_mode
+            ),
+            cache_write_input: price_for_cache_usage(
+              usage.cache_write_input_tokens,
+              prices,
+              :cache_write_input,
+              pricing_mode
+            ),
+            output: price_for_usage(usage.output_tokens, prices, :output, pricing_mode)
+          )
+        end
+
+        private
+
+        def price_for_cache_usage(tokens, prices, key, pricing_mode)
+          return 0.0 unless tokens.positive?
+
+          price_for(prices, key, pricing_mode) || price_for(prices, :input, pricing_mode)
+        end
+
+        def price_for_usage(tokens, prices, key, pricing_mode)
+          tokens.positive? ? price_for(prices, key, pricing_mode) : 0.0
+        end
+
+        def price_for(prices, key, pricing_mode)
+          mode = normalized_pricing_mode(pricing_mode)
+          return prices[key] unless mode
+
+          prices[:"#{mode}_#{key}"] || prices[key]
+        end
+
+        def normalized_pricing_mode(value)
+          return nil if value.nil?
+
+          mode = value.to_s.strip
+          return nil if mode.empty? || mode == "standard"
+
+          mode
+        end
+      end
+    end
+  end
+end