orizu 0.5.19 → 0.5.21

package/vendor/gepa-python/src/gepa/proposer/base.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from dataclasses import dataclass, field
+from typing import Any, Generic, Protocol
+
+from gepa.core.data_loader import DataId
+from gepa.core.state import GEPAState
+
+
+@dataclass
+class SubsampleEvaluation:
+    """Result of evaluating a candidate on the subsample minibatch.
+
+    Captures all outputs from ``adapter.evaluate``:
+
+    - ``scores``: per-example numeric scores.
+    - ``outputs``: per-example raw outputs (e.g. the model's response text).
+    - ``objective_scores``: optional per-example multi-objective score dicts.
+    - ``trajectories``: optional per-example execution traces (captured for both
+      before and after evaluations).
+    """
+
+    scores: list[float]
+    outputs: list[Any] = field(default_factory=list)
+    objective_scores: list[dict[str, float]] | None = None
+    trajectories: list[Any] | None = None
+
+
+@dataclass
+class CandidateProposal(Generic[DataId]):
+    candidate: dict[str, str]
+    parent_program_ids: list[int]
+    # Optional mini-batch / subsample info
+    subsample_indices: list[DataId] | None = None
+    subsample_scores_before: list[float] | None = None
+    subsample_scores_after: list[float] | None = None
+    # Rich evaluation data (superset of scores — includes outputs, objective_scores, trajectories)
+    eval_before: SubsampleEvaluation | None = None
+    eval_after: SubsampleEvaluation | None = None
+    # Free-form metadata for logging/trace
+    tag: str = ""
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class ProposeNewCandidate(Protocol[DataId]):
+    """
+    Strategy that receives the current optimizer state and proposes a new candidate or returns None.
+    It may compute subsample evaluations, set trace fields in state, etc.
+    The engine will handle acceptance and full eval unless the strategy already did those and encoded in metadata.
+    """
+
+    def propose(self, state: GEPAState[Any, DataId]) -> CandidateProposal | None: ...

package/vendor/gepa-python/src/gepa/proposer/merge.py

@@ -0,0 +1,403 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+import math
+import random
+from collections.abc import Callable, Iterable, Sequence
+from copy import deepcopy
+
+from gepa.core.adapter import Candidate, DataInst, RolloutOutput
+from gepa.core.callbacks import (
+    EvaluationEndEvent,
+    EvaluationStartEvent,
+    GEPACallback,
+    notify_callbacks,
+)
+from gepa.core.data_loader import DataId, DataLoader
+from gepa.core.state import GEPAState, ObjectiveScores, ProgramIdx
+from gepa.gepa_utils import find_dominator_programs
+from gepa.logging.logger import LoggerProtocol
+from gepa.proposer.base import CandidateProposal, ProposeNewCandidate
+
+AncestorLog = tuple[int, int, int]
+MergeDescription = tuple[int, int, tuple[int, ...]]
+MergeAttempt = tuple[Candidate, ProgramIdx, ProgramIdx, ProgramIdx] | None
+
+
+def does_triplet_have_desirable_predictors(
+    program_candidates: Sequence[Candidate],
+    ancestor: ProgramIdx,
+    id1: ProgramIdx,
+    id2: ProgramIdx,
+) -> bool:
+    found_predictors: list[tuple[int, int]] = []
+    pred_names = list(program_candidates[ancestor].keys())
+    for pred_idx, pred_name in enumerate(pred_names):
+        pred_anc = program_candidates[ancestor][pred_name]
+        pred_id1 = program_candidates[id1][pred_name]
+        pred_id2 = program_candidates[id2][pred_name]
+        if (pred_anc == pred_id1 or pred_anc == pred_id2) and pred_id1 != pred_id2:
+            same_as_ancestor_id = 1 if pred_anc == pred_id1 else 2
+            found_predictors.append((pred_idx, same_as_ancestor_id))
+
+    return len(found_predictors) > 0
+
+
+def filter_ancestors(
+    i: ProgramIdx,
+    j: ProgramIdx,
+    common_ancestors: Iterable[ProgramIdx],
+    merges_performed: tuple[list[AncestorLog], list[MergeDescription]],
+    agg_scores: Sequence[float],
+    program_candidates: Sequence[Candidate],
+) -> list[ProgramIdx]:
+    filtered_ancestors: list[ProgramIdx] = []
+    for ancestor in common_ancestors:
+        if (i, j, ancestor) in merges_performed[0]:
+            continue
+
+        if agg_scores[ancestor] > agg_scores[i] or agg_scores[ancestor] > agg_scores[j]:
+            continue
+
+        if not does_triplet_have_desirable_predictors(program_candidates, ancestor, i, j):
+            continue
+
+        filtered_ancestors.append(ancestor)
+    return filtered_ancestors
+
+
+def find_common_ancestor_pair(
+    rng: random.Random,
+    parent_list: Sequence[Sequence[int | None]],
+    program_indexes: Sequence[int],
+    merges_performed: tuple[list[AncestorLog], list[MergeDescription]],
+    agg_scores: Sequence[float],
+    program_candidates: Sequence[Candidate],
+    max_attempts: int = 10,
+) -> tuple[int, int, int] | None:
+    def get_ancestors(node: int, ancestors_found: set[int]) -> list[int]:
+        parents = parent_list[node]
+        for parent in parents:
+            if parent is not None and parent not in ancestors_found:
+                ancestors_found.add(parent)
+                get_ancestors(parent, ancestors_found)
+
+        return list(ancestors_found)
+
+    for _ in range(max_attempts):
+        if len(program_indexes) < 2:
+            return None
+        i, j = rng.sample(list(program_indexes), 2)
+        if i == j:
+            continue
+
+        if j < i:
+            i, j = j, i
+
+        ancestors_i = get_ancestors(i, set())
+        ancestors_j = get_ancestors(j, set())
+
+        if j in ancestors_i or i in ancestors_j:
+            # If one is an ancestor of the other, we cannot merge them
+            continue
+
+        common_ancestors = set(ancestors_i) & set(ancestors_j)
+        common_ancestors = filter_ancestors(i, j, common_ancestors, merges_performed, agg_scores, program_candidates)
+        if common_ancestors:
+            # Select a random common ancestor
+            common_ancestor = rng.choices(
+                list(common_ancestors),
+                k=1,
+                weights=[agg_scores[ancestor] for ancestor in common_ancestors],
+            )[0]
+            return (i, j, common_ancestor)
+
+    return None
+
+
+def sample_and_attempt_merge_programs_by_common_predictors(
+    agg_scores: Sequence[float],
+    rng: random.Random,
+    merge_candidates: Sequence[int],
+    merges_performed: tuple[list[AncestorLog], list[MergeDescription]],
+    program_candidates: Sequence[Candidate],
+    parent_program_for_candidate: Sequence[Sequence[int | None]],
+    has_val_support_overlap: Callable[[ProgramIdx, ProgramIdx], bool] | None = None,
+    max_attempts: int = 10,
+) -> MergeAttempt:
+    if len(merge_candidates) < 2:
+        return None
+    if len(parent_program_for_candidate) < 3:
+        return None
+
+    for _ in range(max_attempts):
+        ids_to_merge = find_common_ancestor_pair(
+            rng,
+            parent_program_for_candidate,
+            list(merge_candidates),
+            merges_performed=merges_performed,
+            agg_scores=agg_scores,
+            program_candidates=program_candidates,
+            max_attempts=max_attempts,
+        )
+        if ids_to_merge is None:
+            continue
+        id1, id2, ancestor = ids_to_merge
+
+        if (id1, id2, ancestor) in merges_performed[0]:
+            continue
+        assert agg_scores[ancestor] <= agg_scores[id1], "Ancestor should not be better than its descendants"
+        assert agg_scores[ancestor] <= agg_scores[id2], "Ancestor should not be better than its descendants"
+        assert id1 != id2, "Cannot merge the same program"
+
+        # Now we have a common ancestor, which is outperformed by both its descendants
+
+        new_program: Candidate = deepcopy(program_candidates[ancestor])
+
+        new_prog_desc: tuple[ProgramIdx, ...] = ()
+
+        pred_names = set(program_candidates[ancestor].keys())
+        assert pred_names == set(program_candidates[id1].keys()) == set(program_candidates[id2].keys()), (
+            "Predictors should be the same across all programs"
+        )
+        for pred_name in pred_names:
+            pred_anc = program_candidates[ancestor][pred_name]
+            pred_id1 = program_candidates[id1][pred_name]
+            pred_id2 = program_candidates[id2][pred_name]
+            if (pred_anc == pred_id1 or pred_anc == pred_id2) and pred_id1 != pred_id2:
+                # We have a predictor that is the same as one of its ancestors, so we can update it with the other
+                same_as_ancestor_id = 1 if pred_anc == pred_id1 else 2
+                new_value_idx = id2 if same_as_ancestor_id == 1 else id1
+                new_program[pred_name] = program_candidates[new_value_idx][pred_name]
+                new_prog_desc = (*new_prog_desc, new_value_idx)
+            elif pred_anc != pred_id1 and pred_anc != pred_id2:
+                # Both predictors are different from  the ancestor, and it is difficult to decide which one gives the benefits
+                # We randomly select one of the descendants to update the predictor
+                # The probability of selecting is proportional to the agg_scores of the descendants
+                # prog_to_get_instruction_from = id1 if (rng.random() < (agg_scores[id1] / (agg_scores[id1] + agg_scores[id2]))) else id2
+                prog_to_get_instruction_from = (
+                    id1
+                    if agg_scores[id1] > agg_scores[id2]
+                    else (id2 if agg_scores[id2] > agg_scores[id1] else rng.choice([id1, id2]))
+                )
+                new_program[pred_name] = program_candidates[prog_to_get_instruction_from][pred_name]
+                new_prog_desc = (*new_prog_desc, prog_to_get_instruction_from)
+            elif pred_id1 == pred_id2:
+                # Either both predictors are the same, or both are different from the ancestor
+                # If both are different from the ancestor, we should use the new predictor, so selecting either one of the descendants is fine
+                # If both are same as the ancesor, again selecting any one of the descendants is fine
+                # So let's select id1
+                new_program[pred_name] = program_candidates[id1][pred_name]
+                new_prog_desc = (*new_prog_desc, id1)
+            else:  # pragma: no cover - defensive
+                raise AssertionError("Unexpected case in predictor merging logic")
+
+        if (id1, id2, new_prog_desc) in merges_performed[1]:
+            # This triplet has already been merged, so we skip it
+            continue
+
+        if has_val_support_overlap and not has_val_support_overlap(id1, id2):
+            # Not enough overlapping validation support for candidates
+            continue
+
+        merges_performed[1].append((id1, id2, new_prog_desc))
+
+        return new_program, id1, id2, ancestor
+
+    return None
+
+
+class MergeProposer(ProposeNewCandidate[DataId]):
+    """
+    Implements merge flow that combines compatible descendants of a common ancestor.
+
+    - Find merge candidates among Pareto front dominators
+    - Attempt a merge via sample_and_attempt_merge_programs_by_common_predictors
+    - Subsample eval on valset-driven selected indices
+    - Return proposal if merge's subsample score >= max(parents)
+    The engine handles full eval + adding to state.
+    """
+
+    def __init__(
+        self,
+        logger: LoggerProtocol,
+        valset: DataLoader[DataId, DataInst],
+        evaluator: Callable[
+            [list[DataInst], dict[str, str]],
+            tuple[list[RolloutOutput], list[float], Sequence[ObjectiveScores] | None],
+        ],
+        use_merge: bool,
+        max_merge_invocations: int,
+        val_overlap_floor: int = 5,
+        rng: random.Random | None = None,
+        callbacks: list[GEPACallback] | None = None,
+    ):
+        self.logger = logger
+        self.valset = valset
+        self.evaluator = evaluator
+        self.use_merge = use_merge
+        self.max_merge_invocations = max_merge_invocations
+        self.rng = rng if rng is not None else random.Random(0)
+        self.callbacks = callbacks
+
+        if val_overlap_floor <= 0:
+            raise ValueError("val_overlap_floor should be a positive integer")
+        self.val_overlap_floor = val_overlap_floor
+        # Internal counters matching original behavior
+        self.merges_due = 0
+        self.total_merges_tested = 0
+        self.merges_performed: tuple[list[AncestorLog], list[MergeDescription]] = ([], [])
+
+        # Toggle controlled by engine: set True when last iter found new program
+        self.last_iter_found_new_program = False
+
+    def schedule_if_needed(self) -> None:
+        if self.use_merge and self.total_merges_tested < self.max_merge_invocations:
+            self.merges_due += 1
+
+    def select_eval_subsample_for_merged_program(
+        self,
+        scores1: dict[DataId, float],
+        scores2: dict[DataId, float],
+        num_subsample_ids: int = 5,
+    ) -> list[DataId]:
+        common_ids = list(set(scores1.keys()) & set(scores2.keys()))
+
+        p1 = [idx for idx in common_ids if scores1[idx] > scores2[idx]]
+        p2 = [idx for idx in common_ids if scores2[idx] > scores1[idx]]
+        p3 = [idx for idx in common_ids if idx not in p1 and idx not in p2]
+
+        n_each = max(1, math.ceil(num_subsample_ids / 3))
+        selected: list[DataId] = []
+        for bucket in (p1, p2, p3):
+            if len(selected) >= num_subsample_ids:
+                break
+            available = [idx for idx in bucket if idx not in selected]
+            take = min(len(available), n_each, num_subsample_ids - len(selected))
+            if take > 0:
+                selected += self.rng.sample(available, k=take)
+
+        remaining = num_subsample_ids - len(selected)
+        if remaining > 0:
+            unused = [idx for idx in common_ids if idx not in selected]
+            if len(unused) >= remaining:
+                selected += self.rng.sample(unused, k=remaining)
+            elif common_ids:
+                selected += self.rng.choices(common_ids, k=remaining)
+
+        return selected[:num_subsample_ids]
+
+    def propose(self, state: GEPAState[RolloutOutput, DataId]) -> CandidateProposal[DataId] | None:
+        i = state.i + 1
+        state.full_program_trace[-1]["invoked_merge"] = True
+
+        # Only attempt when scheduled by engine and after a new program in last iteration
+        if not (self.use_merge and self.last_iter_found_new_program and self.merges_due > 0):
+            self.logger.log(f"Iteration {i}: No merge candidates scheduled")
+            return None
+
+        pareto_front_programs = state.get_pareto_front_mapping()
+
+        tracked_scores: Sequence[float] = getattr(
+            state, "per_program_tracked_scores", state.program_full_scores_val_set
+        )
+        merge_candidates = find_dominator_programs(pareto_front_programs, list(tracked_scores))
+
+        def has_val_support_overlap(id1: ProgramIdx, id2: ProgramIdx) -> bool:
+            common_ids = set(state.prog_candidate_val_subscores[id1].keys()) & set(
+                state.prog_candidate_val_subscores[id2].keys()
+            )
+            return len(common_ids) >= self.val_overlap_floor
+
+        merge_output = sample_and_attempt_merge_programs_by_common_predictors(
+            agg_scores=list(tracked_scores),
+            rng=self.rng,
+            merge_candidates=merge_candidates,
+            merges_performed=self.merges_performed,
+            program_candidates=state.program_candidates,
+            parent_program_for_candidate=state.parent_program_for_candidate,
+            has_val_support_overlap=has_val_support_overlap,
+        )
+
+        if merge_output is None:
+            self.logger.log(f"Iteration {i}: No merge candidates found")
+            return None
+
+        new_program, id1, id2, ancestor = merge_output
+        state.full_program_trace[-1]["merged"] = True
+        state.full_program_trace[-1]["merged_entities"] = (id1, id2, ancestor)
+        self.merges_performed[0].append((id1, id2, ancestor))
+        self.logger.log(f"Iteration {i}: Merged programs {id1} and {id2} via ancestor {ancestor}")
+
+        subsample_ids = self.select_eval_subsample_for_merged_program(
+            state.prog_candidate_val_subscores[id1],
+            state.prog_candidate_val_subscores[id2],
+        )
+        if not subsample_ids:
+            self.logger.log(
+                f"Iteration {i}: Skipping merge of {id1} and {id2} due to insufficient overlapping val coverage"
+            )
+            return None
+
+        assert set(subsample_ids).issubset(state.prog_candidate_val_subscores[id1].keys())
+        assert set(subsample_ids).issubset(state.prog_candidate_val_subscores[id2].keys())
+        id1_sub_scores = [state.prog_candidate_val_subscores[id1][k] for k in subsample_ids]
+        id2_sub_scores = [state.prog_candidate_val_subscores[id2][k] for k in subsample_ids]
+        state.full_program_trace[-1]["subsample_ids"] = subsample_ids
+
+        mini_devset = self.valset.fetch(subsample_ids)
+
+        # Notify evaluation start for merged candidate
+        notify_callbacks(
+            self.callbacks,
+            "on_evaluation_start",
+            EvaluationStartEvent(
+                iteration=i,
+                candidate_idx=None,
+                batch_size=len(mini_devset),
+                capture_traces=False,
+                parent_ids=[id1, id2],
+                inputs=mini_devset,
+                is_seed_candidate=False,
+            ),
+        )
+
+        outputs_by_id, scores_by_id, objective_by_id, actual_evals_count = state.cached_evaluate_full(
+            new_program, subsample_ids, self.valset.fetch, self.evaluator
+        )
+        new_sub_scores = [scores_by_id[eid] for eid in subsample_ids]
+        outputs = [outputs_by_id[eid] for eid in subsample_ids]
+
+        notify_callbacks(
+            self.callbacks,
+            "on_evaluation_end",
+            EvaluationEndEvent(
+                iteration=i,
+                candidate_idx=None,
+                scores=new_sub_scores,
+                has_trajectories=False,
+                parent_ids=[id1, id2],
+                outputs=outputs,
+                trajectories=None,
+                objective_scores=[objective_by_id[eid] for eid in subsample_ids] if objective_by_id else None,
+                is_seed_candidate=False,
+            ),
+        )
+
+        state.full_program_trace[-1]["id1_subsample_scores"] = id1_sub_scores
+        state.full_program_trace[-1]["id2_subsample_scores"] = id2_sub_scores
+        state.full_program_trace[-1]["new_program_subsample_scores"] = new_sub_scores
+
+        # Count evals via hook mechanism
+        state.increment_evals(actual_evals_count)
+
+        # Acceptance will be evaluated by engine (>= max(parents))
+        return CandidateProposal(
+            candidate=new_program,
+            parent_program_ids=[id1, id2],
+            subsample_indices=subsample_ids,
+            subsample_scores_before=[sum(id1_sub_scores), sum(id2_sub_scores)],
+            subsample_scores_after=new_sub_scores,
+            tag="merge",
+            metadata={"ancestor": ancestor},
+        )

package/vendor/gepa-python/src/gepa/proposer/reflective_mutation/base.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from dataclasses import dataclass
+from typing import Any, ClassVar, Mapping, Protocol, runtime_checkable
+
+from gepa.core.adapter import Trajectory
+from gepa.core.state import GEPAState
+
+
+@runtime_checkable
+class CandidateSelector(Protocol):
+    def select_candidate_idx(self, state: GEPAState) -> int: ...
+
+
+class ReflectionComponentSelector(Protocol):
+    def __call__(
+        self,
+        state: GEPAState,
+        trajectories: list[Trajectory],
+        subsample_scores: list[float],
+        candidate_idx: int,
+        candidate: dict[str, str],
+    ) -> list[str]: ...
+
+
+class LanguageModel(Protocol):
+    def __call__(self, prompt: str | list[dict[str, Any]]) -> str: ...
+
+
+@dataclass
+class Signature:
+    prompt_template: ClassVar[str]
+    input_keys: ClassVar[list[str]]
+    output_keys: ClassVar[list[str]]
+
+    @classmethod
+    def prompt_renderer(cls, input_dict: Mapping[str, Any]) -> str | list[dict[str, Any]]:
+        raise NotImplementedError
+
+    @classmethod
+    def output_extractor(cls, lm_out: str) -> dict[str, str]:
+        raise NotImplementedError
+
+    @classmethod
+    def run(cls, lm: LanguageModel, input_dict: Mapping[str, Any]) -> dict[str, str]:
+        full_prompt = cls.prompt_renderer(input_dict)
+        lm_res = lm(full_prompt)
+        lm_out = lm_res.strip()
+        return cls.output_extractor(lm_out)
+
+    @classmethod
+    def run_with_metadata(
+        cls, lm: LanguageModel, input_dict: Mapping[str, Any]
+    ) -> tuple[dict[str, str], str | list[dict[str, Any]], str]:
+        """Like ``run()``, but also returns the rendered prompt and raw LM output.
+
+        Returns:
+            A tuple of (extracted_output, rendered_prompt, raw_lm_output).
+        """
+        full_prompt = cls.prompt_renderer(input_dict)
+        lm_res = lm(full_prompt)
+        lm_out = lm_res.strip()
+        return cls.output_extractor(lm_out), full_prompt, lm_out