mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/logging/utils.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+
+from mantisdk.algorithm.gepa.lib.core.adapter import DataInst
+from mantisdk.algorithm.gepa.lib.core.data_loader import DataId
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState, ValsetEvaluation
+from mantisdk.algorithm.gepa.lib.strategies.eval_policy import EvaluationPolicy
+
+
+def log_detailed_metrics_after_discovering_new_program(
+    logger,
+    gepa_state: GEPAState,
+    new_program_idx,
+    valset_evaluation: ValsetEvaluation,
+    objective_scores,
+    experiment_tracker,
+    linear_pareto_front_program_idx,
+    valset_size: int,
+    val_evaluation_policy: EvaluationPolicy[DataId, DataInst],
+    log_individual_valset_scores_and_programs: bool = False,
+):
+    # best_prog_per_agg_val_score = idxmax(gepa_state.program_full_scores_val_set)
+    best_prog_per_agg_val_score = val_evaluation_policy.get_best_program(gepa_state)
+    best_score_on_valset = val_evaluation_policy.get_valset_score(best_prog_per_agg_val_score, gepa_state)
+
+    # avg, coverage = gepa_state.get_program_average_val_subset(new_program_idx)
+    valset_score = val_evaluation_policy.get_valset_score(new_program_idx, gepa_state)
+    valset_scores = valset_evaluation.scores_by_val_id
+    coverage = len(valset_scores)
+    logger.log(
+        f"Iteration {gepa_state.i + 1}: Valset score for new program: {valset_score}"
+        f" (coverage {coverage} / {valset_size})"
+    )
+
+    agg_valset_score_new_program = val_evaluation_policy.get_valset_score(new_program_idx, gepa_state)
+
+    logger.log(f"Iteration {gepa_state.i + 1}: Val aggregate for new program: {agg_valset_score_new_program}")
+    logger.log(f"Iteration {gepa_state.i + 1}: Individual valset scores for new program: {valset_scores}")
+    if objective_scores:
+        logger.log(f"Iteration {gepa_state.i + 1}: Objective aggregate scores for new program: {objective_scores}")
+    logger.log(f"Iteration {gepa_state.i + 1}: New valset pareto front scores: {gepa_state.pareto_front_valset}")
+    if gepa_state.objective_pareto_front:
+        logger.log(f"Iteration {gepa_state.i + 1}: Objective pareto front scores: {gepa_state.objective_pareto_front}")
+
+    pareto_scores = list(gepa_state.pareto_front_valset.values())
+    assert all(score > float("-inf") for score in pareto_scores), (
+        "Should have at least one valid score per validation example"
+    )
+    assert len(pareto_scores) > 0
+    pareto_avg = sum(pareto_scores) / len(pareto_scores)
+
+    logger.log(f"Iteration {gepa_state.i + 1}: Valset pareto front aggregate score: {pareto_avg}")
+    logger.log(
+        f"Iteration {gepa_state.i + 1}: Updated valset pareto front programs: {gepa_state.program_at_pareto_front_valset}"
+    )
+    if gepa_state.program_at_pareto_front_objectives:
+        logger.log(
+            f"Iteration {gepa_state.i + 1}: Updated objective pareto front programs: {gepa_state.program_at_pareto_front_objectives}"
+        )
+    logger.log(
+        f"Iteration {gepa_state.i + 1}: Best valset aggregate score so far: {max(gepa_state.program_full_scores_val_set)}"
+    )
+    logger.log(
+        f"Iteration {gepa_state.i + 1}: Best program as per aggregate score on valset: {best_prog_per_agg_val_score}"
+    )
+    logger.log(f"Iteration {gepa_state.i + 1}: Best score on valset: {best_score_on_valset}")
+    logger.log(f"Iteration {gepa_state.i + 1}: Linear pareto front program index: {linear_pareto_front_program_idx}")
+    logger.log(f"Iteration {gepa_state.i + 1}: New program candidate index: {new_program_idx}")
+
+    metrics = {
+        "iteration": gepa_state.i + 1,
+        "new_program_idx": new_program_idx,
+        "valset_pareto_front_agg": pareto_avg,
+        "valset_pareto_front_programs": {k: list(v) for k, v in gepa_state.program_at_pareto_front_valset.items()},
+        "best_valset_agg_score": best_score_on_valset,
+        "linear_pareto_front_program_idx": linear_pareto_front_program_idx,
+        "best_program_as_per_agg_score_valset": best_prog_per_agg_val_score,
+        "best_score_on_valset": best_score_on_valset,
+        "val_evaluated_count_new_program": coverage,
+        "val_total_count": valset_size,
+        "val_program_average": valset_score,
+    }
+    if log_individual_valset_scores_and_programs:
+        metrics.update(
+            {
+                "valset_pareto_front_scores": dict(gepa_state.pareto_front_valset),
+                "individual_valset_score_new_program": dict(valset_scores),
+            }
+        )
+    if objective_scores:
+        metrics["objective_scores_new_program"] = dict(objective_scores)
+    if valset_evaluation.objective_scores_by_val_id:
+        metrics["objective_scores_by_val_new_program"] = {
+            val_id: dict(scores) for val_id, scores in valset_evaluation.objective_scores_by_val_id.items()
+        }
+    if gepa_state.objective_pareto_front:
+        metrics["objective_pareto_front_scores"] = dict(gepa_state.objective_pareto_front)
+        metrics["objective_pareto_front_programs"] = {
+            k: list(v) for k, v in gepa_state.program_at_pareto_front_objectives.items()
+        }
+
+    experiment_tracker.log_metrics(metrics, step=gepa_state.i + 1)

mantisdk/algorithm/gepa/lib/proposer/base.py

@@ -0,0 +1,31 @@
+# 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 mantisdk.algorithm.gepa.lib.core.data_loader import DataId
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState
+
+
+@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
+    # 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: ...

mantisdk/algorithm/gepa/lib/proposer/merge.py

@@ -0,0 +1,357 @@
+# 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 mantisdk.algorithm.gepa.lib.core.adapter import Candidate, DataInst, RolloutOutput
+from mantisdk.algorithm.gepa.lib.core.data_loader import DataId, DataLoader
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState, ObjectiveScores, ProgramIdx
+from mantisdk.algorithm.gepa.lib.gepa_utils import find_dominator_programs
+from mantisdk.algorithm.gepa.lib.logging.logger import LoggerProtocol
+from mantisdk.algorithm.gepa.lib.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,
+    ):
+        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)
+
+        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
+
+        new_sub_scores, actual_evals_count = state.cached_evaluate(
+            new_program, subsample_ids, self.valset.fetch, self.evaluator
+        )
+        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
+        state.total_num_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},
+        )

mantisdk/algorithm/gepa/lib/proposer/reflective_mutation/base.py

@@ -0,0 +1,49 @@
+# 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 mantisdk.algorithm.gepa.lib.core.adapter import Trajectory
+from mantisdk.algorithm.gepa.lib.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) -> 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:
+        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_out = lm(full_prompt).strip()
+        return cls.output_extractor(lm_out)