ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/batch.rb

@@ -0,0 +1,432 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # A Batch is a provider-side batch of requests: chats awaiting a
+  # response (or texts awaiting embeddings) go in together, answers come
+  # back at batch prices, typically within hours. Persist the id, pick the
+  # batch back up from any process, and collect the results once
+  # processing ends.
+  #
+  #   chats = documents.map do |doc|
+  #     RubyLLM.chat(model: "claude-haiku-4-5").ask_later(doc.text)
+  #   end
+  #   batch = RubyLLM.batch(chats)
+  #   batch.id                # => "msgbatch_01EhcDuvb5XfWqcdJArbsfNX"
+  #   batch.refresh.complete? # => false, check back later
+  #   batch.messages          # the responses, in submission order
+  #
+  class Batch
+    include Support::Inspectable
+
+    AWAITING_ROLES = %i[user tool].freeze # :nodoc:
+
+    # The provider's batch id. Persist it to load the batch again later
+    # from any process with ::find.
+    attr_reader :id
+
+    # The provider-neutral lifecycle status: +:pending+, +:succeeded+,
+    # +:failed+, or +:cancelled+. Refreshed by #refresh.
+    attr_reader :status
+
+    # The provider-reported status string, such as "in_progress".
+    # Refreshed by #refresh.
+    attr_reader :raw_status
+
+    # The provider-reported request tallies by state, or +nil+ when the
+    # provider does not report them.
+    attr_reader :request_counts
+
+    # The submitted Chat objects in order, or +nil+ when the batch was
+    # loaded by id via ::find or holds embedding requests.
+    attr_reader :chats
+
+    # The submitted EmbeddingRequest objects in order, or +nil+ when the
+    # batch was loaded by id via ::find or holds chats.
+    attr_reader :requests
+
+    # The normalized outcome of each collected request, in submission order.
+    # Values are +:succeeded+, +:failed+, or +:cancelled+.
+    attr_reader :statuses
+
+    class << self
+      # Submits chats or embedding requests to their shared provider as a
+      # batch and returns a new Batch. Accepts a single Chat or an array.
+      # Every chat must be awaiting the model (see Chat#ask_later), and all
+      # requests must use the same provider.
+      #
+      #   chats = tickets.map do |ticket|
+      #     RubyLLM.chat(model: "claude-haiku-4-5").ask_later(ticket.body)
+      #   end
+      #   batch = RubyLLM::Batch.submit(chats)
+      #   batch.status     # => :pending
+      #   batch.raw_status # => "in_progress"
+      #
+      # Raises ArgumentError if the batch is empty, mixes providers, mixes
+      # chats with embedding requests, or includes a chat that is not
+      # awaiting the model.
+      def submit(chats)
+        records = wrap_records(chats)
+        return submit_embeddings(records) if records.any?(EmbeddingRequest)
+
+        submit_chats(records)
+      end
+
+      # Returns a Batch reflecting the provider's current state for +id+.
+      # Use it to pick a batch back up from any process.
+      #
+      #   batch = RubyLLM::Batch.find("msgbatch_01EhcDuvb5XfWqcdJArbsfNX",
+      #                               provider: :anthropic)
+      #   batch.complete? # => true
+      #
+      # Pass +context:+ to use a Context in place of the global
+      # configuration. Raises ArgumentError if +provider+ is not given.
+      def find(id, provider: nil, context: nil)
+        config = context&.config || RubyLLM.config
+        persisted = config.batch_store&.fetch(id, provider:, context:)
+        return persisted if persisted
+
+        unless provider
+          raise ArgumentError, 'Provider must be specified to find a batch that is not persisted by RubyLLM'
+        end
+
+        provider = Provider.resolve!(provider).new(config)
+        raise Error, "#{provider.slug} doesn't support batch requests" unless provider.batches?
+
+        new(provider:, store: config.batch_store, **provider.find_batch(id))
+      end
+
+      private
+
+      def submit_chats(records)
+        chats = normalize_chats(records)
+
+        provider = shared_provider(chats)
+        payload = { provider: provider.slug, provider_class: provider.class.display_name, requests: chats.size }
+        RubyLLM.instrument('batch.ruby_llm', payload, config: provider.config) do |event|
+          requests = chats.each_with_index.map do |chat, index|
+            { custom_id: index.to_s, model: chat.model.id, payload: chat.render }
+          end
+          store = provider.config.batch_store
+          batch = new(provider:, chats:, store:, **provider.create_batch(requests))
+          store&.persist(batch, records)
+          event[:batch_id] = batch.id
+          batch
+        end
+      end
+
+      def submit_embeddings(requests)
+        unless requests.all?(EmbeddingRequest)
+          raise ArgumentError, 'A batch takes chats or embedding requests, not both'
+        end
+
+        provider = shared_provider(requests)
+        payload = { provider: provider.slug, provider_class: provider.class.display_name, requests: requests.size }
+        RubyLLM.instrument('batch.ruby_llm', payload, config: provider.config) do |event|
+          lines = requests.each_with_index.map do |request, index|
+            { custom_id: index.to_s, model: request.model.id, payload: request.render, text: request.text }
+          end
+          batch = new(provider:, requests:, store: provider.config.batch_store, **provider.create_batch(lines))
+          event[:batch_id] = batch.id
+          batch
+        end
+      end
+
+      def wrap_records(records)
+        return [records] if records.respond_to?(:to_llm)
+
+        case records
+        when Chat, EmbeddingRequest then [records]
+        else Array(records)
+        end
+      end
+
+      def normalize_chats(records)
+        normalized = records.map { |chat| chat.respond_to?(:to_llm) ? chat.to_llm : chat }
+        raise ArgumentError, 'Cannot submit an empty batch' if normalized.empty?
+
+        unless normalized.all? { |chat| awaiting_model?(chat) }
+          raise ArgumentError,
+                'Every chat in a batch must be awaiting the model; stage one with ask_later, or run_tools first'
+        end
+
+        normalized
+      end
+
+      def awaiting_model?(chat)
+        !chat.complete? && AWAITING_ROLES.include?(chat.messages.last&.role)
+      end
+
+      def shared_provider(chats)
+        slugs = chats.map { |chat| chat.provider.slug }.uniq
+        raise ArgumentError, "A batch takes one provider per submission, got: #{slugs.join(', ')}" if slugs.size > 1
+
+        provider = chats.first.provider
+        raise Error, "#{provider.slug} doesn't support batch requests" unless provider.batches?
+
+        provider
+      end
+    end
+
+    def initialize(provider:, chats: nil, requests: nil, batch_protocol: nil, store: nil, **attributes) # :nodoc:
+      @provider = provider
+      @chats = chats
+      @requests = requests
+      @batch_protocol = protocol_name(batch_protocol)
+      @store = store
+      @delivered = {}
+      @statuses = []
+      apply(attributes)
+    end
+
+    # The slug of the provider running the batch, as a String.
+    def provider
+      @provider.slug
+    end
+
+    attr_reader :batch_protocol, :reported_cost # :nodoc:
+
+    # Returns whether the batch has finished processing, as of the last
+    # state fetched from the provider. Never contacts the provider; poll
+    # with #refresh.
+    #
+    #   sleep 60 until batch.refresh.complete?
+    #
+    def complete?
+      @completed
+    end
+
+    # Returns whether the provider completed the batch successfully.
+    def succeeded?
+      status == :succeeded
+    end
+
+    # Returns whether the provider failed or expired the batch.
+    def failed?
+      status == :failed
+    end
+
+    # Returns whether the provider cancelled the batch.
+    def cancelled?
+      status == :cancelled
+    end
+
+    # Re-fetches the batch from the provider, updating #status, #raw_status,
+    # #request_counts, and #complete?. Returns +self+.
+    def refresh
+      apply(@provider.find_batch(id))
+      persist_state
+      self
+    end
+
+    # Asks the provider to cancel the batch and applies the new state.
+    # Requests already processed still return results. Returns +self+.
+    def cancel
+      apply(@provider.cancel_batch(id))
+      persist_state
+      self
+    end
+
+    # Returns the answers in submission order, +nil+ where a request
+    # failed. In a chat batch the answers are Messages, each also appended
+    # to its chat; in an embeddings batch they are Embeddings, each also
+    # hydrated into its request's EmbeddingRequest#result. Fetches results
+    # from the provider; cached once #complete? is true, so collecting
+    # early keeps reading fresh.
+    #
+    #   batch.messages.each do |message|
+    #     puts message.content
+    #   end
+    #
+    def messages
+      return @messages if @messages
+
+      collected = collect_results
+      @messages = collected if @completed
+      collected
+    end
+
+    alias results messages
+
+    # Returns token usage aggregated across the batch's collected responses.
+    def tokens
+      Tokens.aggregate(messages.compact.map(&:tokens))
+    end
+
+    # Returns a Cost for the batch. Uses the provider's reported total when
+    # available, otherwise aggregates collected response costs at batch rates.
+    # The total is +nil+ until the batch ends or when pricing is unknown.
+    def cost
+      return Cost.aggregate([reported_cost], complete: complete?) if reported_cost
+      return Cost.aggregate([], complete: false) unless complete?
+
+      Cost.aggregate(messages.compact.map(&:cost))
+    end
+
+    private
+
+    def persist_state
+      @store&.sync(self)
+    end
+
+    def apply(attributes)
+      @batch_protocol = protocol_name(attributes[:batch_protocol]) if attributes[:batch_protocol]
+      @id = attributes.fetch(:id)
+      @raw_status = attributes.fetch(:raw_status)
+      @completed = attributes.fetch(:completed)
+      @request_counts = attributes[:request_counts]
+      @request_count = attributes[:request_count]
+      @reported_cost = attributes[:reported_cost] if attributes[:reported_cost]
+      @status = @provider.batch_status(@raw_status, completed: @completed, batch_protocol: @batch_protocol)
+    end
+
+    def protocol_name(protocol)
+      return if protocol.nil?
+
+      protocol.is_a?(Module) ? @provider.batch_protocol_name(protocol) : protocol.to_s
+    end
+
+    def collect_results
+      results = @provider.batch_results(id, batch_protocol: @batch_protocol)
+      slots = Array.new(result_slot_count(results))
+
+      results.each do |index, result, failure_status|
+        slots[index] = result
+        deliver(index, result, failure_status)
+      end
+
+      fill_missing_statuses(slots.size) if complete?
+
+      slots
+    end
+
+    def result_slot_count(results)
+      chats&.size || requests&.size || @request_count || ((results.map(&:first).max || -1) + 1)
+    end
+
+    def fill_missing_statuses(size)
+      missing_status = cancelled? ? :cancelled : :failed
+      size.times { |index| statuses[index] ||= missing_status }
+    end
+
+    # Collecting early keeps reading fresh, so a result already delivered
+    # comes back on every later poll: hand each one over once.
+    def deliver(index, result, failure_status)
+      statuses[index] = result ? :succeeded : failure_status
+      return unless result
+
+      if result.is_a?(Embedding)
+        deliver_embedding(index, result)
+      else
+        deliver_message(index, result)
+      end
+    end
+
+    def deliver_embedding(index, embedding)
+      request = requests&.[](index)
+      delivered = @delivered[index]
+      attach_batch_usage(
+        embedding,
+        operation: :embedding,
+        model: request&.model,
+        category: :embeddings,
+        instrument: !request.nil? && !delivered
+      )
+      return if delivered
+
+      @delivered[index] = true
+      request&.result = embedding
+    end
+
+    def deliver_message(index, message)
+      chat = chats&.[](index)
+      delivered = @delivered[index] || (chat && already_in_chat?(chat, message))
+      attach_batch_usage(
+        message,
+        operation: :chat,
+        model: chat&.model,
+        category: :text_tokens,
+        instrument: !chat.nil? && !delivered
+      )
+      return if delivered
+
+      @delivered[index] = true
+      chat&.add_completion(message, record_usage: true)
+    end
+
+    def attach_batch_usage(result, operation:, model:, category:, instrument:)
+      return unless result.ruby_llm_usage_entries.empty?
+
+      model ||= RubyLLM.models.find(result.model, provider: @provider.slug, config: @provider.config)
+      entry = Accounting::Usage::Entry.new(
+        operation:,
+        provider: @provider.slug,
+        model: result.model || model.id,
+        status: :succeeded,
+        tokens: result.tokens,
+        cost: @provider.batch_cost(result.tokens, model:, category:),
+        message: result.is_a?(Message) ? result : nil
+      )
+      result.ruby_llm_usage_entries = [entry]
+      Accounting::Usage.instrument(entry, config: @provider.config) if instrument
+    end
+
+    # A plain answer is the chat's last message once it arrives. A tool-call
+    # answer is not: running its tools adds messages after it, so we match on its
+    # tool-call ids instead.
+    def already_in_chat?(chat, message)
+      if message.tool_call?
+        chat.messages.any? { |m| m.tool_call? && m.tool_calls.keys.intersect?(message.tool_calls.keys) }
+      else
+        !AWAITING_ROLES.include?(chat.messages.last&.role)
+      end
+    end
+
+    module Helpers # :nodoc:
+      private
+
+      def batch_result_index(id)
+        Integer(id)
+      end
+
+      def batch_failure(custom_id, detail, status: 'failed')
+        RubyLLM.logger.warn ["Batch request #{custom_id} #{status}", detail].compact.join(': ')
+        status.to_s.match?(/cancel/i) ? :cancelled : :failed
+      end
+
+      def batch_error_message(line)
+        response = line['response']
+        body = response['body'] if response.is_a?(Hash)
+        body_error = body['error'] if body.is_a?(Hash)
+        response_error = response['error'] if response.is_a?(Hash)
+
+        batch_error_value(line['error']) ||
+          line['error_message'] ||
+          batch_error_value(body_error) ||
+          batch_error_value(response_error)
+      end
+
+      def batch_error_value(error)
+        case error
+        when Hash then error['message']
+        when String then error
+        end
+      end
+
+      def single_batch_model!(requests, provider_name)
+        models = requests.map { |request| request.fetch(:model) }.uniq
+        return models.first if models.one?
+
+        raise Error, "#{provider_name} batch requests must use one model per submission"
+      end
+
+      def batch_payload(request, except: [])
+        excluded = (Array(except) + [:stream]).map(&:to_s)
+        request.fetch(:payload).reject { |key, _| excluded.include?(key.to_s) }
+      end
+    end
+
+    def inspect_attributes # :nodoc:
+      { id: id, status: status, raw_status: raw_status, chats: chats&.count, requests: requests&.count }
+    end
+  end
+end

data/lib/ruby_llm/cached_content.rb

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # A CachedContent is a provider-side prompt cache resource. Create one
+  # with ::create from a long, stable prompt prefix, then attach it to a
+  # chat so later requests read the cached tokens instead of resending
+  # them. This lifecycle requires a provider with managed cache resources;
+  # automatic prompt caching is configured with Chat#with_caching.
+  #
+  #   cache = RubyLLM.cache(big_document, model: 'gemini-3.7-flash', ttl: 3600)
+  #   chat = RubyLLM.chat(model: 'gemini-3.7-flash').with_caching(id: cache)
+  #   chat.ask "What does the document conclude?"
+  #   cache.delete
+  #
+  # Cache names are provider-owned. Persist #provider alongside #name and
+  # pass it back when finding the cache later.
+  class CachedContent
+    include Support::Inspectable
+
+    def inspect_attributes # :nodoc:
+      { name: name, model: model, provider: provider, expires_at: expires_at }
+    end
+
+    # The provider-assigned resource name, such as
+    # <tt>"cachedContents/abc123"</tt>.
+    attr_reader :name
+
+    # The model the cache was created for.
+    attr_reader :model
+
+    # The slug of the provider that stores the cache.
+    attr_reader :provider
+
+    # The Time the provider stored the cache.
+    attr_reader :created_at
+
+    # The Time the provider will delete the cache.
+    attr_reader :expires_at
+
+    # The number of tokens stored in the cache.
+    attr_reader :tokens
+
+    # The raw provider response data for the cache, as a Hash.
+    attr_reader :metadata
+
+    def initialize(name:, **attributes) # :nodoc:
+      @name = name
+      @model = attributes[:model]
+      @provider_instance = attributes[:provider_instance]
+      @provider = attributes[:provider] || @provider_instance&.slug
+      @created_at = attributes[:created_at]
+      @expires_at = attributes[:expires_at]
+      @tokens = attributes[:tokens]
+      @metadata = attributes[:metadata] || {}
+    end
+
+    # Deletes the cache resource from the provider. Returns +self+.
+    def delete
+      @provider_instance.delete_cache(name)
+      self
+    end
+
+    # Extends the cache's lifetime to +ttl:+ seconds from now, given as an
+    # Integer or a provider duration string such as <tt>"600s"</tt>.
+    # Updates #expires_at and returns +self+.
+    #
+    #   cache.renew(ttl: 3600)
+    #
+    def renew(ttl:)
+      refreshed = @provider_instance.extend_cache(name, ttl: ttl)
+      @expires_at = refreshed.expires_at
+      @metadata = refreshed.metadata
+      self
+    end
+
+    # Creates a provider-side prompt cache from +content+ and returns a
+    # CachedContent. Also available as RubyLLM.cache.
+    #
+    #   RubyLLM::CachedContent.create(big_document, model: 'gemini-3.7-flash')
+    #
+    # +content+ is the text to cache; pass file attachments with +with:+
+    # the way Chat#ask accepts them. +instructions:+ caches a system
+    # prompt alongside the content, and +ttl:+ sets the cache lifetime in
+    # seconds (Integer) or as a provider duration string such as
+    # <tt>"300s"</tt>. When +provider:+ is omitted, the model's default
+    # provider is used. The content must exceed the model's minimum
+    # cacheable token count.
+    def self.create(content, model:, ttl: nil, instructions: nil, provider: nil, context: nil, with: nil)
+      config = context&.config || RubyLLM.config
+      model_instance, provider_instance = Models.resolve(model, provider: provider, config: config)
+
+      provider_instance.cache_content(content, model: model_instance, ttl:, instructions:, with:)
+    end
+
+    # Fetches an existing cache resource by +name+ and returns a
+    # CachedContent. When +provider:+ is omitted, the provider of the
+    # configured default model is used.
+    #
+    #   cache = RubyLLM::CachedContent.find("cachedContents/abc123", provider: :gemini)
+    #
+    def self.find(name, provider: nil, context: nil)
+      config = context&.config || RubyLLM.config
+
+      provider_instance = if provider
+                            Provider.resolve!(provider).new(config)
+                          else
+                            Models.resolve(config.default_model, config:).last
+                          end
+      provider_instance.find_cache(name)
+    end
+  end
+end

data/lib/ruby_llm/chat/tool_concurrency.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    module ToolConcurrency # :nodoc: all
+      MODES = %i[threads fibers].freeze
+      Result = Struct.new(:index, :tool_call, :value, :error, keyword_init: true)
+
+      module_function
+
+      def run(mode, tool_calls, on_result: nil, &)
+        case mode
+        when :threads
+          run_with_threads(tool_calls, on_result:, &)
+        when :fibers
+          run_with_fibers(tool_calls, on_result:, &)
+        end
+      end
+
+      def run_with_threads(tool_calls, on_result:, &execute)
+        executor = rails_executor
+        workflow_context = Support::Instrumentation.current_workflow
+        queue = Queue.new
+        threads = tool_calls.each_value.with_index.map do |tool_call, index|
+          thread = Thread.new do
+            Support::Instrumentation.with_workflow(workflow_context) do
+              queue << capture_result(index, tool_call, executor, execute)
+            end
+          end
+          thread.report_on_exception = false
+          thread
+        end
+
+        collect_results(queue, threads.size, on_result:)
+      ensure
+        threads&.each(&:join)
+      end
+
+      def run_with_fibers(tool_calls, on_result:, &execute)
+        begin
+          require 'async'
+          require 'async/queue'
+        rescue LoadError
+          raise LoadError, "The 'async' gem is required for concurrent tool execution with fibers. " \
+                           "Add `gem 'async', '>= 2.0'` to your Gemfile or use `concurrency: :threads`."
+        end
+        if Gem.loaded_specs.fetch('async').version < Gem::Version.new('2.0')
+          raise LoadError, "The 'async' gem version 2.0 or newer is required for concurrent tool execution with fibers."
+        end
+
+        executor = rails_executor
+        workflow_context = Support::Instrumentation.current_workflow
+        Async do |task|
+          queue = Async::Queue.new
+          tasks = tool_calls.each_value.with_index.map do |tool_call, index|
+            task.async do
+              Support::Instrumentation.with_workflow(workflow_context) do
+                queue << capture_result(index, tool_call, executor, execute)
+              end
+            end
+          end
+
+          collect_results(queue, tasks.size, on_result:)
+        ensure
+          tasks&.each(&:wait)
+        end.wait
+      end
+
+      def collect_results(queue, count, on_result:)
+        results = Array.new(count)
+        errors = []
+
+        count.times do
+          result = queue.pop
+          if result.error
+            errors << result.error
+          else
+            results[result.index] = [result.tool_call, result.value]
+            on_result&.call(result.tool_call, result.value)
+          end
+        end
+
+        raise errors.first if errors.any?
+
+        results
+      end
+
+      def capture_result(index, tool_call, rails_executor, execute)
+        tool_call, value = run_tool_call(tool_call, rails_executor, execute)
+        Result.new(index:, tool_call:, value:)
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        Result.new(index:, tool_call:, error: e)
+      end
+
+      def run_tool_call(tool_call, rails_executor, execute)
+        if rails_executor
+          rails_executor.wrap { [tool_call, execute.call(tool_call)] }
+        else
+          [tool_call, execute.call(tool_call)]
+        end
+      end
+
+      def rails_executor
+        defined?(Rails) && Rails.respond_to?(:application) && Rails.application&.executor
+      end
+
+      private_class_method :run_with_threads, :run_with_fibers, :collect_results, :capture_result, :run_tool_call,
+                           :rails_executor
+    end
+  end
+end