dspy 0.28.1 → 0.29.0

data/lib/gepa/core/engine.rb

@@ -0,0 +1,226 @@
+# frozen_string_literal: true
+
+require 'sorbet-runtime'
+
+require_relative 'state'
+require_relative 'result'
+require_relative '../telemetry'
+
+module GEPA
+  module Core
+    class Engine
+      extend T::Sig
+
+      sig do
+        params(
+          evaluator: T.proc.params(dataset: T::Array[T.untyped], candidate: T::Hash[String, String])
+                           .returns([T::Array[T.untyped], T::Array[Float]]),
+          valset: T::Array[T.untyped],
+          seed_candidate: T::Hash[String, String],
+          max_metric_calls: Integer,
+          perfect_score: Float,
+          seed: Integer,
+          reflective_proposer: T.untyped,
+          logger: T.untyped,
+          experiment_tracker: T.untyped,
+          merge_proposer: T.nilable(T.untyped),
+          run_dir: T.nilable(String),
+          track_best_outputs: T::Boolean,
+          display_progress_bar: T::Boolean,
+          telemetry: T.nilable(T.untyped),
+          raise_on_exception: T::Boolean
+        ).void
+      end
+      def initialize(
+        evaluator:,
+        valset:,
+        seed_candidate:,
+        max_metric_calls:,
+        perfect_score:,
+        seed:, # rubocop:disable Lint/UnusedMethodArgument -- kept for parity and future use
+        reflective_proposer:,
+        logger:,
+        experiment_tracker:,
+        merge_proposer: nil,
+        run_dir: nil,
+        track_best_outputs: false,
+        display_progress_bar: false,
+        telemetry: nil,
+        raise_on_exception: true
+      )
+        @run_dir = run_dir
+        @evaluator = evaluator
+        @valset = valset
+        @seed_candidate = seed_candidate
+        @max_metric_calls = max_metric_calls
+        @perfect_score = perfect_score
+        @reflective_proposer = reflective_proposer
+        @merge_proposer = merge_proposer
+        @logger = logger
+        @experiment_tracker = experiment_tracker
+        @track_best_outputs = track_best_outputs
+        @display_progress_bar = display_progress_bar
+        @telemetry = telemetry || GEPA::Telemetry
+        @raise_on_exception = raise_on_exception
+      end
+
+      sig { returns(GEPA::Core::State) }
+      def run
+        with_span('gepa.engine.run', max_metric_calls: @max_metric_calls) do
+          state = GEPA::Core::State.initialize_gepa_state(
+            run_dir: @run_dir,
+            logger: @logger,
+            seed_candidate: @seed_candidate,
+            valset_evaluator: ->(candidate) { full_evaluator(candidate) },
+            track_best_outputs: @track_best_outputs
+          )
+
+          @experiment_tracker.log_metrics({ base_program_full_valset_score: state.program_full_scores_val_set.first }, step: 0)
+
+          if @merge_proposer
+            @merge_proposer.last_iter_found_new_program = false
+          end
+
+          while state.total_num_evals < @max_metric_calls
+            break unless iteration_step(state)
+          end
+
+          state.save(@run_dir)
+          state
+        end
+      end
+
+      private
+
+      sig { params(state: GEPA::Core::State).returns(T::Boolean) }
+      def iteration_step(state)
+        state.i += 1
+        trace_entry = { iteration: state.i }
+        state.full_program_trace << trace_entry
+
+        progress = false
+
+        with_span('gepa.engine.iteration', iteration: state.i) do
+          merge_result = process_merge_iteration(state)
+          case merge_result
+          when :accepted
+            return true
+          when :attempted
+            return false
+          end
+
+          reflective_result = process_reflective_iteration(state)
+          return false if reflective_result == :no_candidate
+          progress = true if reflective_result == :accepted
+        end
+
+        progress
+      rescue StandardError => e
+        @logger.log("Iteration #{state.i}: Exception during optimization: #{e}")
+        @logger.log(e.backtrace&.join("\n"))
+        raise e if @raise_on_exception
+        true
+      end
+
+      sig { params(state: GEPA::Core::State).returns(Symbol) }
+      def process_merge_iteration(state)
+        return :skipped unless @merge_proposer && @merge_proposer.use_merge
+
+        if @merge_proposer.merges_due.positive? && @merge_proposer.last_iter_found_new_program
+          proposal = @merge_proposer.propose(state)
+          @merge_proposer.last_iter_found_new_program = false
+
+          if proposal&.tag == 'merge'
+            parent_sums = Array(proposal.subsample_scores_before).map(&:to_f)
+            new_sum = Array(proposal.subsample_scores_after).map(&:to_f).sum
+
+            if parent_sums.empty?
+              @logger.log("Iteration #{state.i}: Missing parent subscores for merge proposal, skipping")
+              return :handled
+            end
+
+            if new_sum >= parent_sums.max
+              with_span('gepa.engine.full_evaluation', iteration: state.i) do
+                run_full_evaluation(state, proposal.candidate, proposal.parent_program_ids)
+              end
+              @merge_proposer.merges_due -= 1
+              @merge_proposer.total_merges_tested += 1
+              return :accepted
+            else
+              @logger.log(
+                "Iteration #{state.i}: Merge subsample score #{new_sum.round(4)} "\
+                "did not beat parents #{parent_sums.map { |v| v.round(4) }}, skipping"
+              )
+              return :attempted
+            end
+          end
+        end
+
+        @merge_proposer.last_iter_found_new_program = false
+        :skipped
+      end
+
+      sig { params(state: GEPA::Core::State).void }
+      def process_reflective_iteration(state)
+        proposal = @reflective_proposer.propose(state)
+        unless proposal
+          @logger.log("Iteration #{state.i}: Reflective mutation did not propose a new candidate")
+          return :no_candidate
+        end
+
+        before = Array(proposal.subsample_scores_before).map(&:to_f)
+        after = Array(proposal.subsample_scores_after).map(&:to_f)
+        if after.empty? || after.sum <= before.sum
+          @logger.log("Iteration #{state.i}: New subsample score is not better, skipping")
+          return :skipped
+        end
+
+        with_span('gepa.engine.full_evaluation', iteration: state.i) do
+          run_full_evaluation(state, proposal.candidate, proposal.parent_program_ids)
+        end
+
+        if @merge_proposer&.use_merge
+          @merge_proposer.last_iter_found_new_program = true
+          @merge_proposer.schedule_if_needed
+        end
+
+        :accepted
+      end
+
+      sig do
+        params(state: GEPA::Core::State, new_program: T::Hash[String, String], parents: T::Array[Integer]).void
+      end
+      def run_full_evaluation(state, new_program, parents)
+        outputs, scores = full_evaluator(new_program)
+        avg_score = scores.sum / scores.length.to_f
+
+        state.num_full_ds_evals += 1
+        state.total_num_evals += scores.length
+
+        state.update_state_with_new_program(
+          parents,
+          new_program,
+          avg_score,
+          outputs,
+          scores,
+          @run_dir,
+          state.total_num_evals
+        )
+
+        @experiment_tracker.log_metrics({ new_program_full_score: avg_score }, step: state.i)
+      end
+
+      sig { params(candidate: T::Hash[String, String]).returns([T::Array[T.untyped], T::Array[Float]]) }
+      def full_evaluator(candidate)
+        @evaluator.call(@valset, candidate)
+      end
+
+      sig do
+        params(operation: String, attrs: T::Hash[Symbol, T.untyped], block: T.proc.returns(T.untyped)).returns(T.untyped)
+      end
+      def with_span(operation, attrs = {}, &block)
+        @telemetry.with_span(operation, attrs, &block)
+      end
+    end
+  end
+end

data/lib/gepa/core/evaluation_batch.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require 'sorbet-runtime'
+
+module GEPA
+  module Core
+    # Container for evaluating a candidate on a batch.
+    class EvaluationBatch < T::Struct
+      extend T::Sig
+
+      const :outputs, T::Array[T.untyped]
+      const :scores, T::Array[Float]
+      const :trajectories, T.nilable(T::Array[T.untyped])
+
+      sig { override.params(args: T.untyped, kwargs: T.untyped).void }
+      def initialize(*args, **kwargs)
+        super
+        raise ArgumentError, 'outputs and scores length mismatch' unless outputs.length == scores.length
+
+        if trajectories
+          raise ArgumentError, 'trajectories length mismatch' unless trajectories.length == outputs.length
+        end
+      end
+    end
+  end
+end

data/lib/gepa/core/result.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'set'
+require 'sorbet-runtime'
+
+module GEPA
+  module Core
+    # Snapshot of GEPA optimization output with helpers for common queries.
+    class Result < T::Struct
+      extend T::Sig
+
+      const :candidates, T::Array[T::Hash[String, String]]
+      const :parents, T::Array[T::Array[T.nilable(Integer)]]
+      const :val_aggregate_scores, T::Array[Float]
+      const :val_subscores, T::Array[T::Array[Float]]
+      const :per_val_instance_best_candidates, T::Array[T::Array[Integer]]
+      const :discovery_eval_counts, T::Array[Integer]
+      const :best_outputs_valset, T.nilable(T::Array[T::Array[T::Array[T.untyped]]]), default: nil
+      const :total_metric_calls, T.nilable(Integer), default: nil
+      const :num_full_val_evals, T.nilable(Integer), default: nil
+      const :run_dir, T.nilable(String), default: nil
+      const :seed, T.nilable(Integer), default: nil
+
+      sig { returns(Integer) }
+      def num_candidates
+        candidates.length
+      end
+
+      sig { returns(Integer) }
+      def num_val_instances
+        per_val_instance_best_candidates.length
+      end
+
+      sig { returns(Integer) }
+      def best_idx
+        val_aggregate_scores.each_with_index.max_by { |score, _i| score }&.last || 0
+      end
+
+      sig { returns(T::Hash[String, String]) }
+      def best_candidate
+        candidates.fetch(best_idx)
+      end
+
+      sig { returns(T::Hash[Symbol, T.untyped]) }
+      def to_h
+        {
+          candidates: candidates.map(&:dup),
+          parents: parents.map(&:dup),
+          val_aggregate_scores: val_aggregate_scores.dup,
+          val_subscores: val_subscores.map(&:dup),
+          best_outputs_valset: best_outputs_valset&.map { |arr| arr.map(&:dup) },
+          per_val_instance_best_candidates: per_val_instance_best_candidates.map(&:dup),
+          discovery_eval_counts: discovery_eval_counts.dup,
+          total_metric_calls: total_metric_calls,
+          num_full_val_evals: num_full_val_evals,
+          run_dir: run_dir,
+          seed: seed,
+          best_idx: best_idx
+        }
+      end
+
+      sig { returns(String) }
+      def to_json(*_args)
+        JSON.pretty_generate(to_h)
+      end
+
+      sig do
+        params(
+          state: T.untyped,
+          run_dir: T.nilable(String),
+          seed: T.nilable(Integer)
+        ).returns(Result)
+      end
+      def self.from_state(state, run_dir: nil, seed: nil)
+        new(
+          candidates: state.program_candidates.map(&:dup),
+          parents: state.parent_program_for_candidate.map(&:dup),
+          val_aggregate_scores: state.program_full_scores_val_set.map(&:to_f),
+          best_outputs_valset: state.respond_to?(:best_outputs_valset) ? state.best_outputs_valset&.map(&:dup) : nil,
+          val_subscores: state.prog_candidate_val_subscores.map { |scores| scores.map(&:to_f) },
+          per_val_instance_best_candidates: state.program_at_pareto_front_valset.map { |set| set.to_a },
+          discovery_eval_counts: state.num_metric_calls_by_discovery.map(&:to_i),
+          total_metric_calls: state.respond_to?(:total_num_evals) ? state.total_num_evals : nil,
+          num_full_val_evals: state.respond_to?(:num_full_ds_evals) ? state.num_full_ds_evals : nil,
+          run_dir: run_dir,
+          seed: seed
+        )
+      end
+    end
+  end
+end

data/lib/gepa/core/state.rb

@@ -0,0 +1,231 @@
+# frozen_string_literal: true
+
+require 'fileutils'
+require 'json'
+require 'set'
+require 'sorbet-runtime'
+
+require_relative '../utils/pareto'
+require_relative '../telemetry'
+
+module GEPA
+  module Core
+    class State
+      extend T::Sig
+
+      attr_accessor :i, :num_full_ds_evals, :total_num_evals
+      attr_reader :program_candidates,
+                  :parent_program_for_candidate,
+                  :program_full_scores_val_set,
+                  :program_at_pareto_front_valset,
+                  :prog_candidate_val_subscores,
+                  :list_of_named_predictors,
+                  :named_predictor_id_to_update_next_for_program_candidate,
+                  :num_metric_calls_by_discovery,
+                  :full_program_trace,
+                  :per_program_tracked_scores,
+                  :pareto_front_valset,
+                  :best_outputs_valset
+
+      sig do
+        params(
+          seed_candidate: T::Hash[String, String],
+          base_valset_eval_output: [T::Array[T.untyped], T::Array[Float]],
+          track_best_outputs: T::Boolean
+        ).void
+      end
+      def initialize(seed_candidate, base_valset_eval_output, track_best_outputs: false)
+        outputs, scores = base_valset_eval_output
+        raise ArgumentError, 'validation scores must not be empty' if scores.empty?
+
+        valset_base_score = scores.sum / scores.length.to_f
+
+        @program_candidates = [seed_candidate.dup]
+        @program_full_scores_val_set = [valset_base_score]
+        @per_program_tracked_scores = [valset_base_score]
+
+        @pareto_front_valset = scores.dup
+        @parent_program_for_candidate = [[nil]]
+        @program_at_pareto_front_valset = Array.new(scores.length) { Set.new([0]) }
+
+        @list_of_named_predictors = seed_candidate.keys
+        @named_predictor_id_to_update_next_for_program_candidate = [0]
+
+        @prog_candidate_val_subscores = [scores.dup]
+        @num_metric_calls_by_discovery = [0]
+
+        @best_outputs_valset = if track_best_outputs
+          outputs.map { |output| [[0, output]] }
+        end
+
+        @full_program_trace = []
+        @i = -1
+        @num_full_ds_evals = 0
+        @total_num_evals = 0
+      end
+
+      sig { returns(T::Boolean) }
+      def consistent?
+        size = @program_candidates.length
+        raise 'program_full_scores_val_set mismatch' unless @program_full_scores_val_set.length == size
+        raise 'per_program_tracked_scores mismatch' unless @per_program_tracked_scores.length == size
+        raise 'parent_program_for_candidate mismatch' unless @parent_program_for_candidate.length == size
+        raise 'named_predictor_id_to_update mismatch' unless @named_predictor_id_to_update_next_for_program_candidate.length == size
+        raise 'prog_candidate_val_subscores mismatch' unless @prog_candidate_val_subscores.length == size
+        raise 'num_metric_calls mismatch' unless @num_metric_calls_by_discovery.length == size
+        raise 'pareto fronts length mismatch' unless @pareto_front_valset.length == @program_at_pareto_front_valset.length
+
+        @program_at_pareto_front_valset.each do |front|
+          front.each do |idx|
+            raise 'pareto index out of range' unless idx < size
+          end
+        end
+        true
+      end
+
+      sig { params(run_dir: T.nilable(String)).void }
+      def save(run_dir)
+        return if run_dir.nil?
+
+        FileUtils.mkdir_p(run_dir)
+        File.open(File.join(run_dir, 'gepa_state.bin'), 'wb') do |file|
+          data = instance_variables.each_with_object({}) do |ivar, acc|
+            acc[ivar.to_s.delete('@')] = instance_variable_get(ivar)
+          end
+          Marshal.dump(data, file)
+        end
+      end
+
+      sig { params(run_dir: String).returns(State) }
+      def self.load(run_dir)
+        File.open(File.join(run_dir, 'gepa_state.bin'), 'rb') do |file|
+          data = Marshal.load(file)
+          state = allocate
+          data.each { |key, value| state.instance_variable_set("@#{key}", value) }
+          state.consistent?
+          state
+        end
+      end
+
+      sig do
+        params(
+          parent_program_idx: T::Array[Integer],
+          new_program: T::Hash[String, String],
+          valset_score: Float,
+          valset_outputs: T::Array[T.untyped],
+          valset_subscores: T::Array[Float],
+          run_dir: T.nilable(String),
+          num_metric_calls: Integer
+        ).returns([Integer, Integer])
+      end
+      def update_state_with_new_program(
+        parent_program_idx,
+        new_program,
+        valset_score,
+        valset_outputs,
+        valset_subscores,
+        run_dir,
+        num_metric_calls
+      )
+        new_program_idx = @program_candidates.length
+        @program_candidates << new_program.dup
+        @num_metric_calls_by_discovery << num_metric_calls
+
+        max_predictor_id = parent_program_idx.map { |idx| @named_predictor_id_to_update_next_for_program_candidate[idx] }.compact.max
+        @named_predictor_id_to_update_next_for_program_candidate << (max_predictor_id || 0)
+        @parent_program_for_candidate << parent_program_idx.dup
+
+        @prog_candidate_val_subscores << valset_subscores.dup
+        @program_full_scores_val_set << valset_score.to_f
+
+        valset_subscores.each_with_index do |new_score, task_idx|
+          old_score = @pareto_front_valset[task_idx]
+          if new_score > old_score
+            @pareto_front_valset[task_idx] = new_score
+            @program_at_pareto_front_valset[task_idx] = Set.new([new_program_idx])
+            if @best_outputs_valset
+              @best_outputs_valset[task_idx] = [[new_program_idx, valset_outputs[task_idx]]]
+            end
+            write_best_output(run_dir, task_idx, new_program_idx, valset_outputs[task_idx])
+          elsif new_score == old_score
+            @program_at_pareto_front_valset[task_idx].add(new_program_idx)
+            if @best_outputs_valset
+              @best_outputs_valset[task_idx] << [new_program_idx, valset_outputs[task_idx]]
+            end
+          end
+        end
+
+        raise 'valset subscores length mismatch' unless valset_subscores.length == @program_at_pareto_front_valset.length
+
+        @per_program_tracked_scores = @program_full_scores_val_set.dup
+        linear_idx = GEPA::Utils::Pareto.idxmax(@per_program_tracked_scores)
+
+        [new_program_idx, linear_idx]
+      end
+
+      sig do
+        params(
+          eval_output: [T::Array[T.untyped], T::Array[Float]],
+          output_dir: String
+        ).void
+      end
+      def self.write_eval_output_to_directory(eval_output, output_dir)
+        _, scores = eval_output
+        scores.each_with_index do |_score, task_idx|
+          dir = File.join(output_dir, "task_#{task_idx}")
+          FileUtils.mkdir_p(dir)
+          path = File.join(dir, 'iter_0_prog_0.json')
+          File.write(path, JSON.pretty_generate(scores[task_idx]))
+        end
+      end
+
+      sig do
+        params(
+          run_dir: T.nilable(String),
+          logger: T.untyped,
+          seed_candidate: T::Hash[String, String],
+          valset_evaluator: T.proc.params(arg0: T::Hash[String, String]).returns([T::Array[T.untyped], T::Array[Float]]),
+          track_best_outputs: T::Boolean
+        ).returns(State)
+      end
+      def self.initialize_gepa_state(run_dir:, logger:, seed_candidate:, valset_evaluator:, track_best_outputs: false)
+        if run_dir && File.exist?(File.join(run_dir, 'gepa_state.bin')) && File.exist?(File.join(run_dir, 'prog_candidates'))
+          logger.log('Loading gepa state from run dir')
+          return load(run_dir)
+        end
+
+        valset_out = valset_evaluator.call(seed_candidate)
+        if run_dir
+          write_eval_output_to_directory(valset_out, File.join(run_dir, 'generated_best_outputs_valset'))
+        end
+
+        state = new(seed_candidate, valset_out, track_best_outputs: track_best_outputs)
+        state.num_full_ds_evals = 1
+        state.total_num_evals = valset_out.last.length
+        state
+      end
+
+      private
+
+      sig do
+        params(run_dir: T.nilable(String), task_idx: Integer, program_idx: Integer, output: T.untyped).void
+      end
+      def write_best_output(run_dir, task_idx, program_idx, output)
+        return if run_dir.nil?
+
+        dir = File.join(run_dir, 'generated_best_outputs_valset', "task_#{task_idx}")
+        FileUtils.mkdir_p(dir)
+        payload = ensure_jsonable(output)
+        File.write(File.join(dir, "iter_#{@i + 1}_prog_#{program_idx}.json"), JSON.pretty_generate(payload))
+      end
+
+      sig { params(value: T.untyped).returns(T.untyped) }
+      def ensure_jsonable(value)
+        JSON.parse(JSON.generate(value))
+      rescue StandardError
+        GEPA::Utils::Pareto.json_default(value)
+      end
+    end
+  end
+end
+

data/lib/gepa/logging/experiment_tracker.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module GEPA
+  module Logging
+    # Lightweight experiment tracker that records metrics locally and can fan out to user hooks.
+    class ExperimentTracker
+      attr_reader :events
+
+      def initialize(subscribers: [])
+        @subscribers = Array(subscribers)
+        @events = []
+      end
+
+      def with_subscriber(proc = nil, &block)
+        @subscribers << (proc || block)
+        self
+      end
+
+      def initialize_backends; end
+
+      def start_run; end
+
+      def log_metrics(metrics, step: nil)
+        entry = { metrics: symbolize_keys(metrics), step: step }
+        @events << entry
+
+        @subscribers.each do |subscriber|
+          subscriber.call(entry)
+        rescue StandardError => e
+          DSPy.log('gepa.experiment_tracker.error', error: e.message)
+        end
+      end
+
+      def end_run; end
+
+      def active?
+        !@events.empty?
+      end
+
+      def each_event(&block)
+        @events.each(&block)
+      end
+
+      private
+
+      def symbolize_keys(hash)
+        hash.each_with_object({}) do |(k, v), memo|
+          memo[k.to_sym] = v
+        end
+      end
+    end
+  end
+end
+

data/lib/gepa/logging/logger.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require 'forwardable'
+
+module GEPA
+  module Logging
+    # Minimal logger interface used across GEPA components.
+    class Logger
+      extend Forwardable
+
+      def initialize(io: $stdout)
+        @io = io
+      end
+
+      def log(message)
+        write(message)
+      end
+
+      private
+
+      attr_reader :io
+
+      def write(message)
+        io.puts(message)
+        io.flush if io.respond_to?(:flush)
+      end
+    end
+
+    # Logger that fans out messages to multiple IO streams.
+    class CompositeLogger < Logger
+      def initialize(*ios)
+        @ios = ios.flatten
+      end
+
+      def log(message)
+        @ios.each do |io|
+          io.puts(message)
+          io.flush if io.respond_to?(:flush)
+        end
+      end
+    end
+
+    # Logger that captures messages into memory (handy for tests).
+    class BufferingLogger < Logger
+      attr_reader :messages
+
+      def initialize
+        @messages = []
+      end
+
+      def log(message)
+        @messages << message
+      end
+    end
+  end
+end
+

data/lib/gepa/logging.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative 'logging/logger'
+require_relative 'logging/experiment_tracker'
+
+module GEPA
+  module Logging
+  end
+end