squishling 0.1.0 → 0.3.0

data/lib/squishling/invoker.rb

@@ -8,91 +8,199 @@ module Squishling
       present, holds additional state about the caller. Respond only with JSON matching the required schema.
     NOTE
 
+    # The most of a rejected output that is forwarded to the next escalation step.
+    MAX_FORWARDED_CHARS = 4_000
+
+    # One escalation's request: the system prompt, the output Schema, and the first user message. The judge's
+    # role (:judge) skips squish_validate and names the judge in errors and logs.
+    Prompt = Data.define(:instructions, :schema, :input, :role)
+
+    # One of a sampling harness's two samples, between rounds: the steps it may attempt, in order.
+    Sample = Struct.new(:name, :steps, :chat, :rejected, :result)
+
     def initialize(definition, receiver, inputs)
       @definition = definition
       @receiver = receiver
       @inputs = inputs
+      @client = LLMClient.new(definition.label)
+      @output_check = OutputCheck.new(definition, receiver, inputs)
     end
 
+    # Runs the harness (Definition#harness) over the escalation (Definition#escalation_path).
     def call
-      instructions = @definition.instructions(@receiver)
+      purpose = @definition.purpose(@receiver)
       schema = @definition.schema
-      raise ConfigurationError, "#{@definition.label} has no instructions" if instructions.nil? || instructions.empty?
+      raise ConfigurationError, "#{@definition.label} has no purpose" if purpose.nil? || purpose.empty?
       raise ConfigurationError, "#{@definition.label} has no output_schema" unless schema
 
-      chat = build_chat
-      chat.with_instructions("#{instructions}\n\n#{INPUT_NOTE}")
-      chat.with_schema(schema.llm_schema)
-      apply_params(chat)
-
-      response = ask(chat, JSON.generate(payload))
-      attempts = 0
-      loop do
-        attempts += 1
-        data, errors = parse(response.content)
-        errors = schema.validate(data) if errors.empty?
-        return schema.build(data, squished: true) if errors.empty?
-
-        if attempts > Squishling.config.max_retries
-          raise InvalidOutputError.new(errors:, raw: response.content, attempts:)
-        end
+      harness = @definition.harness
+      path = @definition.escalation_path
+      body = payload
+      prompt = Prompt.new(instructions: "#{purpose}\n\n#{INPUT_NOTE}", schema:,
+        input: JSON.generate(body), role: nil)
+      return escalate(path, prompt) unless harness.squishsum?
 
-        response = ask(chat, retry_message(errors))
-      end
+      # The sample steps and the judge are resolved before any sample is requested, so a missing one fails
+      # before costing anything.
+      sample_steps, rest = sampling_plan(path, harness)
+      judge = harness.judged? && Judge.new(definition: @definition, receiver: @receiver, client: @client, harness:,
+        rest:, purpose:, payload: body, escalate: method(:escalate),
+        log_failure: method(:log_failure))
+      sample(sample_steps, prompt, harness, judge)
     end
 
     private
 
-    def build_chat
-      RubyLLM.chat(**chat_options)
-    rescue RubyLLM::ModelNotFoundError, RubyLLM::ConfigurationError => e
-      raise ConfigurationError, "#{@definition.label}: #{e.message}"
+    # Splits the escalation into runs of consecutive identical steps (a step's attempts). A squishsum samples
+    # the first run twice; an ensemble samples the first run, then the second. Returns each sample's steps and
+    # the rest of the escalation, which is where a judge looks for its default step.
+    def sampling_plan(path, harness)
+      runs = path.chunk_while { |step, following| step == following }.to_a
+      if harness.ensemble? && runs.size < 2
+        raise ConfigurationError, "#{@definition.label}: the #{harness.type} harness needs a second escalation " \
+                                  "step to sample (adjacent identical steps count as one step's attempts)"
+      end
+
+      return [runs.first(2), runs.drop(2).flatten(1)] if harness.ensemble?
+
+      [[runs.first, runs.first], runs.drop(1).flatten(1)]
     end
 
-    # RubyLLM validates some settings locally (e.g. an impossible thinking budget for the model)
-    # and raises ArgumentError before any request is sent.
-    def apply_params(chat)
-      Params.apply(chat, @definition.params)
-    rescue ArgumentError => e
-      raise ConfigurationError, "#{@definition.label}: invalid params #{@definition.params.inspect} (#{e.message})"
+    # Makes each attempt of the path until an output passes the schema and the squish_validate check.
+    # Consecutive attempts on the same step continue the same conversation, so the model sees what it got wrong;
+    # a new step starts a fresh chat told about the last rejected output (unless the step sets
+    # forward_rejected: false).
+    def escalate(path, prompt)
+      chat = nil
+      rejected = nil # [raw content, errors] of the last invalid output
+      models = []
+
+      path.each.with_index(1) do |step, attempt|
+        models << step.model
+        last = attempt == path.size
+        if chat && step == path[attempt - 2]
+          message = retry_message(rejected.last)
+        else
+          chat = start_chat(step, prompt)
+          message = rejected && step.forward_rejected ? escalation_message(prompt.input, *rejected) : prompt.input
+        end
+
+        begin
+          response = @client.ask(chat, message, step)
+        rescue LLMError => e
+          squawk(prompt, path, attempt, nil, e)
+          raise if last
+
+          # The failed request may be half-recorded in this chat, so the next attempt starts a fresh one.
+          chat = nil
+          log_failure(prompt.role, attempt, path, e.message)
+          next
+        end
+
+        result, errors = @output_check.call(response.content, prompt)
+        if errors.empty?
+          squawk(prompt, path, attempt, response, nil)
+          return result
+        end
+
+        error = InvalidOutputError.new(errors:, raw: response.content, attempts: attempt, models: models.dup,
+          source: source(prompt.role))
+        squawk(prompt, path, attempt, response, error)
+        raise error if last
+
+        rejected = [response.content, errors]
+        log_failure(prompt.role, attempt, path, errors.join("; "))
+      end
     end
 
-    # Transient HTTP failures are already retried by RubyLLM (config.max_retries); anything that
-    # still fails is surfaced as an LLMError rather than retried again here. A 400 means the request
-    # we built is invalid (an unsupported param, a schema the provider rejects), so it's a setup
-    # mistake: a fallback would otherwise hide it on every call.
-    def ask(chat, message)
-      chat.ask(message)
-    rescue RubyLLM::ConfigurationError, RubyLLM::UnauthorizedError, RubyLLM::ForbiddenError => e
-      raise ConfigurationError, "#{@definition.label}: #{e.class}: #{e.message}"
-    rescue RubyLLM::BadRequestError => e
-      raise ConfigurationError, "#{@definition.label}: the provider rejected the request (#{e.message})#{params_hint}"
-    rescue RubyLLM::Error, Faraday::Error => e
-      raise LLMError, "#{@definition.label}: #{e.class}: #{e.message}"
+    # Two samples, each on its own steps and retried per their attempts the same way escalate retries.
+    # Each round sends the pending samples' requests concurrently; everything else runs on this thread.
+    def sample(sample_steps, prompt, harness, judge)
+      samples = %w[a b].zip(sample_steps).map { |name, steps| Sample.new(name, steps) }
+
+      1.upto(sample_steps.map(&:size).max) do |attempt|
+        pending = samples.reject(&:result)
+        break if pending.empty?
+
+        jobs = pending.map do |sample|
+          step = sample.steps[attempt - 1]
+          message = sample_message(sample, step, prompt)
+          -> { @client.ask(sample.chat, message, step) }
+        end
+        outcomes = LLMClient.concurrently(jobs)
+        # A setup mistake (or anything that isn't a failed request) ends the call, whichever sample hit it.
+        fatal = outcomes.map(&:last).find { |error| error && !error.is_a?(LLMError) }
+        raise fatal if fatal
+
+        pending.zip(outcomes).each { |sample, outcome| record(sample, outcome, prompt, attempt) }
+      end
+
+      settle(samples, harness, judge)
     end
 
-    # Models missing from RubyLLM's registry (e.g. newly released ones) are only usable when a
-    # provider is named, so RubyLLM is told to assume they exist.
-    def chat_options
-      model = @definition.model
-      provider = @definition.provider
-      options = { model:, provider: }.compact
-      options[:assume_model_exists] = true if model && provider && !known_model?(model, provider)
-      options
+    # The next message for a sample: a retry in its chat after invalid output, otherwise a fresh chat.
+    def sample_message(sample, step, prompt)
+      return retry_message(sample.rejected.last) if sample.chat
+
+      sample.chat = start_chat(step, prompt)
+      sample.rejected && step.forward_rejected ? escalation_message(prompt.input, *sample.rejected) : prompt.input
     end
 
-    def known_model?(model, provider)
-      RubyLLM.models.find(model, provider:)
-      true
-    rescue RubyLLM::ModelNotFoundError
-      false
+    def record(sample, (response, error), prompt, attempt)
+      steps = sample.steps
+      last = attempt == steps.size
+      role = "sample #{sample.name}"
+      if error
+        squawk(prompt, steps, attempt, nil, error)
+        raise error if last
+
+        sample.chat = nil
+        return log_failure(role, attempt, steps, error.message)
+      end
+
+      result, errors = @output_check.call(response.content, prompt)
+      if errors.empty?
+        squawk(prompt, steps, attempt, response, nil)
+        return sample.result = result
+      end
+
+      models = steps.first(attempt).map(&:model)
+      invalid = InvalidOutputError.new(errors:, raw: response.content, attempts: attempt, models:)
+      squawk(prompt, steps, attempt, response, invalid)
+      raise invalid if last
+
+      sample.rejected = [response.content, errors]
+      log_failure(role, attempt, steps, errors.join("; "))
     end
 
-    def params_hint
-      params = @definition.params
-      return "" if params.empty?
+    # Agreeing samples are accepted; otherwise the judge, if any, picks one or the call fails.
+    def settle(samples, harness, judge)
+      first, second = samples.map(&:result)
+      return first if agree?(first, second, harness.compare)
+
+      models = samples.map { |sample| sample.steps.first.model }
+      unless judge
+        log_warning("samples disagreed")
+        raise DisagreementError.new(candidates: [first, second], models:)
+      end
+
+      log_warning("samples disagreed, asking the judge (#{judge.step.display_name})")
+      choice, reason, detail = judge.verdict(first, second)
+      return { a: first, b: second }.fetch(choice) unless choice == :neither
 
-      ". Check params #{params.inspect}; reasoning models often reject sampling params such as temperature and top_p"
+      raise DisagreementError.new(candidates: [first, second], verdict: :neither, reason:, detail:,
+        models: models + [judge.step.model])
+    end
+
+    # compare: is the developer's own code, so its exceptions propagate unwrapped.
+    def agree?(first, second, compare)
+      return Schema.jsonify(first.to_h) == Schema.jsonify(second.to_h) unless compare
+
+      @receiver.instance_exec(first, second, **@inputs, &compare) ? true : false
+    end
+
+    def start_chat(step, prompt)
+      @client.chat(step, instructions: prompt.instructions, schema: prompt.schema.llm_schema)
     end
 
     # Only the arguments, the declared squish_context names, and a squish! call's context leave the process.
@@ -126,26 +234,67 @@ module Squishling
       @receiver.instance_variable_get(:"@#{name}")
     end
 
-    # RubyLLM parses structured output itself and leaves the raw string when that fails, so
-    # content is a Hash on success, or a String/nil when the model refused, was cut off, or
-    # ignored the schema.
-    def parse(content)
-      return [nil, ["response was empty"]] if content.nil? || (content.is_a?(String) && content.strip.empty?)
-      return [content, []] unless content.is_a?(String)
+    def retry_message(errors)
+      "Your previous response was rejected:\n- #{errors.join("\n- ")}\nRespond again with corrected JSON only."
+    end
+
+    # The first message to a fresh chat after an earlier model's output was rejected. The rejected output is
+    # model-generated and may cross providers, so it is capped at MAX_FORWARDED_CHARS.
+    def escalation_message(input, raw, errors)
+      previous = rejected_text(raw)
+      previous = "(an empty response)" if previous.strip.empty? || raw.nil?
+      if previous.length > MAX_FORWARDED_CHARS
+        omitted = previous.length - MAX_FORWARDED_CHARS
+        previous = "#{previous[0, MAX_FORWARDED_CHARS]}... [truncated, #{omitted} more characters]"
+      end
+      "#{input}\n\nA previous attempt at this request returned:\n#{previous}\n" \
+        "It was rejected:\n- #{errors.join("\n- ")}\nRespond with corrected JSON only."
+    end
 
-      [JSON.parse(strip_code_fence(content)), []]
-    rescue JSON::ParserError => e
-      [nil, ["response was not valid JSON (#{e.message.lines.first&.strip})"]]
+    # The rejected output as text. Output that was rejected for being unrepresentable in JSON (Infinity, invalid
+    # UTF-8) can't be serialized or measured as is, so it is scrubbed or replaced rather than raising.
+    def rejected_text(raw)
+      return utf8(raw).scrub("?") if raw.is_a?(String)
+
+      JSON.generate(raw)
+    rescue JSON::JSONError
+      "(a response JSON can't represent)"
     end
 
-    # Models without native structured output sometimes wrap JSON in a markdown code fence.
-    def strip_code_fence(text)
-      text[/\A\s*```(?:json)?\s*\n(.*?)\n\s*```\s*\z/m, 1] || text
+    # A binary-tagged string counts every byte as valid, so it is read as UTF-8 before scrubbing.
+    def utf8(text)
+      text.encoding == Encoding::BINARY ? text.dup.force_encoding(Encoding::UTF_8) : text
     end
 
-    def retry_message(errors)
-      "Your previous response did not match the required schema:\n- #{errors.join("\n- ")}\n" \
-        "Respond again with corrected JSON only."
+    # Runs the observability hook, if any, with this attempt's raw output (nil when the call itself failed),
+    # the error that ended it (nil when it was accepted), and what was asked.
+    def squawk(prompt, path, attempt, response, error)
+      hook = @definition.squawk
+      return unless hook
+
+      step = path[attempt - 1]
+      metadata = {
+        label: @definition.label, attempt:, attempts: path.size, final: attempt == path.size,
+        model: response&.model || step.model, provider: step.provider, params: step.params, input: prompt.input,
+        usage: response&.tokens&.to_h
+      }
+      Squawk.call(hook, output: response&.content, metadata:, error:)
+    end
+
+    def source(role)
+      role == :judge ? "Judge" : "LLM"
+    end
+
+    def log_failure(role, attempt, path, reason)
+      step = path[attempt - 1]
+      next_step = path[attempt]
+      action = next_step == step ? "retrying" : "escalating to #{next_step.display_name}"
+      who = role ? "#{role} attempt" : "attempt"
+      log_warning("#{who} #{attempt} of #{path.size} (#{step.display_name}) failed, #{action}: #{reason}")
+    end
+
+    def log_warning(message)
+      Squishling.config.logger&.warn("[Squishling] #{@definition.label} #{message}")
     end
   end
 end

data/lib/squishling/judge.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+
+module Squishling
+  # Picks between a judged harness's two disagreeing samples (judged_squishsum or judged_ensemble), or rejects
+  # both. The judge is the harness's judge: step, or else the escalation's next step after the ones the samples
+  # ran on. A chat judge answers with a strict verdict schema; a :judgment judge (a System One decision model such
+  # as Jev) answers one choice question through RubyLLM.judge, and its pick counts only at or above
+  # min_confidence.
+  class Judge
+    VERDICT_SCHEMA = {
+      "type" => "object",
+      "properties" => {
+        "verdict" => { "type" => "string", "enum" => %w[a b neither] },
+        "reason" => { "type" => "string" }
+      },
+      "required" => %w[verdict reason],
+      "additionalProperties" => false
+    }.freeze
+
+    INPUT_NOTE = <<~NOTE
+      The input is a JSON object. "input" is what the operation received ("arguments" and, when present,
+      "context"); "output_schema" is the format both candidates follow; "candidates" holds "a" and "b".
+      Respond only with JSON matching the required schema: your verdict ("a", "b", or "neither") and a short reason.
+    NOTE
+
+    CHOICES = {
+      a: "Candidate a correctly carries out the purpose for this input (also when both do)",
+      b: "Candidate b correctly carries out the purpose for this input",
+      neither: "Neither candidate is clearly correct"
+    }.freeze
+
+    def initialize(definition:, receiver:, client:, harness:, rest:, purpose:, payload:, escalate:, log_failure:)
+      @definition = definition
+      @receiver = receiver
+      @client = client
+      @harness = harness
+      @purpose = purpose
+      @payload = payload
+      @escalate = escalate
+      @log_failure = log_failure
+      @steps = resolve_steps(rest)
+    end
+
+    # The judge's first attempt, for logs and DisagreementError#models.
+    def step
+      @steps.first
+    end
+
+    # [:a | :b, reason, detail] for the chosen candidate, or [:neither, reason, detail]. The detail is text
+    # Squishling wrote (nil for a chat judge, whose reason is model-written) and is safe for error messages.
+    def verdict(first, second)
+      candidates = { a: Schema.jsonify(first.to_h), b: Schema.jsonify(second.to_h) }
+      judgment? ? judgment(candidates) : chat(candidates)
+    end
+
+    private
+
+    def judgment?
+      @harness.judge&.fetch(:type) == :judgment
+    end
+
+    # A declared judge brings its own model and provider. A chat judge gets the method's generation params
+    # under its own, like any escalation step; a judgment judge sends only its own params, as provider options.
+    def resolve_steps(rest)
+      if (judge = @harness.judge)
+        params = judge[:type] == :judgment ? {} : @definition.params
+        return ModelPath.steps([judge], provider: nil, params:)
+      end
+
+      if rest.empty?
+        ordinal = @harness.ensemble? ? "third" : "second"
+        raise ConfigurationError, "#{@definition.label}: the #{@harness.type} harness needs a judge: or a " \
+                                  "#{ordinal} escalation step to judge with"
+      end
+
+      rest.take_while { |step| step == rest.first }
+    end
+
+    def chat(candidates)
+      prompt = Invoker::Prompt.new(
+        instructions: "#{judge_instructions}\n\nThe operation's purpose:\n<purpose>\n#{@purpose}\n" \
+                      "</purpose>\n\n#{INPUT_NOTE}",
+        schema: Schema.for(VERDICT_SCHEMA), input: JSON.generate(state(candidates)), role: :judge
+      )
+      result = @escalate.call(@steps, prompt)
+      [result[:verdict].to_sym, result[:reason], nil]
+    end
+
+    def judgment(candidates)
+      input = { purpose: @purpose, **state(candidates) }
+      questions = { winner: { type: :choice, instructions: judge_instructions, options: CHOICES } }
+      answer = judge_with_retries(input, questions)[:winner]
+      unless answer.respond_to?(:choice) && answer.respond_to?(:probabilities)
+        raise InvalidOutputError.new(errors: ["the judgment didn't answer the winner question"], source: "Judge")
+      end
+
+      interpret(answer.choice.to_s.to_sym, answer.probabilities)
+    end
+
+    # The threshold is checked on the exact probability; it's rounded only for the reason text.
+    def interpret(choice, probabilities)
+      probability = probabilities.to_h.find { |key, _| key.to_s == choice.to_s }&.last.to_f
+      minimum = @harness.judge[:min_confidence]
+      shown = probability.round(3)
+      picked = %i[a b].include?(choice)
+      if picked && probability >= minimum
+        text = "chosen with probability #{shown}"
+        return [choice, text, text]
+      end
+
+      text =
+        if choice == :neither then "the judge chose neither (probability #{shown})"
+        elsif picked then "the judge chose #{choice} with probability #{shown}, below min_confidence #{minimum}"
+        else "the judge returned an unrecognized choice (probability #{shown})"
+        end
+      [:neither, text, text]
+    end
+
+    # A judgment always matches its questions, so only a failed request (LLMError) moves on to the next attempt.
+    def judge_with_retries(input, questions)
+      @steps.each.with_index(1) do |step, attempt|
+        return @client.judge(input, questions:, step:)
+      rescue LLMError => e
+        raise if attempt == @steps.size
+
+        @log_failure.call(:judge, attempt, @steps, e.message)
+      end
+    end
+
+    def state(candidates)
+      { input: @payload, output_schema: @definition.schema.llm_schema["schema"], candidates: }
+    end
+
+    def judge_instructions
+      value = @harness.judge_instructions || Harness::DEFAULT_JUDGE_INSTRUCTIONS
+      value = @receiver.instance_exec(&value) if value.is_a?(Proc)
+      return value if value.is_a?(String) && !value.strip.empty?
+
+      raise ConfigurationError, "#{@definition.label}: the judge_instructions proc must return a non-blank String"
+    end
+  end
+end

data/lib/squishling/llm_client.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+module Squishling
+  # RubyLLM requests for one squished method, with RubyLLM's errors mapped onto Squishling's taxonomy.
+  class LLMClient
+    # Where RubyLLM's own code lives, to tell its ArgumentErrors from the developer's (see #judge).
+    RUBY_LLM_LIB = File.dirname(RubyLLM.method(:judge).source_location.first)
+    private_constant :RUBY_LLM_LIB
+
+    # Runs each job on its own thread and returns, in order, each job's [value, exception]. Only provider
+    # requests run here: parsing, validation, and the developer's callbacks stay on the caller's thread. Fiber
+    # storage (RubyLLM's usage owner) is inherited by the threads, and RubyLLM's instrumentation context is
+    # carried over the way RubyLLM's own concurrent tool calls do it. Unlike those, the threads don't enter the
+    # Rails executor: the caller already holds it, and a second share can deadlock against a pending code reload.
+    # If the caller is interrupted while waiting (a timeout, Thread#raise), the requests still running are
+    # stopped rather than left running unobserved.
+    def self.concurrently(jobs)
+      return jobs.map { |job| capture(job) } if jobs.size < 2
+
+      context = instrumentation_context
+      threads = jobs.map do |job|
+        thread = Thread.new { with_instrumentation_context(context) { capture(job) } }
+        thread.report_on_exception = false
+        thread
+      end
+      threads.map(&:value)
+    ensure
+      threads&.each { |thread| thread.kill.join if thread.alive? }
+    end
+
+    def self.capture(job)
+      [job.call, nil]
+    rescue Exception => e # rubocop:disable Lint/RescueException -- re-raised on the caller's thread
+      [nil, e]
+    end
+
+    def self.instrumentation_context
+      instrumentation = defined?(RubyLLM::Support::Instrumentation) && RubyLLM::Support::Instrumentation
+      return unless instrumentation.respond_to?(:current_workflow) && instrumentation.respond_to?(:capture_context)
+
+      [instrumentation.current_workflow, instrumentation.capture_context]
+    end
+
+    def self.with_instrumentation_context(context, &)
+      return yield unless context
+
+      instrumentation = RubyLLM::Support::Instrumentation
+      instrumentation.with_workflow(context.first) { instrumentation.with_context(context.last, &) }
+    end
+
+    private_class_method :capture, :instrumentation_context, :with_instrumentation_context
+
+    def initialize(label)
+      @label = label
+    end
+
+    # A fresh chat on one escalation step, with the system prompt, the strict output schema, and its params.
+    def chat(step, instructions:, schema:)
+      chat = build_chat(step)
+      chat.with_instructions(instructions)
+      chat.with_schema(schema)
+      apply_params(chat, step.params)
+      chat
+    end
+
+    # Transient HTTP failures are already retried by RubyLLM (RubyLLM.config.max_retries); anything
+    # that still fails is surfaced as an LLMError, which moves on to the next attempt of the
+    # escalation (or propagates from the last one). A 400 means the request we built is invalid (an
+    # unsupported param, a schema the provider rejects), so it's a setup mistake: escalating or a
+    # fallback would otherwise hide it on every call.
+    def ask(chat, message, step)
+      request(step) { chat.ask(message) }
+    end
+
+    # A System One judgment (RubyLLM.judge) on one step. The step's params are sent as provider options.
+    # RubyLLM raises ArgumentError for a judgment it can't build (e.g. a model on a provider that takes none, or
+    # a provider option its protocol reserves), and a bare RubyLLM::Error without an HTTP response for a provider
+    # that doesn't support judgments at all. An ArgumentError from anywhere else (an instrumentation subscriber
+    # runs inside the call) is the developer's own and propagates as-is.
+    def judge(input, questions:, step:)
+      request(step) do
+        RubyLLM.judge(input, questions:, provider_options: step.params, **model_options(step))
+      rescue RubyLLM::ModelNotFoundError => e
+        raise ConfigurationError, "#{@label}: #{e.message}"
+      rescue ArgumentError => e
+        raise unless raised_by_ruby_llm?(e)
+
+        raise ConfigurationError, "#{@label}: #{e.message}"
+      rescue RubyLLM::Error => e
+        raise unless e.instance_of?(RubyLLM::Error) && e.response.nil?
+
+        raise ConfigurationError, "#{@label}: #{e.message}"
+      end
+    end
+
+    private
+
+    def raised_by_ruby_llm?(error)
+      error.backtrace&.first&.start_with?("#{RUBY_LLM_LIB}/") || false
+    end
+
+    def build_chat(step)
+      RubyLLM.chat(**model_options(step))
+    rescue RubyLLM::ModelNotFoundError, RubyLLM::ConfigurationError => e
+      raise ConfigurationError, "#{@label}: #{e.message}"
+    end
+
+    # RubyLLM validates some settings locally (e.g. an impossible thinking budget for the model)
+    # and raises ArgumentError before any request is sent.
+    def apply_params(chat, params)
+      Params.apply(chat, params)
+    rescue ArgumentError => e
+      raise ConfigurationError, "#{@label}: invalid params #{params.inspect} (#{e.message})"
+    end
+
+    def request(step)
+      yield
+    rescue RubyLLM::ConfigurationError, RubyLLM::UnauthorizedError, RubyLLM::ForbiddenError => e
+      raise ConfigurationError, "#{@label}: #{e.class}: #{e.message}"
+    rescue RubyLLM::BadRequestError => e
+      raise ConfigurationError,
+        "#{@label}: the provider rejected the request (#{e.message})#{params_hint(step.params)}"
+    rescue RubyLLM::Error, Faraday::Error => e
+      raise LLMError, "#{@label}: #{e.class}: #{e.message}"
+    end
+
+    # Models missing from RubyLLM's registry (e.g. newly released ones) are only usable when a
+    # provider is named, so RubyLLM is told to assume they exist.
+    def model_options(step)
+      model = step.model
+      provider = step.provider
+      options = { model:, provider: }.compact
+      options[:assume_model_exists] = true if model && provider && !known_model?(model, provider)
+      options
+    end
+
+    def known_model?(model, provider)
+      RubyLLM.models.find(model, provider:)
+      true
+    rescue RubyLLM::ModelNotFoundError
+      false
+    end
+
+    def params_hint(params)
+      return "" if params.empty?
+
+      ". Check params #{params.inspect}; reasoning models often reject sampling params such as temperature and top_p"
+    end
+  end
+end