smith-agents 0.6.1 → 0.8.0

data/lib/smith/tool/call_allowance.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallAllowance
+      EXHAUSTION_POLICIES = %i[raise complete].freeze
+
+      def self.charge_legacy!(allowance)
+        LegacyCallAllowance.charge!(allowance) { yield if block_given? }
+      end
+
+      def initialize(remaining, on_exhaustion: :raise)
+        invalid = (remaining.is_a?(Integer) && remaining.negative?) ||
+                  (!remaining.is_a?(Integer) && !remaining.is_a?(CallBudget))
+        raise ArgumentError, "tool call allowance must be a non-negative integer" if invalid
+
+        initialize_scope(CallBudget.coerce(remaining), on_exhaustion:, parent: nil)
+      end
+
+      def scope(budget, on_exhaustion: @on_exhaustion)
+        self.class.allocate.tap do |allowance|
+          allowance.__send__(:initialize_scope, CallBudget.coerce(budget), on_exhaustion:, parent: self)
+        end
+      end
+
+      def charge!(tool_name = nil)
+        batch = CallBatch.coerce(exact? ? [tool_name] : 1, exact: exact?)
+
+        synchronize do
+          Thread.handle_interrupt(Object => :never) do
+            raise BudgetExceeded, "agent tool_calls budget exceeded" unless reservable?(batch)
+
+            yield if block_given?
+            consume!(batch)
+          end
+        end
+      end
+
+      def remaining
+        synchronize { lineage.map { _1.__send__(:counter).remaining }.min }
+      end
+
+      def remaining_for(tool_name)
+        name = canonical_tool_name(tool_name)
+        synchronize do
+          exact_scopes = lineage.select { _1.__send__(:budget).exact? }
+          return nil if exact_scopes.empty?
+
+          exact_scopes.map { _1.__send__(:counter).remaining_for(name) }.min
+        end
+      end
+
+      def reserve_batch(tool_names_or_size, ledger: nil)
+        batch = CallBatch.coerce(tool_names_or_size, exact: exact?)
+
+        synchronize do
+          Thread.handle_interrupt(Object => :never) do
+            return unless reservable?(batch)
+
+            workflow_ledger = ledger if ledger&.limits&.key?(:tool_calls)
+            ledger_reservation = workflow_ledger&.reserve!(:tool_calls, batch.size)
+            reservation = CallReservation.new(
+              limit: batch.size,
+              ledger: workflow_ledger,
+              ledger_reservation: ledger_reservation
+            )
+
+            consume!(batch)
+            reservation
+          end
+        end
+      rescue BudgetExceeded
+        nil
+      end
+
+      def complete_on_exhaustion?
+        @on_exhaustion == :complete
+      end
+
+      def exact?
+        lineage.any? { _1.__send__(:budget).exact? }
+      end
+
+      def used?
+        synchronize { counter.used? }
+      end
+
+      def [](key)
+        remaining if key == :remaining
+      end
+
+      protected
+
+      attr_reader :budget
+
+      private
+
+      def initialize_scope(budget, on_exhaustion:, parent:)
+        unless EXHAUSTION_POLICIES.include?(on_exhaustion)
+          raise ArgumentError, "tool call exhaustion policy must be :raise or :complete"
+        end
+
+        @budget = budget
+        @parent = parent
+        @lineage = [self, *Array(parent&.__send__(:lineage))].freeze
+        @mutex = parent ? parent.__send__(:shared_mutex) : Mutex.new
+        @counter = CallAllowanceCounter.new(budget)
+        @on_exhaustion = on_exhaustion
+        self
+      end
+
+      attr_reader :lineage, :counter
+
+      def shared_mutex = @mutex
+
+      def synchronize(&) = @mutex.synchronize(&)
+
+      def canonical_tool_name(value)
+        name = value.to_s
+        raise ArgumentError, "tool name must not be empty" if name.empty?
+
+        name
+      end
+
+      def reservable?(batch)
+        lineage.all? { _1.__send__(:counter).available?(batch) }
+      end
+
+      def consume!(batch)
+        lineage.each { _1.__send__(:counter).consume!(batch) }
+      end
+    end
+  end
+end

data/lib/smith/tool/call_allowance_counter.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallAllowanceCounter
+      def initialize(budget)
+        @budget = budget
+        @remaining = budget.total
+        @remaining_by_tool = budget.tool_limits&.dup
+      end
+
+      attr_reader :remaining
+
+      def remaining_for(tool_name)
+        @remaining_by_tool.fetch(tool_name, 0)
+      end
+
+      def used?
+        remaining < @budget.total
+      end
+
+      def available?(batch)
+        remaining >= batch.size && tool_counts_available?(batch.counts)
+      end
+
+      def consume!(batch)
+        @remaining -= batch.size
+        return unless @budget.exact?
+
+        batch.counts.each { |name, count| @remaining_by_tool[name] -= count }
+      end
+
+      private
+
+      def tool_counts_available?(counts)
+        return true unless @budget.exact?
+        return false unless counts
+
+        counts.all? { |name, count| @remaining_by_tool.fetch(name, 0) >= count }
+      end
+    end
+  end
+end

data/lib/smith/tool/call_batch.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallBatch
+      attr_reader :size, :counts
+
+      def self.coerce(value, exact:)
+        return from_size(value, exact:) if value.is_a?(Integer)
+        raise ArgumentError, "tool call batch size must be a positive integer" unless value.is_a?(Array)
+
+        names = value.map { canonical_tool_name(_1) }
+        validate_size!(names.length)
+        new(size: names.length, counts: names.tally.freeze)
+      end
+
+      def self.from_size(size, exact:)
+        validate_size!(size)
+        raise ArgumentError, "exact tool call allowance requires tool names for batch reservation" if exact
+
+        new(size:, counts: nil)
+      end
+      private_class_method :from_size
+
+      def self.canonical_tool_name(value)
+        name = value.to_s
+        raise ArgumentError, "tool name must not be empty" if name.empty?
+
+        name
+      end
+      private_class_method :canonical_tool_name
+
+      def self.validate_size!(size)
+        return if size.is_a?(Integer) && size.positive?
+
+        raise ArgumentError, "tool call batch size must be a positive integer"
+      end
+      private_class_method :validate_size!
+
+      def initialize(size:, counts:)
+        @size = size
+        @counts = counts
+        freeze
+      end
+    end
+  end
+end

data/lib/smith/tool/call_budget.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require "dry-struct"
+
+require_relative "../types"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallBudget < Dry::Struct
+      OwnedString = Types::String.constructor { |value| value.is_a?(String) ? value.dup.freeze : value }
+      ToolName = OwnedString.constrained(min_size: 1)
+      Limit = Types::Integer.constrained(gteq: 0)
+      PositiveLimit = Types::Integer.constrained(gt: 0)
+      ToolLimits = Types::Hash.map(ToolName, PositiveLimit)
+
+      private_constant :OwnedString, :ToolName, :Limit, :PositiveLimit, :ToolLimits
+
+      attribute :total, Limit
+      attribute :tool_limits, ToolLimits.optional.default(nil)
+
+      def self.coerce(value)
+        return value if value.is_a?(self)
+        return new(total: value) if value.is_a?(Integer)
+
+        raise ArgumentError, "tool call budget must be an integer or Smith::Tool::CallBudget"
+      end
+
+      def initialize(...)
+        super
+        validate_tool_limits!
+        tool_limits&.freeze
+        freeze
+      end
+
+      def exact? = !tool_limits.nil?
+
+      def limit_for(tool_name)
+        tool_limits&.fetch(tool_name.to_s, nil)
+      end
+
+      private
+
+      def validate_tool_limits!
+        return unless tool_limits
+        return if total <= tool_limits.values.sum
+
+        raise ArgumentError, "aggregate tool call allowance exceeds the sum of its per-tool limits"
+      end
+    end
+  end
+end

data/lib/smith/tool/call_reservation.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class CallReservation
+      extend Dry::Initializer
+
+      option :limit
+      option :ledger
+      option :ledger_reservation
+
+      def initialize(...)
+        super
+        @claimed = 0
+        @settled = false
+        @mutex = Mutex.new
+      end
+
+      def claim
+        @mutex.synchronize do
+          return false if @settled || @claimed >= limit
+
+          @claimed += 1
+          true
+        end
+      end
+
+      def settle!
+        Thread.handle_interrupt(Exception => :never) do
+          @mutex.synchronize do
+            return if @settled
+
+            ledger&.reconcile!(ledger_reservation, @claimed)
+            @settled = true
+          end
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/capture.rb

@@ -3,19 +3,75 @@
 module Smith
   class Tool < RubyLLM::Tool
     module Capture
+      CAPTURE_FAILED = Object.new.freeze
+      private_constant :CAPTURE_FAILED
+
       private
 
+      def ensure_capture_ready!
+        return unless self.class.capture_result && self.class.capture_result_strict?
+
+        capture_collector(true)
+      end
+
       def capture_result_if_configured(kwargs, result)
         block = self.class.capture_result
         return unless block
 
-        collector = self.class.current_tool_result_collector
+        capture_with_policy(block, kwargs, result, self.class.capture_result_strict?)
+      end
+
+      def capture_with_policy(block, kwargs, result, strict)
+        collector = capture_collector(strict)
         return unless collector
 
-        captured = block.call(kwargs, result)
-        collector.call({ tool: name.to_s, captured: captured }) if captured
+        captured = captured_value(block, kwargs, result, strict)
+        return if captured.equal?(CAPTURE_FAILED)
+
+        append_capture(collector, captured, strict)
+      end
+
+      def capture_collector(strict)
+        collector = self.class.current_tool_result_collector
+        return collector unless strict
+        return collector if collector.respond_to?(:call)
+
+        reason = collector.nil? ? :collector_missing : :collector_invalid
+        raise tool_capture_failed(reason)
+      end
+
+      def captured_value(block, kwargs, result, strict)
+        captured = call_capture_block(block, kwargs, result, strict)
+        return captured if captured.equal?(CAPTURE_FAILED)
+        return captured unless captured.nil? && strict
+
+        raise tool_capture_failed(:capture_empty)
+      end
+
+      def call_capture_block(block, kwargs, result, strict)
+        block.call(kwargs, result)
       rescue StandardError => e
-        Smith.config.logger&.warn("[Smith] capture_result failed for #{name}: #{e.message}")
+        capture_failure(e, strict, reason: :capture_block_failed)
+      end
+
+      def append_capture(collector, captured, strict)
+        collector.call({ tool: name.to_s, captured: captured }) if strict || captured
+      rescue StandardError => e
+        capture_failure(e, strict, reason: :collector_failed)
+      end
+
+      def capture_failure(error, strict, reason:)
+        return CAPTURE_FAILED.tap { log_capture_failure(error) } unless strict
+
+        raise tool_capture_failed(reason), cause: error
+      end
+
+      def tool_capture_failed(reason)
+        ToolCaptureFailed.for_runtime(tool_name: name, reason:)
+      end
+
+      def log_capture_failure(error)
+        Smith.config.logger&.warn("[Smith] capture_result failed for #{name}: #{error.message}")
       end
     end
   end

data/lib/smith/tool/capture_configuration.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module CaptureConfiguration
+      def capture_result(strict: false, &block)
+        return @capture_result unless block
+
+        raise ArgumentError, "capture_result strict must be true or false" unless [true, false].include?(strict)
+
+        @capture_result_strict = strict
+        @capture_result = block
+      end
+
+      def capture_result_strict?
+        @capture_result_strict == true
+      end
+    end
+  end
+end

data/lib/smith/tool/chat_execution_callbacks.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ChatExecutionCallbacks
+      private
+
+      def execute_tool_with_callbacks(tool_call)
+        batch = execution_batch(tool_call)
+        dispatch_call = batch&.dispatch_for(tool_call)
+        return super unless dispatch_call && batch.request_for(dispatch_call)
+
+        execute_managed_tool_callbacks(batch, dispatch_call) do
+          super(batch.source_for(dispatch_call))
+        end
+      end
+
+      def execute_managed_tool_callbacks(batch, dispatch_call, &block)
+        claim = batch.claim_dispatch!(dispatch_call)
+        with_dispatch_claim(batch, claim) { complete_managed_callbacks(batch, dispatch_call, &block) }
+      rescue StandardError => e
+        handle_managed_callback_failure(batch, dispatch_call, claim, e)
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        record_batch_failure(batch, e)
+        raise
+      end
+
+      def with_dispatch_claim(batch, claim, &)
+        context = batch.context.merge(current_tool_dispatch_claim: claim).freeze
+        Tool::ScopedContext.around(context, &)
+      end
+
+      def complete_managed_callbacks(batch, dispatch_call)
+        result = yield
+        notify_rejected_callback_dispatch(batch, dispatch_call) unless batch.started?(dispatch_call)
+        result
+      end
+
+      def handle_managed_callback_failure(batch, dispatch_call, claim, error)
+        record_batch_failure(batch, error)
+        raise error unless claim
+        raise error if error.is_a?(ToolFailureNotificationFailed)
+
+        failure = invocation_dispatch_failure(batch, dispatch_call, error)
+        notify_invocation_failure(batch, dispatch_call, failure)
+        raise failure
+      end
+
+      def notify_rejected_callback_dispatch(batch, dispatch_call)
+        notify_invocation_failure(
+          batch,
+          dispatch_call,
+          ToolDispatchRejected.new("tool invocation arguments were rejected before execution")
+        )
+      end
+    end
+  end
+end

data/lib/smith/tool/chat_execution_context.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ChatExecutionContext
+      include BoundedCompletionInstallation
+      include ExecutionFailureHandling
+      include ExecutionDispatch
+      include ChatExecutionCallbacks
+      include ExecutionBatchLifecycle
+
+      def self.install(chat)
+        return chat unless chat.respond_to?(:tools) && chat.tools.respond_to?(:values)
+        unless chat.respond_to?(:execute_tool, true)
+          raise Error, "unsupported RubyLLM chat execution interface: missing #execute_tool"
+        end
+
+        chat.extend(self) unless chat.singleton_class < self
+        chat.__send__(:smith_tool_execution_batches)
+        chat
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_authority.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionAuthority
+      THREAD_KEY = :smith_tool_execution_authority
+      private_constant :THREAD_KEY
+
+      def self.current = Thread.current[THREAD_KEY]
+
+      def self.around(tool:, dispatch_claim:, &block)
+        raise ArgumentError, "block required" unless block
+
+        previous = current
+        Thread.handle_interrupt(Object => :never) do
+          Thread.current[THREAD_KEY] = new(tool:, dispatch_claim:)
+          begin
+            Thread.handle_interrupt(Object => :immediate, &block)
+          ensure
+            Thread.current[THREAD_KEY] = previous
+          end
+        end
+      end
+
+      def initialize(tool:, dispatch_claim:)
+        @tool = tool
+        @dispatch_claim = dispatch_claim
+        @available = true
+        @mutex = Mutex.new
+      end
+
+      def claim(tool, dispatch_claim)
+        @mutex.synchronize do
+          return false unless @available && @tool.equal?(tool) && @dispatch_claim.equal?(dispatch_claim)
+
+          @available = false
+          true
+        end
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_authorization.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Smith
+  class Tool < RubyLLM::Tool
+    module ExecutionAuthorization
+      private
+
+      def authorize_tool_execution!
+        claim = self.class.__send__(:current_tool_dispatch_claim)
+        return unless claim
+        return if ExecutionAuthority.current&.claim(self, claim)
+
+        raise ToolExecutionNotAdmitted, "managed Smith tool execution requires its exact admitted dispatch authority"
+      end
+    end
+  end
+end

data/lib/smith/tool/execution_batch.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require "dry-initializer"
+
+module Smith
+  class Tool < RubyLLM::Tool
+    class ExecutionBatch
+      extend Dry::Initializer
+
+      option :context
+      option :tool_calls
+      option :source_calls
+      option :dispatch_calls
+      option :source_calls_by_dispatch
+      option :requests
+      option :targets
+      option :call_admissions, default: proc { {}.compare_by_identity.freeze }
+      option :call_reservation, optional: true
+
+      attr_reader :capture_failures, :fatal_failures, :notification_failures, :terminal_failures
+
+      def initialize(...)
+        super
+        @state = ExecutionBatchState.new(requests:)
+        @mutex = Mutex.new
+        initialize_failure_queues
+        @host_admitted = false
+        @host_admission_required = context.fetch(:current_invocation_batch_admitter) && requests.any?
+      end
+
+      def request_for(tool_call) = requests[tool_call]
+
+      def target_for(tool_call) = targets[tool_call]
+
+      def admission_for(tool_call) = call_admissions[tool_call]
+
+      def dispatch_for(tool_call) = requests.key?(tool_call) ? tool_call : dispatch_calls[tool_call]
+
+      def source_for(tool_call) = source_calls_by_dispatch[tool_call]
+
+      def claim_dispatch!(tool_call) = @state.claim_dispatch!(tool_call)
+
+      def verify_dispatch!(tool_call, current_tool:, claim:)
+        request = requests.fetch(tool_call)
+        target = targets.fetch(tool_call)
+        return if dispatch_admitted? &&
+                  @state.dispatch_claimed?(tool_call, claim) &&
+                  target_unchanged?(current_tool, target, request) &&
+                  call_unchanged?(tool_call, request)
+
+        raise ToolDispatchRejected, "admitted tool invocation changed before dispatch"
+      end
+
+      def mark_started!(tool_call, claim:) = @state.mark_started!(tool_call, claim:)
+
+      def mark_executed!(tool_call, claim:) = @state.mark_executed!(tool_call, claim:)
+
+      def started?(tool_call) = @state.started?(tool_call)
+
+      def mark_host_admitted!
+        @mutex.synchronize { @host_admitted = true }
+      end
+
+      def host_admitted?
+        @mutex.synchronize { @host_admitted }
+      end
+
+      def claim_failure_request(tool_call) = @state.claim_failure_request(tool_call)
+
+      def claim_unsettled_request = @state.claim_unsettled_request
+
+      def complete_failure_notification!(tool_call) = @state.complete_failure_notification!(tool_call)
+
+      def release_failure_notification!(tool_call, state) = @state.release_failure_notification!(tool_call, state)
+
+      def settle! = call_reservation&.settle!
+
+      private
+
+      def initialize_failure_queues
+        @capture_failures = Queue.new
+        @fatal_failures = Queue.new
+        @notification_failures = Queue.new
+        @terminal_failures = Queue.new
+      end
+
+      def dispatch_admitted? = !@host_admission_required || host_admitted?
+
+      def target_unchanged?(current_tool, target, request)
+        current_tool.equal?(target) && current_tool.instance_of?(request.tool_class)
+      end
+
+      def call_unchanged?(tool_call, request)
+        tool_call.id.equal?(request.invocation.tool_call_id) &&
+          tool_call.name.to_s == request.invocation.tool_name &&
+          tool_call.arguments.equal?(request.arguments)
+      end
+    end
+  end
+end