smith-agents 0.13.1 → 0.14.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 7c3bc332fe08b2ff077b3039603ddc8356ac65d5034f47abcf4253c3e167f8c6
-  data.tar.gz: 9eaf0e8548b473da25838206b63e2a5639d7ee4391f7cb6e30980a5b5aa46e67
+  metadata.gz: d24b9d2b16aa4bc15335adf1ac519a2b1979306e78c5adbd9d79d7c4c37b4ae1
+  data.tar.gz: f89dc27bdede7609559dd8b943c6f8a8ca430d8b65384cb2bee098589d9389b5
 SHA512:
-  metadata.gz: 1ff7f020acb2efd11f7143a074d3c65093137be5bc2646e10a668e1e5242b37255db15909f0c99aa3b6073ad73e31a2557f61fa3d2a35d2de716ba122d7deb4e
-  data.tar.gz: fab4fd19cd33c21b6800a8d75740bc62533aa4100a2d9ae5e26bad31ad5a1fa57b4869cce9efc9a351b41366ee5ee5a8cb24c0d21dc8369effdbe5b8fd8e703c
+  metadata.gz: 568b8f7e26296999a607319f95cebd0f6eab757a6bc109449b4a1e53738f74d57bbad20c13bfe02949204ea3e428ce19135b0713b8cb57d0f15c2f4806ec13ec
+  data.tar.gz: 403c1afd783f9a811b8c9a3e1f8ecadef3ab340b4250f75e73ee596104ea5fa9246bdcb964ca01f4ce55fc9b4861a59d528a25625e37b23d43834248024c4602

data/CHANGELOG.md

@@ -4,6 +4,58 @@ All notable changes to Smith are documented in this file.
 
 Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). Versioning follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). Smith is pre-1.0 and under active development; expect occasional contract tightening between minor versions until 1.0.
 
+## [0.14.0] - 2026-11-09
+
+### Added
+
+- `optimize` takes `mode:` (`:sync`, `:batch`, or a callable of the context
+  and attempt) and `on_request_change:`, as `execute` and `fan_out` do. In
+  batch mode each call of the loop, a round's generator and then its
+  evaluator, is a batch request of its own, queued under a wait of its own,
+  and the run waits for its result without holding a worker; the pending
+  batch keeps the loop's `progress` (the round, the call it waits for, the
+  candidate, feedback and last score, and the verdicts so far, as JSON gives
+  them). A delivery runs the step again at its attempt and the loop goes on
+  from that progress: the call it waited for takes its result, nothing
+  before it runs again, and the next call is batched in turn, so each call is
+  made, charged and held on a budget ledger once. The mode callable is asked
+  at each call. A failed or missing result fails its call, and `retry_on`
+  runs the step again with its loop afresh, as any retry of an optimize step
+  does, making and paying for the earlier rounds' calls again. Progress that
+  is not JSON, or that would take a stored run past `persistence_max_bytes`,
+  fails the step before anything is queued. Restore checks the progress (its
+  round, call, last score and kept verdicts) and that it is the progress of
+  the call the run's wait was made for, and refuses any other with
+  `Smith::SerializationError`; lowering `max_rounds` or dropping batch mode
+  therefore strands runs waiting under the old setting. Runtime readiness
+  checks the generator and the evaluator for a batch interface.
+
+### Changed
+
+- A batch wait's results settle their holds on a host's budget ledger before
+  the step runs again, so the request the step holds next (a retry's, or an
+  optimize step's next call) never overlaps them; a ledger with room for
+  one hold no longer refuses that request. Restore refuses a pending batch
+  whose attempt is not the one its wait's token was made for
+  (`Smith::SerializationError`); a run stored by Smith 0.13 restores as it
+  is.
+
+### Upgrading
+
+- Restore checks that a stored batch wait's pending batch is the one its
+  token was made for (its attempt, and an optimize step's round and call).
+  Runs stored by Smith 0.13 pass; a pending batch edited by hand, or moved
+  to another attempt, is refused with `Smith::SerializationError`.
+- A host ledger now sees a delivered result's hold settled before the
+  step's next hold is taken; a ledger that counted both at once no longer
+  needs room for them together.
+- While runs wait on an optimize step's batch, keep its `max_rounds` and its
+  batch mode as they were: restore refuses progress the step could not go
+  on from, so lowering one or dropping the other strands those runs until
+  their waits have drained.
+- `Smith::EXECUTION_SEMANTICS_VERSION` stays `"7"`: a workflow that does not
+  declare batch mode on an optimize step runs and stores as before.
+
 ## [0.13.1] - 2026-11-08
 
 ### Changed

data/README.md

@@ -30,7 +30,7 @@ What you get:
 
 ```ruby
 # Gemfile
-gem "smith-agents", "0.13.1", require: "smith"
+gem "smith-agents", "0.14.0", require: "smith"
 ```
 
 In a Rails app:
@@ -206,7 +206,8 @@ anything is sent, and a step can wait for budget instead.
 register as the model writes it. Treat it as a draft: Smith checks the reply
 only once it is complete. `execute :triager, mode: :batch` sends the request
 through Anthropic's batch API at half price; the run waits, and your app
-delivers the result when the batch ends.
+delivers the result when the batch ends. `fan_out` and `optimize` take the same
+`mode:`; an `optimize` step batches each call of its loop in turn.
 [`docs/AGENTS.md`](docs/AGENTS.md)
 
 ## Untrusted content

data/docs/AGENTS.md

@@ -219,7 +219,7 @@ end
 ```
 
 `mode:` is `:sync` (the default), `:batch`, or a callable of the context and
-the step attempt returning one, on `execute` and `fan_out`; a batch step
+the step attempt returning one, on `execute`, `fan_out` and `optimize`; a batch step
 declares `on_failure`, and cannot stream or be `parallel:`. Smith prepares
 the call as a synchronous one, admits it against the budget, and hands the
 exact request (the provider payload and the call's own headers, never the
@@ -253,3 +253,34 @@ built in, and a batch request is one turn, so a batching agent has no tools;
 runtime readiness reports a missing queue (`:batch_queue_missing`) and an
 agent that cannot batch (`:batch_unavailable`), and graph inspection shows the
 mode.
+
+An `optimize` step in batch mode sends each call of its loop as a batch
+request of its own, a round's generator and then its evaluator, since the
+evaluator reads what the generator wrote: a round is two waits. Each call is
+queued under a wait of its own (its round and call tell the token apart), and
+the run's pending batch keeps the loop's progress (`progress`): the round, the
+call it waits for, the candidate, the feedback and last score, and the
+verdicts so far, as JSON gives them. A candidate or feedback that is not JSON
+fails the step before anything is queued, as does progress that would take a
+stored run past `persistence_max_bytes` (`Smith::PersistedStateTooLarge`). A
+delivery settles its result's hold, then runs the step again at its attempt,
+and the loop goes on from that progress: the call it waited for takes its
+result, nothing before it runs again (`before_eval` included), and the next
+call is batched in turn, so each call is made, charged and held on a budget
+ledger once. The mode callable is asked at each call; a call it gives `:sync`
+goes over the wire and the loop goes on.
+
+A failed or missing result fails its call as on any batch step, and
+`retry_on` (in process, as an optimize step takes no durable retry) runs the
+step again: as on any retry of an optimize step, the loop starts afresh, so
+the new attempt makes, and pays for, the earlier rounds' calls again, under
+the mode the callable gives it (`attempt == 1 ? :batch : :sync` sends the
+whole loop over the wire). A delivery that crashes leaves the stored wait as
+it was: in lax mode the host delivers the results again and the loop goes on
+as it would have; in strict mode the delivery was stored as a step under way,
+so the run raises `Smith::StepInProgressOnRestore`, as any batch step's does,
+and the host decides. Restore refuses progress the step could not go on from
+(a round past `max_rounds`, or a call the run's wait was not made for) with
+`Smith::SerializationError`, so lowering `max_rounds`, or dropping batch mode,
+strands the runs waiting under the old setting: keep both until those waits
+have drained.

data/docs/PATTERNS.md

@@ -403,6 +403,10 @@ value, including a Hash with `accept: true`, is ignored and the evaluator
 judges as before. A deterministic check that makes a candidate unacceptable
 can reject it this way without spending an evaluator call.
 
+`mode: :batch` (or a callable of the context and attempt) sends each call of
+the loop through the provider's batch API, the run waiting for each result
+and going on from where the loop stood: docs/AGENTS.md, Batch mode.
+
 Every round's verdict is an output. Once its loop has run, the step's record
 carries its verdicts under `:evaluations`, in attempt and round order, whether
 the step completed or failed (a step that failed before its loop, or any other

data/lib/smith/version.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
 module Smith
-  VERSION = "0.13.1"
+  VERSION = "0.14.0"
   EXECUTION_SEMANTICS_VERSION = "7"
 end

data/lib/smith/workflow/batch_charges.rb

@@ -16,17 +16,21 @@ module Smith
         Thread.current[:smith_call_reservation] = ledger.reserve_spent!(estimates)
       end
 
-      def charge_untaken_batch_results(transition)
+      def charge_untaken_batch_results(transition, delivered)
         @batch_delivery.untaken.each do |branch_key, entry|
           next unless entry.result["type"] == "succeeded"
 
           Attribution.with(branch_key:) do
-            charge_batch_result(batch_agent_class(transition, branch_key), entry.result.fetch("message"))
+            charge_batch_result(batch_agent_class(transition, branch_key, delivered), entry.result.fetch("message"))
           end
         end
       end
 
-      def batch_agent_class(transition, branch_key)
+      # The agent whose call a result answers: a fan-out branch's, an
+      # optimize step's generator or evaluator (the call its pending batch
+      # `delivered` waited for), or the step's own.
+      def batch_agent_class(transition, branch_key, delivered)
+        return optimizer_batch_agent(transition, delivered.fetch("progress")) if transition.optimized?
         return resolve_agent_class(transition) unless branch_key
 
         resolve_fanout_agent_class(transition, transition.fanout_config.fetch(:branches).fetch(branch_key))

data/lib/smith/workflow/batch_execution.rb

@@ -17,15 +17,19 @@ module Smith
     module BatchExecution
       # Batches end within 24 hours; the wait allows one more.
       WAIT_SECONDS = 25 * 3600
-      # The requests an attempt captured, for the step to queue.
-      Captured = Data.define(:token, :requests, :attempt, :key_generation)
+      # The requests an attempt captured, for the step to queue, and an
+      # optimize step's `progress` (OptimizerBatch).
+      Captured = Data.define(:token, :requests, :attempt, :key_generation, :progress) do
+        def initialize(progress: nil, **) = super
+      end
 
       private
 
       # The attempt sends its requests as a batch: its mode says so and no
-      # result of it is being delivered.
+      # result of it is being delivered. An optimize step batches each of
+      # its loop's calls instead (OptimizerBatch).
       def batch_send?(transition)
-        transition.batches? && !delivering_batch? && batch_mode_for(transition) == :batch
+        transition.batches? && !transition.optimized? && !delivering_batch? && batch_mode_for(transition) == :batch
       end
 
       # A delivery runs only the stored batch wait's step (BatchWait).
@@ -45,10 +49,11 @@ module Smith
 
       # An agent step's batch: its one call.
       def send_agent_batch(transition, agent_class, prepared_input)
-        send_batch(transition, nil => lambda do
+        capture = lambda do
           admit_batch_call { reserve_for_serial(transition, effective_call_ledger) }
           with_agent_context(agent_class) { capture_batch_request(agent_class, prepared_input) }
-        end)
+        end
+        send_batch(transition, { nil => capture })
       end
 
       # A fan-out's batch: each branch's call, captured in its own branch
@@ -82,11 +87,11 @@ module Smith
 
       # The request the agent's call would send (BatchExchange throws it out
       # of the call before anything is sent), and its hold.
-      def capture_batch_request(agent_class, prepared_input)
+      def capture_batch_request(agent_class, prepared_input, **call)
         @batch_capture = nil
         request = catch(Providers::BatchExchange::CAPTURED) do
           @batch_capturing = true
-          invoke_agent(agent_class, prepared_input)
+          invoke_agent(agent_class, prepared_input, **call)
           nil
         ensure
           @batch_capturing = false
@@ -94,13 +99,15 @@ module Smith
         [request || raise(BatchModeUnavailable, "a batch step's agent made no provider request"), @batch_capture.hold]
       end
 
-      def send_batch(transition, captures)
+      # `position` tells apart the waits of one attempt (an optimize step's
+      # calls), `progress` is its loop's.
+      def send_batch(transition, captures, position: [], progress: nil)
         refuse_batch_wait!
-        token = wait_token_for(transition, ["batch", @step_attempt || 1])
+        token = wait_token_for(transition, ["batch", @step_attempt || 1, *position])
         captured = captures.transform_values(&:call)
         hold_batch!(token, captured.transform_values(&:last))
         Captured.new(token:, requests: captured.transform_values(&:first), attempt: @step_attempt || 1,
-                     key_generation: @step_key_generation || 0)
+                     key_generation: @step_key_generation || 0, progress:)
       end
 
       # The run can wait for a batch: a queue takes its requests, it is not
@@ -141,8 +148,11 @@ module Smith
       end
 
       def pending_batch_record(transition, captured, ids)
-        { "transition" => transition.name.to_s, "attempt" => captured.attempt,
-          "key_generation" => captured.key_generation, "requests" => ids.transform_values { Batch.digest(_1) } }.freeze
+        record = { "transition" => transition.name.to_s, "attempt" => captured.attempt,
+                   "key_generation" => captured.key_generation,
+                   "requests" => ids.transform_values { Batch.digest(_1) } }
+        record["progress"] = captured.progress if captured.progress
+        record.freeze
       end
     end
   end

data/lib/smith/workflow/batch_holds.rb

@@ -40,12 +40,13 @@ module Smith
       # Each result settles the hold its request took on a host's budget
       # ledger (BatchExecution), taken by a call or not: a succeeded one
       # reconciles what it cost, any other (or a wait that expired) releases.
-      def settle_batch_holds(transition, token, entries)
+      # `delivered` is the pending batch the results answer.
+      def settle_batch_holds(transition, token, entries, delivered)
         ledger = settling_ledger
         return unless ledger
 
         entries.each do |branch_key, entry|
-          spent = batch_spent(transition, branch_key, entry.result)
+          spent = batch_spent(transition, branch_key, entry.result, delivered)
           Budget::HostLedger.settle(ledger, Batch.custom_id(token, branch_key), spent)
         end
       end
@@ -60,12 +61,12 @@ module Smith
         nil
       end
 
-      def batch_spent(transition, branch_key, result)
+      def batch_spent(transition, branch_key, result, delivered)
         return unless result["type"] == "succeeded"
 
         message = result.fetch("message")
         usage = batch_reply_usage(message)
-        model = batch_model_id(batch_agent_class(transition, branch_key), message)
+        model = batch_model_id(batch_agent_class(transition, branch_key, delivered), message)
         Budget::HostLedger.amount(usage, model:, provider: "anthropic", mode: :batch) if usage
       end
     end

data/lib/smith/workflow/batch_wait.rb

@@ -16,8 +16,10 @@ module Smith
     # by a sibling's failure) is charged all the same (BatchCharges).
     module BatchWait
       PENDING_BATCH_FIELDS = %w[attempt key_generation requests transition].freeze
+      # An optimize step's pending batch also keeps its loop's progress.
+      PENDING_PROGRESS_FIELDS = (PENDING_BATCH_FIELDS + %w[progress]).sort.freeze
       DIGEST = /\A\h{64}\z/
-      private_constant :PENDING_BATCH_FIELDS, :DIGEST
+      private_constant :PENDING_BATCH_FIELDS, :PENDING_PROGRESS_FIELDS, :DIGEST
 
       private
 
@@ -55,17 +57,19 @@ module Smith
 
       # The step runs again at the batch's attempt, each call taking its
       # result. It may enter another wait (a batch sent again, or a durable
-      # retry's backoff); otherwise it counts. A completion that raises
-      # leaves the run waiting.
+      # retry's backoff); otherwise it counts. The results settle their
+      # holds first, so a request the step holds next never overlaps them. A
+      # completion that raises leaves the run waiting.
       def complete_batch_wait(transition, entries, done)
         record = @awaiting
+        delivered = @pending_batch
         @awaiting = nil
         @batch_delivery = Batch::Delivery.new(entries, on_change: transition.on_request_change || :keep)
+        settle_batch_holds(transition, record.fetch("token"), entries, delivered)
         step = with_step_context(transition) do
           Execution.instance_method(:execute_step_body).bind_call(self, transition)
         ensure
-          settle_batch_holds(transition, record.fetch("token"), entries)
-          charge_untaken_batch_results(transition)
+          charge_untaken_batch_results(transition, delivered)
         end
         settle_batch_wait(step)
         return done if wait_pending?(step)
@@ -99,26 +103,49 @@ module Smith
       # A persisted pending batch names the stored batch wait's transition
       # (which WaitRecord checked is still a batch step), an attempt its
       # retry policy allows, a key generation only a declared
-      # `new_idempotency_key_on` moves, and the digest of each request the
-      # step's branches send; a stored batch wait needs one. Anything else
-      # raises SerializationError.
+      # `new_idempotency_key_on` moves, the digest of each request the
+      # step's branches send, and for an optimize step its loop's progress
+      # (OptimizerBatch); a stored batch wait needs one, the one its token
+      # was made for. Anything else raises SerializationError.
       def restore_pending_batch(value)
         return if value.nil? && !batch_waiting?
 
         record = value.is_a?(Hash) ? value.transform_keys(&:to_s) : {}
+        check_pending_batch!(record)
+        record.transform_values(&:freeze).freeze
+      end
+
+      def check_pending_batch!(record)
         raise Smith::SerializationError, "persisted workflow pending_batch is invalid" unless
           well_formed_pending_batch?(record)
+        return if batch_wait_made_for?(record)
 
-        record.transform_values(&:freeze).freeze
+        raise Smith::SerializationError, "persisted workflow pending_batch is not the one its wait was made for"
+      end
+
+      # The wait's token is the one the step makes for the record's attempt
+      # (and an optimize step's round and call), so no stored record moves a
+      # wait to another attempt or call.
+      def batch_wait_made_for?(record)
+        transition = self.class.transition_named(@awaiting["transition"])
+        position = transition.optimized? ? record["progress"].values_at("round", "role") : []
+        wait_token_for(transition, ["batch", record["attempt"], *position]) == @awaiting["token"]
       end
 
       def well_formed_pending_batch?(record)
         transition = self.class.transition_named(@awaiting["transition"]) if batch_waiting?
-        transition && record.keys.sort == PENDING_BATCH_FIELDS && record["transition"] == transition.name.to_s &&
-          allowed_batch_position?(transition, *record.values_at("attempt", "key_generation")) &&
-          batch_digests?(transition, record["requests"])
+        transition && record.keys.sort == pending_batch_fields(transition) &&
+          record["transition"] == transition.name.to_s && pending_batch_parts?(transition, record)
+      end
+
+      def pending_batch_parts?(transition, record)
+        allowed_batch_position?(transition, *record.values_at("attempt", "key_generation")) &&
+          batch_digests?(transition, record["requests"]) &&
+          (!transition.optimized? || optimizer_progress?(transition, record["progress"], record["attempt"]))
       end
 
+      def pending_batch_fields(transition) = transition.optimized? ? PENDING_PROGRESS_FIELDS : PENDING_BATCH_FIELDS
+
       def allowed_batch_position?(transition, attempt, generation)
         config = transition.retry_config
         attempt.is_a?(Integer) && attempt.between?(1, config&.fetch(:attempts) || 1) && generation.is_a?(Integer) &&

data/lib/smith/workflow/evaluator_optimizer.rb

@@ -5,10 +5,13 @@ require "json"
 require_relative "agent_result"
 require_relative "execution_result_snapshot"
 require_relative "optimization_state"
+require_relative "optimizer_batch"
 
 module Smith
   class Workflow
     module EvaluatorOptimizer
+      include OptimizerBatch
+
       private
 
       def execute_optimization_step(transition, prepared_input: nil)
@@ -26,15 +29,21 @@ module Smith
           role: :evaluator
         )
         pend_evaluations
-        run_optimization_loop(state)
+        run_optimization_loop(state, transition)
       end
 
-      def run_optimization_loop(state)
-        state.config[:max_rounds].times do |round|
+      # From round 0, or from the round and call a batch delivery goes on
+      # from (OptimizerBatch).
+      def run_optimization_loop(state, transition)
+        first, call = resume_optimizer_progress(state)
+        ((first || 0)...state.config[:max_rounds]).each do |round|
+          at_evaluator = round == first && call == "evaluator"
           # The overlay scopes every trace and usage fact from this round's
           # generator and evaluator calls to the round that produced them.
           result = Attribution.with(round: round) do
-            IdempotencyKey.__send__(:within, :round, round) { run_optimization_round(state, round) }
+            IdempotencyKey.__send__(:within, :round, round) do
+              run_optimization_round(state, round, transition, at_evaluator:)
+            end
           end
           return result if result
         end
@@ -44,12 +53,13 @@ module Smith
       end
 
       # Each role is a position of its own in the round, so the generator,
-      # the before_eval callback and the evaluator never share a key.
-      def run_optimization_round(state, round)
-        IdempotencyKey.__send__(:within, :generator) { generate_candidate!(state, round) }
-        rejection = IdempotencyKey.__send__(:within, :before_eval) { before_eval_rejection(state) }
-        evaluation = rejection ||
-                     normalize_evaluation(IdempotencyKey.__send__(:within, :evaluator) { evaluate_candidate(state) })
+      # the before_eval callback and the evaluator never share a key. A
+      # round a delivery goes on from at its evaluator has its candidate.
+      def run_optimization_round(state, round, transition, at_evaluator: false)
+        rejection = generate_and_screen(state, round, transition) unless at_evaluator
+        evaluation = rejection || normalize_evaluation(
+          IdempotencyKey.__send__(:within, :evaluator) { evaluate_candidate(state, round, transition) }
+        )
         validate_evaluation_structure!(evaluation)
         validate_evaluation_fields!(evaluation, state.config)
         record_evaluation(state, evaluation, round, rejection ? :before_eval : :evaluator)
@@ -201,15 +211,19 @@ module Smith
         end
       end
 
-      def generate_candidate!(state, round)
+      # The round's candidate, and a before_eval rejection of it, if any.
+      def generate_and_screen(state, round, transition)
+        IdempotencyKey.__send__(:within, :generator) { generate_candidate!(state, round, transition) }
+        IdempotencyKey.__send__(:within, :before_eval) { before_eval_rejection(state) }
+      end
+
+      def generate_candidate!(state, round, transition)
         input = prepare_generator_input(state.prepared_input, round, state.candidate, state.feedback)
-        result = invoke_agent_with_budget(state.generator_class, input)
-        state.candidate = result
+        state.candidate = optimizer_call(transition, state, round, "generator", input)
       end
 
-      def evaluate_candidate(state)
-        input = build_evaluator_input(state)
-        invoke_with_evaluator_schema(state.evaluator_class, state.config[:evaluator_schema], input)
+      def evaluate_candidate(state, round, transition)
+        optimizer_call(transition, state, round, "evaluator", build_evaluator_input(state))
       end
 
       # evaluator_context: :inject_state appends the candidate as a
@@ -248,10 +262,6 @@ module Smith
         evaluation if evaluation.key?(:accept) && evaluation[:accept].equal?(false)
       end
 
-      def invoke_with_evaluator_schema(evaluator_class, schema, input)
-        invoke_agent_with_budget(evaluator_class, input, output_schema: schema)
-      end
-
       def invoke_agent_with_budget(agent_class, prepared_input, output_schema: agent_class.output_schema)
         Thread.current[:smith_last_agent_result] = nil
         clear_failed_billable_attempts

data/lib/smith/workflow/execution.rb

@@ -76,7 +76,7 @@ module Smith
         prepared_input = step_prepared_input(transition)
         return send_agent_batch(transition, agent_class, prepared_input) if batch_send?(transition)
 
-        guarded_agent_output(transition, agent_class, prepared_input)
+        optimizer_batch_output(transition) { guarded_agent_output(transition, agent_class, prepared_input) }
       end
 
       def guarded_agent_output(transition, agent_class, prepared_input)

data/lib/smith/workflow/graph/batch_diagnostic.rb

@@ -27,6 +27,7 @@ module Smith
 
         def agents(transition)
           return transition.fanout_config.fetch(:branches).values if transition.fanout?
+          return transition.optimization_config.values_at(:generator, :evaluator) if transition.optimized?
 
           [transition.agent_name]
         end

data/lib/smith/workflow/optimizer_batch.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Smith
+  class Workflow
+    # Batch mode on an optimize step (`optimize ..., mode: :batch`). Each of
+    # the loop's calls, a round's generator and then its evaluator, is a
+    # batch request of its own, as the evaluator reads what the generator
+    # wrote. The call is captured as an agent step's is (BatchExecution), the
+    # attempt ends, and the run waits for the result, its pending batch
+    # keeping the loop's progress (`progress`: the round, the call it waits for,
+    # the candidate, feedback and last score, and the verdicts so far, as
+    # JSON gives them). A delivery runs the step again at its attempt, and
+    # the loop goes on from that progress: the call it waited for takes its
+    # result, and the next call is batched in turn, so each call is made,
+    # and charged, once. The step's mode is asked at each call.
+    module OptimizerBatch
+      ROLES = %w[generator evaluator].freeze
+      PROGRESS_FIELDS = %w[candidate evaluations feedback last_score role round].freeze
+      RECORD_FIELDS = %w[round source verdict].freeze
+      SOURCES = %w[evaluator before_eval].freeze
+      ATTEMPT = /\A[1-9]\d*\z/
+      CAPTURED = :smith_optimizer_batch
+      private_constant :ROLES, :PROGRESS_FIELDS, :RECORD_FIELDS, :SOURCES, :ATTEMPT, :CAPTURED
+
+      private
+
+      # The step's output, or the call its loop batched, thrown out of the
+      # loop for the step to queue.
+      def optimizer_batch_output(transition, &)
+        return yield unless transition.optimized? && transition.batches?
+
+        catch(CAPTURED, &)
+      end
+
+      # One of the loop's calls (`role` "generator" or "evaluator"): batched
+      # when the step's mode says so and no result of it is being
+      # delivered, else made as ever.
+      def optimizer_call(transition, state, round, role, input)
+        if transition.batches? && !batch_result_due? && batch_mode_for(transition) == :batch
+          throw(CAPTURED, batch_optimizer_call(transition, state, round, role, input))
+        end
+
+        agent_class, output_schema = optimizer_agent(state, role)
+        invoke_agent_with_budget(agent_class, input, output_schema:)
+      end
+
+      def optimizer_agent(state, role)
+        return [state.generator_class, state.generator_class.output_schema] if role == "generator"
+
+        [state.evaluator_class, state.config[:evaluator_schema]]
+      end
+
+      # The call's request, admitted and held as an agent step's is, under
+      # a wait of its own: the round and the call tell its token apart. A
+      # stored run whose progress would take it past persistence_max_bytes
+      # fails here, before anything is queued.
+      def batch_optimizer_call(transition, state, round, role, input)
+        agent_class, output_schema = optimizer_agent(state, role)
+        progress = optimizer_progress(state, round, role)
+        StepCompletion.instance_method(:check_result_storable!).bind_call(self, progress)
+        capture = lambda do
+          admit_batch_call { reserve_serial_budget(effective_call_ledger, agent_budget: agent_class&.budget) }
+          with_agent_context(agent_class) { capture_batch_request(agent_class, input, output_schema:) }
+        end
+        send_batch(transition, { nil => capture }, position: [round, role], progress:)
+      end
+
+      # What a delivery's loop goes on from, as JSON gives it.
+      def optimizer_progress(state, round, role)
+        progress = { "round" => round, "role" => role, "candidate" => state.candidate, "feedback" => state.feedback,
+                     "last_score" => state.last_score, "evaluations" => kept_evaluations }
+        Batch.deep_freeze(JSON.parse(JSON.generate(progress)))
+      rescue JSON::GeneratorError, JSON::NestingError, Encoding::UndefinedConversionError
+        raise WorkflowError, "an optimize step waits for a batch only with JSON values for its candidate and feedback"
+      end
+
+      # Every attempt's verdicts so far, by attempt.
+      def kept_evaluations
+        @pending_evaluations.to_h do |attempt, records|
+          [attempt.to_s, records.map { |record| record.merge(source: record[:source].to_s) }]
+        end
+      end
+
+      # The round and call the loop goes on from, with its state and
+      # verdicts as the pending batch kept them; none unless a result of
+      # this step's loop is being delivered.
+      def resume_optimizer_progress(state)
+        progress = @pending_batch&.fetch("progress", nil) if delivering_batch?
+        return unless progress
+
+        state.candidate, state.feedback, state.last_score =
+          snapshot_value(progress.values_at("candidate", "feedback", "last_score"))
+        restore_kept_evaluations(progress.fetch("evaluations"))
+        state.evaluations = @pending_evaluations.values.last.dup.freeze
+        progress.values_at("round", "role")
+      end
+
+      # Kept verdicts back as the loop records them, in attempt order, the
+      # attempt under way last.
+      def restore_kept_evaluations(kept)
+        restored = kept.to_h do |attempt, records|
+          [attempt.to_i, records.map { kept_record(_1) }]
+        end
+        @pending_evaluations = @pending_evaluations.merge(restored).sort.to_h
+      end
+
+      def kept_record(record)
+        { round: record.fetch("round"), source: record.fetch("source").to_sym,
+          verdict: frozen_verdict(normalize_evaluation(record.fetch("verdict"))) }.freeze
+      end
+
+      # The agent whose call an optimize step's pending batch waited for.
+      def optimizer_batch_agent(transition, progress)
+        role = progress.fetch("role").to_sym
+        resolve_registered_agent!(transition.optimization_config.fetch(role), workflow_class: self.class,
+                                                                              transition_name: transition.name, role:)
+      end
+
+      # A stored pending batch's `progress` (BatchWait): a call of a
+      # round the step runs, a last score that is a number or none, and kept
+      # verdicts of attempts up to the batch's.
+      def optimizer_progress?(transition, progress, attempt)
+        progress.is_a?(Hash) && progress.keys.sort == PROGRESS_FIELDS && optimizer_position?(transition, progress) &&
+          (progress["last_score"].nil? || progress["last_score"].is_a?(Numeric)) &&
+          kept_evaluations?(progress["evaluations"], attempt)
+      end
+
+      def optimizer_position?(transition, progress)
+        round = progress["round"]
+        ROLES.include?(progress["role"]) && round.is_a?(Integer) &&
+          round.between?(0, transition.optimization_config.fetch(:max_rounds) - 1)
+      end
+
+      def kept_evaluations?(kept, attempt)
+        kept.is_a?(Hash) && kept.all? do |key, records|
+          key.is_a?(String) && ATTEMPT.match?(key) && key.to_i <= attempt && records.is_a?(Array) &&
+            records.all? { kept_record?(_1) }
+        end
+      end
+
+      def kept_record?(record)
+        record.is_a?(Hash) && record.keys.sort == RECORD_FIELDS && record["round"].is_a?(Integer) &&
+          !record["round"].negative? && SOURCES.include?(record["source"]) && record["verdict"].is_a?(Hash)
+      end
+    end
+  end
+end

data/lib/smith/workflow/transition.rb

@@ -125,6 +125,8 @@ module Smith
         commit_execution_primitive!("workflow")
       end
 
+      # `mode:` and `on_request_change:` as on `execute`: in batch mode each
+      # of the loop's calls is a batch request of its own (OptimizerBatch).
       def optimize(generator:, evaluator:, max_rounds:, evaluator_schema:,
                    improvement_threshold: nil,
                    evaluator_context: nil,
@@ -133,16 +135,19 @@ module Smith
                    on_converged: :raise,
                    on_threshold: :raise,
                    input: nil,
-                   output_to: nil)
+                   output_to: nil,
+                   mode: nil,
+                   on_request_change: nil)
         validate_execution_primitive_conflict!("optimize")
         assign_step_io!(input, output_to)
+        assign_batch_mode!(mode, on_request_change)
         validate_optimize_controls!(generator, evaluator, max_rounds, evaluator_schema)
         validate_optimize_exit_modes!(on_exhaustion: on_exhaustion, on_converged: on_converged,
                                       on_threshold: on_threshold)
         validate_optimize_evaluator_context!(evaluator_context)
         validate_optimize_before_eval!(before_eval)
 
-        config = {
+        @optimization_config = {
           generator: normalize_agent_reference!(generator, "optimizer generator"),
           evaluator: normalize_agent_reference!(evaluator, "optimizer evaluator"),
           max_rounds: max_rounds,
@@ -153,8 +158,6 @@ module Smith
           on_converged: on_converged,
           on_threshold: on_threshold
         }.freeze
-
-        @optimization_config = config
         commit_execution_primitive!("optimize")
       end
 
@@ -174,8 +177,7 @@ module Smith
       def fan_out(branches:, input: nil, output_to: nil, mode: nil, on_request_change: nil)
         validate_execution_primitive_conflict!("fan_out")
         assign_step_io!(input, output_to)
-        @batch_mode = batch_mode!(mode, stream: nil, parallel: false)
-        @on_request_change = request_change!(on_request_change)
+        assign_batch_mode!(mode, on_request_change)
 
         normalized = normalize_fanout_branches!(branches)
         @fanout_branch_lookup = normalized.to_h do |key, agent|
@@ -323,6 +325,13 @@ module Smith
         mode
       end
 
+      # `mode:` and `on_request_change:` of a step that has no stream and no
+      # parallel copies (`fan_out`, `optimize`).
+      def assign_batch_mode!(mode, on_request_change)
+        @batch_mode = batch_mode!(mode, stream: nil, parallel: false)
+        @on_request_change = request_change!(on_request_change)
+      end
+
       # What a batch result does when the step renders another request by
       # the time it comes: :keep it (the default), :resend the request over
       # the wire, or :fail the step with BatchRequestChanged.

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: smith-agents
 version: !ruby/object:Gem::Version
-  version: 0.13.1
+  version: 0.14.0
 platform: ruby
 authors:
 - Samuel Ralak
@@ -522,6 +522,7 @@ files:
 - lib/smith/workflow/message_value_normalizer.rb
 - lib/smith/workflow/nested_execution.rb
 - lib/smith/workflow/optimization_state.rb
+- lib/smith/workflow/optimizer_batch.rb
 - lib/smith/workflow/orchestration_state.rb
 - lib/smith/workflow/orchestrator_worker.rb
 - lib/smith/workflow/parallel.rb