ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/chat/tool_concurrency.rb

@@ -8,38 +8,92 @@ module RubyLLM
 
       module_function
 
-      def run(mode, tool_calls, on_result: nil, &)
+      # Runs +tool_calls+ in rounds split by the calls +exclusive+ picks
+      # out. A round with a result +paused+ picks out stops the rounds
+      # after it, which wait for that call to finish.
+      def run(mode, tool_calls, exclusive: nil, paused: nil, on_result: nil, config: nil, &)
+        rounds = rounds(tool_calls, exclusive)
         case mode
         when :threads
-          run_with_threads(tool_calls, on_result:, &)
+          run_with_threads(rounds, paused:, on_result:, config:, &)
         when :fibers
-          run_with_fibers(tool_calls, on_result:, &)
+          run_with_fibers(rounds, paused:, on_result:, config:, &)
         end
       end
 
-      def run_with_threads(tool_calls, on_result:, &execute)
+      # Groups the calls in the order the model gave them, so an exclusive
+      # call waits for the calls before it and the calls after it wait for
+      # it, as with a fair read-write lock.
+      def rounds(tool_calls, exclusive)
+        indexed = tool_calls.each_value.with_index.to_a
+        return [indexed] unless exclusive
+
+        indexed.chunk_while { |(first, _), (second, _)| !exclusive.call(first) && !exclusive.call(second) }.to_a
+      end
+
+      def run_with_threads(rounds, paused:, on_result:, config:, &execute)
         caller_context = execution_context
         workflow_context = Support::Instrumentation.current_workflow
-        instrumentation_context = Support::Instrumentation.capture_context
-        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
-              Support::Instrumentation.with_context(instrumentation_context) do
-                queue << capture_result(index, tool_call, caller_context, execute)
-              end
+        instrumentation_context = Support::Instrumentation.capture_context(config:)
+        results = Array.new(rounds.sum(&:size))
+        rounds.each do |round|
+          queue = Queue.new
+          threads = []
+          round.each do |tool_call, index|
+            thread = Thread.new do
+              result = capture_result(index, tool_call, caller_context, execute,
+                                      workflow_context:, instrumentation_context:)
+            ensure
+              queue << (result || Result.new(index:, tool_call:, error: thread_exited_error(tool_call)))
             end
+            thread.report_on_exception = false
+            threads << thread
           end
-          thread.report_on_exception = false
-          thread
+
+          collect_results(queue, threads.size, results, on_result:)
+          break if paused?(round, results, paused)
+        ensure
+          threads&.each(&:join)
         end
+        results
+      end
 
-        collect_results(queue, threads.size, on_result:)
+      def run_with_fibers(rounds, paused:, on_result:, config:, &execute)
+        require_async
+        caller_context = execution_context
+        contexts = { workflow_context: Support::Instrumentation.current_workflow,
+                     instrumentation_context: Support::Instrumentation.capture_context(config:) }
+        results = Array.new(rounds.sum(&:size))
+        # Inside a reactor, Sync runs the block in the calling fiber, so
+        # results persist with the caller's execution state and connection.
+        Sync do |task|
+          isolated(caller_context) do
+            rounds.each do |round|
+              run_fiber_round(task, round, results, caller_context, contexts, on_result:, &execute)
+              break if paused?(round, results, paused)
+            end
+          end
+        end
+        results
+      end
+
+      def run_fiber_round(task, round, results, caller_context, contexts, on_result:, &execute)
+        queue = Async::Queue.new
+        tasks = []
+        round.each do |tool_call, index|
+          worker = task.async { queue << capture_result(index, tool_call, caller_context, execute, **contexts) }
+          tasks << worker
+        end
+
+        collect_results(queue, tasks.size, results, on_result:)
+      rescue Async::Stop
+        tasks&.each(&:stop)
+        raise
       ensure
-        threads&.each(&:join)
+        tasks&.each(&:wait)
       end
 
-      def run_with_fibers(tool_calls, on_result:, &execute)
+      def require_async
         begin
           require 'async'
           require 'async/queue'
@@ -47,37 +101,16 @@ module RubyLLM
           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
+        return if Gem.loaded_specs.fetch('async').version >= Gem::Version.new('2.0')
 
-        caller_context = execution_context
-        workflow_context = Support::Instrumentation.current_workflow
-        instrumentation_context = Support::Instrumentation.capture_context
-        # Inside a reactor, Sync runs the block in the calling fiber, so
-        # results persist with the caller's execution state and connection.
-        Sync do |task|
-          isolated(caller_context) do
-            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
-                  Support::Instrumentation.with_context(instrumentation_context) do
-                    queue << capture_result(index, tool_call, caller_context, execute)
-                  end
-                end
-              end
-            end
+        raise LoadError, "The 'async' gem version 2.0 or newer is required for concurrent tool execution with fibers."
+      end
 
-            collect_results(queue, tasks.size, on_result:)
-          ensure
-            tasks&.each(&:wait)
-          end
-        end
+      def paused?(round, results, paused)
+        paused && round.any? { |_tool_call, index| paused.call(results[index].last) }
       end
 
-      def collect_results(queue, count, on_result:)
-        results = Array.new(count)
+      def collect_results(queue, count, results, on_result:)
         errors = []
 
         count.times do
@@ -91,17 +124,24 @@ module RubyLLM
         end
 
         raise errors.first if errors.any?
-
-        results
       end
 
-      def capture_result(index, tool_call, caller_context, execute)
-        value = isolated(caller_context) { execute.call(tool_call) }
+      def capture_result(index, tool_call, caller_context, execute, workflow_context:, instrumentation_context:)
+        value = nil
+        Support::Instrumentation.with_workflow(workflow_context) do
+          Support::Instrumentation.with_context(instrumentation_context) do
+            value = isolated(caller_context) { execute.call(tool_call) }
+          end
+        end
         Result.new(index:, tool_call:, value:)
       rescue Exception => e # rubocop:disable Lint/RescueException
         Result.new(index:, tool_call:, error: e)
       end
 
+      def thread_exited_error(tool_call)
+        Error.new("Tool #{tool_call.name} exited its thread without a result")
+      end
+
       # Rails gives a thread, and a fiber under fiber isolation, execution
       # state of its own: the block runs as its own unit of work, so its
       # connections return to the pool. A fiber sharing the caller's state
@@ -128,7 +168,8 @@ module RubyLLM
         defined?(Rails) && Rails.respond_to?(:application) && Rails.application&.executor
       end
 
-      private_class_method :run_with_threads, :run_with_fibers, :collect_results, :capture_result, :isolated,
+      private_class_method :rounds, :paused?, :run_with_threads, :run_with_fibers, :run_fiber_round, :require_async,
+                           :collect_results, :capture_result, :thread_exited_error, :isolated,
                            :execution_context, :rails_executor
     end
   end

data/lib/ruby_llm/chat/tool_search.rb

@@ -3,6 +3,11 @@
 module RubyLLM
   class Chat
     # Deferred tool registration and the tool set rendered for the request.
+    # A provider with tool search of its own gets every tool, the deferred
+    # ones marked. Any other gets the tools that are not deferred, a
+    # +tool_search+ tool, and the deferred tools the conversation has loaded
+    # with it or called, in the order it first did, so every request repeats
+    # the one before it and adds to its end.
     module ToolSearch # :nodoc:
       private
 
@@ -10,7 +15,7 @@ module RubyLLM
         names = @tool_deferrals.filter_map { |name, deferred| name if deferred }
         mcp.each do |server|
           explicit = @mcp_deferrals[server.name.to_sym]
-          server.tools.each do |tool|
+          MCP::Outages.tools(server).each do |tool|
             names << tool.name.to_sym if explicit.nil? ? server.defers?(tool) : explicit
           end
         end
@@ -19,8 +24,38 @@ module RubyLLM
 
       def request_tools
         deferred = deferred_tool_names
+        return searched_request_tools(deferred) if searching_tools_here?(deferred)
+
         tools.to_h { |name, tool| [name, deferred.include?(name) ? Tool::Deferred.new(tool) : tool] }
       end
+
+      def searched_request_tools(deferred)
+        all = tools
+        loaded = loaded_tool_names(deferred).select { |name| all.key?(name) }
+        all.except(*deferred).merge(all.slice(*loaded))
+      end
+
+      def tool_search_tool
+        return unless searching_tools_here?(deferred_tool_names)
+
+        @tool_search_tool ||= Tools::Search.new(-> { deferred_tools })
+      end
+
+      def searching_tools_here?(deferred)
+        deferred.any? && !@provider.tool_search?(@model, protocol: @protocol)
+      end
+
+      def loaded_tool_names(deferred)
+        results = messages.select(&:tool_result?).to_h { |message| [message.tool_call_id, message.content] }
+        calls = messages.select(&:tool_call?).flat_map { |message| message.tool_calls.to_a }
+        calls.flat_map { |tool_call| names_loaded_by(tool_call, results) }.uniq & deferred
+      end
+
+      def names_loaded_by(tool_call, results)
+        return [tool_call.qualified_name.to_sym] unless tool_call.qualified_name == Tools::Search.tool_name
+
+        Tools::Search.loaded_names(results[tool_call.id]).map(&:to_sym)
+      end
     end
   end
 end

data/lib/ruby_llm/chat/tool_sequences.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Tool lookup by namespace and the ordered batches of client toolsets.
+    # When the provider runs a namespace's calls in one response in order,
+    # each call waits for the ones before it, and once one fails the rest
+    # are answered with the provider's skipped-call error instead of
+    # running.
+    module ToolSequences # :nodoc:
+      private
+
+      def tool_for(call)
+        tool = tools[call.qualified_name.to_sym]
+        tool if tool && tool.namespace == call.namespace && tool.name.to_s == call.name.to_s
+      end
+
+      def partition_tool_sequences(pending)
+        executable = {}
+        denied = {}
+        pending.each do |id, call|
+          decision = if skipped_result(call)
+                       true
+                     elsif !input_pending?(call)
+                       execution_decision(call)
+                     end
+          (decision ? executable : denied)[id] = call unless decision.nil?
+        end
+        [executable, denied]
+      end
+
+      def ordered?(call)
+        !skipped_tool_call_error(call).nil?
+      end
+
+      def skipped_result(call)
+        error = skipped_tool_call_error(call)
+        { error: } if error && sequence_failed?(call)
+      end
+
+      def skipped_tool_call_error(call)
+        @provider.skipped_tool_call_error(call, model: @model, protocol: @protocol)
+      end
+
+      def sequence_blocked?(call)
+        earlier_sequence_calls(call).any? { |candidate| tool_result_for(candidate).nil? }
+      end
+
+      # An exclusive call waits for the calls before it, and the calls after
+      # it wait for it, also across rounds: a call held back by one that
+      # waits on an approval or a Task stays pending. Calls in +running+
+      # run in this round, in order, so they hold nothing back.
+      def exclusive_blocked?(call, running = {})
+        response = messages.reverse.find { |message| message.tool_calls.key?(call.id) }
+        return false unless response
+
+        exclusive = exclusive_call?(call)
+        response.tool_calls.take_while { |candidate| candidate.id != call.id }.any? do |earlier|
+          (exclusive || exclusive_call?(earlier)) && !running.key?(earlier.id) && tool_result_for(earlier).nil?
+        end
+      end
+
+      def exclusive_call?(call)
+        !call.remote? && tool_for(call)&.exclusive? == true
+      end
+
+      def sequence_failed?(call)
+        earlier_sequence_calls(call).any? { |candidate| tool_result_for(candidate)&.tool_error_message }
+      end
+
+      def earlier_sequence_calls(call)
+        return [] unless ordered?(call)
+
+        response = messages.reverse.find { |message| message.tool_calls.key?(call.id) }
+        response.tool_calls.take_while { |candidate| candidate.id != call.id }
+                .select { |candidate| candidate.namespace == call.namespace }
+      end
+
+      def tool_result_for(call)
+        messages.reverse.find { |message| message.tool_result? && message.tool_call_id == call.id }
+      end
+    end
+  end
+end