ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/support/instrumentation.rb

@@ -35,8 +35,9 @@ module RubyLLM
         @mutex.synchronize { @subscribers = (@subscribers - [subscriber]).freeze }
       end
 
-      def capture_context
-        @subscribers.filter_map do |subscriber|
+      def capture_context(config: nil)
+        instrumenter = (config || RubyLLM.config).instrumenter
+        (@subscribers | [instrumenter]).filter_map do |subscriber|
           [subscriber, subscriber.capture_context] if subscriber.respond_to?(:capture_context)
         end
       end
@@ -63,15 +64,15 @@ module RubyLLM
       private_class_method :dispatch
 
       def current_workflow
-        Thread.current[WORKFLOW_CONTEXT_KEY]
+        Fiber[WORKFLOW_CONTEXT_KEY]
       end
 
       def with_workflow(context)
         previous_context = current_workflow
-        Thread.current[WORKFLOW_CONTEXT_KEY] = context
+        Fiber[WORKFLOW_CONTEXT_KEY] = context
         yield
       ensure
-        Thread.current[WORKFLOW_CONTEXT_KEY] = previous_context
+        Fiber[WORKFLOW_CONTEXT_KEY] = previous_context
       end
     end
   end

data/lib/ruby_llm/support/rescuable.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Support
+    # The +rescue_from+ macro Agent, Tool, and MCP share: handlers named
+    # with +with:+ or given as a block, searched in reverse declaration
+    # order from the class up through its ancestors when an exception is
+    # raised, and run on the instance whose work raised.
+    module Rescuable # :nodoc:
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      # The class side of Rescuable.
+      module ClassMethods # :nodoc:
+        def rescue_from(*exception_classes, with: nil, &block)
+          raise ArgumentError, 'rescue_from needs a handler: pass with: or a block' unless with || block
+          raise ArgumentError, 'rescue_from takes with: or a block, not both' if with && block
+
+          exception_classes.flatten.each do |exception_class|
+            rescue_handlers << [rescue_handler_key(exception_class), with || block]
+          end
+        end
+
+        def rescue_handlers
+          @rescue_handlers ||= []
+        end
+
+        # Returns the handler for +exception+, searched from this class up
+        # through +within+, or through every ancestor without it.
+        def rescue_handler_for(exception, within: nil)
+          rescue_handlers.reverse_each do |class_name, handler|
+            exception_class = rescue_handler_class(class_name)
+            return handler if exception_class && exception.is_a?(exception_class)
+          end
+          return if equal?(within)
+
+          superclass.rescue_handler_for(exception, within:) if superclass.respond_to?(:rescue_handler_for)
+        end
+
+        def rescue_handlers?
+          rescue_handlers.any? || (superclass.respond_to?(:rescue_handlers?) && superclass.rescue_handlers?)
+        end
+
+        private
+
+        def rescue_handler_key(exception_class)
+          case exception_class
+          when ->(klass) { klass.is_a?(Module) && !(klass.is_a?(Class) && (klass <=> StandardError).nil?) }
+            exception_class.name || exception_class
+          when String then exception_class
+          else raise ArgumentError, "#{exception_class.inspect} is not a StandardError class or its name"
+          end
+        end
+
+        def rescue_handler_class(class_name)
+          return class_name if class_name.is_a?(Module)
+
+          Object.const_get(class_name)
+        rescue NameError
+          nil
+        end
+      end
+
+      # Runs the ::rescue_from handler matching +exception+ and returns its
+      # value, or re-raises when no handler matches.
+      def rescue_with_handler(exception)
+        handler = self.class.rescue_handler_for(exception)
+        raise exception unless handler
+
+        run_rescue_handler(handler, exception)
+      end
+
+      private
+
+      def run_rescue_handler(handler, exception)
+        return instance_exec(exception, &handler) unless handler.is_a?(Symbol)
+
+        handler_method = method(handler)
+        handler_method.arity.zero? ? handler_method.call : handler_method.call(exception)
+      end
+    end
+  end
+end

data/lib/ruby_llm/support/utils.rb

@@ -12,9 +12,24 @@ module RubyLLM
 
       module_function
 
-      def hash_get(hash, key)
-        symbol_key = key.to_sym
-        hash.key?(symbol_key) ? hash[symbol_key] : hash[key.to_s]
+      def cache_directory
+        home = Dir.home
+        host_os = RbConfig::CONFIG['host_os']
+
+        if host_os.match?(/darwin/i)
+          File.join(home, 'Library', 'Caches', 'RubyLLM')
+        elsif host_os.match?(/mswin|mingw|cygwin/i)
+          local_app_data = ENV.fetch('LOCALAPPDATA', nil)
+          local_app_data = File.join(home, 'AppData', 'Local') if local_app_data.to_s.empty?
+          File.join(local_app_data, 'RubyLLM', 'Cache')
+        else
+          xdg_cache = ENV.fetch('XDG_CACHE_HOME', nil)
+          xdg_cache = File.join(home, '.cache') if xdg_cache.to_s.empty?
+          File.join(xdg_cache, 'ruby_llm')
+        end
+      rescue ArgumentError
+        # Dir.home raises when the home directory cannot be determined.
+        nil
       end
 
       def strip_schema_metadata(schema)
@@ -28,6 +43,14 @@ module RubyLLM
         name.gsub(UNDERSCORE_BOUNDARY, '_').downcase
       end
 
+      # The path Rails autoloads a class from, which follows the application's
+      # inflections when Active Support is loaded.
+      def file_path(class_name)
+        return ActiveSupport::Inflector.underscore(class_name) if defined?(ActiveSupport::Inflector)
+
+        underscore(class_name.gsub('::', '/'))
+      end
+
       def to_safe_array(item)
         case item
         when Array then item

data/lib/ruby_llm/task.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  # A tool call that finishes later. Return one from Tool#execute and a
+  # chat pauses the call instead of waiting for it, the way it pauses for
+  # an approval. Chat#pending_tasks returns the tasks it waits on, and each
+  # Chat#complete hands every one back to its tool's Tool#resume, which
+  # returns the call's result once it has one:
+  #
+  #   class RequestSignature < RubyLLM::Tool
+  #     description "Sends a contract for signature and reports when it is signed"
+  #
+  #     def execute(contract_id:)
+  #       envelope = Signatures.send_for_signature(Contract.find(contract_id))
+  #       RubyLLM::Task.new(status: :input_required, data: { "envelope_id" => envelope.id })
+  #     end
+  #
+  #     def resume(task, **)
+  #       envelope = Signatures.find(task.data["envelope_id"])
+  #       envelope.signed? ? "Signed on #{envelope.signed_on}" : task
+  #     end
+  #   end
+  #
+  # A task is saved with its tool call as #to_h, so a Rails chat can pause
+  # in one process and resume in another. MCP::Task is the task of an MCP
+  # tool call that a server runs in the background.
+  class Task
+    include Support::Inspectable
+    include Support::DeprecatedJobPredicates
+
+    FINISHED = %i[completed failed cancelled].freeze
+    STATUSES = %i[working input_required].concat(FINISHED).freeze
+    private_constant :FINISHED, :STATUSES
+
+    # The task's state as a Symbol: +:working+, +:input_required+,
+    # +:completed+, +:failed+, or +:cancelled+.
+    attr_reader :status
+
+    # What the task is doing, such as how far it got, or +nil+.
+    attr_reader :status_message
+
+    # How many seconds to wait before checking on the task again, or +nil+.
+    attr_reader :poll_interval
+
+    # The tool's own state for the task, a Hash with String keys that
+    # serializes to JSON.
+    attr_reader :data
+
+    # The ToolCall paused on the task, when it came from a chat.
+    attr_reader :tool_call
+
+    # Returns the Task that +state+, the Hash #to_h returned, describes.
+    def self.load(state, tool_call: nil)
+      task = state.fetch('task')
+      new(status: task.fetch('status'), status_message: task['status_message'],
+          poll_interval: task['poll_interval'], data: task.fetch('data', {}), tool_call:)
+    end
+
+    # Returns a task with +status+, +:working+ unless the call waits on a
+    # person, +:input_required+. +data+ is what Tool#resume needs to carry
+    # the call on, as a Hash that serializes to JSON. Raises ArgumentError
+    # for a status other than the five #status lists.
+    def initialize(status: :working, status_message: nil, poll_interval: nil, data: {}, tool_call: nil)
+      @status = status.to_sym
+      raise ArgumentError, "Unknown task status: #{status.inspect}" unless STATUSES.include?(@status)
+
+      @status_message = status_message
+      @poll_interval = poll_interval
+      @data = Support::Utils.deep_stringify_keys(Support::Utils.deep_dup(data))
+      @tool_call = tool_call
+    end
+
+    # Returns +true+ when the task waits on a person, so checking on it
+    # changes nothing until they act.
+    def input_required?
+      status == :input_required
+    end
+
+    # Returns +true+ once the task completed, failed, or was cancelled.
+    def complete?
+      FINISHED.include?(status)
+    end
+
+    # Returns +true+ when the task finished with a result.
+    def succeeded?
+      status == :completed
+    end
+
+    # Returns +true+ when the task failed with an error.
+    def failed?
+      status == :failed
+    end
+
+    # Returns +true+ when the task was cancelled before it finished.
+    def cancelled?
+      status == :cancelled
+    end
+
+    # Checks on the task once and returns +self+. Subclasses that can check
+    # on their work override it; the base implementation does nothing.
+    def refresh
+      self
+    end
+
+    # Stops the task and returns +self+. A chat calls it when it is
+    # cancelled while the call waits. Subclasses that can stop their work
+    # override it; the base implementation does nothing.
+    def cancel
+      self
+    end
+
+    # Returns the task as a Hash that serializes to JSON, which ::load turns
+    # back into the task.
+    def to_h
+      task = { 'status' => status.to_s, 'status_message' => status_message,
+               'poll_interval' => poll_interval, 'data' => (data unless data.empty?) }
+      { 'task' => task.compact }
+    end
+
+    private
+
+    def inspect_attributes
+      { status:, status_message: }
+    end
+  end
+end

data/lib/ruby_llm/test_helper.rb

@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require 'ruby_llm'
+
+module RubyLLM
+  # Lets a test script what the model replies, with no provider and no API
+  # key. Replies go through the real conversation loop, so streaming, tools,
+  # persistence, cost, and #cancel all behave as they do in production.
+  # Require it in your test helper, and include it like ActiveJob::TestHelper:
+  #
+  #   require "ruby_llm/test_helper"
+  #
+  #   class SupportAgentTest < ActiveSupport::TestCase
+  #     include RubyLLM::TestHelper
+  #
+  #     test "answers" do
+  #       stub_replies "Hello there", cost: 0.01
+  #
+  #       reply = SupportAgent.new.ask("Hi")
+  #
+  #       assert_equal "Hello there", reply.content
+  #       assert_equal 0.01, reply.cost.total
+  #     end
+  #   end
+  #
+  # It also works in RSpec: <tt>config.include RubyLLM::TestHelper</tt>. A
+  # chat's request reaches the stub. Other operations, such as RubyLLM.embed,
+  # still call their provider.
+  module TestHelper
+    # The scripted replies, shared across threads because a system test's
+    # server and its jobs run in threads of their own.
+    module Script # :nodoc:
+      @lock = Mutex.new
+      @replies = nil
+
+      class << self
+        def stub(replies, options)
+          @lock.synchronize { @replies = [replies, options] }
+        end
+
+        def clear
+          @lock.synchronize { @replies = nil }
+        end
+
+        def active?
+          @lock.synchronize { !@replies.nil? }
+        end
+
+        # The next reply. The last one repeats, so a test that doesn't count
+        # requests only names one.
+        def next_reply
+          @lock.synchronize do
+            replies, options = @replies
+            [replies.size > 1 ? replies.shift : replies.first, options]
+          end
+        end
+      end
+    end
+
+    # Answers requests from the script, once #stub_replies has been called.
+    module StubbedProvider # :nodoc:
+      # Skips the API key check, which some providers define for themselves.
+      def initialize(config)
+        if Script.active?
+          @config = config
+          @connection = Transport::Connection.new(self, config)
+        else
+          super
+        end
+      end
+
+      def complete(messages, usage_recorder: nil, progress_recorder: nil, **options, &)
+        return super unless Script.active?
+
+        reply, settings = Script.next_reply
+        raise reply if reply.is_a?(Exception)
+
+        tool_call = reply if reply.is_a?(ToolCall)
+        text = tool_call ? '' : reply.to_s
+        stream(text, settings, options[:model], &) if block_given?
+        Message.new(role: :assistant, content: text, model_id: options[:model].id,
+                    tool_calls: tool_call && { tool_call.id => tool_call }, input_tokens: 1, output_tokens: 1,
+                    reported_cost: settings[:cost])
+      end
+
+      private
+
+      def stream(text, settings, model)
+        text.scan(/\S+\s*|\s+/).each do |word|
+          sleep settings[:interval] if settings[:interval]
+          yield Chunk.new(role: :assistant, content: word, model_id: model.id)
+        end
+      end
+    end
+
+    Provider.prepend(StubbedProvider)
+
+    def self.included(base) # :nodoc:
+      base.after { Script.clear } if base.respond_to?(:after)
+    end
+
+    def after_teardown # :nodoc:
+      super
+    ensure
+      Script.clear
+    end
+
+    # Makes every chat answer with +replies+ until the test ends, in order,
+    # repeating the last, instead of calling a provider. A String streams
+    # word by word, #tool_call asks the conversation loop to run a tool, and
+    # an exception such as RubyLLM::ServerError fails the request.
+    # +cost:+ is what each reply reports, in dollars. +interval:+ is the
+    # seconds between words, so a test can press Stop while a reply streams.
+    #
+    #   stub_replies "Hello there", cost: 0.01
+    #   stub_replies tool_call(:weather, city: "Berlin"), "It is sunny."
+    #   stub_replies "A long reply", interval: 0.1
+    #
+    def stub_replies(*replies, cost: nil, interval: nil)
+      Script.stub(replies, { cost:, interval: })
+    end
+
+    # Returns a reply that calls the tool +name+ with +arguments+, for
+    # #stub_replies.
+    def tool_call(name, **arguments)
+      ToolCall.new(id: "call_#{SecureRandom.hex(4)}", name: name.to_s, arguments:)
+    end
+  end
+end

data/lib/ruby_llm/thinking.rb

@@ -134,14 +134,14 @@ module RubyLLM
         enabled == false || effort == :none
       end
 
-      def resolve(model)
+      def resolve(model, strict: true)
         return self if @intent.nil?
 
         options = Controls.new(model).public_send(@intent)
         return nil if options == :already_on
-        raise ArgumentError, resolution_error(model) unless options
+        raise ArgumentError, resolution_error(model) if strict && !options
 
-        self.class.new(**options)
+        self.class.new(**(options || { enabled: @intent == :enable }))
       end
 
       private

data/lib/ruby_llm/tokens.rb

@@ -12,11 +12,23 @@ module RubyLLM
   #   response.tokens.cache_read  # prompt cache reads
   #   response.tokens.cache_write # prompt cache writes
   #
+  # Operations a provider bills in other units count those units here too,
+  # so Cost prices them the same way:
+  #
+  #   RubyLLM.speak("Hello").tokens.characters             # => 5
+  #   RubyLLM.transcribe("call.mp3").tokens.audio_seconds  # => 42.3
+  #   RubyLLM.rerank(query, documents).tokens.search_units # => 1
+  #
   class Tokens
     include Support::Inspectable
 
     READERS = %i[
-      input output cache_read cache_write thinking cache_write_by_ttl server_tool_use reported_cost
+      input output cache_read cache_write thinking characters audio_seconds search_units cache_write_by_ttl
+      server_tool_use reported_cost
+    ].freeze # :nodoc:
+
+    SUMMED = %i[
+      input output cache_read cache_write thinking characters audio_seconds search_units reported_cost
     ].freeze # :nodoc:
 
     # The number of standard (non-cached) input tokens, or +nil+ if the
@@ -50,6 +62,20 @@ module RubyLLM
     # does not report them.
     attr_reader :thinking
 
+    # The number of input characters the provider bills, such as the text
+    # sent to a text-to-speech model, or +nil+ if the operation is not
+    # billed by character.
+    attr_reader :characters
+
+    # The seconds of audio the provider bills, such as the length of a
+    # transcribed recording, or +nil+ if the operation is not billed by
+    # duration. May be fractional.
+    attr_reader :audio_seconds
+
+    # The number of search units the provider bills, such as Cohere's
+    # rerank searches, or +nil+ if the operation is not billed by search.
+    attr_reader :search_units
+
     # How many times each provider-executed tool ran, as a Hash of String
     # keys and Integer counts, or +nil+ if the provider reported none.
     # Web searches are <tt>"web_search_requests"</tt> on every provider
@@ -72,20 +98,26 @@ module RubyLLM
     # even when every count is +nil+.
     #
     #   tokens = RubyLLM::Tokens.new(input: 100, output: 20)
-    #   RubyLLM.models.find("gpt-5.6").cost_for(tokens).total
+    #   RubyLLM.models.find("gpt-6.1-sol").cost_for(tokens).total
     #
+    # +characters:+, +audio_seconds:+, and +search_units:+ count usage
+    # billed in those units;
     # +cache_write_by_ttl:+ splits the cache writes by lifetime, keyed as
     # #cache_write_by_ttl describes;
     # +server_tool_use:+ counts the provider-executed tools that ran, keyed
     # as #server_tool_use describes;
     # +reported_cost:+ holds the provider's total price in US dollars.
     def initialize(input: nil, output: nil, cache_read: nil, cache_write: nil, thinking: nil,
+                   characters: nil, audio_seconds: nil, search_units: nil,
                    cache_write_by_ttl: nil, server_tool_use: nil, reported_cost: nil)
       @input = input
       @output = output
       @cache_read = cache_read
       @cache_write = cache_write
       @thinking = thinking
+      @characters = characters
+      @audio_seconds = audio_seconds
+      @search_units = search_units
       @cache_write_by_ttl = positive_counts(cache_write_by_ttl)
       @server_tool_use = positive_counts(server_tool_use)
       @reported_cost = reported_cost
@@ -95,10 +127,9 @@ module RubyLLM
     # no attempt reported it.
     def self.aggregate(tokens)
       tokens = Array(tokens).compact
-      return new if tokens.empty?
       return tokens.first if tokens.one?
 
-      values = %i[input output cache_read cache_write thinking reported_cost].to_h do |component|
+      values = SUMMED.to_h do |component|
         reported = tokens.filter_map { |usage| usage.public_send(component) }
         [component, reported.empty? ? nil : reported.sum]
       end
@@ -114,20 +145,22 @@ module RubyLLM
       end)
     end
 
-    def self.aggregate_counts(tokens, reader) # :nodoc:
-      reported = tokens.filter_map(&reader)
-      return nil if reported.empty?
+    def without_reported_cost # :nodoc:
+      self.class.new(**(READERS - [:reported_cost]).to_h { |reader| [reader, public_send(reader)] })
+    end
 
-      reported.each_with_object({}) do |counters, total|
-        counters.each { |key, count| total[key] = total.fetch(key, 0) + count.to_i }
+    def self.aggregate_counts(tokens, reader) # :nodoc:
+      tokens.filter_map(&reader).each_with_object({}) do |counters, total|
+        counters.each { |key, count| total[key] = total.fetch(key, 0) + (count.is_a?(Float) ? count : count.to_i) }
       end
     end
     private_class_method :aggregate_counts
 
     # Returns the counts as a hash with keys +:input_tokens+,
     # +:output_tokens+, +:cache_read_tokens+, +:cache_write_tokens+, and
-    # +:thinking_tokens+, omitting +nil+ counts. Includes
-    # +:cache_write_tokens_by_ttl+ and +:server_tool_use+ when reported.
+    # +:thinking_tokens+, omitting +nil+ counts. Includes +:characters+,
+    # +:audio_seconds+, +:search_units+, +:cache_write_tokens_by_ttl+, and
+    # +:server_tool_use+ when reported.
     # The provider's reported cost is available separately
     # through #reported_cost.
     #
@@ -141,6 +174,9 @@ module RubyLLM
         cache_read_tokens: cache_read,
         cache_write_tokens: cache_write,
         thinking_tokens: thinking,
+        characters: characters,
+        audio_seconds: audio_seconds,
+        search_units: search_units,
         cache_write_tokens_by_ttl: cache_write_by_ttl,
         server_tool_use: server_tool_use
       }.compact
@@ -153,6 +189,9 @@ module RubyLLM
         cache_read: cache_read,
         cache_write: cache_write,
         thinking: thinking,
+        characters: characters,
+        audio_seconds: audio_seconds,
+        search_units: search_units,
         cache_write_by_ttl: cache_write_by_ttl,
         server_tool_use: server_tool_use
       }.compact
@@ -161,7 +200,8 @@ module RubyLLM
     private
 
     def positive_counts(counts)
-      used = counts.to_h.to_h { |key, count| [key.to_s, count.to_i] }.select { |_, count| count.positive? }
+      used = counts.to_h.to_h { |key, count| [key.to_s, count.is_a?(Float) ? count : count.to_i] }
+                   .select { |_, count| count.positive? }
       used unless used.empty?
     end
   end