mantisdk 0.1.0__py3-none-any.whl

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

@@ -0,0 +1,176 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from mantisdk.algorithm.gepa.lib.core.adapter import DataInst, GEPAAdapter, RolloutOutput, Trajectory
+from mantisdk.algorithm.gepa.lib.core.data_loader import DataId, DataLoader, ensure_loader
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState
+from mantisdk.algorithm.gepa.lib.proposer.base import CandidateProposal, ProposeNewCandidate
+from mantisdk.algorithm.gepa.lib.proposer.reflective_mutation.base import (
+    CandidateSelector,
+    LanguageModel,
+    ReflectionComponentSelector,
+)
+from mantisdk.algorithm.gepa.lib.strategies.batch_sampler import BatchSampler
+from mantisdk.algorithm.gepa.lib.strategies.instruction_proposal import InstructionProposalSignature
+
+
+class ReflectiveMutationProposer(ProposeNewCandidate[DataId]):
+    """
+    Implements current reflective mutation flow:
+    - Select candidate via selector
+    - Select minibatch via sampler
+    - capture_traces_and_eval -> trajectories, subsample_scores
+    - skip if all scores==perfect and skip_perfect_score
+    - reflection + mutate -> new candidate
+    - evaluate new candidate on same minibatch -> new_subsample_scores
+    - Return proposal if improved; else None
+    """
+
+    def __init__(
+        self,
+        logger: Any,
+        trainset: list[DataInst] | DataLoader[DataId, DataInst],
+        adapter: GEPAAdapter[DataInst, Trajectory, RolloutOutput],
+        candidate_selector: CandidateSelector,
+        module_selector: ReflectionComponentSelector,
+        batch_sampler: BatchSampler[DataId, DataInst],
+        perfect_score: float,
+        skip_perfect_score: bool,
+        experiment_tracker: Any,
+        reflection_lm: LanguageModel | None = None,
+        reflection_prompt_template: str | None = None,
+    ):
+        self.logger = logger
+        self.trainset = ensure_loader(trainset)
+        self.adapter = adapter
+        self.candidate_selector = candidate_selector
+        self.module_selector = module_selector
+        self.batch_sampler = batch_sampler
+        self.perfect_score = perfect_score
+        self.skip_perfect_score = skip_perfect_score
+        self.experiment_tracker = experiment_tracker
+        self.reflection_lm = reflection_lm
+
+        InstructionProposalSignature.validate_prompt_template(reflection_prompt_template)
+        self.reflection_prompt_template = reflection_prompt_template
+
+    def propose_new_texts(
+        self,
+        candidate: dict[str, str],
+        reflective_dataset: Mapping[str, Sequence[Mapping[str, Any]]],
+        components_to_update: list[str],
+    ) -> dict[str, str]:
+        if self.adapter.propose_new_texts is not None:
+            return self.adapter.propose_new_texts(candidate, reflective_dataset, components_to_update)
+
+        if self.reflection_lm is None:
+            raise ValueError("reflection_lm must be provided when adapter.propose_new_texts is None.")
+        new_texts: dict[str, str] = {}
+        for name in components_to_update:
+            # Gracefully handle cases where a selected component has no data in reflective_dataset
+            if name not in reflective_dataset or not reflective_dataset.get(name):
+                self.logger.log(f"Component '{name}' is not in reflective dataset. Skipping.")
+                continue
+
+            base_instruction = candidate[name]
+            dataset_with_feedback = reflective_dataset[name]
+            new_texts[name] = InstructionProposalSignature.run(
+                lm=self.reflection_lm,
+                input_dict={
+                    "current_instruction_doc": base_instruction,
+                    "dataset_with_feedback": dataset_with_feedback,
+                    "prompt_template": self.reflection_prompt_template,
+                },
+            )["new_instruction"]
+        return new_texts
+
+    def propose(self, state: GEPAState) -> CandidateProposal | None:
+        i = state.i + 1
+
+        curr_prog_id = self.candidate_selector.select_candidate_idx(state)
+        curr_prog = state.program_candidates[curr_prog_id]
+        state.full_program_trace[-1]["selected_program_candidate"] = curr_prog_id
+        self.logger.log(
+            f"Iteration {i}: Selected program {curr_prog_id} score: {state.program_full_scores_val_set[curr_prog_id]}"
+        )
+
+        self.experiment_tracker.log_metrics({"iteration": i, "selected_program_candidate": curr_prog_id}, step=i)
+
+        subsample_ids = self.batch_sampler.next_minibatch_ids(self.trainset, state)
+        state.full_program_trace[-1]["subsample_ids"] = subsample_ids
+        minibatch = self.trainset.fetch(subsample_ids)
+
+        # 1) Evaluate current program with traces
+        # Note: We don't use cache for capture_traces=True evaluations since we need fresh traces for reflection
+        eval_curr = self.adapter.evaluate(minibatch, curr_prog, capture_traces=True)
+        state.total_num_evals += len(subsample_ids)
+        state.full_program_trace[-1]["subsample_scores"] = eval_curr.scores
+
+        # Update cache with current program evaluation results (for future reuse when capture_traces=False)
+        if state.evaluation_cache is not None:
+            objective_scores_list = list(eval_curr.objective_scores) if eval_curr.objective_scores else None
+            state.evaluation_cache.put_batch(
+                curr_prog, subsample_ids, eval_curr.outputs, eval_curr.scores, objective_scores_list
+            )
+
+        if not eval_curr.trajectories or len(eval_curr.trajectories) == 0:
+            self.logger.log(f"Iteration {i}: No trajectories captured. Skipping.")
+            return None
+
+        if self.skip_perfect_score and all(s >= self.perfect_score for s in eval_curr.scores):
+            self.logger.log(f"Iteration {i}: All subsample scores perfect. Skipping.")
+            return None
+
+        self.experiment_tracker.log_metrics({"subsample_score": sum(eval_curr.scores)}, step=i)
+
+        # 2) Decide which predictors to update
+        predictor_names_to_update = self.module_selector(
+            state, eval_curr.trajectories, eval_curr.scores, curr_prog_id, curr_prog
+        )
+
+        # 3) Build reflective dataset and propose texts
+        try:
+            reflective_dataset = self.adapter.make_reflective_dataset(curr_prog, eval_curr, predictor_names_to_update)
+            new_texts = self.propose_new_texts(curr_prog, reflective_dataset, predictor_names_to_update)
+            for pname, text in new_texts.items():
+                self.logger.log(f"Iteration {i}: Proposed new text for {pname}: {text}")
+            self.experiment_tracker.log_metrics(
+                {f"new_instruction_{pname}": text for pname, text in new_texts.items()}, step=i
+            )
+        except Exception as e:
+            self.logger.log(f"Iteration {i}: Exception during reflection/proposal: {e}")
+            import traceback
+
+            self.logger.log(traceback.format_exc())
+            return None
+
+        # 4) Create candidate, evaluate on same minibatch (no need to capture traces)
+        new_candidate = curr_prog.copy()
+        for pname, text in new_texts.items():
+            assert pname in new_candidate, f"{pname} missing in candidate"
+            new_candidate[pname] = text
+
+        def evaluator(b, c):
+            r = self.adapter.evaluate(b, c, capture_traces=False)
+            return r.outputs, r.scores, list(r.objective_scores) if r.objective_scores else None
+
+        new_scores, actual_evals_count = state.cached_evaluate(
+            new_candidate, subsample_ids, self.trainset.fetch, evaluator
+        )
+        state.total_num_evals += actual_evals_count
+        state.full_program_trace[-1]["new_subsample_scores"] = new_scores
+
+        new_sum = sum(new_scores)
+        self.experiment_tracker.log_metrics({"new_subsample_score": new_sum}, step=i)
+
+        return CandidateProposal(
+            candidate=new_candidate,
+            parent_program_ids=[curr_prog_id],
+            subsample_indices=subsample_ids,
+            subsample_scores_before=eval_curr.scores,
+            subsample_scores_after=new_scores,
+            tag="reflective_mutation",
+        )

mantisdk/algorithm/gepa/lib/strategies/batch_sampler.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+import random
+from collections import Counter
+from typing import Protocol
+
+from mantisdk.algorithm.gepa.lib.core.adapter import DataInst
+from mantisdk.algorithm.gepa.lib.core.data_loader import DataId, DataLoader
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState
+
+
+class BatchSampler(Protocol[DataId, DataInst]):
+    def next_minibatch_ids(self, loader: DataLoader[DataId, DataInst], state: GEPAState) -> list[DataId]: ...
+
+
+class EpochShuffledBatchSampler(BatchSampler[DataId, DataInst]):
+    """
+    Mirrors the original batching logic:
+    - Shuffle ids each epoch
+    - Pad to minibatch size with least frequent ids
+    - Deterministic via state.rng1
+    """
+
+    def __init__(self, minibatch_size: int, rng: random.Random | None = None):
+        self.minibatch_size = minibatch_size
+        self.shuffled_ids: list[DataId] = []
+        self.epoch = -1
+        self.id_freqs = Counter()
+        self.last_trainset_size = 0
+        if rng is None:
+            self.rng = random.Random(0)
+        else:
+            self.rng = rng
+
+    def _update_shuffled(self, loader: DataLoader[DataId, DataInst]):
+        all_ids = list(loader.all_ids())
+        trainset_size = len(loader)
+        self.last_trainset_size = trainset_size
+
+        if trainset_size == 0:
+            self.shuffled_ids = []
+            self.id_freqs = Counter()
+            return
+
+        self.shuffled_ids = list(all_ids)
+        self.rng.shuffle(self.shuffled_ids)
+        self.id_freqs = Counter(self.shuffled_ids)
+
+        mod = trainset_size % self.minibatch_size
+        num_to_pad = (self.minibatch_size - mod) if mod != 0 else 0
+        if num_to_pad > 0:
+            for _ in range(num_to_pad):
+                selected_id = self.id_freqs.most_common()[::-1][0][0]
+                self.shuffled_ids.append(selected_id)
+                self.id_freqs[selected_id] += 1
+
+    def next_minibatch_ids(self, loader: DataLoader[DataId, DataInst], state: GEPAState) -> list[DataId]:
+        trainset_size = len(loader)
+        if trainset_size == 0:
+            raise ValueError("Cannot sample a minibatch from an empty loader.")
+
+        base_idx = state.i * self.minibatch_size
+        curr_epoch = 0 if self.epoch == -1 else base_idx // max(len(self.shuffled_ids), 1)
+
+        needs_refresh = not self.shuffled_ids or trainset_size != self.last_trainset_size or curr_epoch > self.epoch
+        if needs_refresh:
+            self.epoch = curr_epoch
+            self._update_shuffled(loader)
+
+        assert len(self.shuffled_ids) >= self.minibatch_size
+        assert len(self.shuffled_ids) % self.minibatch_size == 0
+
+        base_idx = base_idx % len(self.shuffled_ids)
+        end_idx = base_idx + self.minibatch_size
+        assert end_idx <= len(self.shuffled_ids)
+        return self.shuffled_ids[base_idx:end_idx]

mantisdk/algorithm/gepa/lib/strategies/candidate_selector.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+import random
+
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState
+from mantisdk.algorithm.gepa.lib.gepa_utils import idxmax, select_program_candidate_from_pareto_front
+from mantisdk.algorithm.gepa.lib.proposer.reflective_mutation.base import CandidateSelector
+
+
+class ParetoCandidateSelector(CandidateSelector):
+    def __init__(self, rng: random.Random | None):
+        if rng is None:
+            self.rng = random.Random(0)
+        else:
+            self.rng = rng
+
+    def select_candidate_idx(self, state: GEPAState) -> int:
+        assert len(state.program_full_scores_val_set) == len(state.program_candidates)
+        return select_program_candidate_from_pareto_front(
+            state.get_pareto_front_mapping(),
+            state.per_program_tracked_scores,
+            self.rng,
+        )
+
+
+class CurrentBestCandidateSelector(CandidateSelector):
+    def __init__(self):
+        pass
+
+    def select_candidate_idx(self, state: GEPAState) -> int:
+        assert len(state.program_full_scores_val_set) == len(state.program_candidates)
+        return idxmax(state.program_full_scores_val_set)
+
+
+class EpsilonGreedyCandidateSelector(CandidateSelector):
+    def __init__(self, epsilon: float, rng: random.Random | None):
+        assert 0.0 <= epsilon <= 1.0
+        self.epsilon = epsilon
+        if rng is None:
+            self.rng = random.Random(0)
+        else:
+            self.rng = rng
+
+    def select_candidate_idx(self, state: GEPAState) -> int:
+        assert len(state.program_full_scores_val_set) == len(state.program_candidates)
+        if self.rng.random() < self.epsilon:
+            return self.rng.randint(0, len(state.program_candidates) - 1)
+        else:
+            return idxmax(state.program_full_scores_val_set)

mantisdk/algorithm/gepa/lib/strategies/component_selector.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+
+from mantisdk.algorithm.gepa.lib.core.adapter import Trajectory
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState
+from mantisdk.algorithm.gepa.lib.proposer.reflective_mutation.base import ReflectionComponentSelector
+
+
+class RoundRobinReflectionComponentSelector(ReflectionComponentSelector):
+    def __call__(
+        self,
+        state: GEPAState,
+        trajectories: list[Trajectory],
+        subsample_scores: list[float],
+        candidate_idx: int,
+        candidate: dict[str, str],
+    ) -> list[str]:
+        pid = state.named_predictor_id_to_update_next_for_program_candidate[candidate_idx]
+        state.named_predictor_id_to_update_next_for_program_candidate[candidate_idx] = (pid + 1) % len(
+            state.list_of_named_predictors
+        )
+        name = state.list_of_named_predictors[pid]
+        return [name]
+
+
+class AllReflectionComponentSelector(ReflectionComponentSelector):
+    def __call__(
+        self,
+        state: GEPAState,
+        trajectories: list[Trajectory],
+        subsample_scores: list[float],
+        candidate_idx: int,
+        candidate: dict[str, str],
+    ) -> list[str]:
+        return list(candidate.keys())

mantisdk/algorithm/gepa/lib/strategies/eval_policy.py

@@ -0,0 +1,64 @@
+"""Validation evaluation policy protocols and helpers."""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Protocol, runtime_checkable
+
+from mantisdk.algorithm.gepa.lib.core.data_loader import DataId, DataInst, DataLoader
+from mantisdk.algorithm.gepa.lib.core.state import GEPAState, ProgramIdx
+
+
+@runtime_checkable
+class EvaluationPolicy(Protocol[DataId, DataInst]):  # type: ignore
+    """Strategy for choosing validation ids to evaluate and identifying best programs for validation instances."""
+
+    @abstractmethod
+    def get_eval_batch(
+        self, loader: DataLoader[DataId, DataInst], state: GEPAState, target_program_idx: ProgramIdx | None = None
+    ) -> list[DataId]:
+        """Select examples for evaluation for a program"""
+        ...
+
+    @abstractmethod
+    def get_best_program(self, state: GEPAState) -> ProgramIdx:
+        """Return "best" program given all validation results so far across candidates"""
+        ...
+
+    @abstractmethod
+    def get_valset_score(self, program_idx: ProgramIdx, state: GEPAState) -> float:
+        """Return the score of the program on the valset"""
+        ...
+
+
+class FullEvaluationPolicy(EvaluationPolicy[DataId, DataInst]):
+    """Policy that evaluates all validation instances every time."""
+
+    def get_eval_batch(
+        self, loader: DataLoader[DataId, DataInst], state: GEPAState, target_program_idx: ProgramIdx | None = None
+    ) -> list[DataId]:
+        """Always return the full ordered list of validation ids."""
+        return list(loader.all_ids())
+
+    def get_best_program(self, state: GEPAState) -> ProgramIdx:
+        """Pick the program whose evaluated validation scores achieve the highest average."""
+        best_idx, best_score, best_coverage = -1, float("-inf"), -1
+        for program_idx, scores in enumerate(state.prog_candidate_val_subscores):
+            coverage = len(scores)
+            avg = sum(scores.values()) / coverage if coverage else float("-inf")
+            if avg > best_score or (avg == best_score and coverage > best_coverage):
+                best_score = avg
+                best_idx = program_idx
+                best_coverage = coverage
+        return best_idx
+
+    def get_valset_score(self, program_idx: ProgramIdx, state: GEPAState) -> float:
+        """Return the score of the program on the valset"""
+        return state.get_program_average_val_subset(program_idx)[0]
+
+
+__all__ = [
+    "DataLoader",
+    "EvaluationPolicy",
+    "FullEvaluationPolicy",
+]

mantisdk/algorithm/gepa/lib/strategies/instruction_proposal.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+import re
+from collections.abc import Mapping, Sequence
+from typing import Any, ClassVar
+
+from mantisdk.algorithm.gepa.lib.proposer.reflective_mutation.base import Signature
+
+
+class InstructionProposalSignature(Signature):
+    default_prompt_template = """I provided an assistant with the following instructions to perform a task for me:
+```
+<curr_instructions>
+```
+
+The following are examples of different task inputs provided to the assistant along with the assistant's response for each of them, and some feedback on how the assistant's response could be better:
+```
+<inputs_outputs_feedback>
+```
+
+Your task is to write a new instruction for the assistant.
+
+Read the inputs carefully and identify the input format and infer detailed task description about the task I wish to solve with the assistant.
+
+Read all the assistant responses and the corresponding feedback. Identify all niche and domain specific factual information about the task and include it in the instruction, as a lot of it may not be available to the assistant in the future. The assistant may have utilized a generalizable strategy to solve the task, if so, include that in the instruction as well.
+
+Provide the new instructions within ``` blocks."""
+
+    input_keys: ClassVar[list[str]] = ["current_instruction_doc", "dataset_with_feedback", "prompt_template"]
+    output_keys: ClassVar[list[str]] = ["new_instruction"]
+
+    @classmethod
+    def validate_prompt_template(cls, prompt_template: str | None) -> None:
+        if prompt_template is None:
+            return
+        missing_placeholders = [
+            placeholder
+            for placeholder in ("<curr_instructions>", "<inputs_outputs_feedback>")
+            if placeholder not in prompt_template
+        ]
+        if missing_placeholders:
+            raise ValueError(
+                f"Missing placeholder(s) in prompt template: {', '.join(missing_placeholders)}"
+            )
+
+    @classmethod
+    def prompt_renderer(cls, input_dict: Mapping[str, Any]) -> str:
+        current_instruction = input_dict.get("current_instruction_doc")
+        if not isinstance(current_instruction, str):
+            raise TypeError("current_instruction_doc must be a string")
+
+        dataset = input_dict.get("dataset_with_feedback")
+        if not isinstance(dataset, Sequence) or isinstance(dataset, (str, bytes)):
+            raise TypeError("dataset_with_feedback must be a sequence of records")
+        def format_samples(samples):
+            def render_value(value, level=3):
+                # level controls markdown header depth (###, ####, etc.)
+                if isinstance(value, dict):
+                    s = ""
+                    for k, v in value.items():
+                        s += f"{'#' * level} {k}\n"
+                        s += render_value(v, min(level + 1, 6))
+                    if not value:
+                        s += "\n"
+                    return s
+                elif isinstance(value, list | tuple):
+                    s = ""
+                    for i, item in enumerate(value):
+                        s += f"{'#' * level} Item {i + 1}\n"
+                        s += render_value(item, min(level + 1, 6))
+                    if not value:
+                        s += "\n"
+                    return s
+                else:
+                    return f"{str(value).strip()}\n\n"
+
+            def convert_sample_to_markdown(sample, examplenum):
+                s = f"# Example {examplenum}\n"
+                for key, val in sample.items():
+                    s += f"## {key}\n"
+                    s += render_value(val, level=3)
+                return s
+
+            return "\n\n".join(convert_sample_to_markdown(sample, i + 1) for i, sample in enumerate(samples))
+
+        prompt_template = input_dict.get("prompt_template")
+        if prompt_template is None:
+            prompt_template = cls.default_prompt_template
+
+        cls.validate_prompt_template(prompt_template)
+
+        prompt = prompt_template.replace("<curr_instructions>", current_instruction)
+        prompt = prompt.replace("<inputs_outputs_feedback>", format_samples(dataset))
+
+        return prompt
+
+    @classmethod
+    def output_extractor(cls, lm_out: str) -> dict[str, str]:
+        def extract_instruction_text() -> str:
+            # Find the first and last backtick positions (if any)
+            start = lm_out.find("```") + 3
+            end = lm_out.rfind("```")
+
+            # Handle if the first and last backticks are the same or overlap
+            if start >= end:
+                # Handle incomplete blocks
+                stripped = lm_out.strip()
+                if stripped.startswith("```"):
+                    # Remove opening ``` and optional language specifier
+                    match = re.match(r"^```\S*\n?", lm_out)
+                    if match:
+                        return lm_out[match.end() :].strip()
+                elif stripped.endswith("```"):
+                    # Remove closing ```
+                    return stripped[:-3].strip()
+                return stripped
+
+            # Skip optional language specifier
+            content = lm_out[start:end]
+            match = re.match(r"^\S*\n", content)
+            if match:
+                content = content[match.end() :]
+
+            return content.strip()
+
+        return {"new_instruction": extract_instruction_text()}

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

@@ -0,0 +1,10 @@
+from .stop_condition import (
+    CompositeStopper,
+    FileStopper,
+    MaxMetricCallsStopper,
+    NoImprovementStopper,
+    ScoreThresholdStopper,
+    SignalStopper,
+    StopperProtocol,
+    TimeoutStopCondition,
+)