smith-agents 0.11.0 → 0.12.0

data/lib/smith/batch/message_shape.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Smith
+  module Batch
+    # Whether a succeeded result's message has the shape of an Anthropic
+    # Messages reply, as the provider's parser reads it: `content` an Array of
+    # blocks, each a Hash with a String `type` (a text block's `text` a
+    # String), and `usage` a Hash whose token counts are non-negative
+    # Integers. Anything else is refused before the step runs.
+    module MessageShape
+      COUNTS = %w[input_tokens output_tokens].freeze
+      OPTIONAL_COUNTS = %w[cache_read_input_tokens cache_creation_input_tokens].freeze
+      private_constant :COUNTS, :OPTIONAL_COUNTS
+
+      module_function
+
+      def valid?(message)
+        message.is_a?(Hash) && message["content"].is_a?(Array) && message["content"].all? { block?(_1) } &&
+          usage?(message["usage"])
+      end
+
+      def block?(block)
+        return false unless block.is_a?(Hash) && block["type"].is_a?(String)
+
+        block["type"] != "text" || block["text"].is_a?(String)
+      end
+
+      def usage?(usage)
+        usage.is_a?(Hash) && COUNTS.all? { count?(usage[_1]) } &&
+          OPTIONAL_COUNTS.all? { usage[_1].nil? || count?(usage[_1]) }
+      end
+
+      def count?(value) = value.is_a?(Integer) && !value.negative?
+    end
+  end
+end

data/lib/smith/batch.rb

@@ -0,0 +1,166 @@
+# frozen_string_literal: true
+
+require "digest"
+require "json"
+
+require_relative "diagnostic_text"
+require_relative "batch/message_shape"
+
+module Smith
+  # Batch mode (`execute Agent, mode: :batch`, and the same on `fan_out`): a
+  # step's provider request is sent through the host's batch queue instead of
+  # over the wire, and the step waits for its result. Smith renders exactly
+  # the request a synchronous call would post, hands it to
+  # `Smith.config.batch_queue.enqueue(custom_id:, await_token:, request:,
+  # execution_key:)` and waits on `:batch_result`; the host submits batches
+  # (Smith::Providers::Anthropic::Batches), collects their results, and
+  # resumes the run with every result of its wait. Smith never submits or
+  # polls, and waits for the results without sleeping (an in-process
+  # `retry_on` backoff sleeps as on any step; `durable: true` waits instead).
+  module Batch
+    EVENT = "batch_result"
+    RESULT_TYPES = %w[succeeded errored canceled expired].freeze
+    MISSING = "missing"
+    MAX_REQUEST_BYTES = 32 * 1024 * 1024
+    MAX_RESULTS_BYTES = 64 * 1024 * 1024
+    INVALID_REQUEST = "invalid_request_error"
+    CREDENTIAL_HEADERS = %w[authorization x-api-key x-goog-api-key].freeze
+
+    module_function
+
+    # The request a batch carries for one call: the provider, the payload it
+    # would post, and the call's own headers (never credentials: the
+    # connection's stay in it, and a call's own are refused), JSON-shaped,
+    # bounded and frozen.
+    def request(provider:, params:, headers:)
+      if headers.to_h.keys.any? { CREDENTIAL_HEADERS.include?(_1.to_s.downcase) }
+        raise WorkflowError, "a batch request cannot carry credentials in its headers"
+      end
+
+      encoded = JSON.generate("provider" => provider.to_s, "params" => params, "headers" => headers.to_h)
+      if encoded.bytesize > MAX_REQUEST_BYTES
+        raise WorkflowError, "a batch request is #{encoded.bytesize} bytes, past the #{MAX_REQUEST_BYTES} Smith sends"
+      end
+
+      deep_freeze(JSON.parse(encoded))
+    end
+
+    # What identifies a request's content: a SHA-256 of its payload and
+    # headers as JSON gives them, so the request a delivery answers can be
+    # told from one the step would render now.
+    def digest(request) = digest_of(request.fetch("params"), request.fetch("headers"))
+
+    # Whether a call renders a request other than the one sent (`digest`).
+    # Nothing is sent, so nothing in it is refused; one JSON cannot give
+    # counts as changed.
+    def changed?(digest, params, headers)
+      digest_of(params, headers.to_h) != digest
+    rescue JSON::GeneratorError, JSON::NestingError, Encoding::UndefinedConversionError
+      true
+    end
+
+    def digest_of(params, headers) = Digest::SHA256.hexdigest(JSON.generate([params, headers]))
+
+    # A step's request id: its wait token, and for a fan-out branch a key of
+    # the token and the branch (64 hex characters, a valid custom_id).
+    def custom_id(token, branch_key = nil)
+      branch_key.nil? ? token : Digest::SHA256.hexdigest("#{token}/#{branch_key}")
+    end
+
+    # The delivered results, checked: one per expected id, each a result of
+    # a known type, keys as JSON gives them. Raises ArgumentError otherwise.
+    def results!(payload, expected_ids)
+      results = stringified(payload)
+      unless results.keys.sort == expected_ids.sort
+        raise ArgumentError, "a batch delivery must hold exactly the results of its #{expected_ids.length} requests"
+      end
+
+      results.each_value { |result| check_result!(result) }
+      deep_freeze(results)
+    end
+
+    # The message of a succeeded result, or the error a failed one raises:
+    # BatchRequestInvalid for a request the provider refused, else
+    # BatchResultFailed (retryable).
+    def message!(result)
+      case result.fetch("type")
+      when "succeeded" then result.fetch("message")
+      when "errored" then raise_errored!(result)
+      else raise BatchResultFailed.new(kind: result.fetch("type"))
+      end
+    end
+
+    def raise_errored!(result)
+      error = result["error"]
+      type = error.dig("error", "type") || error["type"] if error.is_a?(Hash)
+      raise BatchRequestInvalid if type == INVALID_REQUEST
+
+      raise BatchResultFailed.new(kind: "errored", error_type: type.is_a?(String) ? type : nil)
+    end
+
+    def check_result!(result)
+      raise ArgumentError, "a batch result must be a Hash with a type" unless result.is_a?(Hash)
+
+      type = result["type"]
+      raise ArgumentError, "a batch result's type must be one of #{RESULT_TYPES.join(", ")}" unless
+        RESULT_TYPES.include?(type) || type == MISSING
+      return unless type == "succeeded" && !MessageShape.valid?(result["message"])
+
+      raise ArgumentError, "a succeeded batch result must carry a message of the provider's reply shape"
+    end
+
+    def stringified(payload)
+      raise ArgumentError, "a batch delivery must be a Hash of request ids to results" unless payload.is_a?(Hash)
+
+      encoded = JSON.generate(payload)
+      raise ArgumentError, "a batch delivery is past #{MAX_RESULTS_BYTES} bytes" if encoded.bytesize > MAX_RESULTS_BYTES
+
+      JSON.parse(encoded)
+    rescue JSON::GeneratorError, JSON::NestingError, Encoding::UndefinedConversionError
+      raise ArgumentError, "a batch delivery must be JSON-shaped"
+    end
+
+    def deep_freeze(value)
+      case value
+      when Hash then value.each_value { deep_freeze(_1) }.freeze
+      when Array then value.each { deep_freeze(_1) }.freeze
+      else value.freeze
+      end
+    end
+
+    # The results a resumed step takes, one per branch (nil for a step
+    # without branches), each with the digest of the request it answers, and
+    # what a result does when its call renders another request
+    # (`on_change`). Each is taken once: a follow-up call the step makes
+    # after it (a repair) goes over the wire. `untaken` is what no call took
+    # (a branch cancelled by a sibling's failure) or set aside (a result its
+    # call would not use), paid for all the same.
+    class Delivery
+      Entry = Struct.new(:result, :digest)
+
+      attr_reader :on_change
+
+      def initialize(entries_by_branch, on_change: :keep)
+        @entries = entries_by_branch.dup
+        @set_aside = {}
+        @on_change = on_change
+        @lock = Mutex.new
+      end
+
+      def take(branch_key) = @lock.synchronize { @entries.delete(branch_key) }
+
+      def for?(branch_key) = @lock.synchronize { @entries.key?(branch_key) }
+
+      def set_aside(branch_key, entry) = @lock.synchronize { @set_aside[branch_key] = entry }
+
+      def untaken
+        @lock.synchronize do
+          @set_aside.merge(@entries).tap do
+            @entries.clear
+            @set_aside.clear
+          end
+        end
+      end
+    end
+  end
+end

data/lib/smith/budget/amount.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Smith
+  module Budget
+    # What a provider request may spend (the hold a host's budget ledger is
+    # asked to admit) or did spend (what it is reconciled with): its input
+    # tokens (uncached, cache reads and cache writes together), its output
+    # tokens (thinking included), and its cost from the pricing catalog, nil
+    # when the model has no price there.
+    Amount = Data.define(:input_tokens, :output_tokens, :cost) do
+      def initialize(input_tokens:, output_tokens:, cost:)
+        unless [input_tokens, output_tokens].all? { _1.is_a?(Integer) && !_1.negative? }
+          raise ArgumentError, "a budget amount's token counts must be non-negative Integers"
+        end
+        unless cost.nil? || (cost.is_a?(Numeric) && cost.finite? && !cost.negative?)
+          raise ArgumentError, "a budget amount's cost must be nil or a non-negative number"
+        end
+
+        super
+      end
+
+      def tokens = input_tokens + output_tokens
+    end
+  end
+end

data/lib/smith/budget/continuation.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Smith
+  module Budget
+    # The rest of a tool loop that paused for budget: the messages a step's
+    # call added before a host's budget ledger refused a later round (the
+    # model's tool calls, their results, its thinking and its signatures),
+    # kept with the step's retry so that the retry continues the loop from
+    # them instead of running its tools again. JSON-shaped and at most
+    # MAX_BYTES; a loop whose messages cannot be kept so (an attachment in
+    # them, or too many bytes) does not pause. A response memo
+    # (Smith::ResponseMemo) keeps a call's messages the same way, with each
+    # response's `usage`: its token counts, stop reason and message id.
+    module Continuation
+      MAX_BYTES = 1024 * 1024
+      ROLES = %w[assistant tool user].freeze
+      TOKENS = %w[input_tokens output_tokens cached_tokens cache_creation_tokens thinking_tokens].freeze
+
+      module_function
+
+      # `kept` (an earlier pause's) followed by `messages`, or nil; with
+      # `usage`, each response's usage too (its facts from `facts`, by
+      # message, for a streamed response, whose facts its stream gave).
+      def dump(kept, messages, usage: false, facts: {})
+        dumped = messages.map { message(_1, usage: (facts[_1] || {} if usage)) || (return nil) }
+        list = JSON.parse(JSON.generate([*kept, *dumped]))
+        list if JSON.generate(list).bytesize <= MAX_BYTES
+      rescue JSON::GeneratorError, JSON::NestingError, Encoding::UndefinedConversionError
+        nil
+      end
+
+      def message(message, usage: nil)
+        content = content_of(message.content)
+        return if content == :unkept
+
+        { "role" => message.role.to_s, "content" => content, "tool_call_id" => message.tool_call_id,
+          "model_id" => message.model_id, "tool_calls" => message.tool_calls&.values&.map { tool_call(_1) },
+          "thinking" => thinking_of(message.thinking), "usage" => (usage_of(message, usage) if usage) }.compact
+      end
+
+      def usage_of(message, streamed)
+        tokens = message.respond_to?(:tokens) ? message.tokens&.to_h : nil
+        facts = Providers::ResponseFacts.from_message(message).merge(streamed)
+        usage = { "tokens" => tokens&.transform_keys(&:to_s), "stop_reason" => facts[:stop_reason],
+                  "id" => facts[:provider_message_id] }.compact
+        usage unless usage.empty?
+      end
+
+      def thinking_of(thinking) = thinking && { "text" => thinking.text, "signature" => thinking.signature }.compact
+
+      def content_of(content)
+        case content
+        when nil, String then content
+        when RubyLLM::Content::Raw then { "raw" => content.value }
+        # A structured reply, as RubyLLM parses one under a schema.
+        when Hash, Array then { "json" => content }
+        else :unkept
+        end
+      end
+
+      def tool_call(call)
+        { "id" => call.id, "name" => call.name, "arguments" => call.arguments,
+          "thought_signature" => call.thought_signature }.compact
+      end
+
+      # The messages a kept continuation holds, for a chat; ArgumentError if
+      # it is not one `dump` makes.
+      def load(list)
+        raise ArgumentError, "a budget continuation must be a list of messages" unless valid?(list)
+
+        list.map { message_from(_1) }
+      end
+
+      def message_from(entry)
+        content = entry["content"]
+        message = RubyLLM::Message.new(
+          role: entry["role"].to_sym, content: content(content), tool_call_id: entry["tool_call_id"],
+          model_id: entry["model_id"], tool_calls: tool_calls(entry["tool_calls"]),
+          thinking: thinking(entry["thinking"]), **kept_usage(entry["usage"])
+        )
+        # RubyLLM's constructor reads a Hash as attachments.
+        message.content = content["json"] if structured?(content)
+        message
+      end
+
+      # A kept response's token counts, and its facts as a raw body
+      # Providers::ResponseFacts reads.
+      def kept_usage(usage)
+        return {} unless usage
+
+        { **usage.fetch("tokens", {}).transform_keys(&:to_sym),
+          raw: Providers::ResponseFacts.kept(stop_reason: usage["stop_reason"], id: usage["id"]) }
+      end
+
+      def valid?(list)
+        list.is_a?(Array) && !list.empty? && JSON.generate(list).bytesize <= MAX_BYTES && list.all? { entry?(_1) }
+      rescue JSON::GeneratorError, JSON::NestingError
+        false
+      end
+
+      def entry?(entry)
+        entry.is_a?(Hash) && ROLES.include?(entry["role"]) && content?(entry["content"]) &&
+          entry.values_at("tool_call_id", "model_id").all? { optional_string?(_1) } &&
+          calls?(entry["tool_calls"]) && thinking?(entry["thinking"]) && usage?(entry["usage"])
+      end
+
+      def usage?(usage)
+        return true if usage.nil?
+        return false unless usage.is_a?(Hash) && (usage.keys - %w[tokens stop_reason id]).empty?
+
+        tokens?(usage.fetch("tokens", {})) && optional_string?(usage["stop_reason"]) && optional_string?(usage["id"])
+      end
+
+      def tokens?(tokens)
+        tokens.is_a?(Hash) && (tokens.keys - TOKENS).empty? && tokens.values.all? { _1.is_a?(Integer) && !_1.negative? }
+      end
+
+      def calls?(calls) = calls.nil? || (calls.is_a?(Array) && calls.all? { call?(_1) })
+
+      def thinking?(thinking) = thinking.nil? || (thinking.is_a?(Hash) && thinking.values.all?(String))
+
+      def content?(content)
+        content.nil? || content.is_a?(String) || (content.is_a?(Hash) && [["raw"], ["json"]].include?(content.keys))
+      end
+
+      def structured?(content) = content.is_a?(Hash) && content.key?("json")
+
+      def call?(call)
+        call.is_a?(Hash) && call["id"].is_a?(String) && call["name"].is_a?(String) && call["arguments"].is_a?(Hash) &&
+          optional_string?(call["thought_signature"])
+      end
+
+      def optional_string?(value) = value.nil? || value.is_a?(String)
+
+      def content(value)
+        return value unless value.is_a?(Hash)
+
+        structured?(value) ? nil : RubyLLM::Content::Raw.new(value.fetch("raw"))
+      end
+
+      def tool_calls(calls)
+        calls&.to_h do |call|
+          [call["id"], RubyLLM::ToolCall.new(id: call["id"], name: call["name"], arguments: call["arguments"],
+                                             thought_signature: call["thought_signature"])]
+        end
+      end
+
+      def thinking(value) = value && RubyLLM::Thinking.new(text: value["text"], signature: value["signature"])
+    end
+  end
+end

data/lib/smith/budget/host_ledger.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+require_relative "amount"
+
+module Smith
+  module Budget
+    # The protocol of a host's budget ledger (a workflow's `budget_ledger`),
+    # and how Smith calls it.
+    #
+    # `reserve!(amount, key:, expires_at:)` admits a provider request's hold
+    # (a Budget::Amount) when it returns true, and refuses it otherwise.
+    # `reconcile!(key, amount)` records what the request spent, in place of
+    # its hold if it had one: a key may be reconciled that was never held
+    # (a later round of a tool loop, or a result that came for a hold that
+    # expired). `release!(key)` says it spent nothing. A hold not settled by
+    # `expires_at` (its process died during the request) counts for nothing
+    # after it; Smith runs no timer. A key may come again after a crash: a
+    # reserve of it replaces its hold, a release ends its hold, and its
+    # reconcile is final (a ledger records it once). Fan-out branches call
+    # the ledger from their own threads.
+    #
+    # A refusal raises BudgetUnavailable before the request is sent. An
+    # error from building the ledger or from `reserve!` becomes
+    # BudgetLedgerError (the hold, if `reserve!` took it, is released), as
+    # nothing was sent. An error from `reconcile!` or `release!` is logged
+    # by class and the step goes on: the money is spent, and a ledger that
+    # could not record it must not cost the answer.
+    module HostLedger
+      METHODS = %i[reserve! reconcile! release!].freeze
+      INHERITED = :smith_inherited_budget_ledger
+      SCOPE_STRING_BYTES = 256
+      private_constant :INHERITED, :SCOPE_STRING_BYTES
+
+      module_function
+
+      def check!(ledger)
+        return ledger if ledger.nil? || METHODS.all? { ledger.respond_to?(_1) }
+
+        raise WorkflowError, "a budget_ledger must respond to #{METHODS.join(", ")}"
+      end
+
+      # The run's ledger from its builder and the scope it was created with.
+      def build(builder, scope)
+        check!(builder.call(scope))
+      rescue WorkflowError
+        raise
+      rescue StandardError => e
+        raise BudgetLedgerError, "the budget ledger could not be built (#{e.class})"
+      end
+
+      # Admits the hold until `seconds` (Smith.config.budget_hold_ttl, plus a
+      # batch's wait) from now.
+      def admit!(ledger, amount, key, seconds: Smith.config.budget_hold_ttl)
+        raise BudgetUnavailable unless ledger.reserve!(amount, key:, expires_at: Time.now.utc + seconds) == true
+      rescue BudgetUnavailable
+        raise
+      rescue StandardError => e
+        settle(ledger, key, nil)
+        raise BudgetLedgerError, "the budget ledger could not reserve (#{e.class})"
+      end
+
+      def settle(ledger, key, amount)
+        amount ? ledger.reconcile!(key, amount) : ledger.release!(key)
+      rescue StandardError => e
+        Smith.config.logger&.error("Smith budget ledger #{amount ? "reconcile!" : "release!"} error: #{e.class}")
+        nil
+      end
+
+      # What a response's usage cost, as an Amount, priced as Smith prices
+      # usage; a count that is not a non-negative Integer counts as 0.
+      def amount(usage, model:, provider:, mode: :sync)
+        counts = [usage.input_tokens, usage.output_tokens, usage.cache_read_tokens, usage.cache_write_tokens]
+        input, output, reads, writes = counts.map { _1.is_a?(Integer) && _1.positive? ? _1 : 0 }
+        Amount.new(
+          input_tokens: input + reads + writes, output_tokens: output,
+          cost: Pricing.compute_cost(model:, provider:, input_tokens: input, output_tokens: output,
+                                     cache_read_tokens: reads, cache_write_tokens: writes, mode:)
+        )
+      end
+
+      # The scope a run's ledger is built from: the values its `keys` have in
+      # the run's trusted context when it is created, JSON scalars (a String
+      # of at most SCOPE_STRING_BYTES bytes), kept with the run so that what
+      # a step writes later never chooses another ledger.
+      def scope!(keys, context, untrusted)
+        keys.to_h do |key|
+          raise WorkflowError, "budget_ledger scope key #{key.inspect} is untrusted" if untrusted.include?(key)
+
+          [key, scope_value!(key, context.is_a?(Hash) ? context[key] : nil)]
+        end.freeze
+      end
+
+      # A stored scope, if it names the declared keys; nil (capture it
+      # again) if the declaration changed.
+      def stored_scope(keys, stored)
+        return unless stored.is_a?(Hash) && same_keys?(keys, stored)
+
+        stored = stored.transform_keys(&:to_s)
+        keys.to_h { |key| [key, scope_value!(key, stored[key.to_s])] }.freeze
+      rescue WorkflowError
+        raise SerializationError, "persisted workflow budget_ledger_scope is invalid"
+      end
+
+      def same_keys?(keys, stored) = stored.keys.map(&:to_s).sort == keys.map(&:to_s).sort
+
+      def scope_value!(key, value)
+        case value
+        when nil, true, false, Integer then value
+        when Float then value.finite? ? value : invalid_scope!(key)
+        when String then value.bytesize <= SCOPE_STRING_BYTES ? value.dup.freeze : invalid_scope!(key)
+        else invalid_scope!(key)
+        end
+      end
+
+      def invalid_scope!(key)
+        raise WorkflowError, "budget_ledger scope key #{key.inspect} must hold a JSON scalar " \
+                             "(a String of at most #{SCOPE_STRING_BYTES} bytes)"
+      end
+
+      # A nested workflow built in the block charges `source` (the enclosing
+      # run's ledger, as a callable) unless it names a ledger of its own.
+      def inheriting(source)
+        outer = Thread.current[INHERITED]
+        Thread.current[INHERITED] = source
+        yield
+      ensure
+        Thread.current[INHERITED] = outer
+      end
+
+      def take_inherited
+        source = Thread.current[INHERITED]
+        Thread.current[INHERITED] = nil
+        source
+      end
+    end
+  end
+end

data/lib/smith/budget/ledger.rb

@@ -100,14 +100,14 @@ module Smith
         @representation.internalize(external)
       end
 
-      def create_reservation(amounts)
+      def create_reservation(amounts, checked: true)
         token = Object.new.freeze
         receipt = Reservation.new(
           ledger_identity: @identity,
           token:,
           amounts: @representation.externalize(amounts)
         )
-        next_state = @state_transition.reserve(@publication.state, amounts)
+        next_state = @state_transition.reserve(@publication.state, amounts, checked:)
         @publication.add(@public_state_validator.call(next_state), token:, amounts:)
         receipt
       end
@@ -136,3 +136,5 @@ module Smith
     end
   end
 end
+
+require_relative "spent_reservation"

data/lib/smith/budget/ledger_state_transition.rb

@@ -11,12 +11,9 @@ module Smith
 
       option :limits
 
-      def reserve(state, amounts)
+      def reserve(state, amounts, checked: true)
         DecimalContext.call do
-          amounts.each do |key, amount|
-            requested = state.consumed[key] + state.reserved[key] + amount
-            raise BudgetExceeded unless requested.finite? && requested <= limits.fetch(key)
-          end
+          amounts.each { |key, amount| admit!(state, key, amount, checked) }
 
           LedgerState.new(consumed: state.consumed, reserved: add(state.reserved, amounts))
         end
@@ -49,6 +46,11 @@ module Smith
 
       private
 
+      def admit!(state, key, amount, checked)
+        requested = state.consumed[key] + state.reserved[key] + amount
+        raise BudgetExceeded unless requested.finite? && (!checked || requested <= limits.fetch(key))
+      end
+
       def add(current, amounts)
         merge(current, amounts) { |value, amount| value + amount }
       end