ruby_llm 1.14.1 → 1.15.0

data/lib/ruby_llm/chat.rb

@@ -30,6 +30,7 @@ module RubyLLM
         tool_call: nil,
         tool_result: nil
       }
+      @callbacks = Hash.new { |callbacks, name| callbacks[name] = [] }
     end
 
     def ask(message = nil, with: nil, &)
@@ -112,31 +113,47 @@ module RubyLLM
       self
     end
 
-    def on_new_message(&block)
-      @on[:new_message] = block
-      self
+    def on_new_message(&)
+      set_legacy_callback(:new_message, :on_new_message, :before_message, &)
     end
 
-    def on_end_message(&block)
-      @on[:end_message] = block
-      self
+    def on_end_message(&)
+      set_legacy_callback(:end_message, :on_end_message, :after_message, &)
     end
 
-    def on_tool_call(&block)
-      @on[:tool_call] = block
-      self
+    def on_tool_call(&)
+      set_legacy_callback(:tool_call, :on_tool_call, :before_tool_call, &)
     end
 
-    def on_tool_result(&block)
-      @on[:tool_result] = block
-      self
+    def on_tool_result(&)
+      set_legacy_callback(:tool_result, :on_tool_result, :after_tool_result, &)
+    end
+
+    def before_message(&)
+      add_callback(:before_message, &)
+    end
+
+    def after_message(&)
+      add_callback(:after_message, &)
+    end
+
+    def before_tool_call(&)
+      add_callback(:before_tool_call, &)
+    end
+
+    def after_tool_result(&)
+      add_callback(:after_tool_result, &)
     end
 
     def each(&)
       messages.each(&)
     end
 
-    def complete(&) # rubocop:disable Metrics/PerceivedComplexity
+    def cost
+      Cost.aggregate(messages.map(&:cost))
+    end
+
+    def complete(&)
       response = @provider.complete(
         messages,
         tools: @tools,
@@ -150,7 +167,7 @@ module RubyLLM
         &wrap_streaming_block(&)
       )
 
-      @on[:new_message]&.call unless block_given?
+      run_callbacks(:before_message, :new_message) unless block_given?
 
       if @schema && response.content.is_a?(String) && !response.tool_call?
         begin
@@ -161,7 +178,7 @@ module RubyLLM
       end
 
       add_message response
-      @on[:end_message]&.call(response)
+      run_callbacks(:after_message, :end_message, response)
 
       if response.tool_call?
         handle_tool_calls(response, &)
@@ -221,28 +238,52 @@ module RubyLLM
       sanitized.empty? ? 'response' : sanitized
     end
 
+    def add_callback(name, &block)
+      @callbacks[name] << block if block
+      self
+    end
+
+    def set_legacy_callback(name, legacy_name, additive_name, &block)
+      warn_legacy_callback_deprecation(legacy_name, additive_name) if block
+
+      @on[name] = block
+      self
+    end
+
+    def warn_legacy_callback_deprecation(legacy_name, additive_name)
+      RubyLLM.logger.warn(
+        "`#{legacy_name}` is deprecated and will be removed in RubyLLM 2.0. " \
+        "Use `#{additive_name}` instead."
+      )
+    end
+
+    def run_callbacks(name, legacy_name, *args)
+      @callbacks[name].each { |callback| callback.call(*args) }
+      @on[legacy_name]&.call(*args)
+    end
+
     def wrap_streaming_block(&block)
       return nil unless block_given?
 
-      @on[:new_message]&.call
+      run_callbacks(:before_message, :new_message)
 
       proc do |chunk|
         block.call chunk
       end
     end
 
-    def handle_tool_calls(response, &) # rubocop:disable Metrics/PerceivedComplexity
+    def handle_tool_calls(response, &)
       halt_result = nil
 
       response.tool_calls.each_value do |tool_call|
-        @on[:new_message]&.call
-        @on[:tool_call]&.call(tool_call)
+        run_callbacks(:before_message, :new_message)
+        run_callbacks(:before_tool_call, :tool_call, tool_call)
         result = execute_tool tool_call
-        @on[:tool_result]&.call(result)
+        run_callbacks(:after_tool_result, :tool_result, result)
         tool_payload = result.is_a?(Tool::Halt) ? result.content : result
         content = content_like?(tool_payload) ? tool_payload : tool_payload.to_s
         message = add_message role: :tool, content:, tool_call_id: tool_call.id
-        @on[:end_message]&.call(message)
+        run_callbacks(:after_message, :end_message, message)
 
         halt_result = result if result.is_a?(Tool::Halt)
       end

data/lib/ruby_llm/cost.rb

@@ -0,0 +1,224 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # Represents the cost of token usage for a model response.
+  class Cost
+    COMPONENTS = %i[input output cache_read cache_write thinking].freeze
+    PER_MILLION = 1_000_000.0
+
+    attr_reader :tokens, :model, :category
+
+    def self.aggregate(costs)
+      costs = costs.compact.select(&:tokens?)
+      return new(amounts: {}, has_tokens: false) if costs.empty?
+
+      missing = COMPONENTS.select do |component|
+        costs.any? { |cost| cost.missing?(component) }
+      end
+
+      amounts = COMPONENTS.to_h do |component|
+        [component, missing.include?(component) ? nil : aggregate_component(costs, component)]
+      end
+
+      new(amounts:, missing:, has_tokens: true)
+    end
+
+    # rubocop:disable Metrics/ParameterLists
+    def initialize(tokens: nil, model: nil, amounts: nil, missing: [], has_tokens: nil, category: :text_tokens,
+                   input_details: nil)
+      @tokens = tokens
+      @model = normalize_model(model)
+      @amounts = amounts
+      @missing = missing
+      @has_tokens = has_tokens
+      @category = category.to_sym
+      @input_details = input_details
+    end
+    # rubocop:enable Metrics/ParameterLists
+
+    def input
+      amount_for(:input)
+    end
+
+    def output
+      amount_for(:output)
+    end
+
+    def cache_read
+      amount_for(:cache_read)
+    end
+
+    def cache_write
+      amount_for(:cache_write)
+    end
+
+    def thinking
+      amount_for(:thinking)
+    end
+
+    alias reasoning thinking
+
+    alias cached_input cache_read
+    alias cache_creation cache_write
+
+    def total
+      return nil unless tokens?
+      return nil if COMPONENTS.any? { |component| missing?(component) }
+
+      costs = COMPONENTS.filter_map { |component| public_send(component) }
+      return nil if costs.empty?
+
+      costs.sum
+    end
+
+    def to_h
+      {
+        input: input,
+        output: output,
+        cache_read: cache_read,
+        cache_write: cache_write,
+        thinking: thinking,
+        total: total
+      }.compact
+    end
+
+    def tokens?
+      return @has_tokens unless @has_tokens.nil?
+
+      COMPONENTS.any? { |component| !tokens_for(component).nil? }
+    end
+
+    def missing?(component)
+      return @missing.include?(component) if aggregate?
+      return image_input_missing? if component == :input && detailed_image_input?
+      return false if component == :thinking && !thinking_priced_separately?
+
+      tokens = tokens_for(component)
+      tokens.to_i.positive? && price_for(component).nil?
+    end
+
+    private_class_method def self.aggregate_component(costs, component)
+      values = costs.filter_map { |cost| cost.public_send(component) }
+      values.empty? ? nil : values.sum
+    end
+
+    private
+
+    def amount_for(component)
+      return @amounts[component] if aggregate?
+      return image_input_amount if component == :input && detailed_image_input?
+
+      token_count = tokens_for(component)
+      return nil if token_count.nil?
+
+      token_count = token_count.to_i
+      return 0.0 if token_count.zero?
+
+      price = price_for(component)
+      return nil unless price
+
+      token_count * price / PER_MILLION
+    end
+
+    def aggregate?
+      !@amounts.nil?
+    end
+
+    def tokens_for(component)
+      return unless tokens
+
+      case component
+      when :input
+        tokens.input
+      when :output
+        tokens.output
+      when :cache_read
+        tokens.cache_read
+      when :cache_write
+        tokens.cache_write
+      when :thinking
+        tokens.thinking if thinking_priced_separately?
+      end
+    end
+
+    def price_for(component)
+      case component
+      when :input
+        text_pricing.input
+      when :output
+        output_pricing.output
+      when :cache_read
+        text_pricing.cache_read_input
+      when :cache_write
+        text_pricing.cache_write_input
+      when :thinking
+        text_pricing.reasoning_output
+      end
+    end
+
+    def text_pricing
+      model&.pricing&.text_tokens || RubyLLM::Model::PricingCategory.new
+    end
+
+    def image_pricing
+      model&.pricing&.images || RubyLLM::Model::PricingCategory.new
+    end
+
+    def output_pricing
+      image_cost? && image_pricing.output ? image_pricing : text_pricing
+    end
+
+    def image_cost?
+      %i[image images].include?(category)
+    end
+
+    def detailed_image_input?
+      image_cost? && @input_details.is_a?(Hash) && image_input_parts.any? { |_, tokens, _| !tokens.nil? }
+    end
+
+    def image_input_amount
+      return nil if image_input_missing?
+
+      image_input_parts.filter_map do |_, token_count, price|
+        next if token_count.nil? || token_count.to_i.zero?
+
+        token_count.to_i * price / PER_MILLION
+      end.sum
+    end
+
+    def image_input_missing?
+      image_input_parts.any? do |_, token_count, price|
+        token_count.to_i.positive? && price.nil?
+      end
+    end
+
+    def image_input_parts
+      [
+        [:text, input_detail('text_tokens'), text_pricing.input],
+        [:image, input_detail('image_tokens'), image_pricing.input || text_pricing.input]
+      ]
+    end
+
+    def input_detail(key)
+      @input_details[key] || @input_details[key.to_sym]
+    end
+
+    def thinking_priced_separately?
+      reasoning_price = text_pricing.reasoning_output
+      return false unless reasoning_price
+
+      output_price = text_pricing.output
+      output_price.nil? || reasoning_price != output_price
+    end
+
+    def normalize_model(model)
+      return RubyLLM.models.find(model.to_s) if model.is_a?(String) || model.is_a?(Symbol)
+      return model.to_llm if model.respond_to?(:to_llm)
+      return model if model.respond_to?(:pricing)
+
+      nil
+    rescue ModelNotFoundError
+      nil
+    end
+  end
+end

data/lib/ruby_llm/image.rb

@@ -3,14 +3,15 @@
 module RubyLLM
   # Represents a generated image from an AI model.
   class Image
-    attr_reader :url, :data, :mime_type, :revised_prompt, :model_id
+    attr_reader :url, :data, :mime_type, :revised_prompt, :model_id, :usage
 
-    def initialize(url: nil, data: nil, mime_type: nil, revised_prompt: nil, model_id: nil)
+    def initialize(url: nil, data: nil, mime_type: nil, revised_prompt: nil, model_id: nil, usage: {}) # rubocop:disable Metrics/ParameterLists
       @url = url
       @data = data
       @mime_type = mime_type
       @revised_prompt = revised_prompt
       @model_id = model_id
+      @usage = usage
     end
 
     def base64?
@@ -36,14 +37,46 @@ module RubyLLM
                    provider: nil,
                    assume_model_exists: false,
                    size: '1024x1024',
-                   context: nil)
+                   context: nil,
+                   with: nil,
+                   mask: nil,
+                   params: {})
       config = context&.config || RubyLLM.config
       model ||= config.default_image_model
       model, provider_instance = Models.resolve(model, provider: provider, assume_exists: assume_model_exists,
                                                        config: config)
       model_id = model.id
 
-      provider_instance.paint(prompt, model: model_id, size:)
+      provider_instance.paint(prompt, model: model_id, size:, with:, mask:, params:)
+    end
+
+    def tokens
+      @tokens ||= Tokens.build(
+        input: usage_value('input_tokens'),
+        output: usage_value('output_tokens')
+      )
+    end
+
+    def cost
+      Cost.new(tokens:, model: model_info, category: :images, input_details: input_tokens_details)
+    end
+
+    def model_info
+      return unless model_id
+
+      @model_info ||= RubyLLM.models.find(model_id)
+    rescue ModelNotFoundError
+      nil
+    end
+
+    private
+
+    def input_tokens_details
+      usage_value('input_tokens_details')
+    end
+
+    def usage_value(key)
+      usage[key] || usage[key.to_sym]
     end
   end
 end

data/lib/ruby_llm/message.rb

@@ -64,6 +64,14 @@ module RubyLLM
       tokens&.cache_creation
     end
 
+    def cache_read_tokens
+      tokens&.cache_read
+    end
+
+    def cache_write_tokens
+      tokens&.cache_write
+    end
+
     def thinking_tokens
       tokens&.thinking
     end
@@ -72,6 +80,10 @@ module RubyLLM
       tokens&.thinking
     end
 
+    def cost(model: nil)
+      Cost.new(tokens:, model: model || model_info)
+    end
+
     def to_h
       {
         role: role,
@@ -88,6 +100,14 @@ module RubyLLM
       super - [:@raw]
     end
 
+    def model_info
+      return unless model_id
+
+      @model_info ||= RubyLLM.models.find(model_id)
+    rescue ModelNotFoundError
+      nil
+    end
+
     private
 
     def normalize_content(content, role:, tool_calls:)

data/lib/ruby_llm/model/info.rb

@@ -77,6 +77,23 @@ module RubyLLM
         pricing.text_tokens.output
       end
 
+      def cache_read_input_price_per_million
+        pricing.text_tokens.cache_read_input
+      end
+
+      def cache_write_input_price_per_million
+        pricing.text_tokens.cache_write_input
+      end
+
+      alias cached_input_price_per_million cache_read_input_price_per_million
+      alias cache_creation_input_price_per_million cache_write_input_price_per_million
+
+      def cost_for(tokens)
+        tokens = tokens.tokens if tokens.respond_to?(:tokens)
+
+        Cost.new(tokens:, model: self)
+      end
+
       def provider_class
         RubyLLM::Provider.resolve provider
       end

data/lib/ruby_llm/model/pricing_category.rb

@@ -19,10 +19,21 @@ module RubyLLM
         standard&.output_per_million
       end
 
-      def cached_input
-        standard&.cached_input_per_million
+      def cache_read_input
+        standard&.cache_read_input_per_million || standard&.cached_input_per_million
       end
 
+      def cache_write_input
+        standard&.cache_write_input_per_million || standard&.cache_creation_input_per_million
+      end
+
+      def reasoning_output
+        standard&.reasoning_output_per_million
+      end
+
+      alias cached_input cache_read_input
+      alias cache_creation_input cache_write_input
+
       def [](key)
         key == :batch ? batch : standard
       end