llm_cost_tracker 0.5.2 → 0.6.0

data/lib/llm_cost_tracker/integrations/anthropic.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -16,11 +17,11 @@ module LlmCostTracker
 
         def patch_targets
           [
-            patch_target("Anthropic::Resources::Messages", with: MessagesPatch, methods: :create),
+            patch_target("Anthropic::Resources::Messages", with: MessagesPatch, methods: %i[create stream stream_raw]),
             patch_target(
               "Anthropic::Resources::Beta::Messages",
               with: MessagesPatch,
-              methods: :create,
+              methods: %i[create stream stream_raw],
               optional: true
             )
           ]
@@ -64,6 +65,28 @@ module LlmCostTracker
             ObjectReader.nested(usage, :output_tokens_details, :reasoning_tokens)
           )
         end
+
+        def track_stream(stream, collector:)
+          return stream unless active?
+
+          StreamTracker.wrap(
+            stream,
+            collector: collector,
+            active: -> { active? },
+            finish: ->(errored:) { finish_stream(collector, errored: errored) }
+          )
+        end
+
+        def stream_collector(request)
+          LlmCostTracker::StreamCollector.new(
+            provider: "anthropic",
+            model: request[:model] || request["model"]
+          )
+        end
+
+        def finish_stream(collector, errored:)
+          record_safely { collector.finish!(errored: errored) }
+        end
       end
 
       module MessagesPatch
@@ -78,6 +101,22 @@ module LlmCostTracker
           )
           message
         end
+
+        def stream(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
+          LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
+        end
+
+        def stream_raw(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Anthropic.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Anthropic.stream_collector(request)
+          LlmCostTracker::Integrations::Anthropic.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Anthropic.track_stream(stream, collector: collector)
+        end
       end
     end
   end

data/lib/llm_cost_tracker/integrations/openai.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "stream_tracker"
 
 module LlmCostTracker
   module Integrations
@@ -16,8 +17,16 @@ module LlmCostTracker
 
         def patch_targets
           [
-            patch_target("OpenAI::Resources::Responses", with: ResponsesPatch, methods: :create),
-            patch_target("OpenAI::Resources::Chat::Completions", with: ChatCompletionsPatch, methods: :create)
+            patch_target(
+              "OpenAI::Resources::Responses",
+              with: ResponsesPatch,
+              methods: %i[create stream stream_raw retrieve_streaming]
+            ),
+            patch_target(
+              "OpenAI::Resources::Chat::Completions",
+              with: ChatCompletionsPatch,
+              methods: %i[create stream_raw]
+            )
           ]
         end
 
@@ -70,6 +79,28 @@ module LlmCostTracker
         def regular_input_tokens(input_tokens, cache_read)
           [ObjectReader.integer(input_tokens) - cache_read.to_i, 0].max
         end
+
+        def track_stream(stream, collector:)
+          return stream unless active?
+
+          StreamTracker.wrap(
+            stream,
+            collector: collector,
+            active: -> { active? },
+            finish: ->(errored:) { finish_stream(collector, errored: errored) }
+          )
+        end
+
+        def stream_collector(request)
+          LlmCostTracker::StreamCollector.new(
+            provider: "openai",
+            model: request[:model] || request["model"]
+          )
+        end
+
+        def finish_stream(collector, errored:)
+          record_safely { collector.finish!(errored: errored) }
+        end
       end
 
       module ResponsesPatch
@@ -84,6 +115,31 @@ module LlmCostTracker
           )
           response
         end
+
+        def stream(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
+
+        def stream_raw(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
+
+        def retrieve_streaming(response_id, *args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          collector.provider_response_id = response_id
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
       end
 
       module ChatCompletionsPatch
@@ -98,6 +154,14 @@ module LlmCostTracker
           )
           response
         end
+
+        def stream_raw(*args, **kwargs)
+          request = LlmCostTracker::Integrations::Openai.request_params(args, kwargs)
+          collector = LlmCostTracker::Integrations::Openai.stream_collector(request)
+          LlmCostTracker::Integrations::Openai.enforce_budget!
+          stream = super
+          LlmCostTracker::Integrations::Openai.track_stream(stream, collector: collector)
+        end
       end
     end
   end

data/lib/llm_cost_tracker/integrations/registry.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "monitor"
+
+require_relative "../errors"
 require_relative "openai"
 require_relative "anthropic"
 require_relative "ruby_llm"
@@ -7,14 +10,22 @@ require_relative "ruby_llm"
 module LlmCostTracker
   module Integrations
     module Registry
-      INTEGRATIONS = {
+      DEFAULT_INTEGRATIONS = {
         openai: Openai,
         anthropic: Anthropic,
         ruby_llm: RubyLlm
       }.freeze
+      MUTEX = Monitor.new
 
       module_function
 
+      def register(name, integration)
+        key = name.to_sym
+        validate_integration!(integration)
+        MUTEX.synchronize { @integrations = integrations.merge(key => integration).freeze }
+        integration
+      end
+
       def install!(names = LlmCostTracker.configuration.instrumented_integrations)
         normalize(names).each { |name| fetch(name).install }
       end
@@ -30,13 +41,32 @@ module LlmCostTracker
       end
 
       def fetch(name)
-        INTEGRATIONS.fetch(name.to_sym) do
-          message = "Unknown integration: #{name.inspect}. Use one of: #{INTEGRATIONS.keys.join(', ')}"
+        integrations.fetch(name.to_sym) do
+          message = "Unknown integration: #{name.inspect}. Use one of: #{names.join(', ')}"
           raise LlmCostTracker::Error, message
         end
       end
+
+      def names
+        integrations.keys
+      end
+
+      def reset!
+        MUTEX.synchronize { @integrations = DEFAULT_INTEGRATIONS.dup.freeze }
+      end
+
+      def integrations
+        @integrations || MUTEX.synchronize { @integrations ||= DEFAULT_INTEGRATIONS.dup.freeze }
+      end
+
+      def validate_integration!(integration)
+        return if integration.respond_to?(:install) && integration.respond_to?(:status)
+
+        raise ArgumentError, "integration must respond to install and status"
+      end
     end
 
+    def self.register(name, integration) = Registry.register(name, integration)
     def self.install! = Registry.install!
     def self.checks = Registry.checks
   end

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,82 +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_value_expression(key)))
-    end
-
-    def self.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 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_value_label(value)
-      value.nil? || value == "" ? "(untagged)" : value.to_s
-    end
-
-    def self.tag_value_expression(key, table_name: quoted_table_name)
-      key = validated_tag_key(key)
-      column = "#{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
-
-    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/parsers/openai_usage.rb

@@ -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

data/lib/llm_cost_tracker/period_grouping.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "active_record_adapter"
+
 module LlmCostTracker
   module PeriodGrouping
     PERIOD_FORMATS = {
@@ -34,10 +36,9 @@ module LlmCostTracker
       column = period_column_expression(column)
       formats = PERIOD_FORMATS.fetch(period)
 
-      case connection.adapter_name
-      when /postgres/i
+      if ActiveRecordAdapter.postgresql?(connection)
         postgres_period_expression(period, column, formats)
-      when /mysql/i
+      elsif ActiveRecordAdapter.mysql?(connection)
         "DATE_FORMAT(#{column}, #{connection.quote(formats.fetch(:mysql))})"
       else
         "strftime(#{connection.quote(formats.fetch(:sqlite))}, #{column})"

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