squishling 0.2.0 → 0.3.0

data/lib/squishling/invoker.rb

@@ -8,23 +8,69 @@ 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
 
-    # Makes each attempt of the escalation (Definition#escalation_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.
+    # Runs the harness (Definition#harness) over the escalation (Definition#escalation_path).
     def call
-      @instructions = @definition.instructions(@receiver)
-      @schema = @definition.schema
-      raise ConfigurationError, "#{@definition.label} has no instructions" if @instructions.nil? || @instructions.empty?
-      raise ConfigurationError, "#{@definition.label} has no output_schema" unless @schema
+      purpose = @definition.purpose(@receiver)
+      schema = @definition.schema
+      raise ConfigurationError, "#{@definition.label} has no purpose" if purpose.nil? || purpose.empty?
+      raise ConfigurationError, "#{@definition.label} has no output_schema" unless schema
 
+      harness = @definition.harness
       path = @definition.escalation_path
-      input = JSON.generate(payload)
+      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?
+
+      # 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
+
+    # 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
+
+    # 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 = []
@@ -35,92 +81,126 @@ module Squishling
         if chat && step == path[attempt - 2]
           message = retry_message(rejected.last)
         else
-          chat = start_chat(step)
-          message = rejected ? escalation_message(input, *rejected) : input
+          chat = start_chat(step, prompt)
+          message = rejected && step.forward_rejected ? escalation_message(prompt.input, *rejected) : prompt.input
         end
 
         begin
-          response = ask(chat, message, step)
+          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(attempt, path, e.message)
+          log_failure(prompt.role, attempt, path, e.message)
           next
         end
 
-        result, errors = check(response.content)
-        return result if errors.empty?
+        result, errors = @output_check.call(response.content, prompt)
+        if errors.empty?
+          squawk(prompt, path, attempt, response, nil)
+          return result
+        end
 
-        raise InvalidOutputError.new(errors:, raw: response.content, attempts: attempt, models:) if last
+        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(attempt, path, errors.join("; "))
+        log_failure(prompt.role, attempt, path, errors.join("; "))
       end
     end
 
-    private
+    # 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) }
 
-    def start_chat(step)
-      chat = build_chat(step)
-      chat.with_instructions("#{@instructions}\n\n#{INPUT_NOTE}")
-      chat.with_schema(@schema.llm_schema)
-      apply_params(chat, step.params)
-      chat
-    end
+      1.upto(sample_steps.map(&:size).max) do |attempt|
+        pending = samples.reject(&:result)
+        break if pending.empty?
 
-    def build_chat(step)
-      RubyLLM.chat(**chat_options(step))
-    rescue RubyLLM::ModelNotFoundError, RubyLLM::ConfigurationError => e
-      raise ConfigurationError, "#{@definition.label}: #{e.message}"
-    end
+        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
 
-    # 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, "#{@definition.label}: invalid params #{params.inspect} (#{e.message})"
+        pending.zip(outcomes).each { |sample, outcome| record(sample, outcome, prompt, attempt) }
+      end
+
+      settle(samples, harness, judge)
     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)
-      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(step.params)}"
-    rescue RubyLLM::Error, Faraday::Error => e
-      raise LLMError, "#{@definition.label}: #{e.class}: #{e.message}"
+    # 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
 
-    # 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(step)
-      model = step.model
-      provider = step.provider
-      options = { model:, provider: }.compact
-      options[:assume_model_exists] = true if model && provider && !known_model?(model, provider)
-      options
+    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 known_model?(model, provider)
-      RubyLLM.models.find(model, provider:)
-      true
-    rescue RubyLLM::ModelNotFoundError
-      false
+    # 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
+
+      raise DisagreementError.new(candidates: [first, second], verdict: :neither, reason:, detail:,
+        models: models + [judge.step.model])
     end
 
-    def params_hint(params)
-      return "" if params.empty?
+    # 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
 
-      ". Check params #{params.inspect}; reasoning models often reject sampling params such as temperature and top_p"
+    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.
@@ -154,92 +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)
-
-      [JSON.parse(strip_code_fence(content)), []]
-    rescue JSON::ParserError => e
-      [nil, ["response was not valid JSON (#{e.message.lines.first&.strip})"]]
-    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
-    end
-
-    # Parses and validates a response: [typed result, []] when it passes, [nil or result, errors] otherwise.
-    def check(content)
-      data, errors = parse(content)
-      errors = @schema.validate(data) if errors.empty?
-      return [nil, errors] if errors.any?
-
-      result = @schema.build(data, squished: true)
-      [result, validator_errors(result)]
+    def retry_message(errors)
+      "Your previous response was rejected:\n- #{errors.join("\n- ")}\nRespond again with corrected JSON only."
     end
 
-    # Runs the squish_validate check, if any. Exceptions it raises are the user's own and propagate as-is.
-    def validator_errors(result)
-      validator = @definition.validator
-      return [] unless validator
-
-      value = @receiver.instance_exec(result, **@inputs, &validator)
-      case value
-      when nil, true then []
-      when false then ["the output was rejected by squish_validate"]
-      when String, Array then Array(value).flatten.compact.map(&:to_s).reject { |message| message.strip.empty? }
-      else
-        return contract_errors(value) if value.respond_to?(:success?) && value.respond_to?(:errors)
-
-        raise ConfigurationError, "#{@definition.label}: squish_validate must return nil, true, false, a String, " \
-                                  "an Array of Strings, or a validation result, got #{value.class}"
+    # 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
 
-    # A dry-validation style result: errors.to_h is { key => ["message", ...] }, nested for nested keys.
-    def contract_errors(value)
-      return [] if value.success?
+    # 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)
 
-      errors = value.errors
-      errors = errors.to_h if !errors.is_a?(Array) && errors.respond_to?(:to_h)
-      messages = errors.is_a?(Hash) ? flatten_messages(errors) : Array(errors).map(&:to_s)
-      messages.empty? ? ["the output was rejected by squish_validate"] : messages
+      JSON.generate(raw)
+    rescue JSON::JSONError
+      "(a response JSON can't represent)"
     end
 
-    def flatten_messages(errors, prefix = nil)
-      errors.flat_map do |key, value|
-        path = [prefix, key].compact.join(".")
-        next flatten_messages(value, path) if value.is_a?(Hash)
-
-        Array(value).map { |message| path.empty? ? message.to_s : "#{path} #{message}" }
-      end
+    # 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 was rejected:\n- #{errors.join("\n- ")}\nRespond 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
 
-    # The first message to a fresh chat after an earlier model's output was rejected.
-    def escalation_message(input, raw, errors)
-      previous = raw.is_a?(String) ? raw : JSON.generate(raw)
-      previous = "(an empty response)" if previous.strip.empty? || raw.nil?
-      "#{input}\n\nA previous attempt at this request returned:\n#{previous}\n" \
-        "It was rejected:\n- #{errors.join("\n- ")}\nRespond with corrected JSON only."
+    def source(role)
+      role == :judge ? "Judge" : "LLM"
     end
 
-    def log_failure(attempt, path, reason)
+    def log_failure(role, attempt, path, reason)
       step = path[attempt - 1]
       next_step = path[attempt]
-      action = next_step == step ? "retrying" : "escalating to #{model_name(next_step)}"
-      Squishling.config.logger&.warn("[Squishling] #{@definition.label} attempt #{attempt} of #{path.size} " \
-                                     "(#{model_name(step)}) failed, #{action}: #{reason}")
+      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 model_name(step)
-      step.model || "RubyLLM default model"
+    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