ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/cost.rb

@@ -1,76 +1,166 @@
 # frozen_string_literal: true
 
 module RubyLLM
-  # Represents the cost of token usage for a model response.
+  # A Cost prices token usage in US dollars using pricing from the model
+  # registry. Usage entries own accounting events; Message and Chat aggregate
+  # calls, and Model#cost_for prices any token usage against a specific model.
+  #
+  #   response = chat.ask "Summarize Ruby's object model."
+  #   response.cost.total
+  #
+  #   cost = model.cost_for(response.tokens)
+  #   cost.input
+  #   cost.output
+  #
+  # The components are RubyLLM's normalized token buckets: #input, #output,
+  # #cache_read, #cache_write, and #thinking. When the registry lacks
+  # pricing for tokens that were used, the affected component and #total
+  # return +nil+ instead of a false zero.
+  #
+  # When the provider reports the exact cost of a call, #total returns the
+  # reported amount instead of a registry-price estimate, even when
+  # registry pricing is missing.
+  #
+  # Costs are computed when the object is built. Readers do not recalculate
+  # them when registry prices change.
+  # ::aggregate and ::from_h return the same class, so a single call, a
+  # whole chat, and a stored breakdown all read the same way.
   class Cost
-    COMPONENTS = %i[input output cache_read cache_write thinking].freeze
-    PER_MILLION = 1_000_000.0
+    include Support::Inspectable
+
+    COMPONENTS = %i[input output cache_read cache_write thinking].freeze # :nodoc:
+    PER_MILLION = 1_000_000.0 # :nodoc:
+
+    class << self
+      # Combines several costs into one Cost that sums each component.
+      # Ignores +nil+ entries. A component returns +nil+ when pricing was
+      # missing for one of the calls, or when no call has a cost for that
+      # component. Pass <tt>complete: false</tt> when some requests are still
+      # running or their costs are unknown; #total then remains +nil+.
+      #
+      #   cost = RubyLLM::Cost.aggregate(messages.map(&:cost))
+      #   cost.total
+      #
+      def aggregate(costs, complete: true)
+        costs = costs.compact.select(&:tokens?)
+
+        missing = COMPONENTS.select do |component|
+          costs.any? { |cost| cost.missing?(component) }
+        end
+
+        amounts = COMPONENTS.to_h do |component|
+          values = costs.filter_map { |cost| cost.public_send(component) }
+          [component, missing.include?(component) || values.empty? ? nil : values.sum]
+        end
+
+        new(amounts:, missing:, reported: costs.any?, complete:, total: aggregate_total(costs))
+      end
 
-    attr_reader :tokens, :model, :category
+      # Rebuilds a cost from a stored breakdown Hash, as produced by #to_h
+      # and persisted alongside a usage entry. Keys may be Strings or
+      # Symbols. The component readers return the recorded amounts and
+      # #total equals the recorded +:total+.
+      #
+      #   RubyLLM::Cost.from_h({ input: 0.0001, total: 0.0003 }).total
+      #
+      def from_h(hash, tokens: nil)
+        amounts = COMPONENTS.to_h { |component| [component, hash[component] || hash[component.to_s]] }
+        total_recorded = hash.key?(:total) || hash.key?('total')
+        total = hash[:total] || hash['total'] if total_recorded
+        missing = missing_recorded_components(amounts, tokens, total_recorded)
+
+        new(amounts:, missing:, reported: recorded_tokens?(amounts, tokens, total_recorded), total:)
+      end
 
-    def self.aggregate(costs)
-      costs = costs.compact.select(&:tokens?)
-      return new(amounts: {}, has_tokens: false) if costs.empty?
+      private
 
-      missing = COMPONENTS.select do |component|
-        costs.any? { |cost| cost.missing?(component) }
+      def aggregate_total(costs)
+        totals = costs.map(&:total)
+        totals.sum if totals.any? && totals.none?(&:nil?)
       end
 
-      amounts = COMPONENTS.to_h do |component|
-        [component, missing.include?(component) ? nil : aggregate_component(costs, component)]
+      def missing_recorded_components(amounts, tokens, total_recorded)
+        return [] if total_recorded
+        return COMPONENTS.reject { |component| amounts[component] } unless tokens
+
+        COMPONENTS.select do |component|
+          tokens.public_send(component).to_i.positive? && amounts[component].nil?
+        end
       end
 
-      new(amounts:, missing:, has_tokens: true)
+      def recorded_tokens?(amounts, tokens, total_recorded)
+        return true if total_recorded
+        return COMPONENTS.any? { |component| !tokens.public_send(component).nil? } if tokens
+
+        amounts.values.any? { |amount| !amount.nil? }
+      end
     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
+    def initialize(tokens: nil, model: nil, category: :text_tokens, input_details: nil, tier: :standard, # :nodoc:
+                   amounts: nil, missing: nil, reported: nil, complete: true, total: nil)
+      if amounts
+        @amounts = amounts
+        @missing = missing || []
+        @reported = reported
+      else
+        price_tokens(tokens, model, category, input_details, tier)
+        total = reported_total if total.nil?
+        @reported ||= !total.nil?
+      end
+      @complete = complete
+      @total = total
     end
-    # rubocop:enable Metrics/ParameterLists
 
+    # Returns the cost of input tokens in US dollars, or +nil+ when the
+    # token count or its pricing is unavailable.
     def input
-      amount_for(:input)
+      @amounts[:input]
     end
 
+    # Returns the cost of billable output tokens in US dollars, or +nil+
+    # when the token count or its pricing is unavailable.
     def output
-      amount_for(:output)
+      @amounts[:output]
     end
 
+    # Returns the cost of cache-read input tokens in US dollars, or +nil+
+    # when the token count or its pricing is unavailable.
     def cache_read
-      amount_for(:cache_read)
+      @amounts[:cache_read]
     end
 
+    # Returns the cost of cache-write input tokens in US dollars, or +nil+
+    # when the token count or its pricing is unavailable.
     def cache_write
-      amount_for(:cache_write)
+      @amounts[:cache_write]
     end
 
+    # Returns the cost of thinking tokens in US dollars, or +nil+ when
+    # the model does not price reasoning output separately from regular
+    # output or the token count is unavailable. When not priced
+    # separately, thinking tokens are part of #output.
     def thinking
-      amount_for(:thinking)
+      @amounts[:thinking]
     end
 
-    alias reasoning thinking
-
-    alias cached_input cache_read
-    alias cache_creation cache_write
-
+    # Returns the sum of all components in US dollars, or the exact amount
+    # the provider reported when it reported one. Returns +nil+ when there
+    # is no token usage, or when pricing is missing for tokens that were
+    # used and the provider reported no cost.
     def total
+      return nil unless @complete
       return nil unless tokens?
-      return nil if COMPONENTS.any? { |component| missing?(component) }
+      return @total unless @total.nil?
+      return nil if @missing.any?
 
-      costs = COMPONENTS.filter_map { |component| public_send(component) }
-      return nil if costs.empty?
+      amounts = @amounts.values.compact
+      return nil if amounts.empty?
 
-      costs.sum
+      amounts.sum
     end
 
+    # Returns a hash of component costs in US dollars, plus +:total+,
+    # omitting +nil+ values.
     def to_h
       {
         input: input,
@@ -82,30 +172,39 @@ module RubyLLM
       }.compact
     end
 
-    def tokens?
-      return @has_tokens unless @has_tokens.nil?
+    def tokens? # :nodoc:
+      @reported
+    end
 
-      COMPONENTS.any? { |component| !tokens_for(component).nil? }
+    def missing?(component) # :nodoc:
+      @missing.include?(component)
     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?
+    private
 
-      tokens = tokens_for(component)
-      tokens.to_i.positive? && price_for(component).nil?
+    def price_tokens(tokens, model, category, input_details, tier)
+      @tokens = tokens
+      @model = normalize_model(model)
+      @category = category.to_sym
+      @input_details = input_details
+      @tier = tier.to_sym
+      @amounts = COMPONENTS.to_h { |component| [component, amount_for(component)] }
+      @missing = COMPONENTS.select { |component| missing_component?(component) }
+      @reported = COMPONENTS.any? { |component| !tokens_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
+    def reported_total
+      @tokens.reported_cost if @tokens.respond_to?(:reported_cost)
     end
 
-    private
+    def missing_component?(component)
+      return image_input_missing? if component == :input && detailed_image_input?
+      return false if component == :thinking && !thinking_priced_separately?
+
+      tokens_for(component).to_i.positive? && price_for(component).nil?
+    end
 
     def amount_for(component)
-      return @amounts[component] if aggregate?
       return image_input_amount if component == :input && detailed_image_input?
 
       token_count = tokens_for(component)
@@ -120,56 +219,88 @@ module RubyLLM
       token_count * price / PER_MILLION
     end
 
-    def aggregate?
-      !@amounts.nil?
-    end
-
     def tokens_for(component)
-      return unless tokens
+      return unless @tokens
 
       case component
       when :input
-        tokens.input
+        @tokens.input
       when :output
-        tokens.output
+        @tokens.output
       when :cache_read
-        tokens.cache_read
+        @tokens.cache_read
       when :cache_write
-        tokens.cache_write
+        @tokens.cache_write
       when :thinking
-        tokens.thinking if thinking_priced_separately?
+        @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
+      when :input then input_price
+      when :output then output_price
+      when :cache_read then text_price(:cache_read_input_per_million, text_pricing.cache_read_input)
+      when :cache_write then text_price(:cache_write_input_per_million, text_pricing.cache_write_input)
+      when :thinking then text_price(:reasoning_output_per_million, text_pricing.reasoning_output)
       end
     end
 
+    def input_price
+      return text_input_price if @category == :text_tokens || image_cost?
+
+      category_pricing.input || text_input_price
+    end
+
+    def output_price
+      return image_pricing.output || text_output_price if image_cost?
+      return text_output_price if @category == :text_tokens
+
+      category_pricing.output || text_output_price
+    end
+
+    def text_input_price
+      text_price(:input_per_million, text_pricing.input)
+    end
+
+    def text_output_price
+      text_price(:output_per_million, text_pricing.output)
+    end
+
+    def text_price(attribute, standard_price)
+      applicable_text_tier&.public_send(attribute) || standard_price
+    end
+
+    def applicable_text_tier
+      return text_pricing.batch if @tier == :batch
+
+      text_pricing.tier_for(prompt_tokens)
+    end
+
+    def prompt_tokens
+      tokens_for(:input).to_i + tokens_for(:cache_read).to_i + tokens_for(:cache_write).to_i
+    end
+
     def text_pricing
-      model&.pricing&.text_tokens || RubyLLM::Model::PricingCategory.new
+      @model&.pricing&.text_tokens || RubyLLM::Model::PricingCategory.new
     end
 
     def image_pricing
-      model&.pricing&.images || RubyLLM::Model::PricingCategory.new
+      @model&.pricing&.images || RubyLLM::Model::PricingCategory.new
     end
 
-    def output_pricing
-      image_cost? && image_pricing.output ? image_pricing : text_pricing
+    def category_pricing
+      return text_pricing if @category == :text_tokens
+      return image_pricing if image_cost?
+
+      pricing = @model&.pricing
+      return RubyLLM::Model::PricingCategory.new unless pricing.respond_to?(@category)
+
+      pricing.public_send(@category)
     end
 
     def image_cost?
-      %i[image images].include?(category)
+      %i[image images].include?(@category)
     end
 
     def detailed_image_input?
@@ -193,9 +324,10 @@ module RubyLLM
     end
 
     def image_input_parts
+      text_input_price = applicable_text_tier&.input_per_million || text_pricing.input
       [
-        [:text, input_detail('text_tokens'), text_pricing.input],
-        [:image, input_detail('image_tokens'), image_pricing.input || text_pricing.input]
+        [:text, input_detail('text_tokens'), text_input_price],
+        [:image, input_detail('image_tokens'), image_pricing.input || text_input_price]
       ]
     end
 
@@ -204,10 +336,11 @@ module RubyLLM
     end
 
     def thinking_priced_separately?
-      reasoning_price = text_pricing.reasoning_output
+      tier = applicable_text_tier
+      reasoning_price = tier&.reasoning_output_per_million || text_pricing.reasoning_output
       return false unless reasoning_price
 
-      output_price = text_pricing.output
+      output_price = tier&.output_per_million || text_pricing.output
       output_price.nil? || reasoning_price != output_price
     end
 
@@ -220,5 +353,9 @@ module RubyLLM
     rescue ModelNotFoundError
       nil
     end
+
+    def inspect_attributes # :nodoc:
+      to_h
+    end
   end
 end

data/lib/ruby_llm/downloaded_file.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # A DownloadedFile contains a provider file's bytes. Save it with #save
+  # or read the bytes with #to_blob. It is also a String, so parsers and
+  # existing string operations work directly on the result.
+  #
+  #   RubyLLM.download(file.id, provider: :openai).save("report.pdf")
+  #
+  class DownloadedFile < String
+    include Support::Inspectable
+
+    # Returns the raw file bytes as a String.
+    def to_blob
+      to_s
+    end
+
+    # Writes the file bytes to +path+, expanding it first. Returns +path+.
+    #
+    #   file.save("report.pdf")
+    #
+    def save(path)
+      File.binwrite(File.expand_path(path), to_blob)
+      path
+    end
+
+    private
+
+    def inspect_attributes # :nodoc:
+      { byte_size: bytesize }
+    end
+  end
+end

data/lib/ruby_llm/embedding.rb

@@ -1,59 +1,163 @@
 # frozen_string_literal: true
 
 module RubyLLM
-  # Core embedding interface.
+  # An Embedding is the result of turning text into numerical vectors.
+  # RubyLLM.embed returns one:
+  #
+  #   embedding = RubyLLM.embed("Ruby is a programmer's best friend")
+  #   embedding.vectors # => [0.018, -0.027, ...]
+  #
+  # Pass an array to embed several texts in one call on supported models:
+  #
+  #   RubyLLM.embed(["Ruby", "Rails"]).vectors # => [[...], [...]]
+  #
+  # Store the vectors for similarity search. #tokens and #cost report
+  # usage when the provider supplies it.
+  #
   class Embedding
-    attr_reader :vectors, :model, :input_tokens
+    include Support::Inspectable
+    include Accounting::Usage::Result
 
-    def initialize(vectors:, model:, input_tokens: 0)
+    # The embedding vectors. A flat array of floats when a single text
+    # was embedded, an array of such arrays when an array of texts was
+    # embedded.
+    attr_reader :vectors
+
+    # The sparse vectors, for models that return one beside the dense
+    # vector, as a Hash mapping token id to weight. Shaped like #vectors:
+    # one Hash for a single text, an array of them for an array of texts.
+    # +nil+ on models that return only dense vectors.
+    attr_reader :sparse_vectors
+
+    # The id of the model that produced the vectors, as a String.
+    attr_reader :model
+
+    def initialize(vectors:, model:, sparse_vectors: nil, input_tokens: nil, reported_cost: nil) # :nodoc:
       @vectors = vectors
+      @sparse_vectors = sparse_vectors
       @model = model
       @input_tokens = input_tokens
+      @reported_cost = reported_cost
+    end
+
+    # Returns usage aggregated across every provider attempt.
+    def tokens
+      return ruby_llm_usage_tokens unless ruby_llm_usage_entries.empty?
+
+      Tokens.new(input: @input_tokens, reported_cost: @reported_cost)
+    end
+
+    # Returns the embedding cost across every provider attempt.
+    def cost
+      return ruby_llm_usage_cost unless ruby_llm_usage_entries.empty?
+
+      Cost.new(tokens:, model: model_info, category: :embeddings)
+    end
+
+    def model_info # :nodoc:
+      @model_info ||= RubyLLM.models.find(model)
+    rescue ModelNotFoundError
+      nil
     end
 
-    def self.embed(text, # rubocop:disable Metrics/ParameterLists
+    # Generates embeddings for +text+ and returns an Embedding. +text+
+    # may be a single string or, on supported models, an array of strings.
+    # An array produces one vector per string in a single API call.
+    #
+    #   RubyLLM.embed "Ruby is a programmer's best friend"
+    #   RubyLLM.embed ["Ruby", "Python", "JavaScript"]
+    #   RubyLLM.embed "This is a test sentence",
+    #                 model: "text-embedding-3-large",
+    #                 dimensions: 512
+    #   RubyLLM.embed "RubyLLM makes provider APIs feel native to Ruby.",
+    #                 model: "gemini-embedding-001",
+    #                 provider: :vertexai,
+    #                 task_type: "RETRIEVAL_DOCUMENT",
+    #                 title: "RubyLLM docs"
+    #   RubyLLM.embed "The Ruby logo",
+    #                 model: "gemini-embedding-2",
+    #                 with: "logo.png"
+    #
+    # +model:+ selects the embedding model and defaults to the
+    # configured +default_embedding_model+. +provider:+ forces a specific
+    # provider, and <tt>assume_model_exists: true</tt> skips the model
+    # registry check. +context:+ supplies a Context whose configuration
+    # is used instead of the global one. +dimensions:+ requests a
+    # specific vector size on models that support it. +task_type:+ names
+    # the embedding task in the provider's own vocabulary: Vertex AI and
+    # Gemini take values such as <tt>"RETRIEVAL_QUERY"</tt> or
+    # <tt>"RETRIEVAL_DOCUMENT"</tt>, while Bedrock Cohere takes an input
+    # type such as <tt>"search_document"</tt>. +title:+ labels the
+    # document on Vertex AI and Gemini retrieval tasks. Providers that
+    # have no task concept ignore both. +with:+ passes one or more media
+    # attachments (images, audio, video, PDFs) to embed alongside the
+    # text on multimodal embedding models such as Gemini's
+    # gemini-embedding-2; providers without multimodal embeddings raise
+    # UnsupportedAttachmentError. +provider_options:+ takes options
+    # in the provider's request vocabulary and merges them into the
+    # request as-is. +metadata:+ is not sent to the provider; it is
+    # attached to the emitted +embedding.ruby_llm+ instrumentation event.
+    def self.embed(text,
                    model: nil,
                    provider: nil,
                    assume_model_exists: false,
                    context: nil,
-                   dimensions: nil)
+                   dimensions: nil,
+                   task_type: nil,
+                   title: nil,
+                   with: nil,
+                   provider_options: {},
+                   metadata: nil)
       config = context&.config || RubyLLM.config
       model ||= config.default_embedding_model
-      model, provider_instance = Models.resolve(model, provider: provider, assume_exists: assume_model_exists,
+      model, provider_instance = Models.resolve(model, provider: provider, assume_model_exists: assume_model_exists,
                                                        config: config)
       model_id = model.id
+      empty_tokens = Tokens.new
 
       payload = {
         provider: provider_instance.slug,
-        provider_class: provider_instance.class.name,
+        provider_class: provider_instance.class.display_name,
         model: model_id,
         model_info: model,
         input: text,
-        dimensions: dimensions
+        dimensions: dimensions,
+        task_type: task_type,
+        title: title,
+        attachment_count: Attachment.wrap(with).size,
+        provider_options: provider_options,
+        metadata: metadata,
+        tokens: empty_tokens,
+        cost: Cost.new(tokens: empty_tokens, model:, category: :embeddings)
       }
 
       RubyLLM.instrument('embedding.ruby_llm', payload, config: config) do |event|
-        result = provider_instance.embed(text, model: model_id, dimensions:)
+        result = provider_instance.embed(text, model:, dimensions:, task_type:, title:, with:, provider_options:)
         event[:result] = result
         event[:response_model] = result.model
-        event[:input_tokens] = result.input_tokens
+        event[:tokens] = result.tokens
+        event[:cost] = result.cost
         event[:embedding_dimensions] = vector_dimensions(result.vectors)
         event[:embedding_count] = embedding_count(result.vectors)
         result
       end
     end
 
-    def self.vector_dimensions(vectors)
-      return unless vectors.is_a?(Array)
-
+    private_class_method def self.vector_dimensions(vectors) # :nodoc:
       vector = vectors.first.is_a?(Array) ? vectors.first : vectors
-      vector.length if vector.respond_to?(:length)
+      vector.length
     end
 
-    def self.embedding_count(vectors)
-      return unless vectors.is_a?(Array)
-
+    private_class_method def self.embedding_count(vectors) # :nodoc:
       vectors.first.is_a?(Array) ? vectors.size : 1
     end
+
+    def inspect_attributes # :nodoc:
+      if vectors&.first.is_a?(Array)
+        { model: model, count: vectors.length, dimensions: vectors.first.length }
+      else
+        { model: model, dimensions: vectors&.length }
+      end
+    end
   end
 end

data/lib/ruby_llm/embedding_request.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # An EmbeddingRequest is an embedding awaiting a provider-side batch.
+  # RubyLLM.embed_later stages one; submit an array of them with
+  # RubyLLM.batch, and once the batch completes, collecting the results
+  # fills in #result:
+  #
+  #   requests = texts.map { |text| RubyLLM.embed_later(text) }
+  #   batch = RubyLLM.batch(requests)
+  #   # later, once batch.refresh.complete?
+  #   batch.results
+  #   requests.first.result.vectors
+  #
+  class EmbeddingRequest
+    include Support::Inspectable
+
+    # The text staged for embedding.
+    attr_reader :text
+
+    # The Model of the embedding model the request targets.
+    attr_reader :model
+
+    # The requested vector size, or +nil+ for the model's default.
+    attr_reader :dimensions
+
+    # The Provider the request will be submitted through.
+    attr_reader :provider
+
+    # The Embedding, once the batch has completed and its results have
+    # been collected. +nil+ until then, and for requests that failed.
+    # Batch result collection writes it.
+    attr_accessor :result
+
+    def initialize(text, model: nil, provider: nil, dimensions: nil, context: nil) # :nodoc:
+      config = context&.config || RubyLLM.config
+      model ||= config.default_embedding_model
+      @model, @provider = Models.resolve(model, provider:, config:)
+      @text = text
+      @dimensions = dimensions
+    end
+
+    # Returns the request payload this embedding request would send to the
+    # provider, in the provider's wire format.
+    def render
+      @provider.render_embedding(text, model: @model, dimensions: @dimensions)
+    end
+
+    def inspect_attributes # :nodoc:
+      { text: text, model: model.id, dimensions: dimensions, result: result }
+    end
+  end
+end